unionai-oss / pandera

A light-weight, flexible, and expressive statistical data testing library
https://www.union.ai/pandera
MIT License
3.05k stars 281 forks source link

Allow check type #1598

Closed Hgherzog closed 2 weeks ago

Hgherzog commented 3 weeks ago

Is your feature request related to a problem? Please describe. I am unable to pass kwargs to a function with a pandera decorator like 'check_types'. I want to be able to flexibly use keywords when reading in a data. E.g @pa.check_types def load_df(path, **kwargs) -> Dataframe[DfModel]: return pd.read_csv(path, **kwargs) Currently, load_df would have {'kwargs': {'arg1': val}} if we passed in a kwarg to load_df.

Alternatively, I can just unpack the dictionary once inside load_df like kwargs=kwargs['kwargs].

I would like check types to be able to handle passing kwargs down into the function it wraps out of the box

cosmicBboy commented 2 weeks ago

hi @Hgherzog can you expand on what you want supported?

The following seems to work fine:

from io import StringIO
import pandera as pa
import pandas as pd

class DfModel(pa.DataFrameModel):
    a: int
    b: int

@pa.check_types
def load_df(path, **kwargs) -> pa.typing.DataFrame[DfModel]:
    return pd.read_csv(path, **kwargs)

data = StringIO("""a,b,c
1,2,'foo'
3,4,'bar'
""")

print(load_df(data, usecols=["a", "b"]))
"""
   a  b
0  1  2
1  3  4
"""