add pretty repr

Author: MarcoGorelli

pandas/core/col.py

@@ -17,6 +17,27 @@
     from pandas import DataFrame
 
 
+OP_SYMBOLS = {
+    "__add__": "+",
+    "__radd__": "+",
+    "__sub__": "-",
+    "__rsub__": "-",
+    "__mul__": "*",
+    "__rmul__": "*",
+    "__truediv__": "/",
+    "__rtruediv__": "/",
+    "__floordiv__": "//",
+    "__rfloordiv__": "//",
+    "__ge__": ">=",
+    "__gt__": ">",
+    "__le__": "<=",
+    "__lt__": "<",
+    "__eq__": "==",
+    "__ne__": "!=",
+    "__mod__": "%",
+}
+
+
 def parse_args(df: DataFrame, *args: Any) -> tuple[Series]:
     return tuple([x._func(df) if isinstance(x, Expr) else x for x in args])
 
@@ -29,8 +50,11 @@ def parse_kwargs(df: DataFrame, **kwargs: Any) -> dict[Hashable, Series]:
 
 
 class Expr:
-    def __init__(self, func: Callable[[DataFrame], Any]) -> None:
+    def __init__(
+        self, func: Callable[[DataFrame], Any], repr_str: str | None = None
+    ) -> None:
         self._func = func
+        self._repr_str = repr_str
 
     def __call__(self, df: DataFrame) -> Series:
         result = self._func(df)
@@ -43,9 +67,22 @@ def __call__(self, df: DataFrame) -> Series:
         return result
 
     def _with_binary_op(self, op: str, other: Any) -> Expr:
+        op_symbol = OP_SYMBOLS.get(op, op)
+
         if isinstance(other, Expr):
-            return Expr(lambda df: getattr(self._func(df), op)(other._func(df)))
-        return Expr(lambda df: getattr(self._func(df), op)(other))
+            if op.startswith("__r"):
+                repr_str = f"({other._repr_str} {op_symbol} {self._repr_str})"
+            else:
+                repr_str = f"({self._repr_str} {op_symbol} {other._repr_str})"
+            return Expr(
+                lambda df: getattr(self._func(df), op)(other._func(df)), repr_str
+            )
+        else:
+            if op.startswith("__r"):
+                repr_str = f"({other!r} {op_symbol} {self._repr_str})"
+            else:
+                repr_str = f"({self._repr_str} {op_symbol} {other!r})"
+            return Expr(lambda df: getattr(self._func(df), op)(other), repr_str)
 
     # Binary ops
     def __add__(self, other: Any) -> Expr:
@@ -106,7 +143,31 @@ def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
             parsed_kwargs = parse_kwargs(df, **kwargs)
             return getattr(self(df), attr)(*parsed_args, **parsed_kwargs)
 
-        return lambda *args, **kwargs: Expr(lambda df: func(df, *args, **kwargs))
+        def wrapper(*args: Any, **kwargs: Any) -> Expr:
+            # Create a readable representation for method calls
+            args_repr = ", ".join(
+                repr(arg._repr_str if isinstance(arg, Expr) else arg) for arg in args
+            )
+            kwargs_repr = ", ".join(
+                f"{k}={v._repr_str if isinstance(v, Expr) else v!r}"
+                for k, v in kwargs.items()
+            )
+
+            all_args = []
+            if args_repr:
+                all_args.append(args_repr)
+            if kwargs_repr:
+                all_args.append(kwargs_repr)
+
+            args_str = ", ".join(all_args)
+            repr_str = f"{self._repr_str}.{attr}({args_str})"
+
+            return Expr(lambda df: func(df, *args, **kwargs), repr_str)
+
+        return wrapper
+
+    def __repr__(self) -> str:
+        return self._repr_str or "Expr(...)"
 
     # Namespaces
     @property
@@ -135,14 +196,16 @@ def struct(self) -> NamespaceExpr:
 
 
 class NamespaceExpr:
-    def __init__(self, func: Callable[[DataFrame], Any], namespace: str) -> None:
+    def __init__(self, func: Expr, namespace: str) -> None:
         self._func = func
         self._namespace = namespace
 
     def __getattr__(self, attr: str) -> Any:
         if isinstance(getattr(getattr(Series, self._namespace), attr), property):
+            repr_str = f"{self._func._repr_str}.{self._namespace}.{attr}"
             return Expr(
-                lambda df: getattr(getattr(self._func(df), self._namespace), attr)
+                lambda df: getattr(getattr(self._func(df), self._namespace), attr),
+                repr_str,
             )
 
         def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
@@ -152,11 +215,32 @@ def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
                 *parsed_args, **parsed_kwargs
             )
 
-        return lambda *args, **kwargs: Expr(lambda df: func(df, *args, **kwargs))
+        def wrapper(*args: Any, **kwargs: Any) -> Expr:
+            # Create a readable representation for namespace method calls
+            args_repr = ", ".join(
+                repr(arg._repr_str if isinstance(arg, Expr) else arg) for arg in args
+            )
+            kwargs_repr = ", ".join(
+                f"{k}={v._repr_str if isinstance(v, Expr) else v!r}"
+                for k, v in kwargs.items()
+            )
+
+            all_args = []
+            if args_repr:
+                all_args.append(args_repr)
+            if kwargs_repr:
+                all_args.append(kwargs_repr)
+
+            args_str = ", ".join(all_args)
+            repr_str = f"{self._func._repr_str}.{self._namespace}.{attr}({args_str})"
+
+            return Expr(lambda df: func(df, *args, **kwargs), repr_str)
+
+        return wrapper
 
 
 def col(col_name: Hashable) -> Expr:
     if not isinstance(col_name, Hashable):
         msg = f"Expected Hashable, got: {type(col_name)}"
         raise TypeError(msg)
-    return Expr(lambda df: df[col_name])
+    return Expr(lambda df: df[col_name], f"col({col_name!r})")