feat: simplified registry

Author: ryanseq-gyg

README.md

@@ -55,7 +55,7 @@ uv run pytest tests/ --cov=dataframe_expectations
 
 **Basic usage with Pandas:**
 ```python
-from dataframe_expectations.expectations_suite import DataFrameExpectationsSuite
+from dataframe_expectations.suite import DataFrameExpectationsSuite
 import pandas as pd
 
 # Build a suite with expectations
@@ -82,7 +82,7 @@ runner.run(df)
 
 **PySpark example:**
 ```python
-from dataframe_expectations.expectations_suite import DataFrameExpectationsSuite
+from dataframe_expectations.suite import DataFrameExpectationsSuite
 from pyspark.sql import SparkSession
 
 # Initialize Spark session
@@ -116,7 +116,7 @@ runner.run(df)
 
 **Decorator pattern for automatic validation:**
 ```python
-from dataframe_expectations.expectations_suite import DataFrameExpectationsSuite
+from dataframe_expectations.suite import DataFrameExpectationsSuite
 from pyspark.sql import SparkSession
 
 # Initialize Spark session

dataframe_expectations/__init__.py

@@ -1,5 +1,7 @@
 """DataFrame Expectations - A validation library for pandas and PySpark DataFrames."""
 
-__version__ = "0.3.0"
+from importlib.metadata import version
+
+__version__ = version("dataframe-expectations")
 
 __all__ = []

dataframe_expectations/expectations/aggregation/unique.py

@@ -4,14 +4,14 @@
 from pandas import DataFrame as PandasDataFrame
 from pyspark.sql import DataFrame as PySparkDataFrame
 from pyspark.sql import functions as F
-
-from dataframe_expectations.core.types import DataFrameLike, DataFrameType
 from dataframe_expectations.core.aggregation_expectation import (
     DataFrameAggregationExpectation,
 )
 from dataframe_expectations.core.types import (
     ExpectationCategory,
     ExpectationSubcategory,
+    DataFrameLike,
+    DataFrameType,
 )
 from dataframe_expectations.registry import register_expectation
 from dataframe_expectations.core.utils import requires_params

dataframe_expectations/registry.py

@@ -1,5 +1,5 @@
 import re
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Tuple
 
 from dataframe_expectations.core.expectation import DataFrameExpectation
 from dataframe_expectations.core.types import (
@@ -11,12 +11,20 @@
 
 logger = setup_logger(__name__)
 
+# Type alias for registry entry (factory function + metadata)
+FactoryFunction = Callable[..., DataFrameExpectation]
+RegistryEntry = Tuple[FactoryFunction, ExpectationMetadata]
+
 
 class DataFrameExpectationRegistry:
     """Registry for dataframe expectations."""
 
-    _expectations: Dict[str, Callable[..., DataFrameExpectation]] = {}
-    _metadata: Dict[str, ExpectationMetadata] = {}
+    # Primary registry: keyed by suite_method_name for O(1) suite access
+    _registry: Dict[str, RegistryEntry] = {}
+
+    # Secondary index: maps expectation_name -> suite_method_name for O(1) lookups
+    _by_name: Dict[str, str] = {}
+
     _loaded: bool = False
 
     @classmethod
@@ -41,21 +49,31 @@ def register(
         :return: Decorator function.
         """
 
-        def decorator(func: Callable[..., DataFrameExpectation]):
+        def decorator(func: FactoryFunction) -> FactoryFunction:
             expectation_name = name
 
             logger.debug(
                 f"Registering expectation '{expectation_name}' with function {func.__name__}"
             )
 
-            # Check if the name is already registered
-            if expectation_name in cls._expectations:
-                error_message = f"Expectation '{expectation_name}' is already registered."
+            suite_method = suite_method_name or cls._convert_to_suite_method(expectation_name)
+
+            # Check for duplicate suite method name
+            if suite_method in cls._registry:
+                existing_metadata = cls._registry[suite_method][1]
+                error_message = (
+                    f"Suite method '{suite_method}' is already registered by expectation '{existing_metadata.expectation_name}'. "
+                    f"Cannot register '{expectation_name}'."
+                )
                 logger.error(error_message)
                 raise ValueError(error_message)
 
-            # Register factory function
-            cls._expectations[expectation_name] = func
+            # Check for duplicate expectation name
+            if expectation_name in cls._by_name:
+                existing_suite_method = cls._by_name[expectation_name]
+                error_message = f"Expectation '{expectation_name}' is already registered with suite method '{existing_suite_method}'."
+                logger.error(error_message)
+                raise ValueError(error_message)
 
             # Extract params from @requires_params if present
             extracted_params = []
@@ -64,10 +82,8 @@ def decorator(func: Callable[..., DataFrameExpectation]):
                 extracted_params = list(func._required_params)
                 extracted_types = getattr(func, "_param_types", {})
 
-            # Store metadata
-            cls._metadata[expectation_name] = ExpectationMetadata(
-                suite_method_name=suite_method_name
-                or cls._convert_to_suite_method(expectation_name),
+            metadata = ExpectationMetadata(
+                suite_method_name=suite_method,
                 pydoc=pydoc,
                 category=category,
                 subcategory=subcategory,
@@ -78,6 +94,12 @@ def decorator(func: Callable[..., DataFrameExpectation]):
                 expectation_name=expectation_name,
             )
 
+            # Store in primary registry
+            cls._registry[suite_method] = (func, metadata)
+
+            # Store in secondary index
+            cls._by_name[expectation_name] = suite_method
+
             return func
 
         return decorator
@@ -93,8 +115,9 @@ def _convert_to_suite_method(cls, expectation_name: str) -> str:
             ExpectationValueGreaterThan -> expect_value_greater_than
             ExpectationMinRows -> expect_min_rows
         """
-        # Remove 'Expectation' prefix
+
         name = re.sub(r"^Expectation", "", expectation_name)
+
         # Convert CamelCase to snake_case
         snake = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", name)
         snake = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", snake)
@@ -141,21 +164,28 @@ def _load_all_expectations(cls):
     def get_expectation(cls, expectation_name: str, **kwargs) -> DataFrameExpectation:
         """Get an expectation instance by name.
 
+        Note: This method is kept for backward compatibility with tests.
+        The suite uses get_expectation_by_suite_method() for better performance.
+
         :param expectation_name: The name of the expectation.
         :param kwargs: Parameters to pass to the expectation factory function.
         :return: An instance of DataFrameExpectation.
         """
-        cls._ensure_loaded()  # Lazy load expectations
+        cls._ensure_loaded()
         logger.debug(f"Retrieving expectation '{expectation_name}' with arguments: {kwargs}")
-        if expectation_name not in cls._expectations:
+
+        if expectation_name not in cls._by_name:
             available = cls.list_expectations()
             error_message = (
                 f"Unknown expectation '{expectation_name}'. "
                 f"Available expectations: {', '.join(available)}"
             )
             logger.error(error_message)
             raise ValueError(error_message)
-        return cls._expectations[expectation_name](**kwargs)
+
+        suite_method = cls._by_name[expectation_name]
+        factory, metadata = cls._registry[suite_method]
+        return factory(**kwargs)
 
     @classmethod
     def get_metadata(cls, expectation_name: str) -> ExpectationMetadata:
@@ -166,9 +196,13 @@ def get_metadata(cls, expectation_name: str) -> ExpectationMetadata:
         :raises ValueError: If expectation not found.
         """
         cls._ensure_loaded()
-        if expectation_name not in cls._metadata:
+
+        if expectation_name not in cls._by_name:
             raise ValueError(f"No metadata found for expectation '{expectation_name}'")
-        return cls._metadata[expectation_name]
+
+        suite_method = cls._by_name[expectation_name]
+        factory, metadata = cls._registry[suite_method]
+        return metadata
 
     @classmethod
     def get_all_metadata(cls) -> Dict[str, ExpectationMetadata]:
@@ -177,7 +211,35 @@ def get_all_metadata(cls) -> Dict[str, ExpectationMetadata]:
         :return: Dictionary mapping expectation names to their metadata.
         """
         cls._ensure_loaded()
-        return cls._metadata.copy()
+        return {metadata.expectation_name: metadata for _, (_, metadata) in cls._registry.items()}
+
+    @classmethod
+    def get_expectation_by_suite_method(
+        cls, suite_method_name: str, **kwargs
+    ) -> DataFrameExpectation:
+        """Get an expectation instance by suite method name.
+
+        :param suite_method_name: The suite method name (e.g., 'expect_value_greater_than').
+        :param kwargs: Parameters to pass to the expectation factory function.
+        :return: An instance of DataFrameExpectation.
+        :raises ValueError: If suite method not found.
+        """
+        cls._ensure_loaded()
+        logger.debug(
+            f"Retrieving expectation for suite method '{suite_method_name}' with arguments: {kwargs}"
+        )
+
+        if suite_method_name not in cls._registry:
+            available = list(cls._registry.keys())
+            error_message = (
+                f"Unknown suite method '{suite_method_name}'. "
+                f"Available methods: {', '.join(available[:10])}..."
+            )
+            logger.error(error_message)
+            raise ValueError(error_message)
+
+        factory, metadata = cls._registry[suite_method_name]
+        return factory(**kwargs)
 
     @classmethod
     def get_suite_method_mapping(cls) -> Dict[str, str]:
@@ -187,16 +249,19 @@ def get_suite_method_mapping(cls) -> Dict[str, str]:
                  to expectation names (e.g., 'ExpectationValueGreaterThan').
         """
         cls._ensure_loaded()
-        return {meta.suite_method_name: exp_name for exp_name, meta in cls._metadata.items()}
+        return {
+            suite_method: metadata.expectation_name
+            for suite_method, (_, metadata) in cls._registry.items()
+        }
 
     @classmethod
     def list_expectations(cls) -> list:
         """List all registered expectation names.
 
         :return: List of registered expectation names.
         """
-        cls._ensure_loaded()  # Lazy load expectations
-        return list(cls._expectations.keys())
+        cls._ensure_loaded()
+        return [metadata.expectation_name for _, (_, metadata) in cls._registry.items()]
 
     @classmethod
     def remove_expectation(cls, expectation_name: str):
@@ -205,23 +270,25 @@ def remove_expectation(cls, expectation_name: str):
         :param expectation_name: The name of the expectation to remove.
         :raises ValueError: If expectation not found.
         """
-        cls._ensure_loaded()  # Lazy load expectations
+        cls._ensure_loaded()
         logger.debug(f"Removing expectation '{expectation_name}'")
-        if expectation_name in cls._expectations:
-            del cls._expectations[expectation_name]
-            if expectation_name in cls._metadata:
-                del cls._metadata[expectation_name]
-        else:
+
+        if expectation_name not in cls._by_name:
             error_message = f"Expectation '{expectation_name}' not found."
             logger.error(error_message)
             raise ValueError(error_message)
 
+        # Remove from both dictionaries
+        suite_method = cls._by_name[expectation_name]
+        del cls._registry[suite_method]
+        del cls._by_name[expectation_name]
+
     @classmethod
     def clear_expectations(cls):
         """Clear all registered expectations."""
-        logger.debug(f"Clearing {len(cls._expectations)} expectations from the registry")
-        cls._expectations.clear()
-        cls._metadata.clear()
+        logger.debug(f"Clearing {len(cls._registry)} expectations from the registry")
+        cls._registry.clear()
+        cls._by_name.clear()
         cls._loaded = False  # Allow reloading
 
 

dataframe_expectations/suite.py

@@ -108,7 +108,6 @@ def run(
             pyspark_df = cast(PySparkDataFrame, data_frame)
             was_already_cached = pyspark_df.is_cached
 
-            # Cache the DataFrame if it wasn't already cached
             if not was_already_cached:
                 logger.debug("Caching PySpark DataFrame for expectations suite execution")
                 pyspark_df.cache()
@@ -206,7 +205,6 @@ def wrapper(*args, **kwargs):
                 logger.info(f"Validating DataFrame returned from '{f.__name__}'")
                 self.run(data_frame=result)
 
-                # Return the original DataFrame if validation passes
                 return result
 
             return wrapper
@@ -254,42 +252,32 @@ def __getattr__(self, name: str):
         if not name.startswith("expect_"):
             raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
 
-        mapping = DataFrameExpectationRegistry.get_suite_method_mapping()
+        # Create and return the dynamic method - validation happens in _create_expectation_method
+        return self._create_expectation_method(name)
 
-        # Check if this method exists in the registry
-        if name not in mapping:
-            available = list(mapping.keys())
-            raise AttributeError(
-                f"Unknown expectation method '{name}'. "
-                f"Available methods: {', '.join(available[:5])}..."
-            )
-
-        expectation_name = mapping[name]
-
-        # Create and return the dynamic method
-        return self._create_expectation_method(expectation_name, name)
-
-    def _create_expectation_method(self, expectation_name: str, method_name: str):
+    def _create_expectation_method(self, suite_method_name: str):
         """
         Create a dynamic expectation method.
 
-        Returns a closure that captures the expectation_name and self.
+        Returns a closure that captures the suite_method_name and self.
         """
 
         def dynamic_method(**kwargs):
             """Dynamically generated expectation method."""
-            expectation = DataFrameExpectationRegistry.get_expectation(
-                expectation_name=expectation_name, **kwargs
-            )
+            try:
+                expectation = DataFrameExpectationRegistry.get_expectation_by_suite_method(
+                    suite_method_name=suite_method_name, **kwargs
+                )
+            except ValueError as e:
+                raise AttributeError(str(e)) from e
 
             logger.info(f"Adding expectation: {expectation}")
 
-            # Add to internal list
             self.__expectations.append(expectation)
             return self
 
         # Set helpful name for debugging
-        dynamic_method.__name__ = method_name
+        dynamic_method.__name__ = suite_method_name
 
         return dynamic_method