Implement $meta$tree_format()

NEWS.md

@@ -8,7 +8,8 @@
 ## What's changed
 
 - New method `$unnest()` for `LazyFrame` (#397).
-- New method `$sample()` for `DataFrame`.
+- New method `$sample()` for `DataFrame` (#399).
+- New method `$meta$tree_format()` to display an `Expr` as a tree (#401).
 
 # polars 0.8.1
 

R/expr__meta.R

@@ -150,3 +150,29 @@ ExprMeta_has_multiple_outputs = function() {
 ExprMeta_is_regex_projection = function() {
   .pr$Expr$meta_is_regex_projection(self)
 }
+
+#' Format an expression as a tree
+#'
+#' @param return_as_string Return the tree as a character vector? If `FALSE`
+#' (default), the tree is printed in the console.
+#'
+#' @return
+#' If `return_as_string` is `TRUE`, a character vector describing the tree.
+#'
+#' If `return_as_string` is `FALSE`, prints the tree in the console but doesn't
+#' return any value.
+#'
+#' @name ExprMeta_tree_format
+#' @examples
+#' my_expr = (pl$col("foo") * pl$col("bar"))$sum()$over(pl$col("ham")) / 2
+#' my_expr$meta$tree_format()
+
+ExprMeta_tree_format = function(return_as_string = FALSE) {
+  out <- .pr$Expr$meta_tree_format(self) |>
+    unwrap("in $tree_format():")
+  if (isTRUE(return_as_string)) {
+    out
+  } else {
+    cat(out)
+  }
+}

R/extendr-wrappers.R

@@ -909,6 +909,8 @@ Expr$meta_has_multiple_outputs <- function() .Call(wrap__Expr__meta_has_multiple
 
 Expr$meta_is_regex_projection <- function() .Call(wrap__Expr__meta_is_regex_projection, self)
 
+Expr$meta_tree_format <- function() .Call(wrap__Expr__meta_tree_format, self)
+
 Expr$cat_set_ordering <- function(ordering) .Call(wrap__Expr__cat_set_ordering, self, ordering)
 
 Expr$new_count <- function() .Call(wrap__Expr__new_count)

src/rust/src/lazy/dsl.rs

@@ -2318,6 +2318,16 @@ impl Expr {
         self.0.clone().meta().is_regex_projection()
     }
 
+    fn meta_tree_format(&self) -> RResult<String> {
+        let e = self
+            .0
+            .clone()
+            .meta()
+            .into_tree_formatter()
+            .map_err(polars_to_rpolars_err)?;
+        Ok(format!("{e}"))
+    }
+
     //the only cat ns function from dsl.rs
     fn cat_set_ordering(&self, ordering: Robj) -> Result<Expr, String> {
         let ordering = robj_to!(Map, str, ordering, |s| {

src/rust/src/rdataframe/mod.rs

@@ -339,16 +339,19 @@ impl DataFrame {
         value_vars: Robj,
         value_name: Robj,
         variable_name: Robj,
-    ) -> Result<Self, String> {
+    ) -> RResult<Self> {
         let args = MeltArgs {
             id_vars: strings_to_smartstrings(robj_to!(Vec, String, id_vars)?),
             value_vars: strings_to_smartstrings(robj_to!(Vec, String, value_vars)?),
             value_name: robj_to!(Option, String, value_name)?.map(|s| s.into()),
             variable_name: robj_to!(Option, String, variable_name)?.map(|s| s.into()),
             streamable: false,
         };
-        let df = self.0.melt2(args).map_err(|s| s.to_string())?;
-        Ok(DataFrame(df))
+
+        self.0
+            .melt2(args)
+            .map_err(polars_to_rpolars_err)
+            .map(DataFrame)
     }
 
     pub fn pivot_expr(
@@ -360,7 +363,7 @@ impl DataFrame {
         sort_columns: Robj,
         aggregate_expr: Robj,
         separator: Robj,
-    ) -> Result<Self, String> {
+    ) -> RResult<Self> {
         let fun = if robj_to!(bool, maintain_order)? {
             pivot_stable
         } else {
@@ -376,8 +379,8 @@ impl DataFrame {
             robj_to!(Option, PLExpr, aggregate_expr)?,
             robj_to!(Option, str, separator)?,
         )
-        .map_err(|err| err.to_string())
-        .map(|ok| ok.into())
+        .map_err(polars_to_rpolars_err)
+        .map(DataFrame)
     }
 
     pub fn sample_n(

tests/testthat/_snaps/expr_meta.md

@@ -0,0 +1,23 @@
+# meta$tree_format
+
+    Code
+      e$meta$tree_format()
+    Output
+          binary: //    
+
+             |              |       
+
+          lit(2.0)          window      
+
+                            |               |        
+
+                            col(ham)        sum          
+
+                                            |        
+
+                                            binary: *    
+
+                                            |                |       
+
+                                            col(bar)         col(foo)    
+

tests/testthat/test-expr_meta.R

@@ -76,3 +76,9 @@ test_that("meta$is_regex_projection", {
   expect_true(e1$meta$is_regex_projection())
   expect_false(e2$meta$is_regex_projection())
 })
+
+test_that("meta$tree_format", {
+  e = (pl$col("foo") * pl$col("bar"))$sum()$over(pl$col("ham")) / 2
+  expect_true(is.character(e$meta$tree_format(return_as_string = TRUE)))
+  expect_snapshot(e$meta$tree_format())
+})