Cleanup code for Dot display to use more idiomatic rust.

crates/polars-plan/src/dsl/meta.rs

@@ -1,10 +1,11 @@
-use std::fmt::{Binary, Display};
+use std::fmt::Display;
 use std::ops::BitAnd;
 
 use super::*;
 use crate::plans::conversion::is_regex_projection;
 use crate::plans::ir::tree_format::TreeFmtVisitor;
 use crate::plans::visitor::{AexprNode, TreeWalker};
+use crate::prelude::tree_format::TreeFmtVisitorDisplay;
 
 /// Specialized expressions for Categorical dtypes.
 pub struct MetaNameSpace(pub(crate) Expr);
@@ -159,10 +160,13 @@ impl MetaNameSpace {
 
     /// Get a hold to an implementor of the `Display` trait that will format as
     /// the expression as a tree
-    pub fn into_tree_formatter(self) -> PolarsResult<impl Display + Binary> {
+    pub fn into_tree_formatter(self, display_as_dot: bool) -> PolarsResult<impl Display> {
         let mut arena = Default::default();
         let node = to_aexpr(self.0, &mut arena)?;
         let mut visitor = TreeFmtVisitor::default();
+        if display_as_dot {
+            visitor.display = TreeFmtVisitorDisplay::DisplayDot;
+        }
 
         AexprNode::new(node).visit(&mut visitor, &arena)?;
 

crates/polars-plan/src/plans/ir/tree_format.rs

@@ -382,12 +382,20 @@ impl<'a> TreeFmtNode<'a> {
     }
 }
 
+#[derive(Default)]
+pub enum TreeFmtVisitorDisplay {
+    #[default]
+    DisplayText,
+    DisplayDot,
+}
+
 #[derive(Default)]
 pub(crate) struct TreeFmtVisitor {
     levels: Vec<Vec<String>>,
     prev_depth: usize,
     depth: usize,
     width: usize,
+    pub(crate) display: TreeFmtVisitorDisplay,
 }
 
 impl Visitor for TreeFmtVisitor {
@@ -868,59 +876,64 @@ impl fmt::Display for Canvas {
     }
 }
 
-impl fmt::Display for TreeFmtVisitor {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
-        fmt::Debug::fmt(self, f)
-    }
-}
+fn tree_fmt_text(tree: &TreeFmtVisitor, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
+    let tree_view: TreeView<'_> = tree.levels.as_slice().into();
+    let canvas: Canvas = tree_view.into();
+    write!(f, "{canvas}")?;
 
-impl fmt::Debug for TreeFmtVisitor {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
-        let tree_view: TreeView<'_> = self.levels.as_slice().into();
-        let canvas: Canvas = tree_view.into();
-        write!(f, "{canvas}")?;
-
-        Ok(())
-    }
+    Ok(())
 }
 
 // GraphViz Output
 // Create a simple DOT graph String from TreeFmtVisitor
-
-// Adding a non-standard display trait is not the right way to do this.
-// Not sure how to cleanly communicate this function to the higher level which
-// only wants traits.
-impl fmt::Binary for TreeFmtVisitor {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
-        // Build a dot graph as a string
-        let tree_view: TreeView<'_> = self.levels.as_slice().into();
-        let mut relations: Vec<String> = Vec::new();
-
-        // Non-empty cells (nodes) and their connections (edges)
-        for (i, row) in tree_view.matrix.iter().enumerate() {
-            for (j, cell) in row.iter().enumerate() {
-                if !cell.text.is_empty() {
-                    // Add node
-                    let node_label = &cell.text.join("\n");
-                    let node_desc = format!("n{i}{j} [label=\"{node_label}\", ordering=\"out\"]");
-                    relations.push(node_desc);
-
-                    // Add child edges
-                    if i < tree_view.rows.len() - 1 {
-                        // Iter in reversed order to undo the reversed child order when iterating expressions
-                        for child_col in cell.children_columns.iter().rev() {
-                            let next_row = i + 1;
-                            let edge = format!("n{i}{j} -- n{next_row}{child_col}");
-                            relations.push(edge);
-                        }
+fn tree_fmt_dot(tree: &TreeFmtVisitor, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
+    // Build a dot graph as a string
+    let tree_view: TreeView<'_> = tree.levels.as_slice().into();
+    let mut relations: Vec<String> = Vec::new();
+
+    // Non-empty cells (nodes) and their connections (edges)
+    for (i, row) in tree_view.matrix.iter().enumerate() {
+        for (j, cell) in row.iter().enumerate() {
+            if !cell.text.is_empty() {
+                // Add node
+                let node_label = &cell.text.join("\n");
+                let node_desc = format!("n{i}{j} [label=\"{node_label}\", ordering=\"out\"]");
+                relations.push(node_desc);
+
+                // Add child edges
+                if i < tree_view.rows.len() - 1 {
+                    // Iter in reversed order to undo the reversed child order when iterating expressions
+                    for child_col in cell.children_columns.iter().rev() {
+                        let next_row = i + 1;
+                        let edge = format!("n{i}{j} -- n{next_row}{child_col}");
+                        relations.push(edge);
                     }
                 }
             }
         }
+    }
 
-        let graph_str = relations.join("\n    ");
-        let s = format!("graph {{\n    {graph_str}\n}}");
-        write!(f, "{s}")?;
-        Ok(())
+    let graph_str = relations.join("\n    ");
+    let s = format!("graph {{\n    {graph_str}\n}}");
+    write!(f, "{s}")?;
+    Ok(())
+}
+
+fn tree_fmt(tree: &TreeFmtVisitor, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
+    match tree.display {
+        TreeFmtVisitorDisplay::DisplayText => tree_fmt_text(tree, f),
+        TreeFmtVisitorDisplay::DisplayDot => tree_fmt_dot(tree, f),
+    }
+}
+
+impl fmt::Display for TreeFmtVisitor {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
+        tree_fmt(self, f)
+    }
+}
+
+impl fmt::Debug for TreeFmtVisitor {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
+        tree_fmt(self, f)
     }
 }

crates/polars-python/src/expr/meta.rs

@@ -102,23 +102,21 @@ impl PyExpr {
         self.inner.clone().meta()._into_selector().into()
     }
 
-    fn meta_tree_format(&self) -> PyResult<String> {
+    fn compute_tree_format(&self, display_as_dot: bool) -> Result<String, PyErr> {
         let e = self
             .inner
             .clone()
             .meta()
-            .into_tree_formatter()
+            .into_tree_formatter(display_as_dot)
             .map_err(PyPolarsErr::from)?;
         Ok(format!("{e}"))
     }
 
+    fn meta_tree_format(&self) -> PyResult<String> {
+        self.compute_tree_format(false)
+    }
+
     fn meta_show_graph(&self) -> PyResult<String> {
-        let e = self
-            .inner
-            .clone()
-            .meta()
-            .into_tree_formatter()
-            .map_err(PyPolarsErr::from)?;
-        Ok(format!("{e:b}"))
+        self.compute_tree_format(true)
     }
 }