Add return types on functions

Author: timsaucer

pyproject.toml

@@ -89,8 +89,6 @@ ignore = [
     "PLR0913", # Allow many arguments in function definition
     "PD901",   # Allow variable name df
     # TODO: Enable all of the following, but this PR is getting too large already
-#    "ANN001",
-    "ANN202",
     "PTH",
     "N812",
     "INP001",
@@ -127,7 +125,7 @@ extend-allowed-calls = ["lit", "datafusion.lit"]
     "PLR0913",
     "PT004",
 ]
-"examples/*" = ["D", "W505", "E501", "T201", "S101", "PLR2004", "ANN001"]
+"examples/*" = ["D", "W505", "E501", "T201", "S101", "PLR2004", "ANN001", "ANN202"]
 "dev/*" = ["D", "E", "T", "S", "PLR", "C", "SIM", "UP", "EXE", "N817", "ERA001", "ANN001"]
 "benchmarks/*" = ["D", "F", "T", "BLE", "FURB", "PLR", "E", "TD", "TRY", "S", "SIM", "EXE", "UP", "ERA001", "ANN001"]
 "docs/*" = ["D"]

python/datafusion/__init__.py

@@ -124,15 +124,15 @@ def literal(value: Any) -> Expr:
     return Expr.literal(value)
 
 
-def string_literal(value: str):
+def string_literal(value: str) -> Expr:
     """Create a UTF8 literal expression.
 
     It differs from `literal` which creates a UTF8view literal.
     """
     return Expr.string_literal(value)
 
 
-def str_lit(value: str):
+def str_lit(value: str) -> Expr:
     """Alias for `string_literal`."""
     return string_literal(value)
 

python/datafusion/user_defined.py

@@ -218,13 +218,13 @@ def _decorator(
             volatility: Volatility | str,
             name: str | None = None,
         ) -> Callable:
-            def decorator(func: Callable):
+            def decorator(func: Callable) -> Callable:
                 udf_caller = ScalarUDF.udf(
                     func, input_types, return_type, volatility, name
                 )
 
                 @functools.wraps(func)
-                def wrapper(*args: Any, **kwargs: Any):
+                def wrapper(*args: Any, **kwargs: Any) -> Callable:
                     return udf_caller(*args, **kwargs)
 
                 return wrapper