feat: Support top-level pl.col autocompletion for iPython

py-polars/docs/source/reference/expressions/col.rst

@@ -11,7 +11,7 @@ See the class documentation below for examples and further documentation.
 -----
 
 .. currentmodule:: polars.functions.col
-.. autoclass:: ColumnFactory
+.. autoclass:: Col
    :members: __call__, __getattr__
    :noindex:
    :autosummary:

py-polars/polars/functions/col.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Any, Iterable, Protocol, cast
+from typing import TYPE_CHECKING, Any, Iterable
 
 from polars._utils.wrap import wrap_expr
 from polars.datatypes import is_polars_dtype
@@ -67,25 +67,7 @@ def _create_col(
         raise TypeError(msg)
 
 
-# appease lint by casting `col` with a protocol that conforms to the factory interface
-class Column(Protocol):
-    def __call__(
-        self,
-        name: str | PolarsDataType | Iterable[str] | Iterable[PolarsDataType],
-        *more_names: str | PolarsDataType,
-    ) -> Expr: ...
-
-    def __getattr__(self, name: str) -> Expr: ...
-
-
-# handle attribute lookup on the metaclass (we use the factory uninstantiated)
-class ColumnFactoryMeta(type):
-    def __getattr__(self, name: str) -> Expr:
-        return _create_col(name)
-
-
-# factory that creates columns using `col("name")` or `col.name` syntax
-class ColumnFactory(metaclass=ColumnFactoryMeta):
+class Col:
     """
     Create Polars column expressions.
 
@@ -142,8 +124,8 @@ class ColumnFactory(metaclass=ColumnFactoryMeta):
     └─────┴─────┴─────┘
     """
 
-    def __new__(  # type: ignore[misc]
-        cls,
+    def __call__(
+        self,
         name: str | PolarsDataType | Iterable[str] | Iterable[PolarsDataType],
         *more_names: str | PolarsDataType,
     ) -> Expr:
@@ -291,14 +273,6 @@ def __new__(  # type: ignore[misc]
         """
         return _create_col(name, *more_names)
 
-    # appease sphinx; we actually use '__new__'
-    def __call__(
-        self,
-        name: str | PolarsDataType | Iterable[str] | Iterable[PolarsDataType],
-        *more_names: str | PolarsDataType,
-    ) -> Expr:
-        return _create_col(name, *more_names)
-
     def __getattr__(self, name: str) -> Expr:
         """
         Create a column expression using attribute syntax.
@@ -331,7 +305,11 @@ def __getattr__(self, name: str) -> Expr:
         │ 6   │
         └─────┘
         """
-        return getattr(type(self), name)
+        # For autocomplete to work with iPython
+        if name.startswith("__wrapped__"):
+            return getattr(type(self), name)
+
+        return _create_col(name)
 
 
-col = cast(Column, ColumnFactory)
+col: Col = Col()