refactor(pyspark): restructure pyspark components

pandera/api/base/schema.py

@@ -10,6 +10,8 @@
 from abc import ABC
 from typing import Any, Dict, Optional, Tuple, Type, Union
 
+from typing_extensions import Self
+
 from pandera.backends.base import BaseSchemaBackend
 from pandera.dtypes import DataType
 from pandera.errors import BackendNotFoundError
@@ -139,3 +141,24 @@
 
     def __setstate__(self, state):
         self.__dict__ = state
+
+    def set_name(self, name: str) -> Self:
+        """Used to set or modify the name of a base model object.
+
+        :param str name: the name of the column object
+
+        """
+        self.name = name
+        return self
+
+    def strategy(
+        self, *, size: Optional[int] = None, n_regex_columns: int = 1
+    ):
+        """Create a data synthesis strategy."""
+        raise NotImplementedError
+
+    def example(
+        self, size: Optional[int] = None, n_regex_columns: int = 1
+    ) -> Any:
+        """Generate an example of this data model specification."""
+        raise NotImplementedError

pandera/api/dataframe/model.py

@@ -21,10 +21,20 @@
     cast,
 )
 
+from typing_extensions import Self
+
 from pandera.api.base.model import BaseModel
 from pandera.api.base.schema import BaseSchema
 from pandera.api.checks import Check
-from pandera.api.dataframe.model_components import (
+from pandera.api.parsers import Parser
+from pandera.engines import PYDANTIC_V2
+from pandera.errors import SchemaInitError
+from pandera.import_utils import strategy_import_error
+from pandera.typing import AnnotationInfo
+from pandera.typing.common import DataFrameBase
+from pandera.utils import docstring_substitution
+
+from .model_components import (
     CHECK_KEY,
     DATAFRAME_CHECK_KEY,
     DATAFRAME_PARSER_KEY,
@@ -36,27 +46,19 @@
     FieldParserInfo,
     ParserInfo,
 )
-from pandera.api.dataframe.model_config import BaseConfig
-from pandera.api.parsers import Parser
-from pandera.engines import PYDANTIC_V2
-from pandera.errors import SchemaInitError
-from pandera.import_utils import strategy_import_error
-from pandera.typing import AnnotationInfo
-from pandera.typing.common import DataFrameBase
-from pandera.utils import docstring_substitution
-
-if PYDANTIC_V2:
-    from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
-    from pydantic_core import core_schema
+from .model_config import BaseConfig
 
 try:
     from typing_extensions import get_type_hints
 except ImportError:  # pragma: no cover
     from typing import get_type_hints  # type: ignore
 
+if PYDANTIC_V2:
+    from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
+    from pydantic_core import core_schema
+
 
 TDataFrame = TypeVar("TDataFrame")
-TDataFrameModel = TypeVar("TDataFrameModel", bound="DataFrameModel")
 TSchema = TypeVar("TSchema", bound=BaseSchema)
 
 _CONFIG_KEY = "Config"
@@ -133,11 +135,9 @@
     __root_parsers__: List[Parser] = []
 
     @docstring_substitution(validate_doc=BaseSchema.validate.__doc__)
-    def __new__(cls, *args, **kwargs) -> DataFrameBase[TDataFrameModel]:  # type: ignore [misc]
+    def __new__(cls, *args, **kwargs) -> DataFrameBase[Self]:  # type: ignore [misc]
         """%(validate_doc)s"""
-        return cast(
-            DataFrameBase[TDataFrameModel], cls.validate(*args, **kwargs)
-        )
+        return cast(DataFrameBase[Self], cls.validate(*args, **kwargs))
 
     def __init_subclass__(cls, **kwargs):
         """Ensure :class:`~pandera.api.dataframe.model_components.FieldInfo` instances."""
@@ -163,9 +163,9 @@
         cls.__config__, cls.__extras__ = cls._collect_config_and_extras()
 
     def __class_getitem__(
-        cls: Type[TDataFrameModel],
+        cls: Type[Self],
         item: Union[Type[Any], Tuple[Type[Any], ...]],
-    ) -> Type[TDataFrameModel]:
+    ) -> Type[Self]:
         """Parameterize the class's generic arguments with the specified types"""
         if not hasattr(cls, "__parameters__"):
             raise TypeError(
@@ -181,9 +181,7 @@
                 f"Expected {len(__parameters__)} generic arguments but found {len(item)}"
             )
         if (cls, item) in GENERIC_SCHEMA_CACHE:
