pydeflate 2.1.3__py3-none-any.whl → 2.2.0__py3-none-any.whl

pydeflate/core/source.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass, field
 
 import pandas as pd
 
+from pydeflate.exceptions import ConfigurationError, DataSourceError
 from pydeflate.sources.common import AvailableDeflators
 from pydeflate.sources.dac import read_dac
 from pydeflate.sources.imf import read_weo
@@ -10,6 +11,12 @@ from pydeflate.sources.world_bank import read_wb, read_wb_lcu_ppp, read_wb_usd_p
 
 @dataclass
 class Source:
+    """Base class for data sources implementing SourceProtocol.
+
+    This class handles loading data from external sources, caching,
+    and validation. It implements the SourceProtocol interface.
+    """
+
     name: str
     reader: callable
     update: bool = False
@@ -17,26 +24,100 @@ class Source:
     _idx = ["pydeflate_year", "pydeflate_entity_code", "pydeflate_iso3"]
 
     def __post_init__(self):
-        self.data = self.reader(self.update)
+        """Load and validate data after initialization."""
+        try:
+            self.data = self.reader(self.update)
+        except Exception as e:
+            raise DataSourceError(
+                f"Failed to load data: {e}",
+                source=self.name,
+            ) from e
+
         self.validate()
 
     def validate(self):
-        if self.data.empty:
-            raise ValueError(f"No data found for {self.name}")
+        """Validate that source data is properly formatted.
 
-        # check all columns start with pydeflate_
-        if not all(col.startswith("pydeflate_") for col in self.data.columns):
-            raise ValueError(f"Invalid data format for {self.name}")
+        Raises:
+            DataSourceError: If data is empty or improperly formatted
+            SchemaValidationError: If data doesn't match expected schema
+        """
+        if self.data.empty:
+            raise DataSourceError(f"No data found", source=self.name)
+
+        # Check all columns start with pydeflate_
+        invalid_cols = [
+            col for col in self.data.columns if not col.startswith("pydeflate_")
+        ]
+        if invalid_cols:
+            raise DataSourceError(
+                f"Invalid column names (must start with 'pydeflate_'): {invalid_cols}",
+                source=self.name,
+            )
+
+        # Validate schema if available and enabled
+        # Note: Schema validation is currently experimental
+        # Set environment variable PYDEFLATE_ENABLE_VALIDATION=1 to enable
+        import os
+
+        if os.environ.get("PYDEFLATE_ENABLE_VALIDATION") == "1":
+            try:
+                from pydeflate.schemas import validate_source_data
+
+                validate_source_data(self.data, self.name)
+            except ImportError:
+                # Pandera not available, skip schema validation
+                pass
+            except Exception as e:
+                # Schema validation failed, but don't break for now
+                # This allows us to roll out schema validation gradually
+                import logging
+
+                logger = logging.getLogger("pydeflate")
+                logger.debug(f"Schema validation skipped for {self.name}: {e}")
 
     def lcu_usd_exchange(self) -> pd.DataFrame:
+        """Return local currency to USD exchange rates.
+
+        Returns:
+            DataFrame with exchange rate data
+
+        Raises:
+            DataSourceError: If exchange rate data is missing
+        """
+        if "pydeflate_EXCHANGE" not in self.data.columns:
+            raise DataSourceError(
+                "Exchange rate data (pydeflate_EXCHANGE) not available",
+                source=self.name,
+            )
         return self.data.filter(self._idx + ["pydeflate_EXCHANGE"])
 
     def price_deflator(self, kind: AvailableDeflators = "NGDP_D") -> pd.DataFrame:
-
-        if f"pydeflate_{kind}" not in self.data.columns:
-            raise ValueError(f"No deflator data found for {kind} in {self.name} data.")
-
-        return self.data.filter(self._idx + [f"pydeflate_{kind}"])
+        """Return price deflator data for specified kind.
+
+        Args:
+            kind: Type of deflator (e.g., 'NGDP_D', 'CPI')
+
+        Returns:
+            DataFrame with deflator data
+
+        Raises:
+            ConfigurationError: If deflator kind not available for this source
+        """
+        column_name = f"pydeflate_{kind}"
+        if column_name not in self.data.columns:
+            available = [
+                col.replace("pydeflate_", "")
+                for col in self.data.columns
+                if col.startswith("pydeflate_") and col not in self._idx
+            ]
+            raise ConfigurationError(
+                f"Deflator '{kind}' not available for {self.name}. "
+                f"Available deflators: {', '.join(available)}",
+                parameter="kind",
+            )
+
+        return self.data.filter(self._idx + [column_name])
 
 
 class IMF(Source):

