Improve type hinting for udaf and fix one pylance warning

Author: timsaucer

python/datafusion/udf.py

@@ -21,7 +21,7 @@
 
 import datafusion._internal as df_internal
 from datafusion.expr import Expr
-from typing import Callable, TYPE_CHECKING, TypeVar
+from typing import Callable, TYPE_CHECKING, TypeVar, Type
 from abc import ABCMeta, abstractmethod
 from typing import List, Any, Optional
 from enum import Enum
@@ -167,10 +167,6 @@ def evaluate(self) -> pyarrow.Scalar:
         pass
 
 
-if TYPE_CHECKING:
-    _A = TypeVar("_A", bound=(Callable[..., _R], Accumulator))
-
-
 class AggregateUDF:
     """Class for performing scalar user-defined functions (UDF).
 
@@ -181,9 +177,9 @@ class AggregateUDF:
     def __init__(
         self,
         name: str | None,
-        accumulator: _A,
+        accumulator: Type[Accumulator],
         input_types: list[pyarrow.DataType],
-        return_type: _R,
+        return_type: pyarrow.DataType,
         state_type: list[pyarrow.DataType],
         volatility: Volatility | str,
         arguments: list[Any],
@@ -214,9 +210,9 @@ def __call__(self, *args: Expr) -> Expr:
 
     @staticmethod
     def udaf(
-        accum: _A,
+        accum: Type[Accumulator],
         input_types: pyarrow.DataType | list[pyarrow.DataType],
-        return_type: _R,
+        return_type: pyarrow.DataType,
         state_type: list[pyarrow.DataType],
         volatility: Volatility | str,
         arguments: Optional[list[Any]] = None,