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

How to use _needs_pyarrow method in pandera

Best Python code snippet using pandera_python

test_from_to_format_conversions.py

Source:test_from_to_format_conversions.py

...55    return pd.DataFrame({"str_col": ["a"], "int_col": [1]})56def invalid_input_dataframe() -> pd.DataFrame:57    """Invalid data under the InSchema* models defined above."""58    return pd.DataFrame({"str_col": ["a"]})59def _needs_pyarrow(schema) -> bool:60    return (61        schema62        in {63            InSchemaParquet,64            InSchemaFeather,65            OutSchemaParquet,66            OutSchemaFeather,67        }68        and not pandas_engine.PYARROW_INSTALLED69    )70@pytest.mark.parametrize(71    "schema,to_fn,buf_cls",72    [73        [InSchemaCsv, lambda df, x: df.to_csv(x, index=None), io.StringIO],74        [InSchemaDict, lambda df: df.to_dict(orient="records"), None],75        [76            InSchemaJson,77            lambda df, x: df.to_json(x, orient="records"),78            io.StringIO,79        ],80        [InSchemaJson, lambda df: df.to_json(orient="records"), None],81        [InSchemaFeather, lambda df, x: df.to_feather(x), io.BytesIO],82        [InSchemaParquet, lambda df, x: df.to_parquet(x), io.BytesIO],83        [InSchemaPickle, lambda df, x: df.to_pickle(x), io.BytesIO],84    ],85)86def test_from_format(schema, to_fn, buf_cls):87    """88    Test that check_types-guarded function reads data from source serialization89    format.90    """91    @pa.check_types92    def fn(df: pa.typing.DataFrame[schema]):93        return df94    for df, invalid in [95        (mock_dataframe(), False),96        (invalid_input_dataframe(), True),97    ]:98        buf = None if buf_cls is None else buf_cls()99        if _needs_pyarrow(schema):100            with pytest.raises(ImportError):101                to_fn(df, *([buf] if buf else []))102        else:103            arg = to_fn(df, *([buf] if buf else []))104            if buf:105                if buf.closed:106                    pytest.skip(107                        "skip test for older pandas versions where to_pickle "108                        "closes user-provided buffers: "109                        "https://github.com/pandas-dev/pandas/issues/35679"110                    )111                buf.seek(0)112                arg = buf113            if invalid:114                with pytest.raises(pa.errors.SchemaError):115                    fn(arg)116                return117            out = fn(arg)118            assert df.equals(out)119@pytest.mark.parametrize(120    "schema,from_fn,buf_cls",121    [122        [OutSchemaCsv, pd.read_csv, io.StringIO],123        [OutSchemaDict, pd.DataFrame, None],124        [OutSchemaJson, lambda x: pd.read_json(x, orient="records"), None],125        [OutSchemaFeather, pd.read_feather, io.BytesIO],126        [OutSchemaParquet, pd.read_parquet, io.BytesIO],127        [OutSchemaPickle, pd.read_pickle, io.BytesIO],128    ],129)130def test_to_format(schema, from_fn, buf_cls):131    """132    Test that check_types-guarded function writes data to source serialization133    format.134    """135    @pa.check_types136    def fn(df: pa.typing.DataFrame[InSchema]) -> pa.typing.DataFrame[schema]:137        return df.assign(float_col=1.1)138    @pa.check_types139    def invalid_fn(140        df: pa.typing.DataFrame[InSchema],141    ) -> pa.typing.DataFrame[schema]:142        return df143    df = mock_dataframe()144    if _needs_pyarrow(schema):145        with pytest.raises((ImportError)):146            fn(df)147        return148    try:149        out = fn(df)150    except IOError:151        pytest.skip(152            f"pandas=={pd.__version__} automatically closes the buffer, for "153            "more details see: "154            "https://github.com/pandas-dev/pandas/issues/35679"155        )156    if buf_cls and not isinstance(out, buf_cls):157        out = buf_cls(out)158    out_df = from_fn(out)...

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.