Merge pull request #20 from vertti/improve-types

Improve types

Author: vertti

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 All notable changes to this project will be documented in this file.
 
+## 0.11.0
+
+- Update function parameter types for better type safety
+- Fix missing return statement in df_log decorator
+- Added stricter mypy type checking settings
+
 ## 0.10.1
 
 - Built and published with UV. No functional changes

README.md

@@ -28,6 +28,7 @@ Like type hints for DataFrames, Daffy helps you catch structural mismatches earl
 - Works with both Pandas and Polars DataFrames
 - Project-wide configuration via pyproject.toml
 - Integrated logging for DataFrame structure inspection
+- Enhanced type annotations for improved IDE and type checker support
 
 ## Documentation
 

daffy/config.py

@@ -1,12 +1,12 @@
 """Configuration handling for DAFFY."""
 
 import os
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import tomli
 
 
-def load_config() -> dict:
+def load_config() -> Dict[str, Any]:
     """
     Load daffy configuration from pyproject.toml.
 
@@ -61,7 +61,7 @@ def find_config_file() -> Optional[str]:
 _config_cache = None
 
 
-def get_config() -> dict:
+def get_config() -> Dict[str, Any]:
     """
     Get the daffy configuration, loading it if necessary.
 

daffy/decorators.py

@@ -4,17 +4,26 @@
 import logging
 import re
 from functools import wraps
-from typing import Any, Callable, Dict, List, Optional, Pattern, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Pattern, Tuple, TypeVar, Union
 
 import pandas as pd
 import polars as pl
 
+# Import fully qualified types to satisfy disallow_any_unimported
+from pandas import DataFrame as PandasDataFrame
+from polars import DataFrame as PolarsDataFrame
+
 from daffy.config import get_strict
 
-# New type definition to support regex patterns
-RegexColumnDef = Tuple[str, Pattern]  # Tuple of (pattern_str, compiled_pattern)
-ColumnsDef = Union[List, Dict, List[Union[str, RegexColumnDef]]]
-DataFrameType = Union[pd.DataFrame, pl.DataFrame]
+# Type variables for preserving return types
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+# Improved type definitions to support regex patterns
+RegexColumnDef = Tuple[str, Pattern[str]]  # Tuple of (pattern_str, compiled_pattern)
+ColumnsDef = Union[List[Union[str, RegexColumnDef]], Dict[str, Any]]
+DataFrameType = Union[PandasDataFrame, PolarsDataFrame]
 
 
 def _is_regex_pattern(column: Any) -> bool:
