Test your AI Agents with the all-new Agent to Agent Testing Platform.Learn More

How to use custom_check_teardown method in pandera

Best Python code snippet using pandera_python

test_extensions.py

Source:test_extensions.py

1# pylint: disable=no-member,redefined-outer-name,unused-argument2# pylint: disable=unused-variable3"""Unit tests for pandera API extensions."""4from typing import Any, Optional, Union5import pandas as pd6import pytest7import pandera as pa8import pandera.strategies as st9from pandera import DataType, extensions10from pandera.checks import Check11def test_custom_checks_in_dir(extra_registered_checks):12    """Ensures that autocomplete works with registered custom checks."""13    assert "no_param_check" in dir(pa.Check)14@pytest.mark.parametrize(15    "data",16    [17        pd.Series([10, 10, 10]),18        pd.DataFrame([[10, 10, 10], [10, 10, 10]]),19    ],20)21def test_register_vectorized_custom_check(22    custom_check_teardown: None, data: Union[pd.Series, pd.DataFrame]23) -> None:24    """Test registering a vectorized custom check."""25    @extensions.register_check_method(26        statistics=["val"],27        supported_types=(pd.Series, pd.DataFrame),28        check_type="vectorized",29    )30    def custom_check(pandas_obj, *, val):31        return pandas_obj == val32    check = Check.custom_check(val=10)33    check_result = check(data)34    assert check_result.check_passed35    for kwargs in [36        {"element_wise": True},37        {"element_wise": False},38        {"groupby": "column"},39        {"groups": ["group1", "group2"]},40    ]:41        with pytest.warns(UserWarning):42            Check.custom_check(val=10, **kwargs)43    with pytest.raises(44        ValueError,45        match="method with name 'custom_check' already defined",46    ):47        # pylint: disable=function-redefined48        @extensions.register_check_method(statistics=["val"])49        def custom_check(pandas_obj, val):  # noqa50            return pandas_obj != val51@pytest.mark.parametrize(52    "data",53    [54        pd.Series([10, 10, 10]),55        pd.DataFrame([[10, 10, 10], [10, 10, 10]]),56    ],57)58def test_register_element_wise_custom_check(59    custom_check_teardown: None, data: Union[pd.Series, pd.DataFrame]60) -> None:61    """Test registering an element-wise custom check."""62    @extensions.register_check_method(63        statistics=["val"],64        supported_types=(pd.Series, pd.DataFrame),65        check_type="element_wise",66    )67    def custom_check(element, *, val):68        return element == val69    check = Check.custom_check(val=10)70    check_result = check(data)71    assert check_result.check_passed72    for kwargs in [73        {"element_wise": True},74        {"element_wise": False},75        {"groupby": "column"},76        {"groups": ["group1", "group2"]},77    ]:78        with pytest.warns(UserWarning):79            Check.custom_check(val=10, **kwargs)80    with pytest.raises(81        ValueError,82        match="Element-wise checks should support DataFrame and Series "83        "validation",84    ):85        @extensions.register_check_method(86            supported_types=pd.Series,87            check_type="element_wise",88        )89        def invalid_custom_check(*args):90            pass91def test_register_custom_groupby_check(custom_check_teardown: None) -> None:92    """Test registering a custom groupby check."""93    @extensions.register_check_method(94        statistics=["group_a", "group_b"],95        supported_types=(pd.Series, pd.DataFrame),96        check_type="groupby",97    )98    def custom_check(dict_groups, *, group_a, group_b):99        """100        Test that the mean values in group A is larger than that of group B.101        Note that this function can handle groups of both dataframes and102        series.103        """104        return (105            dict_groups[group_a].values.mean()106            > dict_groups[group_b].values.mean()107        )108    # column groupby check109    data_column_check = pd.DataFrame(110        {111            "col1": [20, 20, 10, 10],112            "col2": list("aabb"),113        }114    )115    schema_column_check = pa.DataFrameSchema(116        {117            "col1": pa.Column(118                int,119                Check.custom_check(group_a="a", group_b="b", groupby="col2"),120            ),121            "col2": pa.Column(str),122        }123    )124    assert isinstance(schema_column_check(data_column_check), pd.DataFrame)125    # dataframe groupby check126    data_df_check = pd.DataFrame(127        {128            "col1": [20, 20, 10, 10],129            "col2": [30, 30, 5, 5],130            "col3": [10, 10, 1, 1],131        },132        index=pd.Index(list("aabb"), name="my_index"),133    )134    schema_df_check = pa.DataFrameSchema(135        columns={136            "col1": pa.Column(int),137            "col2": pa.Column(int),138            "col3": pa.Column(int),139        },140        index=pa.Index(str, name="my_index"),141        checks=Check.custom_check(142            group_a="a", group_b="b", groupby="my_index"143        ),144    )145    assert isinstance(schema_df_check(data_df_check), pd.DataFrame)146    for kwargs in [{"element_wise": True}, {"element_wise": False}]:147        with pytest.warns(UserWarning):148            Check.custom_check(val=10, **kwargs)149@pytest.mark.parametrize(150    "supported_types",151    [152        1,153        10.0,154        "foo",155        {"foo": "bar"},156        {1: 10},157        ["foo", "bar"],158        [1, 10],159        ("foo", "bar"),160        (1, 10),161    ],162)163def test_register_check_invalid_supported_types(supported_types: Any) -> None:164    """Test that TypeError is raised on invalid supported_types arg."""165    with pytest.raises(TypeError):166        @extensions.register_check_method(supported_types=supported_types)167        def custom_check(*args, **kwargs):168            pass169@pytest.mark.skipif(170    not st.HAS_HYPOTHESIS, reason='needs "strategies" module dependencies'171)172def test_register_check_with_strategy(custom_check_teardown: None) -> None:173    """Test registering a custom check with a data generation strategy."""174    import hypothesis  # pylint: disable=import-outside-toplevel,import-error175    def custom_ge_strategy(176        pandas_dtype: DataType,177        strategy: Optional[st.SearchStrategy] = None,178        *,179        min_value: Any,180    ) -> st.SearchStrategy:181        if strategy is None:182            return st.pandas_dtype_strategy(183                pandas_dtype,184                min_value=min_value,185                exclude_min=False,186            )187        return strategy.filter(lambda x: x > min_value)188    @extensions.register_check_method(189        statistics=["min_value"], strategy=custom_ge_strategy190    )191    def custom_ge_check(pandas_obj, *, min_value):192        return pandas_obj >= min_value193    check = Check.custom_ge_check(min_value=0)194    strat = check.strategy(pa.Int)195    with pytest.warns(hypothesis.errors.NonInteractiveExampleWarning):196        assert strat.example() >= 0197def test_schema_model_field_kwarg(custom_check_teardown: None) -> None:198    """Test that registered checks can be specified in a Field."""199    # pylint: disable=missing-class-docstring,too-few-public-methods200    @extensions.register_check_method(201        statistics=["val"],202        supported_types=(pd.Series, pd.DataFrame),203        check_type="vectorized",204    )205    def custom_gt(pandas_obj, val):206        return pandas_obj > val207    @extensions.register_check_method(208        statistics=["min_value", "max_value"],209        supported_types=(pd.Series, pd.DataFrame),210        check_type="vectorized",211    )212    def custom_in_range(pandas_obj, min_value, max_value):213        return (min_value <= pandas_obj) & (pandas_obj <= max_value)214    class Schema(pa.SchemaModel):215        """Schema that uses registered checks in Field."""216        col1: pa.typing.Series[int] = pa.Field(custom_gt=100)217        col2: pa.typing.Series[float] = pa.Field(218            custom_in_range={"min_value": -10, "max_value": 10}219        )220        class Config:221            coerce = True222    data = pd.DataFrame(223        {224            "col1": [101, 1000, 2000],225            "col2": [-5.0, 0.0, 6.0],226        }227    )228    Schema.validate(data)229    for invalid_data in [230        pd.DataFrame({"col1": [0], "col2": [-10.0]}),231        pd.DataFrame({"col1": [1000], "col2": [-100.0]}),232    ]:233        with pytest.raises(pa.errors.SchemaError):234            Schema.validate(invalid_data)235def test_register_before_schema_definitions() -> None:236    """Test that custom checks need to be registered before use."""237    # pylint: disable=missing-class-docstring,too-few-public-methods238    # pylint: disable=function-redefined239    with pytest.raises(240        pa.errors.SchemaInitError,241        match="custom check 'custom_eq' is not available",242    ):243        class Schema1(pa.SchemaModel):244            col: pa.typing.Series[int] = pa.Field(custom_eq=1)245    with pytest.raises(AttributeError):246        pa.Check.custom_eq(1)247    @extensions.register_check_method(statistics=["val"])248    def custom_eq(pandas_obj, val):249        return pandas_obj == val250    class Schema2(pa.SchemaModel):  # noqa F811251        col: pa.typing.Series[int] = pa.Field(custom_eq=1)...

checks_fixtures.py

Source:checks_fixtures.py

...6import pandera as pa7import pandera.extensions as pa_ext8__all__ = "custom_check_teardown", "extra_registered_checks"9@pytest.fixture(scope="function")10def custom_check_teardown() -> Generator[None, None, None]:11    """Remove all custom checks after execution of each pytest function."""12    yield13    for check_name in list(pa.Check.REGISTERED_CUSTOM_CHECKS):14        del pa.Check.REGISTERED_CUSTOM_CHECKS[check_name]15@pytest.fixture(scope="function")16def extra_registered_checks() -> Generator[None, None, None]:17    """temporarily registers custom checks onto the Check class"""18    # pylint: disable=unused-variable19    with mock.patch(20        "pandera.Check.REGISTERED_CUSTOM_CHECKS", new_callable=dict21    ):22        # register custom checks here23        @pa_ext.register_check_method()24        def no_param_check(_: pd.DataFrame) -> bool:...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.