Can't pickle function with @check_types decorator

[bmwilly]

Describe the bug

Functions cannot be pickled if they are decorated with @pa.check_types.

• I have checked that this issue has not already been reported.
  • @pa.check gives pickle error when using along with pandarallel #1217 is possibly related, but it's old, unclear, and the issue author is not being responsive.
• I have confirmed this bug exists on the latest version of pandera.
• (optional) I have confirmed this bug exists on the main branch of pandera.

Note: Please read this guide detailing how to provide the necessary information for us to reproduce your bug.

Code Sample, a copy-pastable example

import pickle
from pathlib import Path

import pandas as pd
import pandera as pa
from pandera.typing import DataFrame, Series


pa.__version__


# Define a schema using pandera
class Schema(pa.DataFrameModel):
    column1: Series[int]
    column2: Series[float]
    column3: Series[str]


# Use the check_types decorator to validate input and output
@pa.check_types
def process_data_and_check_types(df: DataFrame[Schema]) -> DataFrame[Schema]:
    # Example processing: add a new column
    df["column4"] = df["column1"] + df["column2"]
    return df


# Create a sample DataFrame
data = {"column1": [1, 2, 3], "column2": [0.1, 0.2, 0.3], "column3": ["a", "b", "c"]}
df = pd.DataFrame(data)

# Process the DataFrame
processed_df = process_data_and_check_types(df)


# Same function without checking types
def process_data(df: DataFrame[Schema]) -> DataFrame[Schema]:
    # Example processing: add a new column
    df["column4"] = df["column1"] + df["column2"]
    return df


# Can only serialize function without decorator
path = Path("/tmp/tmp.pkl")
with path.open("wb") as f:
    pickle.dump(process_data, f)

# This will raise an error
with path.open("wb") as f:
    pickle.dump(process_data_and_check_types, f)

Expected behavior

Ability to pickle a function decorated with @pa.check_types.

Desktop (please complete the following information):

• OS: macOS 15.1
• Version: pandera 0.20.4

Screenshots

Error from above code:

{
	"name": "NotImplementedError",
	"message": "object proxy must define __reduce_ex__()",
	"stack": "---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[6], line 2
      1 with path.open(\"wb\") as f:
----> 2     pickle.dump(process_data_and_check_types, f)

NotImplementedError: object proxy must define __reduce_ex__()"
}

Additional context

This functionality is required to use a variety of ML frameworks (such as MLFlow and Metaflow) that pickle artifacts under the hood in order to log them.