@@ -30,9 +39,9 @@ def _match_column_with_regex(column_pattern: RegexColumnDef, df_columns: List[st
     return [col for col in df_columns if pattern.match(col)]
 
 
-def _compile_regex_patterns(columns: List) -> List:
+def _compile_regex_patterns(columns: List[Any]) -> List[Union[str, RegexColumnDef]]:
     """Compile regex patterns in the column list."""
-    result = []
+    result: List[Union[str, RegexColumnDef]] = []
     for col in columns:
         if isinstance(col, str) and col.startswith("r/") and col.endswith("/"):
             # Pattern is in the format "r/pattern/"
@@ -97,7 +106,9 @@ def _check_columns(df: DataFrameType, columns: ColumnsDef, strict: bool) -> None
             raise AssertionError(f"DataFrame contained unexpected column(s): {', '.join(extra_columns)}")
 
 
-def df_out(columns: Optional[ColumnsDef] = None, strict: Optional[bool] = None) -> Callable:
+def df_out(
+    columns: Optional[ColumnsDef] = None, strict: Optional[bool] = None
+) -> Callable[[Callable[..., DataFrameType]], Callable[..., DataFrameType]]:
     """Decorate a function that returns a Pandas or Polars DataFrame.
 
     Document the return value of a function. The return value will be validated in runtime.
@@ -109,12 +120,12 @@ def df_out(columns: Optional[ColumnsDef] = None, strict: Optional[bool] = None)
             If None, uses the value from [tool.daffy] strict setting in pyproject.toml.
 
     Returns:
-        Callable: Decorated function
+        Callable: Decorated function with preserved DataFrame return type
     """
 
-    def wrapper_df_out(func: Callable) -> Callable:
+    def wrapper_df_out(func: Callable[..., DataFrameType]) -> Callable[..., DataFrameType]:
         @wraps(func)
-        def wrapper(*args: str, **kwargs: Any) -> Any:
+        def wrapper(*args: Any, **kwargs: Any) -> DataFrameType:
             result = func(*args, **kwargs)
             assert isinstance(result, pd.DataFrame) or isinstance(result, pl.DataFrame), (
                 f"Wrong return type. Expected DataFrame, got {type(result)}"
@@ -128,7 +139,7 @@ def wrapper(*args: str, **kwargs: Any) -> Any:
     return wrapper_df_out
 
 
-def _get_parameter(func: Callable, name: Optional[str] = None, *args: str, **kwargs: Any) -> DataFrameType:
+def _get_parameter(func: Callable[..., Any], name: Optional[str] = None, *args: Any, **kwargs: Any) -> Any:
     if not name:
         if len(args) > 0:
             return args[0]
@@ -144,7 +155,9 @@ def _get_parameter(func: Callable, name: Optional[str] = None, *args: str, **kwa
     return kwargs[name]
 
 
-def df_in(name: Optional[str] = None, columns: Optional[ColumnsDef] = None, strict: Optional[bool] = None) -> Callable:
+def df_in(
+    name: Optional[str] = None, columns: Optional[ColumnsDef] = None, strict: Optional[bool] = None
+) -> Callable[[Callable[..., R]], Callable[..., R]]:
     """Decorate a function parameter that is a Pandas or Polars DataFrame.
 
     Document the contents of an input parameter. The parameter will be validated in runtime.
@@ -157,12 +170,12 @@ def df_in(name: Optional[str] = None, columns: Optional[ColumnsDef] = None, stri
             If None, uses the value from [tool.daffy] strict setting in pyproject.toml.
 
     Returns:
-        Callable: Decorated function
+        Callable: Decorated function with preserved return type
     """
 
-    def wrapper_df_in(func: Callable) -> Callable:
+    def wrapper_df_in(func: Callable[..., R]) -> Callable[..., R]:
         @wraps(func)
-        def wrapper(*args: str, **kwargs: Any) -> Any:
+        def wrapper(*args: Any, **kwargs: Any) -> R:
             df = _get_parameter(func, name, *args, **kwargs)
             assert isinstance(df, pd.DataFrame) or isinstance(df, pl.DataFrame), (
                 f"Wrong parameter type. Expected DataFrame, got {type(df).__name__} instead."
@@ -203,7 +216,7 @@ def _log_output(level: int, func_name: str, df: Any, include_dtypes: bool) -> No
         )
 
 
-def df_log(level: int = logging.DEBUG, include_dtypes: bool = False) -> Callable:
+def df_log(level: int = logging.DEBUG, include_dtypes: bool = False) -> Callable[[Callable[..., T]], Callable[..., T]]:
     """Decorate a function that consumes or produces a Pandas DataFrame or both.
 
     Logs the columns of the consumed and/or produced DataFrame.
@@ -213,15 +226,16 @@ def df_log(level: int = logging.DEBUG, include_dtypes: bool = False) -> Callable
         include_dtypes (bool, optional): When set to True, will log also the dtypes of each column. Defaults to False.
 
     Returns:
-        Callable: Decorated function.
+        Callable: Decorated function with preserved return type.
     """
 
-    def wrapper_df_log(func: Callable) -> Callable:
+    def wrapper_df_log(func: Callable[..., T]) -> Callable[..., T]:
         @wraps(func)
-        def wrapper(*args: str, **kwargs: Any) -> Any:
+        def wrapper(*args: Any, **kwargs: Any) -> T:
             _log_input(level, func.__name__, _get_parameter(func, None, *args, **kwargs), include_dtypes)
             result = func(*args, **kwargs)
             _log_output(level, func.__name__, result, include_dtypes)
+            return result  # Added missing return statement
 
         return wrapper
 

mypy.ini

@@ -1,12 +1,30 @@
 [mypy]
 python_version = 3.10
-warn_return_any = True
+# Enable all strict options
+strict = True
+# The strict flag enables:
+# - disallow_untyped_calls
+# - disallow_untyped_defs
+# - disallow_incomplete_defs
+# - check_untyped_defs
+# - disallow_untyped_decorators
+# - no_implicit_optional
+# - warn_redundant_casts
+# - warn_unused_ignores
+# - warn_no_return
+# - warn_return_any
+# - warn_unreachable
+# - disallow_any_generics
+# Additional strict options not included in strict flag
+warn_incomplete_stub = True
+warn_redundant_casts = True
 warn_unused_configs = True
-ignore_missing_imports = True
-disallow_untyped_calls = True
-disallow_untyped_defs = True
-disallow_incomplete_defs = True
-check_untyped_defs = True
 warn_unused_ignores = True
-warn_redundant_casts = True
-no_implicit_optional = True
+implicit_reexport = False
+strict_optional = True
+strict_equality = True
+
+# Relax rules for tests
+[mypy-tests.*]
+disallow_any_unimported = False
+disallow_any_decorated = False

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "daffy"
-version = "0.10.1"
+version = "0.11.0"
 description = "Function decorators for Pandas and Polars Dataframe column name and data type validation"
 authors = [
  { name="Janne Sinivirta", email="[email protected]" },
@@ -42,6 +42,7 @@ dev = [
     "coverage[toml]>=7.3.2",
     "pydocstyle>=6.3.0",
     "ruff>=0.9.1",
+    "pandas-stubs>=2.2.2.240807",
 ]
 
 [build-system]
@@ -83,4 +84,4 @@ line-length = 120
 target-version = "py39"
 
 [tool.ruff.lint]
-select = ["F", "E", "W", "I", "N"]
+select = ["F", "E", "W", "I", "N"]

tests/test_decorators.py

@@ -1,5 +1,3 @@
-from typing import Any
-
 import pandas as pd
 
 from daffy import df_in, df_out
@@ -8,8 +6,10 @@
 def test_decorator_combinations(basic_pandas_df: pd.DataFrame, extended_pandas_df: pd.DataFrame) -> None:
     @df_in(columns=["Brand", "Price"])
     @df_out(columns=["Brand", "Price", "Year"])
-    def test_fn(my_input: Any) -> Any:
+    def test_fn(my_input: pd.DataFrame) -> pd.DataFrame:
         my_input["Year"] = list(extended_pandas_df["Year"])
         return my_input
 
-    pd.testing.assert_frame_equal(extended_pandas_df, test_fn(basic_pandas_df.copy()))
+    result = test_fn(basic_pandas_df.copy())
+    assert isinstance(result, pd.DataFrame)  # Confirm the type for mypy
+    pd.testing.assert_frame_equal(extended_pandas_df, result)

tests/test_df_log.py

@@ -60,7 +60,7 @@ def test_fn(foo_df: pd.DataFrame) -> pd.DataFrame:
 
 def test_log_df_with_dtypes_polars(basic_polars_df: pl.DataFrame, mocker: MockerFixture) -> None:
     @df_log(include_dtypes=True)
-    def test_fn(foo_df: pd.DataFrame) -> pd.DataFrame:
+    def test_fn(foo_df: pl.DataFrame) -> pl.DataFrame:
         return basic_polars_df
 
     mock_log = mocker.patch("daffy.decorators.logging.log")

tests/test_df_out.py

@@ -9,7 +9,7 @@
 
 
 def test_wrong_return_type() -> None:
-    @df_out()
+    @df_out()  # type: ignore[arg-type]
     def test_fn() -> int:
         return 1
 
@@ -86,7 +86,10 @@ def test_fn(my_input: Any) -> Any:
         return my_input
 
     assert list(basic_pandas_df.columns) == ["Brand", "Price"]  # For sanity
-    pd.testing.assert_frame_equal(extended_pandas_df, test_fn(basic_pandas_df.copy()))
+    result = test_fn(basic_pandas_df.copy())
+    # Type check to ensure we get a pandas DataFrame before comparing
+    assert isinstance(result, pd.DataFrame)
+    pd.testing.assert_frame_equal(extended_pandas_df, result)
 
 
 def test_regex_column_pattern_in_output(basic_pandas_df: pd.DataFrame) -> None: