feat: restructure codebase with core/ module and explicit imports

Author: ryanseq-gyg

dataframe_expectations/__init__.py

@@ -1,16 +1,5 @@
-from enum import Enum
-from typing import Union
+"""DataFrame Expectations - A validation library for pandas and PySpark DataFrames."""
 
-from pandas import DataFrame as PandasDataFrame
-from pyspark.sql import DataFrame as PySparkDataFrame
+__version__ = "0.3.0"
 
-DataFrameLike = Union[PySparkDataFrame, PandasDataFrame]
-
-
-class DataFrameType(str, Enum):
-    """
-    Enum for DataFrame types.
-    """
-
-    PANDAS = "pandas"
-    PYSPARK = "pyspark"
+__all__ = []

dataframe_expectations/core/__init__.py

@@ -0,0 +1,3 @@
+"""Core base classes and interfaces for DataFrame expectations."""
+
+__all__ = []

…/expectations/aggregation_expectation.py …ctations/core/aggregation_expectation.pydataframe_expectations/expectations/aggregation_expectation.py renamed to dataframe_expectations/core/aggregation_expectation.py dataframe_expectations/expectations/aggregation_expectation.py renamed to dataframe_expectations/core/aggregation_expectation.py

@@ -1,8 +1,8 @@
 from abc import abstractmethod
 from typing import List, Union
 
-from dataframe_expectations import DataFrameLike, DataFrameType
-from dataframe_expectations.expectations import DataFrameExpectation
+from dataframe_expectations.core.types import DataFrameLike, DataFrameType
+from dataframe_expectations.core.expectation import DataFrameExpectation
 from dataframe_expectations.result_message import (
     DataFrameExpectationFailureMessage,
     DataFrameExpectationResultMessage,

…tions/expectations/column_expectation.py …_expectations/core/column_expectation.pydataframe_expectations/expectations/column_expectation.py renamed to dataframe_expectations/core/column_expectation.py dataframe_expectations/expectations/column_expectation.py renamed to dataframe_expectations/core/column_expectation.py

@@ -1,7 +1,7 @@
 from typing import Callable
 
-from dataframe_expectations import DataFrameLike, DataFrameType
-from dataframe_expectations.expectations import DataFrameExpectation
+from dataframe_expectations.core.types import DataFrameLike, DataFrameType
+from dataframe_expectations.core.expectation import DataFrameExpectation
 from dataframe_expectations.result_message import (
     DataFrameExpectationFailureMessage,
     DataFrameExpectationResultMessage,

dataframe_expectations/core/expectation.py

@@ -0,0 +1,110 @@
+from abc import ABC, abstractmethod
+from typing import cast
+
+from pandas import DataFrame as PandasDataFrame
+from pyspark.sql import DataFrame as PySparkDataFrame
+
+# Import the connect DataFrame type for Spark Connect
+try:
+    from pyspark.sql.connect.dataframe import DataFrame as PySparkConnectDataFrame
+except ImportError:
+    # Fallback for older PySpark versions that don't have connect
+    PySparkConnectDataFrame = None  # type: ignore[misc,assignment]
+
+from dataframe_expectations.core.types import DataFrameLike, DataFrameType
+from dataframe_expectations.result_message import (
+    DataFrameExpectationResultMessage,
+)
+
+
+class DataFrameExpectation(ABC):
+    """
+    Base class for DataFrame expectations.
+    """
+
+    def get_expectation_name(self) -> str:
+        """
+        Returns the class name as the expectation name.
+        """
+        return type(self).__name__
+
+    @abstractmethod
+    def get_description(self) -> str:
+        """
+        Returns a description of the expectation.
+        """
+        raise NotImplementedError(
+            f"description method must be implemented for {self.__class__.__name__}"
+        )
+
+    def __str__(self):
+        """
+        Returns a string representation of the expectation.
+        """
+        return f"{self.get_expectation_name()} ({self.get_description()})"
+
+    @classmethod
+    def infer_data_frame_type(cls, data_frame: DataFrameLike) -> DataFrameType:
+        """
+        Infer the DataFrame type based on the provided DataFrame.
+        """
+        if isinstance(data_frame, PandasDataFrame):
+            return DataFrameType.PANDAS
+        elif isinstance(data_frame, PySparkDataFrame):
+            return DataFrameType.PYSPARK
+        elif PySparkConnectDataFrame is not None and isinstance(
+            data_frame, PySparkConnectDataFrame
+        ):
+            return DataFrameType.PYSPARK
+        else:
+            raise ValueError(f"Unsupported DataFrame type: {type(data_frame)}")
+
+    def validate(self, data_frame: DataFrameLike, **kwargs):
+        """
+        Validate the DataFrame against the expectation.
+        """
+        data_frame_type = self.infer_data_frame_type(data_frame)
+
+        if data_frame_type == DataFrameType.PANDAS:
+            return self.validate_pandas(data_frame=data_frame, **kwargs)
+        elif data_frame_type == DataFrameType.PYSPARK:
+            return self.validate_pyspark(data_frame=data_frame, **kwargs)
+        else:
+            raise ValueError(f"Unsupported DataFrame type: {data_frame_type}")
+
+    @abstractmethod
+    def validate_pandas(
+        self, data_frame: DataFrameLike, **kwargs
+    ) -> DataFrameExpectationResultMessage:
+        """
+        Validate a pandas DataFrame against the expectation.
+        """
+        raise NotImplementedError(
+            f"validate_pandas method must be implemented for {self.__class__.__name__}"
+        )
+
+    @abstractmethod
+    def validate_pyspark(
+        self, data_frame: DataFrameLike, **kwargs
+    ) -> DataFrameExpectationResultMessage:
+        """
+        Validate a PySpark DataFrame against the expectation.
+        """
+        raise NotImplementedError(
+            f"validate_pyspark method must be implemented for {self.__class__.__name__}"
+        )
+
+    @classmethod
+    def num_data_frame_rows(cls, data_frame: DataFrameLike) -> int:
+        """
+        Count the number of rows in the DataFrame.
+        """
+        data_frame_type = cls.infer_data_frame_type(data_frame)
+        if data_frame_type == DataFrameType.PANDAS:
+            # Cast to PandasDataFrame since we know it's a Pandas DataFrame at this point
+            return len(cast(PandasDataFrame, data_frame))
+        elif data_frame_type == DataFrameType.PYSPARK:
+            # Cast to PySparkDataFrame since we know it's a PySpark DataFrame at this point
+            return cast(PySparkDataFrame, data_frame).count()
+        else:
+            raise ValueError(f"Unsupported DataFrame type: {data_frame_type}")

dataframe_expectations/core/types.py

@@ -0,0 +1,57 @@
+"""Core types, enums, and data models for dataframe-expectations."""
+
+from enum import Enum
+from typing import Any, Dict, Union
+
+from pandas import DataFrame as PandasDataFrame
+from pydantic import BaseModel, ConfigDict, Field
+from pyspark.sql import DataFrame as PySparkDataFrame
+
+# Type aliases
+DataFrameLike = Union[PySparkDataFrame, PandasDataFrame]
+
+
+class DataFrameType(str, Enum):
+    """Enum for DataFrame types."""
+
+    PANDAS = "pandas"
+    PYSPARK = "pyspark"
+
+
+class ExpectationCategory(str, Enum):
+    """Categories for expectations."""
+
+    COLUMN_EXPECTATIONS = "Column Expectations"
+    COLUMN_AGGREGATION_EXPECTATIONS = "Column Aggregation Expectations"
+    DATAFRAME_AGGREGATION_EXPECTATIONS = "DataFrame Aggregation Expectations"
+
+
+class ExpectationSubcategory(str, Enum):
+    """Subcategory of expectations."""
+
+    ANY_VALUE = "Any Value"
+    NUMERICAL = "Numerical"
+    STRING = "String"
+    UNIQUE = "Unique"
+
+
+class ExpectationMetadata(BaseModel):
+    """Metadata for a registered expectation."""
+
+    suite_method_name: str = Field(
+        ..., description="Method name in ExpectationsSuite (e.g., 'expect_value_greater_than')"
+    )
+    pydoc: str = Field(..., description="Human-readable description of the expectation")
+    category: ExpectationCategory = Field(..., description="Category (e.g., 'Column Expectations')")
+    subcategory: ExpectationSubcategory = Field(
+        ..., description="Subcategory (e.g., 'Numerical', 'String')"
+    )
+    params_doc: Dict[str, str] = Field(..., description="Documentation for each parameter")
+    params: list = Field(default_factory=list, description="List of required parameter names")
+    param_types: Dict[str, Any] = Field(
+        default_factory=dict, description="Type hints for parameters"
+    )
+    factory_func_name: str = Field(..., description="Name of the factory function")
+    expectation_name: str = Field(..., description="Name of the expectation class")
+
+    model_config = ConfigDict(frozen=True)  # Make model immutable

…frame_expectations/expectations/utils.py dataframe_expectations/core/utils.pydataframe_expectations/expectations/utils.py renamed to dataframe_expectations/core/utils.py dataframe_expectations/expectations/utils.py renamed to dataframe_expectations/core/utils.py

@@ -1,7 +1,7 @@
 from functools import wraps
 from typing import Any, Callable, Dict, Optional, Tuple, Type, Union, get_args
 
-from dataframe_expectations.expectations import DataFrameExpectation
+from dataframe_expectations.core.expectation import DataFrameExpectation
 
 
 def requires_params(

dataframe_expectations/expectations/__init__.py

@@ -11,7 +11,7 @@
     # Fallback for older PySpark versions that don't have connect
     PySparkConnectDataFrame = None  # type: ignore[misc,assignment]
 
-from dataframe_expectations import DataFrameLike, DataFrameType
+from dataframe_expectations.core.types import DataFrameLike, DataFrameType
 from dataframe_expectations.result_message import (
     DataFrameExpectationResultMessage,
 )

dataframe_expectations/expectations/aggregation/__init__.py

@@ -0,0 +1,3 @@
+"""Aggregation expectations."""
+
+__all__ = []

…n_expectations/any_value_expectations.py …ns/expectations/aggregation/any_value.pydataframe_expectations/expectations/aggregation_expectations/any_value_expectations.py renamed to dataframe_expectations/expectations/aggregation/any_value.py dataframe_expectations/expectations/aggregation_expectations/any_value_expectations.py renamed to dataframe_expectations/expectations/aggregation/any_value.py

@@ -4,16 +4,16 @@
 from pyspark.sql import DataFrame as PySparkDataFrame
 from pyspark.sql import functions as F
 
-from dataframe_expectations import DataFrameLike, DataFrameType
-from dataframe_expectations.expectations.aggregation_expectation import (
+from dataframe_expectations.core.types import DataFrameLike, DataFrameType
+from dataframe_expectations.core.aggregation_expectation import (
     DataFrameAggregationExpectation,
 )
-from dataframe_expectations.expectations.expectation_registry import (
+from dataframe_expectations.registry import (
     ExpectationCategory,
     ExpectationSubcategory,
     register_expectation,
 )
-from dataframe_expectations.expectations.utils import requires_params
+from dataframe_expectations.core.utils import requires_params
 from dataframe_expectations.result_message import (
     DataFrameExpectationFailureMessage,
     DataFrameExpectationResultMessage,