pydeflate/deflate/deflators.py

@@ -3,7 +3,7 @@ from functools import wraps
 import pandas as pd
 
 from pydeflate.core.api import BaseDeflate
-from pydeflate.core.source import DAC, WorldBank, IMF
+from pydeflate.core.source import DAC, IMF, WorldBank
 
 
 def _generate_docstring(source_name: str, price_kind: str) -> str:

pydeflate/deflate/legacy_deflate.py

@@ -4,7 +4,7 @@ import pandas as pd
 from pandas.util._decorators import deprecate_kwarg
 
 from pydeflate.core.api import BaseDeflate
-from pydeflate.core.source import DAC, WorldBank, IMF
+from pydeflate.core.source import DAC, IMF, WorldBank
 
 
 @deprecate_kwarg(old_arg_name="method", new_arg_name="deflator_method")

pydeflate/exceptions.py

@@ -0,0 +1,166 @@
+"""Custom exception hierarchy for pydeflate.
+
+This module defines specific exception types that allow users to handle
+different failure modes appropriately (e.g., retry on network errors,
+fail fast on validation errors).
+"""
+
+from __future__ import annotations
+
+
+class PydeflateError(Exception):
+    """Base exception for all pydeflate errors.
+
+    All exceptions raised by pydeflate inherit from this class,
+    making it easy to catch all pydeflate-specific errors.
+    """
+
+    pass
+
+
+class DataSourceError(PydeflateError):
+    """Raised when there's an issue with a data source.
+
+    This is a base class for all data source related errors.
+    """
+
+    def __init__(self, message: str, source: str | None = None):
+        """Initialize DataSourceError.
+
+        Args:
+            message: Description of the error
+            source: Name of the data source (e.g., 'IMF', 'World Bank')
+        """
+        self.source = source
+        super().__init__(f"[{source}] {message}" if source else message)
+
+
+class NetworkError(DataSourceError):
+    """Raised when network operations fail.
+
+    This typically indicates a transient error that might succeed on retry.
+    """
+
+    pass
+
+
+class SchemaValidationError(DataSourceError):
+    """Raised when data doesn't match expected schema.
+
+    This indicates a problem with the data structure, either from:
+    - External API changes
+    - Corrupted downloaded data
+    - User input with wrong columns/types
+    """
+
+    def __init__(
+        self,
+        message: str,
+        source: str | None = None,
+        expected_schema: dict | None = None,
+        actual_schema: dict | None = None,
+    ):
+        """Initialize SchemaValidationError.
+
+        Args:
+            message: Description of validation failure
+            source: Name of the data source
+            expected_schema: Expected column types/names
+            actual_schema: Actual column types/names found
+        """
+        self.expected_schema = expected_schema
+        self.actual_schema = actual_schema
+        super().__init__(message, source)
+
+
+class CacheError(PydeflateError):
+    """Raised when cache operations fail.
+
+    Examples:
+    - Unable to write to cache directory
+    - Corrupted cache files
+    - Lock file acquisition timeout
+    """
+
+    def __init__(self, message: str, cache_path: str | None = None):
+        """Initialize CacheError.
+
+        Args:
+            message: Description of cache error
+            cache_path: Path to the cache file/directory involved
+        """
+        self.cache_path = cache_path
+        super().__init__(
+            f"Cache error at {cache_path}: {message}" if cache_path else message
+        )
+
+
+class ConfigurationError(PydeflateError):
+    """Raised when configuration parameters are invalid.
+
+    Examples:
+    - Invalid currency code
+    - Base year out of range
+    - Missing required columns in user data
+    - Conflicting parameter combinations
+    """
+
+    def __init__(self, message: str, parameter: str | None = None):
+        """Initialize ConfigurationError.
+
+        Args:
+            message: Description of configuration issue
+            parameter: Name of the problematic parameter
+        """
+        self.parameter = parameter
+        super().__init__(
+            f"Invalid configuration for '{parameter}': {message}"
+            if parameter
+            else message
+        )
+
+
+class MissingDataError(PydeflateError):
+    """Raised when required deflator or exchange data is unavailable.
+
+    This occurs when:
+    - Requested country/year combination has no data in the source
+    - Data gaps in historical records
+    - Future years beyond available estimates
+    """
+
+    def __init__(
+        self,
+        message: str,
+        missing_entities: dict[str, list[int]] | None = None,
+    ):
+        """Initialize MissingDataError.
+
+        Args:
+            message: Description of missing data
+            missing_entities: Dict mapping entity codes to missing years
+        """
+        self.missing_entities = missing_entities
+        super().__init__(message)
+
+
+class PluginError(PydeflateError):
+    """Raised when plugin registration or loading fails.
+
+    Examples:
+    - Plugin doesn't implement required protocol
+    - Plugin name conflicts with existing source
+    - Plugin initialization fails
+    """
+
+    def __init__(self, message: str, plugin_name: str | None = None):
+        """Initialize PluginError.
+
+        Args:
+            message: Description of plugin error
+            plugin_name: Name of the plugin that failed
+        """
+        self.plugin_name = plugin_name
+        super().__init__(
+            f"Plugin '{plugin_name}' error: {message}" if plugin_name else message
+        )