-            return typing.cast(
-                Type[TDataFrameModel], GENERIC_SCHEMA_CACHE[(cls, item)]
-            )
+            return typing.cast(Type[Self], GENERIC_SCHEMA_CACHE[(cls, item)])
 
         param_dict: Dict[TypeVar, Type[Any]] = dict(zip(__parameters__, item))
         extra: Dict[str, Any] = {"__annotations__": {}}
@@ -276,18 +274,18 @@
     @classmethod
     @docstring_substitution(validate_doc=BaseSchema.validate.__doc__)
     def validate(
-        cls: Type[TDataFrameModel],
+        cls: Type[Self],
         check_obj: TDataFrame,
         head: Optional[int] = None,
         tail: Optional[int] = None,
         sample: Optional[int] = None,
         random_state: Optional[int] = None,
         lazy: bool = False,
         inplace: bool = False,
-    ) -> DataFrameBase[TDataFrameModel]:
+    ) -> DataFrameBase[Self]:
         """%(validate_doc)s"""
         return cast(
-            DataFrameBase[TDataFrameModel],
+            DataFrameBase[Self],
             cls.to_schema().validate(
                 check_obj, head, tail, sample, random_state, lazy, inplace
             ),
@@ -296,7 +294,7 @@
     # TODO: add docstring_substitution using generic class
     @classmethod
     @strategy_import_error
-    def strategy(cls: Type[TDataFrameModel], **kwargs):
+    def strategy(cls: Type[Self], **kwargs):
         """Create a ``hypothesis`` strategy for generating a DataFrame.
 
         :param size: number of elements to generate
@@ -309,17 +307,15 @@
     @classmethod
     @strategy_import_error
     def example(
-        cls: Type[TDataFrameModel],
+        cls: Type[Self],
         **kwargs,
-    ) -> DataFrameBase[TDataFrameModel]:
+    ) -> DataFrameBase[Self]:
         """Generate an example of a particular size.
 
         :param size: number of elements in the generated DataFrame.
         :returns: DataFrame object.
         """
-        return cast(
-            DataFrameBase[TDataFrameModel], cls.to_schema().example(**kwargs)
-        )
+        return cast(DataFrameBase[Self], cls.to_schema().example(**kwargs))
 
     @classmethod
     def _get_model_attrs(cls) -> Dict[str, Any]:
@@ -572,9 +568,7 @@
         raise NotImplementedError
 
     @classmethod
-    def empty(
-        cls: Type[TDataFrameModel], *_args
-    ) -> DataFrameBase[TDataFrameModel]:
+    def empty(cls: Type[Self], *_args) -> DataFrameBase[Self]:
         """Create an empty DataFrame instance."""
         raise NotImplementedError
 

pandera/api/pandas/array.py

@@ -81,7 +81,7 @@ def strategy(self, *, size=None):
             size=size,
         )
 
-    def example(self, size=None) -> TDataObject:
+    def example(self, size=None, n_regex_columns: int = 1) -> TDataObject:
         """Generate an example of a particular size.
 
         :param size: number of elements in the generated array.
@@ -272,7 +272,11 @@ def validate(  # type: ignore [override]
             )
         return cast(pd.Series, validated_obj)
 
-    def example(self, size=None) -> pd.Series:
+    def example(
+        self,
+        size=None,
+        n_regex_columns: int = 1,  # pylint: disable=unused-argument
+    ) -> pd.Series:
         """Generate an example of a particular size.
 
         :param size: number of elements in the generated Series.

pandera/api/pandas/components.py

@@ -192,7 +192,11 @@ def strategy_component(self):
             name=self.name,
         )
 
-    def example(self, size=None) -> pd.DataFrame:
+    def example(
+        self,
+        size=None,
+        n_regex_columns: int = 1,  # pylint: disable=unused-argument
+    ) -> pd.DataFrame:
         """Generate an example of a particular size.
 
         :param size: number of elements in the generated Index.
@@ -236,7 +240,12 @@ def __eq__(self, other):
     ###########################
 
     @strategy_import_error
-    def strategy(self, *, size: Optional[int] = None):
+    def strategy(
+        self,
+        *,
+        size: Optional[int] = None,
+        n_regex_columns: int = 1,  # pylint: disable=unused-argument
+    ):
         """Create a ``hypothesis`` strategy for generating an Index.
 
         :param size: number of elements to generate.
@@ -265,7 +274,11 @@ def strategy_component(self):
             name=self.name,
         )
 
-    def example(self, size: Optional[int] = None) -> pd.Index:
+    def example(
+        self,
+        size: Optional[int] = None,
+        n_regex_columns: int = 1,  # pylint: disable=unused-argument
+    ) -> pd.Index:
         """Generate an example of a particular size.
 
         :param size: number of elements in the generated Index.