feat(ir): support pretty printing arbitrary traversable objects (ibis-project#9043)

This enables pretty formatting the dereference mappings and also the
replacements mappings we use in rewrites for better inspection. Also
moved the irrelevant logic out from `format.py`.

Author: kszucs

ibis/common/typing.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import inspect
 import re
 import sys
 from abc import abstractmethod
@@ -259,3 +260,21 @@ class Coercible(Abstract):
     @classmethod
     @abstractmethod
     def __coerce__(cls, value: Any, **kwargs: Any) -> Self: ...
+
+
+def get_defining_frame(obj):
+    """Locate the outermost frame where `obj` is defined."""
+    for frame_info in inspect.stack()[::-1]:
+        for var in frame_info.frame.f_locals.values():
+            if obj is var:
+                return frame_info.frame
+    raise ValueError(f"No defining frame found for {obj}")
+
+
+def get_defining_scope(obj, types=None):
+    """Get variables in the scope where `expr` is first defined."""
+    frame = get_defining_frame(obj)
+    scope = {**frame.f_globals, **frame.f_locals}
+    if types is not None:
+        scope = {k: v for k, v in scope.items() if isinstance(v, types)}
+    return scope

ibis/expr/format.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import functools
-import inspect
 import itertools
 import textwrap
 import types
@@ -13,8 +12,8 @@
 import ibis
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
-import ibis.expr.types as ir
 from ibis import util
+from ibis.common.graph import Node
 
 _infix_ops = {
     # comparison operations
@@ -147,57 +146,31 @@ def inline_args(fields, prefer_positional=False):
     return ", ".join(f"{k}={v}" for k, v in fields.items())
 
 
-def get_defining_frame(expr):
-    """Locate the outermost frame where `expr` is defined."""
-    for frame_info in inspect.stack()[::-1]:
-        for var in frame_info.frame.f_locals.values():
-            if isinstance(var, ir.Expr) and expr.equals(var):
-                return frame_info.frame
-    raise ValueError(f"No defining frame found for {expr}")
-
-
-def get_defining_scope(expr):
-    """Get variables in the scope where `expr` is first defined."""
-    frame = get_defining_frame(expr)
-    scope = {**frame.f_globals, **frame.f_locals}
-    return {k: v for k, v in scope.items() if isinstance(v, ir.Expr)}
-
-
 class Rendered(str):
     def __repr__(self):
         return self
 
 
 @public
-def pretty(expr: ops.Node | ir.Expr, scope: Optional[dict[str, ir.Expr]] = None) -> str:
+def pretty(node: Node, scope: Optional[dict[str, Node]] = None) -> str:
     """Pretty print an expression.
 
     Parameters
     ----------
-    expr
-        The expression to pretty print.
+    node
+        The graph node to pretty print.
     scope
         A dictionary of expression to name mappings used to intermediate
-        assignments.
-        If not provided the names of the expressions will either be
-        - the variable name in the defining scope if
-          `ibis.options.repr.show_variables` is enabled
-        - generated names like `r0`, `r1`, etc. otherwise
+        assignments. If not provided aliases will be generated for each
+        relation.
 
     Returns
     -------
     str
         A pretty printed representation of the expression.
     """
-    if isinstance(expr, ir.Expr):
-        node = expr.op()
-    elif isinstance(expr, ops.Node):
-        node = expr
-    else:
-        raise TypeError(f"Expected an expression or a node, got {type(expr)}")
-
-    if scope is None and ibis.options.repr.show_variables:
-        scope = get_defining_scope(expr)
+    if not isinstance(node, Node):
+        raise TypeError(f"Expected a graph node, got {type(node)}")
 
     refs = {}
     refcnt = itertools.count()

ibis/expr/tests/snapshots/test_format/test_arbitrary_traversables_are_supported/repr.txt

@@ -0,0 +1,9 @@
+r0 := UnboundTable: t
+  a int64
+
+MyNode
+  obj:
+    r0.a
+  children:
+    r0.a
+    r0.a + 1

ibis/expr/tests/test_format.py

@@ -10,6 +10,7 @@
 import ibis.expr.operations as ops
 import ibis.legacy.udf.vectorized as udf
 from ibis import util
+from ibis.common.graph import Node as Traversable
 from ibis.expr.format import fmt, pretty
 
 
@@ -465,3 +466,26 @@ class ValueList(ops.Node):
     result = pretty(vl)
 
     snapshot.assert_match(result, "repr.txt")
+
+
+def test_arbitrary_traversables_are_supported(snapshot):
+    class MyNode(Traversable):
+        __slots__ = ("obj", "children")
+        __argnames__ = ("obj", "children")
+
+        def __init__(self, obj, children):
+            self.obj = obj.op()
+            self.children = tuple(child.op() for child in children)
+
+        @property
+        def __args__(self):
+            return self.obj, self.children
+
+        def __hash__(self):
+            return hash((self.__class__, self.obj, self.children))
+
+    t = ibis.table([("a", "int64")], name="t")
+    node = MyNode(t.a, [t.a, t.a + 1])
+    result = pretty(node)
+
+    snapshot.assert_match(result, "repr.txt")

ibis/expr/types/core.py

@@ -16,8 +16,10 @@
 from ibis.common.exceptions import IbisError, TranslationError
 from ibis.common.grounds import Immutable
 from ibis.common.patterns import Coercible, CoercionError
+from ibis.common.typing import get_defining_scope
 from ibis.config import _default_backend
 from ibis.config import options as opts
+from ibis.expr.format import pretty
 from ibis.expr.types.pretty import to_rich
 from ibis.util import experimental
 
@@ -44,7 +46,6 @@ def _repr_mimebundle_(self, *args, **kwargs):
         return bundle
 
 
-# TODO(kszucs): consider to subclass from Annotable with a single _arg field
 @public
 class Expr(Immutable, Coercible):
     """Base expression class."""
@@ -53,9 +54,11 @@ class Expr(Immutable, Coercible):
     _arg: ops.Node
 
     def _noninteractive_repr(self) -> str:
-        from ibis.expr.format import pretty
-
-        return pretty(self)
+        if ibis.options.repr.show_variables:
+            scope = get_defining_scope(self, types=Expr)
+        else:
+            scope = None
+        return pretty(self.op(), scope=scope)
 
     def _interactive_repr(self) -> str:
         console = Console(force_terminal=False)