Repository: unionai-oss/pandera (Python)

Issue:
SeriesSchema should validate index
Currently, the `SeriesSchema` object doesn't validate the index of the schema. The purpose of this task is to extend the `__init__` signature of `SeriesSchema` to take an `index` argument, which would take a `pa.Index` or `pa.MultiIndex`. In the `validate` / `__call__` call, the index should be checked.


Discussion:

Patch:
diff --git a/pandera/schemas.py b/pandera/schemas.py
index 20faeb843..b6ae8a9a0 100644
--- a/pandera/schemas.py
+++ b/pandera/schemas.py
@@ -990,6 +990,38 @@ def __eq__(self, other):
 class SeriesSchema(SeriesSchemaBase):
     """Series validator."""
 
+    def __init__(
+            self,
+            pandas_dtype: PandasDtypeInputTypes = None,
+            checks: CheckList = None,
+            index=None,
+            nullable: bool = False,
+            allow_duplicates: bool = True,
+            coerce: bool = False,
+            name: str = None) -> None:
+        """Initialize series schema base object.
+
+        :param pandas_dtype: datatype of the column. If a string is specified,
+            then assumes one of the valid pandas string values:
+            http://pandas.pydata.org/pandas-docs/stable/basics.html#dtypes
+        :param checks: If element_wise is True, then callable signature should
+            be:
+
+            ``Callable[Any, bool]`` where the ``Any`` input is a scalar element
+            in the column. Otherwise, the input is assumed to be a
+            pandas.Series object.
+        :type checks: callable
+        :param index: specify the datatypes and properties of the index.
+        :param nullable: Whether or not column can contain null values.
+        :type nullable: bool
+        :param allow_duplicates:
+        :type allow_duplicates: bool
+        """
+        super().__init__(
+            pandas_dtype, checks, nullable, allow_duplicates, coerce, name
+        )
+        self.index = index
+
     @property
     def _allow_groupby(self) -> bool:
         """Whether the schema or schema component allows groupby operations."""
@@ -1052,8 +1084,16 @@ def validate(
         if self.coerce:
             check_obj = self.coerce_dtype(check_obj)
 
+        if self.index is not None and (self.index.coerce or self.coerce):
+            check_obj.index = self.index.coerce_dtype(check_obj.index)
+
+        # validate index
+        if self.index:
+            self.index(check_obj)
+
         return super(SeriesSchema, self).validate(
-            check_obj, head, tail, sample, random_state, lazy)
+            check_obj, head, tail, sample, random_state, lazy
+        )
 
     def __call__(
             self,


Test patch:
diff --git a/tests/test_schemas.py b/tests/test_schemas.py
index aefb10593..84a1dc9e2 100644
--- a/tests/test_schemas.py
+++ b/tests/test_schemas.py
@@ -175,6 +175,33 @@ def test_series_schema_multiple_validators():
         schema.validate(pd.Series([1, 5, 20, 50]))
 
 
+@pytest.mark.parametrize("coerce", [True, False])
+def test_series_schema_with_index(coerce):
+    """Test SeriesSchema with Index and MultiIndex components."""
+    schema_with_index = SeriesSchema(
+        pandas_dtype=Int,
+        index=Index(Int, coerce=coerce),
+    )
+    validated_series = schema_with_index(pd.Series([1, 2, 3], index=[1, 2, 3]))
+    assert isinstance(validated_series, pd.Series)
+
+    schema_with_multiindex = SeriesSchema(
+        pandas_dtype=Int,
+        index=MultiIndex([
+            Index(Int, coerce=coerce),
+            Index(String, coerce=coerce),
+        ])
+    )
+    multi_index = pd.MultiIndex.from_arrays(
+        [[0, 1, 2], ["foo", "bar", "foo"]],
+    )
+    validated_series_multiindex = schema_with_multiindex(
+        pd.Series([1, 2, 3], index=multi_index)
+    )
+    assert isinstance(validated_series_multiindex, pd.Series)
+    assert (validated_series_multiindex.index == multi_index).all()
+
+
 class SeriesGreaterCheck:
     # pylint: disable=too-few-public-methods
     """Class creating callable objects to check if series elements exceed a