Repository: unionai-oss/pandera

Issue:
`pa.check_types` doesn't seem to work with Union of multiple pandera types
**Describe the bug**
A clear and concise description of what the bug is.

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the latest version of pandera.
- [ ] (optional) I have confirmed this bug exists on the master branch of pandera.

**Note**: Please read [this guide](https://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports) detailing how to provide the necessary information for us to reproduce your bug.

#### Code Sample, a copy-pastable example

```python
from typing import Union, List
import pytest
import pandas as pd
import pandera as pa
from pandera.typing import DataFrame, Series

class Model1(pa.SchemaModel):
    col1: Series[int]

class Model2(pa.SchemaModel):
    col2: Series[str]

@pytest.fixture
def good_dataframe():
    return pd.DataFrame({'col1': [1, 2, 3], 'col2': ['four', 'five', 'six']})

@pytest.fixture
def bad_dataframe():
    return pd.DataFrame({'col3': ['this', 'should', 'not', 'work']})

@pa.check_types
def my_function(df: Union[DataFrame[Model1], DataFrame[Model2]]) -> List[str]:
    return df.columns

def test_good_df(good_dataframe):
    cols = my_function(good_dataframe)
    assert sorted(cols) == sorted(['col1', 'col2'])

def test_bad_df(bad_dataframe):
    with pytest.raises(pa.errors.SchemaError):
        my_function(bad_dataframe)
```

#### Expected behavior
Because I am using `pa.check_types` on the `my_function` function above, I was expecting a pa.errors.SchemaError to be raised when I tried to pass in a dataframe to `my_function` that didn't conform to either of the SchemaModels that I've defined as being an either-or input to `my_function`.  However, it says that no SchemaError is raised by that function when I pass in `bad_dataframe`.

It seems like `Union` wipes out the ability of `pa.check_types` to work properly?  I couldn't find anything in the documentation that suggests this shouldn't be possible, but maybe I missed something?

#### Desktop (please complete the following information):

 - OS: MacOS 12.5.1
 - Pandera Version 0.13.3

#### Screenshots
Output when ran `pytest pandera_bug_test.py` with the above code copy-pasted into it.
<img width="915" alt="image" src="https://user-images.githubusercontent.com/15696062/197363319-61921b8a-7b66-4b36-bc64-845d49b6843a.png">



Discussion:
I played around with it a little and I can get a `pydantic.ValidationError` raised if I use `@pa.check_types(with_pydantic=True)`, but I'm not sure why it doesn't behave as I'd expect with just `pandera`.
thanks for reporting this bug @kr-hansen, yeah currently `check_types` doesn't know anything about `Union` (or other generics like `Optional` for that matter)

> I played around with it a little and I can get a pydantic.ValidationError raised if I use

Was this error what you were expecting?

In any case, the `check_types` function will need to be updated to handle `Union`s of `SchemaModel`s, right now it's just skipping it entirely: https://github.com/unionai-oss/pandera/blob/master/pandera/decorators.py#L535-L536

Gonna tag this with `help wanted`, would you be able to help with a contribution here @kr-hansen ?

Basically need to implement the logic you articulated here
> I was expecting a pa.errors.SchemaError to be raised when I tried to pass in a dataframe to my_function that didn't conform to either of the SchemaModels that I've defined as being an either-or input to my_function

So it would iterate through the schema models and break early if the current schema model passes, else continue to try to validate the data on the subsequent schemas in the `Union`. Still not sure what the error message should look like for this... any thoughts?
@cosmicBboy So the error in my test case outlined above is that I was expecting a `pandera.errors.SchemaError`.  If I write a similar test case without a Union in the input and pass it a non-conforming schema as input, that is the type of error I get in that case as `check_types` doesn't properly validate the input.  This works without requiring relying on `pydantic`.

I might look at trying to put together a contribution here sometime in the next week or so.  I'll note I haven't worked a ton with development on typing but am happy to learn and give it a shot.  What you outlined seems pretty straightforward that I could take a stab at it.  Do you guys have a CLA process I could kickstart with my company so that doesn't slow this down?

As for the error message, I'd assume something similar to what I get when it isn't a union.  Something like `error in check_types decorator of function 'function_name': Schema was not one of {union_list}`.  I assume that this last part following the colon would need to be encoded as `msg` somewhere to get passed in [here](https://github.com/unionai-oss/pandera/blob/d2f94165d17601743943d311eec19d363e651135/pandera/decorators.py#L77-L101)?  Or would that get added in a different portion?
> Do you guys have a CLA process I could kickstart with my company so that doesn't slow this down?

what's CLA?

> I'll note I haven't worked a ton with development on typing but am happy to learn and give it a shot

Also happy to help guide you through!

#### Steps

1. Add a conditional [here](https://github.com/unionai-oss/pandera/blob/master/pandera/decorators.py#L535-L536) that checks for `Union` types. For now, I think we should only support `Union[pandera.typing.DataFrame[S1], pandera.typing.DataFrame[S2]]`  types and not something like `Union[pandera.typing.DataFrame[S1], dict]`.
2. Generalize [`annotated_schema_models`](https://github.com/unionai-oss/pandera/blob/master/pandera/decorators.py#L539) by making it `Dict[str, Iterable[Tuple[SchemaModel, AnnotationInfo]]]` instead of `Dict[str, Tuple[SchemaModel, AnnotationInfo]]`.
3. Update [`_check_arg`](https://github.com/unionai-oss/pandera/blob/master/pandera/decorators.py#L541-L591) so that it works with the `Iterable[Tuple[SchemaModel, AnnotationInfo]]` 
4. `_check_arg` should be happy as soon as one of the schemas in the Union type passes validation. A `SchemaError/SchemaErrors` exception (depending on whether `lazy=True`) should be raised if none of the Schemas are valid.
5. Update [`_handle_schema_error`](https://github.com/unionai-oss/pandera/blob/master/pandera/decorators.py#L77-L101) to handle multiple `SchemaError/SchemaErrors` objects.

For (5) I'll leave it up to you to play around with what error message is the most useful: would it help to report the failure cases across all of the schemas in the Union type? If yes, will need to think harder about what the `SchemaError(s).failure_cases` attribute will contain: perhaps a `Dict[str, pd.DataFrame]` mapping schema name to failure case?
Sounds great.  A CLA is a [Contributor License Agreement](https://en.wikipedia.org/wiki/Contributor_License_Agreement) which we've typically had to fill out regarding IP/licensing stuff when contributing to Open Source projects, especially those hosted supported by other companies.  Typically just to clarify/formalize that the contributions to the project are still owned/maintained by the Open Source project itself.

I'll take a stab at the steps you outlined above and try to get a PR up on it.
Hi Kyle, I do the legal stuff at Union.  Do you think you should sign this, or should an officer in your company?  (We're using this site to gen up the agreement:  http://www.harmonyagreements.org).  We can do it either way, but it starts out saying that the person signing the agreement claims to have the right to bind the company legally.  
Thanks @georgesnelling.  I'll chat with someone on my legal team specifically about this scenario and respond back here.  I know our company has had this happen on Open Source contributions before, so I figured I'd ask ahead to help the ball get rolling on it.  I'll respond after I chat with someone at my company early next week.
Sounds good. Just so you know, it is not required on our side, but we are happy to configure one for you and / or your company if that is y'all's policy.   Pandera's license is MIT, which is about as unrestrictive as you can get.  
Hi Kyle, I think the information I gave you may be out-of-date.  We have to huddle on our side.  Likely we will add an built-in CLA into the checkin flow, but we have a little work left to do to put that in place.  

Resolving patch:
diff --git a/pandera/decorators.py b/pandera/decorators.py
--- a/pandera/decorators.py
+++ b/pandera/decorators.py
@@ -9,6 +9,7 @@
     Any,
     Callable,
     Dict,
+    Iterable,
     List,
     NoReturn,
     Optional,
@@ -24,6 +25,7 @@
 from pydantic import validate_arguments
 
 from . import errors, schemas
+from .error_handlers import SchemaErrorHandler
 from .inspection_utils import (
     is_classmethod_from_meta,
     is_decorated_classmethod,
@@ -84,6 +86,27 @@ def _handle_schema_error(
 ) -> NoReturn:
     """Reraise schema validation error with decorator context.
 
+    :param fn: check the DataFrame or Series input of this function.
+    :param schema: dataframe/series schema object
+    :param arg_df: dataframe/series we are validating.
+    :param schema_error: original exception.
+    :raises SchemaError: when ``DataFrame`` violates built-in or custom
+        checks.
+    """
+    raise _parse_schema_error(
+        decorator_name, fn, schema, arg_df, schema_error
+    ) from schema_error
+
+
+def _parse_schema_error(
+    decorator_name,
+    fn: Callable,
+    schema: Union[schemas.DataFrameSchema, schemas.SeriesSchema],
+    arg_df: pd.DataFrame,
+    schema_error: errors.SchemaError,
+) -> NoReturn:
+    """Parse schema validation error with decorator context.
+
     :param fn: check the DataFrame or Series input of this function.
     :param schema: dataframe/series schema object
     :param arg_df: dataframe/series we are validating.
@@ -95,14 +118,14 @@ def _handle_schema_error(
     if isinstance(fn, types.MethodType):
         func_name = fn.__self__.__class__.__name__ + "." + func_name
     msg = f"error in {decorator_name} decorator of function '{func_name}': {schema_error}"
-    raise errors.SchemaError(
+    return errors.SchemaError(  # type: ignore[misc]
         schema,
         arg_df,
         msg,
         failure_cases=schema_error.failure_cases,
         check=schema_error.check,
         check_index=schema_error.check_index,
-    ) from schema_error
+    )
 
 
 def check_input(
@@ -524,6 +547,7 @@ def check_types(
     lazy: bool = False,
     inplace: bool = False,
 ) -> Callable:
+    # pylint: disable=too-many-statements
     """Validate function inputs and output based on type annotations.
 
     See the :ref:`User Guide <schema_models>` for more.
@@ -558,66 +582,110 @@ def check_types(
         )
 
     # Front-load annotation parsing
-    annotated_schema_models: Dict[str, Tuple[SchemaModel, AnnotationInfo]] = {}
+    annotated_schema_models: Dict[
+        str,
+        Iterable[Tuple[Union[SchemaModel, None], Union[AnnotationInfo, None]]],
+    ] = {}
     for arg_name_, annotation in typing.get_type_hints(wrapped).items():
         annotation_info = AnnotationInfo(annotation)
         if not annotation_info.is_generic_df:
-            continue
+            # pylint: disable=comparison-with-callable
+            if annotation_info.origin == Union:
+                annotation_model_pairs = []
+                for annot in annotation_info.args:  # type: ignore[union-attr]
+                    sub_annotation_info = AnnotationInfo(annot)
+                    if not sub_annotation_info.is_generic_df:
+                        continue
+
+                    schema_model = cast(SchemaModel, sub_annotation_info.arg)
+                    annotation_model_pairs.append(
+                        (schema_model, sub_annotation_info)
+                    )
+            else:
+                continue
+        else:
+            schema_model = cast(SchemaModel, annotation_info.arg)
+            annotation_model_pairs = [(schema_model, annotation_info)]
 
-        schema_model = cast(SchemaModel, annotation_info.arg)
-        annotated_schema_models[arg_name_] = (schema_model, annotation_info)
+        annotated_schema_models[arg_name_] = annotation_model_pairs
 
     def _check_arg(arg_name: str, arg_value: Any) -> Any:
         """
         Validate function's argument if annoted with a schema, else
         pass-through.
         """
-        schema_model, annotation_info = annotated_schema_models.get(
-            arg_name, (None, None)
+        annotation_model_pairs = annotated_schema_models.get(
+            arg_name, [(None, None)]
         )
 
-        if schema_model is None:
+        if not annotation_model_pairs:
             return arg_value
 
-        if (
-            annotation_info
-            and not (annotation_info.optional and arg_value is None)
-            # the pandera.schema attribute should only be available when
-            # schema.validate has been called in the DF. There's probably
-            # a better way of doing this
-        ):
-            config = schema_model.__config__
-            data_container_type = annotation_info.origin
-            schema = schema_model.to_schema()
-
-            if data_container_type and config and config.from_format:
-                arg_value = data_container_type.from_format(arg_value, config)
+        error_handler = SchemaErrorHandler(lazy=True)
+        for schema_model, annotation_info in annotation_model_pairs:
+            if schema_model is None:
+                return arg_value
 
             if (
-                arg_value.pandera.schema is None
-                # don't re-validate a dataframe that contains the same exact
-                # schema
-                or arg_value.pandera.schema != schema
+                annotation_info
+                and not (annotation_info.optional and arg_value is None)
+                # the pandera.schema attribute should only be available when
+                # schema.validate has been called in the DF. There's probably
+                # a better way of doing this
             ):
-                try:
-                    arg_value = schema.validate(
-                        arg_value,
-                        head,
-                        tail,
-                        sample,
-                        random_state,
-                        lazy,
-                        inplace,
+                config = schema_model.__config__
+                data_container_type = annotation_info.origin
+                schema = schema_model.to_schema()
+
+                if data_container_type and config and config.from_format:
+                    arg_value = data_container_type.from_format(
+                        arg_value, config
                     )
-                except errors.SchemaError as e:
-                    _handle_schema_error(
-                        "check_types", wrapped, schema, arg_value, e
+
+                # Don't do checks if value is still a built-in type
+                if isinstance(
+                    arg_value, (int, str, bool, float, dict, list, tuple, set)
+                ):
+                    return arg_value
+
+                if (
+                    arg_value.pandera.schema is None
+                    # don't re-validate a dataframe that contains the same exact
+                    # schema
+                    or arg_value.pandera.schema != schema
+                ):
+                    try:
+                        arg_value = schema.validate(
+                            arg_value,
+                            head,
+                            tail,
+                            sample,
+                            random_state,
+                            lazy,
+                            inplace,
+                        )
+                    except errors.SchemaError as e:
+                        error_handler.collect_error(
+                            "invalid_type",
+                            _parse_schema_error(
+                                "check_types", wrapped, schema, arg_value, e
+                            ),
+                        )
+                        continue
+
+                if data_container_type and config and config.to_format:
+                    arg_value = data_container_type.to_format(
+                        arg_value, config
                     )
 
-            if data_container_type and config and config.to_format:
-                arg_value = data_container_type.to_format(arg_value, config)
+                return arg_value
 
-            return arg_value
+        if error_handler.collected_errors:
+            if len(error_handler.collected_errors) == 1:
+                raise error_handler.collected_errors[0]["error"]  # type: ignore[misc]
+            raise errors.SchemaErrors(
+                schema, error_handler.collected_errors, arg_value
+            )
 
     sig = inspect.signature(wrapped)
 
diff --git a/pandera/typing/common.py b/pandera/typing/common.py
--- a/pandera/typing/common.py
+++ b/pandera/typing/common.py
@@ -208,6 +208,7 @@ class AnnotationInfo:  # pylint:disable=too-few-public-methods
 
     Attributes:
         origin: The non-parameterized generic class.
+        args: All generic types for accessing as an iterable.
         arg: The first generic type (SchemaModel does not support more than
             1 argument).
         literal: Whether the annotation is a literal.
@@ -247,6 +248,7 @@ def _parse_annotation(self, raw_annotation: Type) -> None:
         self.origin = typing_inspect.get_origin(raw_annotation)
         # Replace empty tuple returned from get_args by None
         args = typing_inspect.get_args(raw_annotation) or None
+        self.args = args
         self.arg = args[0] if args else args
 
         self.metadata = getattr(self.arg, "__metadata__", None)


Test patch:
diff --git a/tests/core/test_decorators.py b/tests/core/test_decorators.py
--- a/tests/core/test_decorators.py
+++ b/tests/core/test_decorators.py
@@ -514,6 +514,12 @@ class OnlyZeroesSchema(SchemaModel):  # pylint:disable=too-few-public-methods
     a: Series[int] = Field(eq=0)
 
 
+class OnlyOnesSchema(SchemaModel):  # pylint:disable=too-few-public-methods
+    """Schema with a single column containing ones."""
+
+    a: Series[int] = Field(eq=1)
+
+
 def test_check_types_arguments() -> None:
     """Test that check_types forwards key-words arguments to validate."""
     df = pd.DataFrame({"a": [0, 0]})
@@ -632,7 +638,6 @@ def transform(df: DataFrame[InSchema]):
         transform(df)
     except errors.SchemaError as exc:
         assert exc.schema == InSchema.to_schema()
-        assert exc.data.equals(df)
 
 
 def test_check_types_error_output() -> None:
@@ -655,7 +660,6 @@ def transform_derived(
         transform_derived(df)
     except errors.SchemaError as exc:
         assert exc.schema == DerivedOutSchema.to_schema()
-        assert exc.data.equals(df)
 
     df = pd.DataFrame({"a": [1]}, index=["1"])
 
@@ -672,7 +676,6 @@ def transform(df: DataFrame[InSchema]) -> DataFrame[OutSchema]:
         transform(df)
     except errors.SchemaError as exc:
         assert exc.schema == OutSchema.to_schema()
-        assert exc.data.equals(df)
 
 
 def test_check_types_optional_out() -> None:
@@ -884,6 +887,73 @@ def static_method(
         instance.static_method(df1=in2, df2=in1)
 
 
+def test_check_types_union_args() -> None:
+    """Test that the @check_types decorator works with
+    typing.Union[pandera.typing.DataFrame[S1], pandera.typing.DataFrame[S2]] type inputs/outputs"""
+
+    @check_types
+    def validate_union(
+        df: typing.Union[
+            DataFrame[OnlyZeroesSchema],
+            DataFrame[OnlyOnesSchema],
+        ],
+    ) -> typing.Union[DataFrame[OnlyZeroesSchema], DataFrame[OnlyOnesSchema]]:
+        return df
+
+    validate_union(pd.DataFrame({"a": [0, 0]}))
+    validate_union(pd.DataFrame({"a": [1, 1]}))
+
+    with pytest.raises(errors.SchemaErrors):
+        validate_union(pd.DataFrame({"a": [0, 1]}))
+    with pytest.raises(errors.SchemaErrors):
+        validate_union(pd.DataFrame({"a": [2, 2]}))
+
+    @check_types
+    def validate_union_wrong_outputs(
+        df: typing.Union[
+            DataFrame[OnlyZeroesSchema], DataFrame[OnlyOnesSchema]
+        ],
+    ) -> typing.Union[DataFrame[OnlyZeroesSchema], DataFrame[OnlyOnesSchema]]:
+        new_df = df.copy()
+        new_df["a"] = [0, 1]
+        return new_df
+
+    with pytest.raises(errors.SchemaErrors):
+        validate_union_wrong_outputs(pd.DataFrame({"a": [0, 0]}))
+
+
+def test_check_types_non_dataframes() -> None:
+    """Test to skip check_types for non-dataframes"""
+
+    @check_types
+    def only_int_type(val: int) -> int:
+        return val
+
+    @check_types
+    def union_int_str_types(
+        val: typing.Union[int, str]
+    ) -> typing.Union[int, str]:
+        return val
+
+    only_int_type(1)
+    int_val = union_int_str_types(2)
+    str_val = union_int_str_types("2")
+    assert isinstance(int_val, int)
+    assert isinstance(str_val, str)
+
+    @check_types(with_pydantic=True)
+    def union_df_int_types_pydantic_check(
+        val: typing.Union[DataFrame[OnlyZeroesSchema], int]
+    ) -> typing.Union[DataFrame[OnlyZeroesSchema], int]:
+        return val
+
+    union_df_int_types_pydantic_check(pd.DataFrame({"a": [0, 0]}))
+    int_val_pydantic = union_df_int_types_pydantic_check(5)
+    str_val_pydantic = union_df_int_types_pydantic_check("5")  # type: ignore[arg-type]
+    assert isinstance(int_val_pydantic, int)
+    assert isinstance(str_val_pydantic, int)
+
+
 def test_coroutines(event_loop: AbstractEventLoop) -> None:
     # pylint: disable=missing-class-docstring,too-few-public-methods,missing-function-docstring
     class Schema(SchemaModel):