Fixing linting errors

Author: crystalxyz

python/datafusion/udf.py

@@ -19,10 +19,10 @@
 
 from __future__ import annotations
 
+import functools
 from abc import ABCMeta, abstractmethod
 from enum import Enum
 from typing import TYPE_CHECKING, Callable, List, Optional, TypeVar
-import functools
 
 import pyarrow
 
@@ -117,15 +117,17 @@ class udf:
         This class can be used both as a **function** and as a **decorator**.
 
         Usage:
-            - **As a function**: Call `udf(func, input_types, return_type, volatility, name)`.
-            - **As a decorator**: Use `@udf(input_types, return_type, volatility, name)`.
-            In this case, do **not** pass `func` explicitly.
+            - **As a function**: Call `udf(func, input_types, return_type, volatility,
+              name)`.
+            - **As a decorator**: Use `@udf(input_types, return_type, volatility,
+              name)`. In this case, do **not** pass `func` explicitly.
 
         Args:
             func (Callable, optional): **Only needed when calling as a function.**
                 Skip this argument when using `udf` as a decorator.
             input_types (list[pyarrow.DataType]): The data types of the arguments
-                to `func`. This list must be of the same length as the number of arguments.
+                to `func`. This list must be of the same length as the number of
+                arguments.
             return_type (_R): The data type of the return value from the function.
             volatility (Volatility | str): See `Volatility` for allowed values.
             name (Optional[str]): A descriptive name for the function.
@@ -139,7 +141,8 @@ class udf:
             ```
             def double_func(x):
                 return x * 2
-            double_udf = udf(double_func, [pyarrow.int32()], pyarrow.int32(), "volatile", "double_it")
+            double_udf = udf(double_func, [pyarrow.int32()], pyarrow.int32(),
+            "volatile", "double_it")
             ```
 
             **Using `udf` as a decorator:**
@@ -149,8 +152,13 @@ def double_udf(x):
                 return x * 2
             ```
         """
+
         def __new__(cls, *args, **kwargs):
-            if args and callable(args[0]):  
+            """Create a new UDF.
+
+            Trigger UDF function or decorator depending on if the first args is callable
+            """
+            if args and callable(args[0]):
                 # Case 1: Used as a function, require the first parameter to be callable
                 return cls._function(*args, **kwargs)
             else:
@@ -185,23 +193,22 @@ def _decorator(
             input_types: list[pyarrow.DataType],
             return_type: _R,
             volatility: Volatility | str,
-            name: Optional[str] = None
+            name: Optional[str] = None,
         ):
             def decorator(func):
                 udf_caller = ScalarUDF.udf(
-                    func,
-                    input_types, 
-                    return_type, 
-                    volatility,
-                    name
+                    func, input_types, return_type, volatility, name
                 )
-                
+
                 @functools.wraps(func)
                 def wrapper(*args, **kwargs):
                     return udf_caller(*args, **kwargs)
+
                 return wrapper
+
             return decorator
 
+
 class Accumulator(metaclass=ABCMeta):
     """Defines how an :py:class:`AggregateUDF` accumulates values."""
 
@@ -268,8 +275,8 @@ def __call__(self, *args: Expr) -> Expr:
     class udaf:
         """Create a new User-Defined Aggregate Function (UDAF).
 
-        This class allows you to define an **aggregate function** that can be used in data
-        aggregation or window function calls. 
+        This class allows you to define an **aggregate function** that can be used in
+        data aggregation or window function calls.
 
         Usage:
             - **As a function**: Call `udaf(accum, input_types, return_type, state_type,
@@ -279,12 +286,12 @@ class udaf:
             When using `udaf` as a decorator, **do not pass `accum` explicitly**.
 
         **Function example:**
-        
-            If your `:py:class:Accumulator` can be instantiated with no arguments, you can
-            simply pass it's type as `accum`. If you need to pass additional arguments to
-            it's constructor, you can define a lambda or a factory method. During runtime the
-            `:py:class:Accumulator` will be constructed for every instance in which this UDAF is
-            used. The following examples are all valid.
+
+            If your `:py:class:Accumulator` can be instantiated with no arguments, you
+            can simply pass it's type as `accum`. If you need to pass additional
+            arguments to it's constructor, you can define a lambda or a factory method.
+            During runtime the `:py:class:Accumulator` will be constructed for every
+            instance in which this UDAF is used. The following examples are all valid.
             ```
             import pyarrow as pa
             import pyarrow.compute as pc
@@ -308,11 +315,14 @@ def evaluate(self) -> pa.Scalar:
             def sum_bias_10() -> Summarize:
                 return Summarize(10.0)
 
-            udaf1 = udaf(Summarize, pa.float64(), pa.float64(), [pa.float64()], "immutable")
-            udaf2 = udaf(sum_bias_10, pa.float64(), pa.float64(), [pa.float64()], "immutable")
-            udaf3 = udaf(lambda: Summarize(20.0), pa.float64(), pa.float64(), [pa.float64()], "immutable")
+            udaf1 = udaf(Summarize, pa.float64(), pa.float64(), [pa.float64()],
+                "immutable")
+            udaf2 = udaf(sum_bias_10, pa.float64(), pa.float64(), [pa.float64()],
+                "immutable")
+            udaf3 = udaf(lambda: Summarize(20.0), pa.float64(), pa.float64(),
+                [pa.float64()], "immutable")
             ```
-        
+
         **Decorator example:**
             ```
             @udaf(pa.float64(), pa.float64(), [pa.float64()], "immutable")
@@ -321,7 +331,8 @@ def udf4() -> Summarize:
             ```
 
         Args:
-            accum: The accumulator python function. **Only needed when calling as a function. Skip this argument when using `udaf` as a decorator.**
+            accum: The accumulator python function. **Only needed when calling as a
+                function. Skip this argument when using `udaf` as a decorator.**
             input_types: The data types of the arguments to ``accum``.
             return_type: The data type of the return value.
             state_type: The data types of the intermediate accumulation.
@@ -334,13 +345,18 @@ def udf4() -> Summarize:
         """
 
         def __new__(cls, *args, **kwargs):
-            if args and callable(args[0]):  
+            """Create a new UDAF.
+
+            Trigger UDAF function or decorator depending on if the first args is
+            callable
+            """
+            if args and callable(args[0]):
                 # Case 1: Used as a function, require the first parameter to be callable
                 return cls._function(*args, **kwargs)
             else:
                 # Case 2: Used as a decorator with parameters
                 return cls._decorator(*args, **kwargs)
-            
+
         @staticmethod
         def _function(
             accum: Callable[[], Accumulator],
@@ -368,31 +384,27 @@ def _function(
                 state_type=state_type,
                 volatility=volatility,
             )
-            
+
         @staticmethod
         def _decorator(
             input_types: pyarrow.DataType | list[pyarrow.DataType],
             return_type: pyarrow.DataType,
             state_type: list[pyarrow.DataType],
             volatility: Volatility | str,
-            name: Optional[str] = None
+            name: Optional[str] = None,
         ):
             def decorator(accum: Callable[[], Accumulator]):
                 udaf_caller = AggregateUDF.udaf(
-                    accum,
-                    input_types, 
-                    return_type,
-                    state_type,
-                    volatility,
-                    name
+                    accum, input_types, return_type, state_type, volatility, name
                 )
-                
+
                 @functools.wraps(accum)
                 def wrapper(*args, **kwargs):
                     return udaf_caller(*args, **kwargs)
+
                 return wrapper
-            return decorator
 
+            return decorator
 
 
 class WindowEvaluator(metaclass=ABCMeta):