Search code examples
pythonpandaspandera

How to validate a dataframe index using SchemaModel in Pandera

I can validate a DataFrame index using the DataFrameSchema like this:

import pandera as pa

from pandera import Column, DataFrameSchema, Check, Index

schema = DataFrameSchema(
    columns={
        "column1": pa.Column(int),
    },
    index=pa.Index(int, name="index_name"),
)
# raises the error as expected
schema.validate(
    pd.DataFrame({"column1": [1, 2, 3]}, index=pd.Index([1, 2, 3], name="index_incorrect_name")) 
)

Is there a way to do the same using a SchemaModel?

Solution

  • Found an answer in GitHub

    You can use pa.typing.Index to type-annotate an index.

    class Schema(pa.SchemaModel):
    column1: pa.typing.Series[int]
    index_name: pa.typing.Index[int] = pa.Field(check_name=True)

    See how you can validate a MultiIndex index: https://pandera.readthedocs.io/en/stable/schema_models.html#multiindex