From e6391a8b1aedcfa587d56e8ed335c99fbdf6482a Mon Sep 17 00:00:00 2001
From: Vincent Arel-Bundock <vincent.arel-bundock@umontreal.ca>
Date: Fri, 6 Dec 2024 06:01:56 -0500
Subject: [PATCH] nits

---
 DESCRIPTION     |  2 +-
 NEWS.md         |  1 +
 R/comparisons.R | 15 +++++++--------
 R/predictions.R | 39 ++++++++++++++++++++++-----------------
 4 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index 00826813c..8c4277ace 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,6 +1,6 @@
 Package: marginaleffects
 Title: Predictions, Comparisons, Slopes, Marginal Means, and Hypothesis Tests
-Version: 0.24.0.3
+Version: 0.24.0.4
 Authors@R: 
     c(person(given = "Vincent",
              family = "Arel-Bundock",
diff --git a/NEWS.md b/NEWS.md
index bd7fcf463..cc4bed0b4 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -3,6 +3,7 @@
 ## Development
 
 * Version 0.24.0 accidentally removed the "contrast" column from the output object in calls with only one focal predictor. This column is reinstated.
+* Reinstate some attributes lost with `marginaleffects_lean` but necessary for printing.
 * More informative warning for `lme4` and `glmmTMB` models with `re.form=NULL`
 
 Bugs:
diff --git a/R/comparisons.R b/R/comparisons.R
index 148a6b87f..546f66906 100644
--- a/R/comparisons.R
+++ b/R/comparisons.R
@@ -547,15 +547,16 @@ comparisons <- function(model,
     get_marginaleffects_attributes(newdata, include_regex = "^newdata.*class|explicit|matrix|levels"))
 
   # Global option for lean return object
-  lean = getOption("marginaleffects_lean", default = FALSE)
+  lean <- getOption("marginaleffects_lean", default = FALSE)
 
   # Only add (potentially large) attributes if lean is FALSE
+  # extra attributes needed for print method, even with lean return object
+  attr(out, "conf_level") <- conf_level
+  attr(out, "by") <- by
+  attr(out, "lean") <- lean
+  attr(out, "vcov.type") <- vcov.type
   if (isTRUE(lean)) {
-    for (a in setdiff(names(attributes(out)), c("names", "row.names", "class"))) attr(out, a) = NULL
-    ## extra attributes needed for print method, even with lean return object
-    attr(out, "conf_level") <- conf_level
-    attr(out, "by") <- by
-    attr(out, "lean") <- TRUE
+    for (a in setdiff(names(attributes(out)), c("names", "row.names", "class"))) attr(out, a) <- NULL
   } else {
     # other attributes
     attr(out, "newdata") <- newdata
@@ -573,8 +574,6 @@ comparisons <- function(model,
     attr(out, "comparison_label") <- comparison_label
     attr(out, "hypothesis_by") <- hyp_by
     attr(out, "transform_label") <- transform_label
-    attr(out, "conf_level") <- conf_level
-    attr(out, "by") <- by
 
     if (inherits(model, "brmsfit")) {
       insight::check_if_installed("brms")
diff --git a/R/predictions.R b/R/predictions.R
index 7cddb4539..8dd9cdd08 100644
--- a/R/predictions.R
+++ b/R/predictions.R
@@ -543,30 +543,35 @@ predictions <- function(model,
   out <- set_marginaleffects_attributes(out, attr_cache = newdata_attr_cache)
 
   # Global option for lean return object
-  lean = getOption("marginaleffects_lean", default = FALSE)
-
-  # Only add (potentially large) attributes if lean isn't TRUE
+  lean <- getOption("marginaleffects_lean", default = FALSE)
+
+  # Only add (potentially large) attributes if lean is FALSE
+  # extra attributes needed for print method, even with lean return object
+  attr(out, "conf_level") <- conf_level
+  attr(out, "by") <- by
+  attr(out, "lean") <- lean
+  attr(out, "vcov.type") <- vcov.type
   if (isTRUE(lean)) {
-    for (a in setdiff(aa, c("names", "row.names", "class"))) attr(mfx, a) = NULL
-    attr(out, "lean") <- TRUE
+    for (a in setdiff(names(attributes(out)), c("names", "row.names", "class"))) {
+      attr(out, a) <- NULL
+    }
   } else {
     # other attributes
-    attr(out, "model") <- model
-    attr(out, "type") <- type_string
+    attr(out, "newdata") <- newdata
+    attr(out, "call") <- call_attr
+    attr(out, "type") <- type
     attr(out, "model_type") <- class(model)[1]
-    attr(out, "vcov.type") <- get_vcov_label(vcov)
+    attr(out, "model") <- model
+    attr(out, "variables") <- predictors
     attr(out, "jacobian") <- J
-    attr(out, "vcov") <- V
-    attr(out, "newdata") <- newdata
+    attr(out, "vcov") <- vcov
+    attr(out, "vcov.type") <- vcov.type
     attr(out, "weights") <- marginaleffects_wts_internal
-    attr(out, "conf_level") <- conf_level
-    attr(out, "by") <- by
-    attr(out, "call") <- call_attr
-    attr(out, "hypothesis_by") <- hyp_by
-    attr(out, "transform_label") <- names(transform)[1]
+    attr(out, "comparison") <- comparison
     attr(out, "transform") <- transform[[1]]
-    # save newdata for use in recall()
-    attr(out, "newdata") <- newdata
+    attr(out, "comparison_label") <- comparison_label
+    attr(out, "hypothesis_by") <- hyp_by
+    attr(out, "transform_label") <- transform_label
 
     if (inherits(model, "brmsfit")) {
       insight::check_if_installed("brms")