diff --git a/pandera/api/dataframe/container.py b/pandera/api/dataframe/container.py index b341b3af7..c849df0b6 100644 --- a/pandera/api/dataframe/container.py +++ b/pandera/api/dataframe/container.py @@ -823,9 +823,7 @@ def update_index(self, index_name: str, **kwargs) -> Self: tmp_index_cols = schema.index.names schema = schema.reset_index() - schema = schema.update_column(index_name, **kwargs) - schema = schema.set_index(tmp_index_cols) return cast(Self, schema) @@ -857,9 +855,7 @@ def update_indexes(self, update_dict: Dict[str, Dict[str, Any]]) -> Self: tmp_index_cols = schema.index.names schema = schema.reset_index() - schema = schema.update_columns(update_dict) - schema = schema.set_index(tmp_index_cols) return cast(Self, schema) @@ -1098,6 +1094,9 @@ def set_index( nullable=new_schema.columns[col].nullable, unique=new_schema.columns[col].unique, coerce=new_schema.columns[col].coerce, + title=new_schema.columns[col].title, + description=new_schema.columns[col].description, + metadata=new_schema.columns[col].metadata, ) ) diff --git a/pandera/api/pandas/components.py b/pandera/api/pandas/components.py index b70e222e2..22227593d 100644 --- a/pandera/api/pandas/components.py +++ b/pandera/api/pandas/components.py @@ -380,6 +380,11 @@ def names(self): """Get index names in the MultiIndex schema component.""" return [index.name for index in self.indexes] + @property + def named_indexes(self) -> Dict[str, Any]: + """Get named indexes.""" + return {index.name: index for index in self.indexes} + @property def coerce(self): """Whether or not to coerce data types.""" diff --git a/tests/pandas/test_schemas.py b/tests/pandas/test_schemas.py index ce59ca530..1b2c6110d 100644 --- a/tests/pandas/test_schemas.py +++ b/tests/pandas/test_schemas.py @@ -2272,9 +2272,31 @@ def test_update_index(): schema.update_index("does_not_exist", dtype=str) +def test_update_index_with_properties(): + + schema = DataFrameSchema(index=Index(dtype=int, name="idx")) + + assert schema.index.title is None + assert schema.index.description is None + assert schema.index.metadata is None + + schema = schema.update_index( + index_name="idx", + unique=True, + title="badger", + description="pear", + metadata={"thing": 123}, + ) + + assert schema.index.unique + assert schema.index.title == "badger" + assert schema.index.description == "pear" + assert schema.index.metadata == {"thing": 123} + + def test_update_multi_index(): """ - Test that schemas can correctly update a multi_index column via update_column method. + Test that schemas can correctly update a multi_index column via update_index method. """ schema = DataFrameSchema( @@ -2342,9 +2364,12 @@ def test_update_multi_indexes(): schema = schema.update_indexes( { "a": {"dtype": str}, + "b": {"dtype": str}, } ) - multi_idx = pd.MultiIndex.from_arrays([["hello"], [1.0]], names=["a", "b"]) + multi_idx = pd.MultiIndex.from_arrays( + [["hello"], ["1.0"]], names=["a", "b"] + ) df = pd.DataFrame({"c": [1]}, index=multi_idx) assert isinstance(schema.validate(df), pd.DataFrame) @@ -2358,6 +2383,43 @@ def test_update_multi_indexes(): ) +def test_update_multi_indexes_with_properties(): + """ + Test that schemas can correctly update a multi_index column via + update_indexes method with properties. + """ + + schema = DataFrameSchema( + index=MultiIndex( + [Index(dtype=int, name="a"), Index(dtype=float, name="b")] + ), + ) + schema = schema.update_indexes( + { + "a": { + "dtype": str, + "title": "badger", + "description": "pear", + "metadata": {"thing": 123}, + }, + "b": { + "dtype": str, + "title": "badger", + "description": "pear", + "metadata": {"thing": 123}, + }, + } + ) + + assert schema.index.named_indexes["a"].title == "badger" + assert schema.index.named_indexes["a"].description == "pear" + assert schema.index.named_indexes["a"].metadata == {"thing": 123} + + assert schema.index.named_indexes["b"].title == "badger" + assert schema.index.named_indexes["b"].description == "pear" + assert schema.index.named_indexes["b"].metadata == {"thing": 123} + + def test_rename_indexes(): """ Test that schemas can correctly rename an index column via rename_indexes method.