Bugfix: #2058 Check_types for callable (#2069)

* Enhance parameterization for one of the tests #2058

Signed-off-by: Yaakov Bressler <yaakovbressler@gmail.com>
Signed-off-by: ybressler <40807730+ybressler@users.noreply.github.com>

* Test cases which we expect to pass

Signed-off-by: ybressler <40807730+ybressler@users.noreply.github.com>

* Fix failing tests - use different method for determining if type is literal

Signed-off-by: ybressler <40807730+ybressler@users.noreply.github.com>

* Update tests/pandas/test_decorators.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Signed-off-by: Yaakov Bressler <yaakovbressler@gmail.com>
Signed-off-by: ybressler <40807730+ybressler@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: ybressler

pandera/typing/common.py

@@ -2,6 +2,7 @@
 
 import copy
 import inspect
+import typing
 from typing import (  # type: ignore[attr-defined]
     TYPE_CHECKING,
     Any,
@@ -238,7 +239,8 @@ def _parse_annotation(self, raw_annotation: type) -> None:
             self.arg = typing_inspect.get_args(self.arg)[0]
 
         self.metadata = metadata
-        self.literal = typing_inspect.is_literal_type(self.arg)
+
+        self.literal = typing_inspect.get_origin(self.arg) is typing.Literal
 
         if self.literal:
             self.arg = typing_inspect.get_args(self.arg)[0]

tests/pandas/test_decorators.py

@@ -3,6 +3,7 @@
 import asyncio
 import pickle
 import typing
+from contextlib import nullcontext
 
 import numpy as np
 import pandas as pd
@@ -536,42 +537,67 @@ class OnlyOnesSchema(DataFrameModel):
     a: Series[int] = Field(eq=1)
 
 
-def test_check_types_arguments() -> None:
+@pytest.mark.parametrize(
+    ("check_types_args", "df_return", "expected"),
+    [
+        pytest.param(
+            dict(),
+            pd.DataFrame({"a": [0, 0]}),
+            nullcontext(),
+            id="validate entire df (2 rows)",
+        ),
+        pytest.param(
+            dict(head=1),
+            pd.DataFrame({"a": [0, 1]}),
+            nullcontext(),
+            id="validate header (row 1) - while row 2 is bad",
+        ),
+        pytest.param(
+            dict(tail=1),
+            pd.DataFrame({"a": [1, 0]}),
+            nullcontext(),
+            id="validate tail, (row 2) - while row 1 is bad",
+        ),
+        pytest.param(
+            dict(lazy=True),
+            pd.DataFrame({"a": [0, 0]}),
+            nullcontext(),
+            id="validate entire df (2 rows) - lazy mode ",
+        ),
+        pytest.param(
+            dict(lazy=True),
+            pd.DataFrame({"a": [1, 1]}),
+            pytest.raises(
+                errors.SchemaErrors,
+                match=r"DATA",  # error msg is specific for lazy-
+            ),
+            id="1's not allowed in schema - lazy mode",
+        ),
+        pytest.param(
+            dict(),
+            pd.DataFrame({"a": [1, 1]}),
+            pytest.raises(
+                errors.SchemaError,
+                match=r"failed element-wise validator",  # error msg is specific for regular-mode
+            ),
+            id="1's not allowed in schema - regular mode",
+        ),
+    ],
+)
+def test_check_types_arguments(
+    check_types_args: dict, df_return: pd.DataFrame, expected
+) -> None:
     """Test that check_types forwards key-words arguments to validate."""
     df = pd.DataFrame({"a": [0, 0]})
 
-    @check_types()
-    def transform_empty_parenthesis(
+    @check_types(**check_types_args)
+    def transform_with_checks(
         df: DataFrame[OnlyZeroesSchema],
     ) -> DataFrame[OnlyZeroesSchema]:  # pylint: disable=unused-argument
-        return df
-
-    transform_empty_parenthesis(df)  # type: ignore
-
-    @check_types(head=1)
-    def transform_head(
-        df: DataFrame[OnlyZeroesSchema],  # pylint: disable=unused-argument
-    ) -> DataFrame[OnlyZeroesSchema]:
-        return pd.DataFrame({"a": [0, 0]})  # type: ignore
-
-    transform_head(df)  # type: ignore
-
-    @check_types(tail=1)
-    def transform_tail(
-        df: DataFrame[OnlyZeroesSchema],  # pylint: disable=unused-argument
-    ) -> DataFrame[OnlyZeroesSchema]:
-        return pd.DataFrame({"a": [1, 0]})  # type: ignore
-
-    transform_tail(df)  # type: ignore
+        return df_return
 
-    @check_types(lazy=True)
-    def transform_lazy(
-        df: DataFrame[OnlyZeroesSchema],  # pylint: disable=unused-argument
-    ) -> DataFrame[OnlyZeroesSchema]:
-        return pd.DataFrame({"a": [1, 1]})  # type: ignore
-
-    with pytest.raises(errors.SchemaErrors, match=r"DATA"):
-        transform_lazy(df)  # type: ignore
+    with expected:
+        transform_with_checks(df)  # type: ignore
 
 
 def test_check_types_unchanged() -> None:
@@ -754,6 +780,39 @@ def transform(
     assert transform(None) is None
 
 
+@pytest.mark.parametrize(
+    "callable_annotation",
+    [
+        pytest.param(typing.Callable[[None], None], id="no args, no return"),
+        pytest.param(typing.Callable[[None], int], id="no args, returns int"),
+        pytest.param(
+            typing.Callable[..., int], id="no info on args, returns int"
+        ),
+        pytest.param(
+            typing.Callable[..., list[int]],
+            id="no info on args, returns list of int",
+        ),
+        pytest.param(
+            typing.Callable[[typing.Any], int],
+            id="includes info on callable args",
+        ),
+    ],
+)
+def test_check_types_callables(callable_annotation: typing.Callable) -> None:
+    """
+    Ensures `check_types` validates a dataframe, while passing in an additional callable argument
+    """
+
+    class MySchema1(DataFrameModel):
+        a: int
+
+    @check_types
+    def some_transformation(df: MySchema1, f: callable_annotation):  # type: ignore[valid-type]
+        pass
+
+    _ = some_transformation(pd.DataFrame({"a": [1, 2]}), lambda x: 1)
+
+
 def test_check_types_coerce() -> None:
     """Test that check_types return the result of validate."""