Add PyCapsule typing protocol and helper detection

Introduce a _PyCapsule typing protocol to enable type checkers
to recognize PyCapsule-based registrations. Restrict the
AggregateUDF udaf overload to the PyCapsule protocol and
update from_pycapsule to wrap raw capsule inputs using
the internal binding directly.

Author: kosiew

python/datafusion/user_defined.py

@@ -28,6 +28,7 @@
     Callable,
     Optional,
     Protocol,
+    TypeGuard,
     TypeVar,
     cast,
     overload,
@@ -92,6 +93,16 @@ class ScalarUDFExportable(Protocol):
     def __datafusion_scalar_udf__(self) -> object: ...  # noqa: D105
 
 
+class _PyCapsule(Protocol):
+    """Lightweight typing proxy for CPython ``PyCapsule`` objects."""
+
+
+def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]:
+    """Return ``True`` when ``value`` is a CPython ``PyCapsule``."""
+
+    return value.__class__.__name__ == "PyCapsule"
+
+
 class ScalarUDF:
     """Class for performing scalar user-defined functions (UDF).
 
@@ -401,7 +412,7 @@ def udaf(accum: AggregateUDFExportable) -> AggregateUDF: ...
 
     @overload
     @staticmethod
-    def udaf(accum: object) -> AggregateUDF: ...
+    def udaf(accum: _PyCapsule) -> AggregateUDF: ...
 
     @staticmethod
     def udaf(*args: Any, **kwargs: Any):  # noqa: D417, C901
@@ -523,7 +534,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Expr:
 
             return decorator
 
-        if hasattr(args[0], "__datafusion_aggregate_udf__"):
+        if hasattr(args[0], "__datafusion_aggregate_udf__") or _is_pycapsule(args[0]):
             return AggregateUDF.from_pycapsule(args[0])
 
         if args and callable(args[0]):
@@ -533,12 +544,17 @@ def wrapper(*args: Any, **kwargs: Any) -> Expr:
         return _decorator(*args, **kwargs)
 
     @staticmethod
-    def from_pycapsule(func: AggregateUDFExportable | object) -> AggregateUDF:
+    def from_pycapsule(func: AggregateUDFExportable | _PyCapsule) -> AggregateUDF:
         """Create an Aggregate UDF from AggregateUDF PyCapsule object.
 
         This function will instantiate a Aggregate UDF that uses a DataFusion
         AggregateUDF that is exported via the FFI bindings.
         """
+        if _is_pycapsule(func):
+            aggregate = cast(AggregateUDF, object.__new__(AggregateUDF))
+            aggregate._udaf = df_internal.AggregateUDF.from_pycapsule(func)
+            return aggregate
+
         capsule = cast(AggregateUDFExportable, func)
         name = str(capsule.__class__)
         return AggregateUDF(