Enabled ruff rule ANN201

Author: Spaarsh

benchmarks/db-benchmark/groupby-datafusion.py

@@ -37,7 +37,7 @@
 exec(open("./_helpers/helpers.py").read())
 
 
-def ans_shape(batches):
+def ans_shape(batches) -> tuple[int, int]:
     rows, cols = 0, 0
     for batch in batches:
         rows += batch.num_rows
@@ -48,7 +48,7 @@ def ans_shape(batches):
     return rows, cols
 
 
-def execute(df):
+def execute(df) -> list:
     print(df.execution_plan().display_indent())
     return df.collect()
 

benchmarks/db-benchmark/join-datafusion.py

@@ -29,7 +29,7 @@
 exec(open("./_helpers/helpers.py").read())
 
 
-def ans_shape(batches):
+def ans_shape(batches) -> tuple[int, int]:
     rows, cols = 0, 0
     for batch in batches:
         rows += batch.num_rows

benchmarks/tpch/tpch.py

@@ -21,7 +21,7 @@
 from datafusion import SessionContext
 
 
-def bench(data_path, query_path):
+def bench(data_path, query_path) -> None:
     with open("results.csv", "w") as results:
         # register tables
         start = time.time()

dev/release/generate-changelog.py

@@ -24,7 +24,7 @@
 from github import Github
 
 
-def print_pulls(repo_name, title, pulls):
+def print_pulls(repo_name, title, pulls) -> None:
     if len(pulls) > 0:
         print(f"**{title}:**")
         print()
@@ -34,7 +34,7 @@ def print_pulls(repo_name, title, pulls):
         print()
 
 
-def generate_changelog(repo, repo_name, tag1, tag2, version):
+def generate_changelog(repo, repo_name, tag1, tag2, version) -> None:
     # get a list of commits between two tags
     print(f"Fetching list of commits between {tag1} and {tag2}", file=sys.stderr)
     comparison = repo.compare(tag1, tag2)
@@ -154,7 +154,7 @@ def generate_changelog(repo, repo_name, tag1, tag2, version):
     )
 
 
-def cli(args=None):
+def cli(args=None) -> None:
     """Process command line arguments."""
     if not args:
         args = sys.argv[1:]

docs/source/conf.py

@@ -73,7 +73,7 @@
 autoapi_python_class_content = "both"
 
 
-def autoapi_skip_member_fn(app, what, name, obj, skip, options):  # noqa: ARG001
+def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool:  # noqa: ARG001
     skip_contents = [
         # Re-exports
         ("class", "datafusion.DataFrame"),
@@ -93,7 +93,7 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options):  # noqa: ARG001
     return skip
 
 
-def setup(sphinx):
+def setup(sphinx) -> None:
     sphinx.connect("autoapi-skip-member", autoapi_skip_member_fn)
 
 

pyproject.toml

@@ -86,7 +86,7 @@ ignore = [
     # "EM101",
     "PLR0913",
     # "PLR1714",
-    "ANN201",
+    # "ANN201",
     "C400",
     "TRY003",
     "B904",

python/datafusion/__init__.py

@@ -92,17 +92,17 @@
 ]
 
 
-def column(value: str):
+def column(value: str) -> Expr:
     """Create a column expression."""
     return Expr.column(value)
 
 
-def col(value: str):
+def col(value: str) -> Expr:
     """Create a column expression."""
     return Expr.column(value)
 
 
-def literal(value):
+def literal(value) -> Expr:
     """Create a literal expression."""
     return Expr.literal(value)
 
@@ -120,6 +120,6 @@ def str_lit(value):
     return string_literal(value)
 
 
-def lit(value):
+def lit(value) -> Expr:
     """Create a literal expression."""
     return Expr.literal(value)

python/datafusion/context.py

@@ -719,7 +719,7 @@ def from_polars(self, data: polars.DataFrame, name: str | None = None) -> DataFr
 
     # https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
     # is the discussion on how we arrived at adding register_view
-    def register_view(self, name: str, df: DataFrame):
+    def register_view(self, name: str, df: DataFrame) -> None:
         """Register a :py:class: `~datafusion.detaframe.DataFrame` as a view.
 
         Args:

python/datafusion/expr.py

@@ -708,7 +708,7 @@ def get_lower_bound(self) -> WindowFrameBound:
         """Returns starting bound."""
         return WindowFrameBound(self.window_frame.get_lower_bound())
 
-    def get_upper_bound(self):
+    def get_upper_bound(self) -> WindowFrameBound:
         """Returns end bound."""
         return WindowFrameBound(self.window_frame.get_upper_bound())