ENH: Introduce pandas.col

pandas/core/col.py

@@ -51,16 +51,32 @@ def _parse_kwargs(df: DataFrame, **kwargs: Any) -> dict[Hashable, Series]:
     }
 
 
+def _pretty_print_args_kwargs(*args: Any, **kwargs: Any) -> str:
+    inputs_repr = ", ".join(
+        arg._repr_str if isinstance(arg, Expression) else repr(arg) for arg in args
+    )
+    kwargs_repr = ", ".join(
+        f"{k}={v._repr_str if isinstance(v, Expression) else v!r}"
+        for k, v in kwargs.items()
+    )
+
+    all_args = []
+    if inputs_repr:
+        all_args.append(inputs_repr)
+    if kwargs_repr:
+        all_args.append(kwargs_repr)
+
+    return ", ".join(all_args)
+
+
 class Expression:
     """
     Class representing a deferred column.
 
     This is not meant to be instantiated directly. Instead, use :meth:`pandas.col`.
     """
 
-    def __init__(
-        self, func: Callable[[DataFrame], Any], repr_str: str | None = None
-    ) -> None:
+    def __init__(self, func: Callable[[DataFrame], Any], repr_str: str) -> None:
         self._func = func
         self._repr_str = repr_str
 
@@ -138,6 +154,19 @@ def __mod__(self, other: Any) -> Expression:
     def __rmod__(self, other: Any) -> Expression:
         return self._with_binary_op("__rmod__", other)
 
+    def __array_ufunc__(
+        self, ufunc: Callable[..., Any], method: str, *inputs: Any, **kwargs: Any
+    ) -> Expression:
+        def func(df: DataFrame) -> Any:
+            parsed_inputs = _parse_args(df, *inputs)
+            parsed_kwargs = _parse_kwargs(df, *kwargs)
+            return ufunc(*parsed_inputs, **parsed_kwargs)
+
+        args_str = _pretty_print_args_kwargs(*inputs, **kwargs)
+        repr_str = f"{ufunc.__name__}({args_str})"
+
+        return Expression(func, repr_str)
+
     # Everything else
     def __getattr__(self, attr: str, /) -> Any:
         if attr in Series._accessors:
@@ -149,23 +178,7 @@ def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
             return getattr(self(df), attr)(*parsed_args, **parsed_kwargs)
 
         def wrapper(*args: Any, **kwargs: Any) -> Expression:
-            # Create a readable representation for method calls
-            args_repr = ", ".join(
-                repr(arg._repr_str if isinstance(arg, Expression) else arg)
-                for arg in args
-            )
-            kwargs_repr = ", ".join(
-                f"{k}={v._repr_str if isinstance(v, Expression) 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)
+            args_str = _pretty_print_args_kwargs(*args, **kwargs)
             repr_str = f"{self._repr_str}.{attr}({args_str})"
 
             return Expression(lambda df: func(df, *args, **kwargs), repr_str)
@@ -200,25 +213,8 @@ def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
             )
 
         def wrapper(*args: Any, **kwargs: Any) -> Expression:
-            # Create a readable representation for namespace method calls
-            args_repr = ", ".join(
-                repr(arg._repr_str if isinstance(arg, Expression) else arg)
-                for arg in args
-            )
-            kwargs_repr = ", ".join(
-                f"{k}={v._repr_str if isinstance(v, Expression) 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)
+            args_str = _pretty_print_args_kwargs(*args, **kwargs)
             repr_str = f"{self._func._repr_str}.{self._namespace}.{attr}({args_str})"
-
             return Expression(lambda df: func(df, *args, **kwargs), repr_str)
 
         return wrapper

pandas/tests/test_col.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 
+import numpy as np
 import pytest
 
 import pandas as pd
@@ -31,6 +32,8 @@
         (pd.col("a") < 1, [False, False], "(col('a') < 1)"),
         (pd.col("a") <= 1, [True, False], "(col('a') <= 1)"),
         (pd.col("a") == 1, [True, False], "(col('a') == 1)"),
+        (np.log(pd.col("a")), [0.0, 0.6931471805599453], "log(col('a'))"),
+        (np.divide(pd.col("a"), pd.col("a")), [1.0, 1.0], "divide(col('a'), col('a'))"),
     ],
 )
 def test_col_simple(