pydeflate/exchange/exchangers.py

@@ -3,7 +3,7 @@ from functools import wraps
 import pandas as pd
 
 from pydeflate.core.api import BaseExchange
-from pydeflate.core.source import DAC, WorldBank, IMF, WorldBankPPP
+from pydeflate.core.source import DAC, IMF, WorldBank, WorldBankPPP
 
 
 def _generate_docstring(source_name: str) -> str:

pydeflate/plugins.py

@@ -0,0 +1,289 @@
+"""Plugin system for custom data sources.
+
+This module provides a registry-based plugin architecture that allows
+users to register custom data sources without modifying pydeflate's code.
+
+Example:
+    >>> from pydeflate.plugins import register_source
+    >>> from pydeflate.protocols import SourceProtocol
+    >>>
+    >>> @register_source("my_custom_source")
+    >>> class MyCustomSource:
+    ...     def __init__(self, update: bool = False):
+    ...         self.name = "my_custom_source"
+    ...         self.data = load_my_data(update)
+    ...         self._idx = ["pydeflate_year", "pydeflate_entity_code", "pydeflate_iso3"]
+    ...
+    ...     def lcu_usd_exchange(self): ...
+    ...     def price_deflator(self, kind): ...
+    ...     def validate(self): ...
+    >>>
+    >>> # Now use it:
+    >>> result = deflate_with_source("my_custom_source", data=df, ...)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, TypeVar
+
+from pydeflate.exceptions import PluginError
+from pydeflate.protocols import SourceProtocol
+
+# Type variable for source classes
+SourceType = TypeVar("SourceType", bound=type)
+
+
+class SourceRegistry:
+    """Registry for data source implementations.
+
+    This class maintains a mapping of source names to their implementation
+    classes. It provides methods to register, retrieve, and list sources.
+    """
+
+    def __init__(self):
+        """Initialize an empty registry."""
+        self._sources: dict[str, type] = {}
+        self._factories: dict[str, Callable[..., SourceProtocol]] = {}
+
+    def register(
+        self,
+        name: str,
+        source_class: type | None = None,
+        factory: Callable[..., SourceProtocol] | None = None,
+        *,
+        override: bool = False,
+    ) -> None:
+        """Register a data source.
+
+        Args:
+            name: Unique name for this source (e.g., 'IMF', 'WorldBank')
+            source_class: Class that implements SourceProtocol
+            factory: Factory function that returns a SourceProtocol instance
+            override: If True, allow replacing existing sources
+
+        Raises:
+            PluginError: If name already registered and override=False
+            PluginError: If neither source_class nor factory is provided
+            PluginError: If source doesn't implement SourceProtocol
+        """
+        # Validation
+        if source_class is None and factory is None:
+            raise PluginError(
+                "Must provide either source_class or factory",
+                plugin_name=name,
+            )
+
+        if not override and name in self._sources:
+            raise PluginError(
+                f"Source '{name}' already registered. Use override=True to replace.",
+                plugin_name=name,
+            )
+
+        # Check protocol conformance if source_class provided
+        if source_class is not None:
+            if not self._check_protocol_conformance(source_class):
+                raise PluginError(
+                    f"Source class must implement SourceProtocol",
+                    plugin_name=name,
+                )
+            self._sources[name] = source_class
+
+        if factory is not None:
+            self._factories[name] = factory
+
+    def _check_protocol_conformance(self, source_class: type) -> bool:
+        """Check if a class implements SourceProtocol.
+
+        Args:
+            source_class: Class to check
+
+        Returns:
+            True if class implements the protocol
+        """
+        required_methods = ["lcu_usd_exchange", "price_deflator", "validate"]
+        required_attrs = ["name", "data", "_idx"]
+
+        for method in required_methods:
+            if not hasattr(source_class, method):
+                return False
+            if not callable(getattr(source_class, method)):
+                return False
+
+        # Check that class can be instantiated and has required attributes
+        # Note: We can't fully check this without instantiation,
+        # so we just check the class has these as class attributes or in __init__
+        # This is a best-effort check
+        return True
+
+    def get(self, name: str, **kwargs) -> SourceProtocol:
+        """Get a source instance by name.
+
+        Args:
+            name: Name of the registered source
+            **kwargs: Keyword arguments to pass to source constructor
+
+        Returns:
+            Instance of the source
+
+        Raises:
+            PluginError: If source not found
+        """
+        # Try factory first
+        if name in self._factories:
+            try:
+                return self._factories[name](**kwargs)
+            except Exception as e:
+                raise PluginError(
+                    f"Factory function failed: {e}",
+                    plugin_name=name,
+                ) from e
+
+        # Try class
+        if name in self._sources:
+            try:
+                return self._sources[name](**kwargs)
+            except Exception as e:
+                raise PluginError(
+                    f"Source instantiation failed: {e}",
+                    plugin_name=name,
+                ) from e
+
+        raise PluginError(
+            f"Source '{name}' not found. Available sources: {self.list_sources()}",
+            plugin_name=name,
+        )
+
+    def list_sources(self) -> list[str]:
+        """List all registered source names.
+
+        Returns:
+            List of source names
+        """
+        all_names = set(self._sources.keys()) | set(self._factories.keys())
+        return sorted(all_names)
+
+    def is_registered(self, name: str) -> bool:
+        """Check if a source is registered.
+
+        Args:
+            name: Source name to check
+
+        Returns:
+            True if registered
+        """
+        return name in self._sources or name in self._factories
+
+    def unregister(self, name: str) -> None:
+        """Remove a source from the registry.
+
+        Args:
+            name: Source name to remove
+
+        Raises:
+            PluginError: If source not found
+        """
+        if not self.is_registered(name):
+            raise PluginError(
+                f"Cannot unregister '{name}': not found",
+                plugin_name=name,
+            )
+
+        self._sources.pop(name, None)
+        self._factories.pop(name, None)
+
+
+# Global registry instance
+_global_registry = SourceRegistry()
+
+
+def register_source(
+    name: str,
+    *,
+    override: bool = False,
+) -> Callable[[SourceType], SourceType]:
+    """Decorator to register a source class.
+
+    Example:
+        >>> @register_source("my_source")
+        ... class MySource:
+        ...     def __init__(self, update: bool = False):
+        ...         self.name = "my_source"
+        ...         ...
+        ...
+        ...     def lcu_usd_exchange(self): ...
+        ...     def price_deflator(self, kind): ...
+
+    Args:
+        name: Unique name for this source
+        override: Allow replacing existing sources
+
+    Returns:
+        Decorator function
+    """
+
+    def decorator(cls: SourceType) -> SourceType:
+        _global_registry.register(name, source_class=cls, override=override)
+        return cls
+
+    return decorator
+
+
+def get_source(name: str, **kwargs) -> SourceProtocol:
+    """Get a registered source instance.
+
+    Args:
+        name: Name of the source
+        **kwargs: Arguments to pass to source constructor
+
+    Returns:
+        Source instance
+
+    Raises:
+        PluginError: If source not found or instantiation fails
+    """
+    return _global_registry.get(name, **kwargs)
+
+
+def list_sources() -> list[str]:
+    """List all available sources.
+
+    Returns:
+        Sorted list of source names
+    """
+    return _global_registry.list_sources()
+
+
+def is_source_registered(name: str) -> bool:
+    """Check if a source is registered.
+
+    Args:
+        name: Source name
+
+    Returns:
+        True if registered
+    """
+    return _global_registry.is_registered(name)
+
+
+# Pre-register built-in sources
+def _register_builtin_sources():
+    """Register pydeflate's built-in sources."""
+    from pydeflate.core.source import DAC, IMF, WorldBank, WorldBankPPP
+
+    _global_registry.register("IMF", source_class=IMF, override=True)
+    _global_registry.register("World Bank", source_class=WorldBank, override=True)
+    _global_registry.register(
+        "World Bank PPP", source_class=WorldBankPPP, override=True
+    )
+    _global_registry.register("DAC", source_class=DAC, override=True)
+
+    # Aliases for convenience
+    _global_registry.register("imf", source_class=IMF, override=True)
+    _global_registry.register("wb", source_class=WorldBank, override=True)
+    _global_registry.register("worldbank", source_class=WorldBank, override=True)
+    _global_registry.register("dac", source_class=DAC, override=True)
+    _global_registry.register("oecd", source_class=DAC, override=True)
+
+
+# Register built-in sources when module is imported
+_register_builtin_sources()