diff --git a/.github/workflows/pr_code_changes.yaml b/.github/workflows/pr_code_changes.yaml index bdffb04a..960bb42c 100644 --- a/.github/workflows/pr_code_changes.yaml +++ b/.github/workflows/pr_code_changes.yaml @@ -44,6 +44,31 @@ jobs: run: uv pip install -e .[dev] --system - name: Run mypy (informational) run: mypy src/policyengine || echo "::warning::mypy found errors (non-blocking until codebase is clean)" + Python-Compat: + name: Install + smoke-import (py${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + allow-prereleases: true + - name: Install uv + uses: astral-sh/setup-uv@v5 + - name: Install package (no country-model extras) + # `h5py` is used transitively by policyengine.core.scoping_strategy + # but is normally supplied via the [us]/[uk] extras (through + # policyengine-core). Install it directly so the smoke import can + # exercise the wrapper without the country models, which pin + # versions that don't support 3.9/3.10 yet. + run: uv pip install --system . h5py + - name: Smoke-import core modules + run: python -c "import policyengine; from policyengine.core import Dataset, Policy, Simulation; from policyengine.outputs import aggregate, poverty, inequality; print('import OK')" Test: runs-on: macos-latest strategy: diff --git a/.gitignore b/.gitignore index 57a0fc21..3c351eab 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ *.ipynb _build/ .env -**/.DS_Store \ No newline at end of file +**/.DS_Store +build/ diff --git a/changelog.d/support-py39.added.md b/changelog.d/support-py39.added.md new file mode 100644 index 00000000..edb247fc --- /dev/null +++ b/changelog.d/support-py39.added.md @@ -0,0 +1 @@ +Support Python 3.9–3.12 (in addition to 3.13–3.14). PEP 604 `X | Y` annotations (evaluated at runtime by pydantic) are rewritten as `Optional[X]` / `Union[X, Y]`; `StrEnum` (3.11+) is replaced with `class Foo(str, Enum)`; PEP 695 generic class syntax in `core/cache.py` and `core/output.py` is rewritten using `typing.TypeVar` + `typing.Generic`. Ruff and mypy target versions dropped to py39. Requires `policyengine-us==1.602.0+` and `policyengine-uk==2.74.0+` from the `[us]`/`[uk]`/`[dev]` extras to also support 3.9/3.10. diff --git a/pyproject.toml b/pyproject.toml index 94b31711..bc5034b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,8 +11,12 @@ authors = [ {name = "PolicyEngine", email = "hello@policyengine.org"}, ] license = {file = "LICENSE"} -requires-python = ">=3.13" +requires-python = ">=3.9" classifiers = [ + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", ] @@ -73,7 +77,7 @@ filterwarnings = [ [tool.ruff] line-length = 88 -target-version = "py313" +target-version = "py39" extend-exclude = ["*.ipynb"] [tool.ruff.lint] @@ -84,7 +88,16 @@ select = [ "W", # pycodestyle warnings "UP", # pyupgrade ] -ignore = ["E501"] # Ignore line length errors +ignore = [ + "E501", # Ignore line length errors + # The following pyupgrade rules would require Python 3.10+, but we + # support 3.9+. Re-enable these once the 3.9 floor is dropped. + "UP006", # prefer `list` over `List` — OK on 3.9 at runtime, but for + # pydantic models we use typing.List for consistency + "UP007", # prefer `X | Y` over `Union[X, Y]` — needs 3.10+ + "UP035", # `typing.List` is deprecated — same as UP006 + "UP045", # prefer `X | None` over `Optional[X]` — needs 3.10+ +] [tool.ruff.format] quote-style = "double" @@ -93,7 +106,7 @@ skip-magic-trailing-comma = false line-ending = "auto" [tool.mypy] -python_version = "3.13" +python_version = "3.9" warn_return_any = true warn_unused_configs = true ignore_missing_imports = true diff --git a/src/policyengine/core/cache.py b/src/policyengine/core/cache.py index 44de06e3..410301e4 100644 --- a/src/policyengine/core/cache.py +++ b/src/policyengine/core/cache.py @@ -1,5 +1,6 @@ import logging from collections import OrderedDict +from typing import Generic, Optional, TypeVar import psutil @@ -8,15 +9,17 @@ _MEMORY_THRESHOLDS_GB = [8, 16, 32] _warned_thresholds: set[int] = set() +T = TypeVar("T") -class LRUCache[T]: + +class LRUCache(Generic[T]): """Least-recently-used cache with configurable size limit and memory monitoring.""" def __init__(self, max_size: int = 100): self._max_size = max_size self._cache: OrderedDict[str, T] = OrderedDict() - def get(self, key: str) -> T | None: + def get(self, key: str) -> Optional[T]: """Get item from cache, marking it as recently used.""" if key not in self._cache: return None diff --git a/src/policyengine/core/dataset.py b/src/policyengine/core/dataset.py index f10a5d22..27f51d16 100644 --- a/src/policyengine/core/dataset.py +++ b/src/policyengine/core/dataset.py @@ -1,3 +1,4 @@ +from typing import Optional from uuid import uuid4 import numpy as np @@ -76,7 +77,7 @@ class YearData(BaseModel): household: pd.DataFrame class MyDataset(Dataset): - data: YearData | None = None + data: Optional[YearData] = None """ model_config = ConfigDict(arbitrary_types_allowed=True) @@ -84,13 +85,13 @@ class MyDataset(Dataset): id: str = Field(default_factory=lambda: str(uuid4())) name: str description: str - dataset_version: DatasetVersion | None = None + dataset_version: Optional[DatasetVersion] = None filepath: str is_output_dataset: bool = False - tax_benefit_model: TaxBenefitModel | None = None + tax_benefit_model: Optional[TaxBenefitModel] = None year: int - data: BaseModel | None = None + data: Optional[BaseModel] = None def map_to_entity( @@ -98,8 +99,8 @@ def map_to_entity( source_entity: str, target_entity: str, person_entity: str = "person", - columns: list[str] | None = None, - values: np.ndarray | None = None, + columns: Optional[list[str]] = None, + values: Optional[np.ndarray] = None, how: str = "sum", ) -> MicroDataFrame: """Map data from source entity to target entity using join keys. diff --git a/src/policyengine/core/dynamic.py b/src/policyengine/core/dynamic.py index 81ef62b7..d707b9b2 100644 --- a/src/policyengine/core/dynamic.py +++ b/src/policyengine/core/dynamic.py @@ -1,5 +1,6 @@ from collections.abc import Callable from datetime import datetime +from typing import Optional from uuid import uuid4 from pydantic import BaseModel, Field @@ -10,9 +11,9 @@ class Dynamic(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str - description: str | None = None + description: Optional[str] = None parameter_values: list[ParameterValue] = [] - simulation_modifier: Callable | None = None + simulation_modifier: Optional[Callable] = None created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) diff --git a/src/policyengine/core/output.py b/src/policyengine/core/output.py index a4bf969a..e71634ab 100644 --- a/src/policyengine/core/output.py +++ b/src/policyengine/core/output.py @@ -1,4 +1,4 @@ -from typing import TypeVar +from typing import Generic, List, TypeVar import pandas as pd from pydantic import BaseModel, ConfigDict @@ -17,10 +17,10 @@ def run(self): raise NotImplementedError("Subclasses must implement run()") -class OutputCollection[T: "Output"](BaseModel): +class OutputCollection(BaseModel, Generic[T]): """Container for a collection of outputs with their DataFrame representation.""" model_config = ConfigDict(arbitrary_types_allowed=True) - outputs: list[T] + outputs: List[T] dataframe: pd.DataFrame diff --git a/src/policyengine/core/parameter.py b/src/policyengine/core/parameter.py index cd5a2c88..49f2b282 100644 --- a/src/policyengine/core/parameter.py +++ b/src/policyengine/core/parameter.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Optional from uuid import uuid4 from pydantic import BaseModel, Field, PrivateAttr @@ -15,15 +15,15 @@ class Parameter(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str - label: str | None = None - description: str | None = None - data_type: type | None = None + label: Optional[str] = None + description: Optional[str] = None + data_type: Optional[type] = None tax_benefit_model_version: TaxBenefitModelVersion - unit: str | None = None + unit: Optional[str] = None # Lazy loading: store core param ref, build values on demand _core_param: Any = PrivateAttr(default=None) - _parameter_values: list["ParameterValue"] | None = PrivateAttr(default=None) + _parameter_values: Optional[list["ParameterValue"]] = PrivateAttr(default=None) def __init__(self, _core_param: Any = None, **data): super().__init__(**data) diff --git a/src/policyengine/core/parameter_node.py b/src/policyengine/core/parameter_node.py index 9a3e25a0..54d384a5 100644 --- a/src/policyengine/core/parameter_node.py +++ b/src/policyengine/core/parameter_node.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from uuid import uuid4 from pydantic import BaseModel, Field @@ -22,8 +22,8 @@ class ParameterNode(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str = Field(description="Full path of the node (e.g., 'gov.hmrc')") - label: str | None = Field( + label: Optional[str] = Field( default=None, description="Human-readable label (e.g., 'HMRC')" ) - description: str | None = Field(default=None, description="Node description") + description: Optional[str] = Field(default=None, description="Node description") tax_benefit_model_version: "TaxBenefitModelVersion" diff --git a/src/policyengine/core/parameter_value.py b/src/policyengine/core/parameter_value.py index 073cd74b..a51ffeb0 100644 --- a/src/policyengine/core/parameter_value.py +++ b/src/policyengine/core/parameter_value.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional, Union from uuid import uuid4 from pydantic import BaseModel, Field @@ -10,7 +10,7 @@ class ParameterValue(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) - parameter: "Parameter | None" = None - value: float | int | str | bool | list | None = None + parameter: "Optional[Parameter]" = None + value: Optional[Union[float, int, str, bool, list]] = None start_date: datetime - end_date: datetime | None = None + end_date: Optional[datetime] = None diff --git a/src/policyengine/core/policy.py b/src/policyengine/core/policy.py index bfb4ca9e..3860a817 100644 --- a/src/policyengine/core/policy.py +++ b/src/policyengine/core/policy.py @@ -1,5 +1,6 @@ from collections.abc import Callable from datetime import datetime +from typing import Optional from uuid import uuid4 from pydantic import BaseModel, Field @@ -10,9 +11,9 @@ class Policy(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str - description: str | None = None + description: Optional[str] = None parameter_values: list[ParameterValue] = [] - simulation_modifier: Callable | None = None + simulation_modifier: Optional[Callable] = None created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) diff --git a/src/policyengine/core/region.py b/src/policyengine/core/region.py index ebf1f93a..7ff55a64 100644 --- a/src/policyengine/core/region.py +++ b/src/policyengine/core/region.py @@ -6,7 +6,7 @@ 2. Filter from a parent region's dataset (e.g., US places/cities, UK countries) """ -from typing import Literal +from typing import Literal, Optional, Union from pydantic import BaseModel, Field, PrivateAttr @@ -15,7 +15,7 @@ # Region type literals for US and UK USRegionType = Literal["national", "state", "congressional_district", "place"] UKRegionType = Literal["national", "country", "constituency", "local_authority"] -RegionType = USRegionType | UKRegionType +RegionType = Union[USRegionType, UKRegionType] class Region(BaseModel): @@ -46,19 +46,19 @@ class Region(BaseModel): ) # Hierarchy - parent_code: str | None = Field( + parent_code: Optional[str] = Field( default=None, description="Code of parent region (e.g., 'us' for states, 'state/nj' for places in New Jersey)", ) # Dataset configuration - dataset_path: str | None = Field( + dataset_path: Optional[str] = Field( default=None, description="GCS path to dedicated dataset (e.g., 'gs://policyengine-us-data/states/CA.h5')", ) # Scoping strategy (preferred over legacy filter fields) - scoping_strategy: ScopingStrategy | None = Field( + scoping_strategy: Optional[ScopingStrategy] = Field( default=None, description="Strategy for scoping dataset to this region (row filtering or weight replacement)", ) @@ -68,20 +68,20 @@ class Region(BaseModel): default=False, description="True if this region filters from a parent dataset rather than having its own", ) - filter_field: str | None = Field( + filter_field: Optional[str] = Field( default=None, description="Dataset field to filter on (e.g., 'place_fips', 'country')", ) - filter_value: str | None = Field( + filter_value: Optional[str] = Field( default=None, description="Value to match when filtering (defaults to code suffix if not set)", ) # Metadata (primarily for US congressional districts) - state_code: str | None = Field( + state_code: Optional[str] = Field( default=None, description="Two-letter state code (e.g., 'CA', 'NJ')" ) - state_name: str | None = Field( + state_name: Optional[str] = Field( default=None, description="Full state name (e.g., 'California', 'New Jersey')", ) @@ -137,7 +137,7 @@ def add_region(self, region: Region) -> None: self._by_type[region.region_type] = [] self._by_type[region.region_type].append(region) - def get(self, code: str) -> Region | None: + def get(self, code: str) -> Optional[Region]: """Get a region by its code. Args: @@ -159,7 +159,7 @@ def get_by_type(self, region_type: str) -> list[Region]: """ return self._by_type.get(region_type, []) - def get_national(self) -> Region | None: + def get_national(self) -> Optional[Region]: """Get the national-level region. Returns: diff --git a/src/policyengine/core/release_manifest.py b/src/policyengine/core/release_manifest.py index 3106998e..90a09f32 100644 --- a/src/policyengine/core/release_manifest.py +++ b/src/policyengine/core/release_manifest.py @@ -3,6 +3,7 @@ from importlib import import_module from importlib.resources import files from pathlib import Path +from typing import Optional import requests from pydantic import BaseModel, Field @@ -35,14 +36,14 @@ class CompatibleModelPackage(BaseModel): class BuiltWithModelPackage(PackageVersion): - git_sha: str | None = None - data_build_fingerprint: str | None = None + git_sha: Optional[str] = None + data_build_fingerprint: Optional[str] = None class DataBuildInfo(BaseModel): - build_id: str | None = None - built_at: str | None = None - built_with_model_package: BuiltWithModelPackage | None = None + build_id: Optional[str] = None + built_at: Optional[str] = None + built_with_model_package: Optional[BuiltWithModelPackage] = None class ArtifactPathReference(BaseModel): @@ -61,8 +62,8 @@ class DataReleaseArtifact(BaseModel): path: str repo_id: str revision: str - sha256: str | None = None - size_bytes: int | None = None + sha256: Optional[str] = None + size_bytes: Optional[int] = None @property def uri(self) -> str: @@ -80,32 +81,32 @@ class DataReleaseManifest(BaseModel): default_factory=list ) default_datasets: dict[str, str] = Field(default_factory=dict) - build: DataBuildInfo | None = None + build: Optional[DataBuildInfo] = None artifacts: dict[str, DataReleaseArtifact] = Field(default_factory=dict) class DataCertification(BaseModel): compatibility_basis: str certified_for_model_version: str - data_build_id: str | None = None - built_with_model_version: str | None = None - built_with_model_git_sha: str | None = None - data_build_fingerprint: str | None = None - certified_by: str | None = None + data_build_id: Optional[str] = None + built_with_model_version: Optional[str] = None + built_with_model_git_sha: Optional[str] = None + data_build_fingerprint: Optional[str] = None + certified_by: Optional[str] = None class CertifiedDataArtifact(BaseModel): - data_package: PackageVersion | None = None + data_package: Optional[PackageVersion] = None dataset: str uri: str - sha256: str | None = None - build_id: str | None = None + sha256: Optional[str] = None + build_id: Optional[str] = None class CountryReleaseManifest(BaseModel): schema_version: int = 1 - bundle_id: str | None = None - published_at: str | None = None + bundle_id: Optional[str] = None + published_at: Optional[str] = None country_id: str policyengine_version: str model_package: PackageVersion @@ -113,8 +114,8 @@ class CountryReleaseManifest(BaseModel): default_dataset: str datasets: dict[str, ArtifactPathReference] = Field(default_factory=dict) region_datasets: dict[str, ArtifactPathTemplate] = Field(default_factory=dict) - certified_data_artifact: CertifiedDataArtifact | None = None - certification: DataCertification | None = None + certified_data_artifact: Optional[CertifiedDataArtifact] = None + certification: Optional[DataCertification] = None @property def default_dataset_uri(self) -> str: @@ -186,7 +187,7 @@ def _specifier_matches(version: str, specifier: str) -> bool: def certify_data_release_compatibility( country_id: str, runtime_model_version: str, - runtime_data_build_fingerprint: str | None = None, + runtime_data_build_fingerprint: Optional[str] = None, ) -> DataCertification: country_manifest = get_release_manifest(country_id) try: @@ -322,7 +323,7 @@ def resolve_dataset_reference(country_id: str, dataset: str) -> str: def resolve_managed_dataset_reference( country_id: str, - dataset: str | None = None, + dataset: Optional[str] = None, *, allow_unmanaged: bool = False, ) -> str: @@ -414,7 +415,7 @@ def resolve_region_dataset_path( country_id: str, region_type: str, **kwargs: str, -) -> str | None: +) -> Optional[str]: manifest = get_release_manifest(country_id) template = manifest.region_datasets.get(region_type) if template is None: diff --git a/src/policyengine/core/scoping_strategy.py b/src/policyengine/core/scoping_strategy.py index 75449a6d..7d9b5126 100644 --- a/src/policyengine/core/scoping_strategy.py +++ b/src/policyengine/core/scoping_strategy.py @@ -12,7 +12,7 @@ import logging from abc import abstractmethod from pathlib import Path -from typing import Annotated, Literal +from typing import Annotated, Literal, Optional, Union import h5py import numpy as np @@ -201,7 +201,7 @@ def _find_region_index(lookup_df: pd.DataFrame, region_code: str) -> int: ) @staticmethod - def _find_household_id_column(df: pd.DataFrame, entity_name: str) -> str | None: + def _find_household_id_column(df: pd.DataFrame, entity_name: str) -> Optional[str]: """Find the column linking an entity to its household.""" candidates = [ "person_household_id", @@ -219,6 +219,6 @@ def cache_key(self) -> str: ScopingStrategy = Annotated[ - RowFilterStrategy | WeightReplacementStrategy, + Union[RowFilterStrategy, WeightReplacementStrategy], Discriminator("strategy_type"), ] diff --git a/src/policyengine/core/simulation.py b/src/policyengine/core/simulation.py index b9af105d..6456e5bc 100644 --- a/src/policyengine/core/simulation.py +++ b/src/policyengine/core/simulation.py @@ -1,5 +1,6 @@ import logging from datetime import datetime +from typing import Optional from uuid import uuid4 from pydantic import BaseModel, Field, model_validator @@ -21,22 +22,22 @@ class Simulation(BaseModel): created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) - policy: Policy | None = None - dynamic: Dynamic | None = None + policy: Optional[Policy] = None + dynamic: Optional[Dynamic] = None dataset: Dataset = None # Scoping strategy (preferred over legacy filter fields) - scoping_strategy: ScopingStrategy | None = Field( + scoping_strategy: Optional[ScopingStrategy] = Field( default=None, description="Strategy for scoping dataset to a sub-national region", ) # Legacy regional filtering parameters (kept for backward compatibility) - filter_field: str | None = Field( + filter_field: Optional[str] = Field( default=None, description="Household-level variable to filter dataset by (e.g., 'place_fips', 'country')", ) - filter_value: str | None = Field( + filter_value: Optional[str] = Field( default=None, description="Value to match when filtering (e.g., '44000', 'ENGLAND')", ) @@ -61,7 +62,7 @@ def _auto_construct_strategy(self) -> "Simulation": ) return self - output_dataset: Dataset | None = None + output_dataset: Optional[Dataset] = None def run(self): self.tax_benefit_model_version.run(self) @@ -96,7 +97,7 @@ def load(self): self.tax_benefit_model_version.load(self) @property - def release_bundle(self) -> dict[str, str | None]: + def release_bundle(self) -> dict[str, Optional[str]]: bundle = ( self.tax_benefit_model_version.release_bundle if self.tax_benefit_model_version is not None diff --git a/src/policyengine/core/tax_benefit_model.py b/src/policyengine/core/tax_benefit_model.py index 02cb94ef..c2d4e26d 100644 --- a/src/policyengine/core/tax_benefit_model.py +++ b/src/policyengine/core/tax_benefit_model.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from pydantic import BaseModel @@ -8,4 +8,4 @@ class TaxBenefitModel(BaseModel): id: str - description: str | None = None + description: Optional[str] = None diff --git a/src/policyengine/core/tax_benefit_model_version.py b/src/policyengine/core/tax_benefit_model_version.py index f253fc5c..7fb03334 100644 --- a/src/policyengine/core/tax_benefit_model_version.py +++ b/src/policyengine/core/tax_benefit_model_version.py @@ -1,5 +1,5 @@ -from datetime import UTC, datetime -from typing import TYPE_CHECKING +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Optional from uuid import uuid4 from pydantic import BaseModel, Field @@ -22,25 +22,27 @@ class TaxBenefitModelVersion(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) model: TaxBenefitModel version: str - description: str | None = None - created_at: datetime | None = Field(default_factory=lambda: datetime.now(UTC)) + description: Optional[str] = None + created_at: Optional[datetime] = Field( + default_factory=lambda: datetime.now(timezone.utc) + ) variables: list["Variable"] = Field(default_factory=list) parameters: list["Parameter"] = Field(default_factory=list) parameter_nodes: list["ParameterNode"] = Field(default_factory=list) # Region registry for geographic simulations - region_registry: "RegionRegistry | None" = Field( + region_registry: "Optional[RegionRegistry]" = Field( default=None, description="Registry of supported geographic regions" ) - release_manifest: CountryReleaseManifest | None = Field( + release_manifest: Optional[CountryReleaseManifest] = Field( default=None, exclude=True, ) - model_package: PackageVersion | None = Field(default=None) - data_package: PackageVersion | None = Field(default=None) - default_dataset_uri: str | None = Field(default=None) - data_certification: DataCertification | None = Field(default=None) + model_package: Optional[PackageVersion] = Field(default=None) + data_package: Optional[PackageVersion] = Field(default=None) + default_dataset_uri: Optional[str] = Field(default=None) + data_certification: Optional[DataCertification] = Field(default=None) @property def parameter_values(self) -> list["ParameterValue"]: @@ -112,7 +114,7 @@ def get_parameter_node(self, name: str) -> "ParameterNode": f"ParameterNode '{name}' not found in {self.model.id} version {self.version}" ) - def get_region(self, code: str) -> "Region | None": + def get_region(self, code: str) -> "Optional[Region]": """Get a region by its code. Args: @@ -126,7 +128,7 @@ def get_region(self, code: str) -> "Region | None": return self.region_registry.get(code) @property - def release_bundle(self) -> dict[str, str | None]: + def release_bundle(self) -> dict[str, Optional[str]]: manifest_certification = ( self.release_manifest.certification if self.release_manifest is not None diff --git a/src/policyengine/core/trace_tro.py b/src/policyengine/core/trace_tro.py index 52ae7b15..ae31a29e 100644 --- a/src/policyengine/core/trace_tro.py +++ b/src/policyengine/core/trace_tro.py @@ -3,6 +3,7 @@ import hashlib import json from collections.abc import Iterable, Mapping +from typing import Optional from .release_manifest import ( CountryReleaseManifest, @@ -28,7 +29,7 @@ def _hash_object(value: str) -> dict[str, str]: } -def _artifact_mime_type(path_or_uri: str) -> str | None: +def _artifact_mime_type(path_or_uri: str) -> Optional[str]: suffix = path_or_uri.rsplit(".", 1)[-1].lower() if "." in path_or_uri else "" return { "h5": "application/x-hdf5", @@ -53,9 +54,9 @@ def build_trace_tro_from_release_bundle( country_manifest: CountryReleaseManifest, data_release_manifest: DataReleaseManifest, *, - certification: DataCertification | None = None, - bundle_manifest_path: str | None = None, - data_release_manifest_path: str | None = None, + certification: Optional[DataCertification] = None, + bundle_manifest_path: Optional[str] = None, + data_release_manifest_path: Optional[str] = None, ) -> dict: certified_artifact = country_manifest.certified_data_artifact if certified_artifact is None: diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index 60aea9c5..03e53495 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Optional from pydantic import BaseModel @@ -8,13 +8,13 @@ class Variable(BaseModel): id: str name: str - label: str | None = None + label: Optional[str] = None tax_benefit_model_version: TaxBenefitModelVersion entity: str - description: str | None = None + description: Optional[str] = None data_type: type = None - possible_values: list[Any] | None = None + possible_values: Optional[list[Any]] = None default_value: Any = None - value_type: type | None = None - adds: list[str] | None = None - subtracts: list[str] | None = None + value_type: Optional[type] = None + adds: Optional[list[str]] = None + subtracts: Optional[list[str]] = None diff --git a/src/policyengine/outputs/aggregate.py b/src/policyengine/outputs/aggregate.py index 9406a4d7..d014b06c 100644 --- a/src/policyengine/outputs/aggregate.py +++ b/src/policyengine/outputs/aggregate.py @@ -1,10 +1,10 @@ -from enum import StrEnum -from typing import Any +from enum import Enum +from typing import Any, Optional from policyengine.core import Output, Simulation -class AggregateType(StrEnum): +class AggregateType(str, Enum): SUM = "sum" MEAN = "mean" COUNT = "count" @@ -14,23 +14,25 @@ class Aggregate(Output): simulation: Simulation variable: str aggregate_type: AggregateType - entity: str | None = None + entity: Optional[str] = None - filter_variable: str | None = None - filter_variable_eq: Any | None = None - filter_variable_leq: Any | None = None - filter_variable_geq: Any | None = None + filter_variable: Optional[str] = None + filter_variable_eq: Optional[Any] = None + filter_variable_leq: Optional[Any] = None + filter_variable_geq: Optional[Any] = None filter_variable_describes_quantiles: bool = False # Convenient quantile specification (alternative to describes_quantiles) - quantile: int | None = ( + quantile: Optional[int] = ( None # Number of quantiles (e.g., 10 for deciles, 5 for quintiles) ) - quantile_eq: int | None = None # Exact quantile (e.g., 3 for 3rd decile) - quantile_leq: int | None = None # Maximum quantile (e.g., 5 for bottom 5 deciles) - quantile_geq: int | None = None # Minimum quantile (e.g., 9 for top 2 deciles) + quantile_eq: Optional[int] = None # Exact quantile (e.g., 3 for 3rd decile) + quantile_leq: Optional[int] = ( + None # Maximum quantile (e.g., 5 for bottom 5 deciles) + ) + quantile_geq: Optional[int] = None # Minimum quantile (e.g., 9 for top 2 deciles) - result: Any | None = None + result: Optional[Any] = None def run(self): # Convert quantile specification to describes_quantiles format diff --git a/src/policyengine/outputs/change_aggregate.py b/src/policyengine/outputs/change_aggregate.py index e1cd3985..87d2e0d9 100644 --- a/src/policyengine/outputs/change_aggregate.py +++ b/src/policyengine/outputs/change_aggregate.py @@ -1,10 +1,10 @@ -from enum import StrEnum -from typing import Any +from enum import Enum +from typing import Any, Optional from policyengine.core import Output, Simulation -class ChangeAggregateType(StrEnum): +class ChangeAggregateType(str, Enum): COUNT = "count" SUM = "sum" MEAN = "mean" @@ -15,34 +15,36 @@ class ChangeAggregate(Output): reform_simulation: Simulation variable: str aggregate_type: ChangeAggregateType - entity: str | None = None + entity: Optional[str] = None # Filter by absolute change - change_geq: float | None = None # Change >= value (e.g., gain >= 500) - change_leq: float | None = None # Change <= value (e.g., loss <= -500) - change_eq: float | None = None # Change == value + change_geq: Optional[float] = None # Change >= value (e.g., gain >= 500) + change_leq: Optional[float] = None # Change <= value (e.g., loss <= -500) + change_eq: Optional[float] = None # Change == value # Filter by relative change (as decimal, e.g., 0.05 = 5%) - relative_change_geq: float | None = None # Relative change >= value - relative_change_leq: float | None = None # Relative change <= value - relative_change_eq: float | None = None # Relative change == value + relative_change_geq: Optional[float] = None # Relative change >= value + relative_change_leq: Optional[float] = None # Relative change <= value + relative_change_eq: Optional[float] = None # Relative change == value # Filter by another variable (e.g., only count people with age >= 30) - filter_variable: str | None = None - filter_variable_eq: Any | None = None - filter_variable_leq: Any | None = None - filter_variable_geq: Any | None = None + filter_variable: Optional[str] = None + filter_variable_eq: Optional[Any] = None + filter_variable_leq: Optional[Any] = None + filter_variable_geq: Optional[Any] = None filter_variable_describes_quantiles: bool = False # Convenient quantile specification (alternative to describes_quantiles) - quantile: int | None = ( + quantile: Optional[int] = ( None # Number of quantiles (e.g., 10 for deciles, 5 for quintiles) ) - quantile_eq: int | None = None # Exact quantile (e.g., 3 for 3rd decile) - quantile_leq: int | None = None # Maximum quantile (e.g., 5 for bottom 5 deciles) - quantile_geq: int | None = None # Minimum quantile (e.g., 9 for top 2 deciles) + quantile_eq: Optional[int] = None # Exact quantile (e.g., 3 for 3rd decile) + quantile_leq: Optional[int] = ( + None # Maximum quantile (e.g., 5 for bottom 5 deciles) + ) + quantile_geq: Optional[int] = None # Minimum quantile (e.g., 9 for top 2 deciles) - result: Any | None = None + result: Optional[Any] = None def run(self): # Convert quantile specification to describes_quantiles format diff --git a/src/policyengine/outputs/congressional_district_impact.py b/src/policyengine/outputs/congressional_district_impact.py index d8162a6d..4a1d0d90 100644 --- a/src/policyengine/outputs/congressional_district_impact.py +++ b/src/policyengine/outputs/congressional_district_impact.py @@ -1,6 +1,6 @@ """Congressional district impact output class for US policy reforms.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import numpy as np from pydantic import ConfigDict @@ -26,7 +26,7 @@ class CongressionalDistrictImpact(Output): reform_simulation: "Simulation" # Results populated by run() - district_results: list[dict] | None = None + district_results: Optional[list[dict]] = None def run(self) -> None: """Group households by geoid and compute per-district metrics.""" diff --git a/src/policyengine/outputs/constituency_impact.py b/src/policyengine/outputs/constituency_impact.py index 5cee7f4d..60f76e0b 100644 --- a/src/policyengine/outputs/constituency_impact.py +++ b/src/policyengine/outputs/constituency_impact.py @@ -5,7 +5,7 @@ that reweights all households to represent that constituency's demographics. """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import h5py import numpy as np @@ -35,7 +35,7 @@ class ConstituencyImpact(Output): year: str = "2025" # Results populated by run() - constituency_results: list[dict] | None = None + constituency_results: Optional[list[dict]] = None def run(self) -> None: """Load weight matrix and compute per-constituency metrics.""" diff --git a/src/policyengine/outputs/decile_impact.py b/src/policyengine/outputs/decile_impact.py index d3339003..b0f2306e 100644 --- a/src/policyengine/outputs/decile_impact.py +++ b/src/policyengine/outputs/decile_impact.py @@ -1,3 +1,5 @@ +from typing import Optional + import pandas as pd from pydantic import ConfigDict @@ -16,19 +18,19 @@ class DecileImpact(Output): baseline_simulation: Simulation reform_simulation: Simulation income_variable: str = "equiv_hbai_household_net_income" - decile_variable: str | None = None # If set, use pre-computed grouping variable - entity: str | None = None + decile_variable: Optional[str] = None # If set, use pre-computed grouping variable + entity: Optional[str] = None decile: int quantiles: int = 10 # Results populated by run() - baseline_mean: float | None = None - reform_mean: float | None = None - absolute_change: float | None = None - relative_change: float | None = None - count_better_off: float | None = None - count_worse_off: float | None = None - count_no_change: float | None = None + baseline_mean: Optional[float] = None + reform_mean: Optional[float] = None + absolute_change: Optional[float] = None + relative_change: Optional[float] = None + count_better_off: Optional[float] = None + count_worse_off: Optional[float] = None + count_no_change: Optional[float] = None def run(self): """Calculate impact for this specific decile.""" @@ -97,16 +99,16 @@ def run(self): def calculate_decile_impacts( - dataset: Dataset | None = None, - tax_benefit_model_version: TaxBenefitModelVersion | None = None, - baseline_policy: Policy | None = None, - reform_policy: Policy | None = None, - dynamic: Dynamic | None = None, + dataset: Optional[Dataset] = None, + tax_benefit_model_version: Optional[TaxBenefitModelVersion] = None, + baseline_policy: Optional[Policy] = None, + reform_policy: Optional[Policy] = None, + dynamic: Optional[Dynamic] = None, income_variable: str = "equiv_hbai_household_net_income", - entity: str | None = None, + entity: Optional[str] = None, quantiles: int = 10, - baseline_simulation: Simulation | None = None, - reform_simulation: Simulation | None = None, + baseline_simulation: Optional[Simulation] = None, + reform_simulation: Optional[Simulation] = None, ) -> OutputCollection[DecileImpact]: """Calculate decile-by-decile impact of a reform. diff --git a/src/policyengine/outputs/inequality.py b/src/policyengine/outputs/inequality.py index 8656dc65..4b16f7a9 100644 --- a/src/policyengine/outputs/inequality.py +++ b/src/policyengine/outputs/inequality.py @@ -1,7 +1,7 @@ """Inequality analysis output types.""" -from enum import StrEnum -from typing import Any +from enum import Enum +from typing import Any, Optional, Union import numpy as np import pandas as pd @@ -10,7 +10,7 @@ from policyengine.core import Output, Simulation -class USInequalityPreset(StrEnum): +class USInequalityPreset(str, Enum): """Preset configurations for US inequality analysis.""" STANDARD = "standard" @@ -86,21 +86,21 @@ class Inequality(Output): simulation: Simulation income_variable: str entity: str = "household" - weight_multiplier_variable: str | None = None - equivalization_variable: str | None = None + weight_multiplier_variable: Optional[str] = None + equivalization_variable: Optional[str] = None equivalization_power: float = 0.0 # Optional demographic filters - filter_variable: str | None = None - filter_variable_eq: Any | None = None - filter_variable_leq: Any | None = None - filter_variable_geq: Any | None = None + filter_variable: Optional[str] = None + filter_variable_eq: Optional[Any] = None + filter_variable_leq: Optional[Any] = None + filter_variable_geq: Optional[Any] = None # Results populated by run() - gini: float | None = None - top_10_share: float | None = None - top_1_share: float | None = None - bottom_50_share: float | None = None + gini: Optional[float] = None + top_10_share: Optional[float] = None + top_1_share: Optional[float] = None + bottom_50_share: Optional[float] = None def run(self): """Calculate inequality metrics.""" @@ -235,10 +235,10 @@ def run(self): def calculate_uk_inequality( simulation: Simulation, income_variable: str = UK_INEQUALITY_INCOME_VARIABLE, - filter_variable: str | None = None, - filter_variable_eq: Any | None = None, - filter_variable_leq: Any | None = None, - filter_variable_geq: Any | None = None, + filter_variable: Optional[str] = None, + filter_variable_eq: Optional[Any] = None, + filter_variable_leq: Optional[Any] = None, + filter_variable_geq: Optional[Any] = None, ) -> Inequality: """Calculate inequality metrics for a UK simulation. @@ -269,11 +269,11 @@ def calculate_uk_inequality( def calculate_us_inequality( simulation: Simulation, income_variable: str = US_INEQUALITY_INCOME_VARIABLE, - preset: USInequalityPreset | str = USInequalityPreset.STANDARD, - filter_variable: str | None = None, - filter_variable_eq: Any | None = None, - filter_variable_leq: Any | None = None, - filter_variable_geq: Any | None = None, + preset: Union[USInequalityPreset, str] = USInequalityPreset.STANDARD, + filter_variable: Optional[str] = None, + filter_variable_eq: Optional[Any] = None, + filter_variable_leq: Optional[Any] = None, + filter_variable_geq: Optional[Any] = None, ) -> Inequality: """Calculate inequality metrics for a US simulation. diff --git a/src/policyengine/outputs/intra_decile_impact.py b/src/policyengine/outputs/intra_decile_impact.py index e2b01243..b91a04e2 100644 --- a/src/policyengine/outputs/intra_decile_impact.py +++ b/src/policyengine/outputs/intra_decile_impact.py @@ -15,6 +15,8 @@ household_weight) so they reflect the share of people, not households. """ +from typing import Optional + import numpy as np import pandas as pd from pydantic import ConfigDict @@ -41,17 +43,17 @@ class IntraDecileImpact(Output): baseline_simulation: Simulation reform_simulation: Simulation income_variable: str = "household_net_income" - decile_variable: str | None = None # If set, use pre-computed grouping + decile_variable: Optional[str] = None # If set, use pre-computed grouping entity: str = "household" decile: int # 1-10 for individual deciles quantiles: int = 10 # Results populated by run() - lose_more_than_5pct: float | None = None - lose_less_than_5pct: float | None = None - no_change: float | None = None - gain_less_than_5pct: float | None = None - gain_more_than_5pct: float | None = None + lose_more_than_5pct: Optional[float] = None + lose_less_than_5pct: Optional[float] = None + no_change: Optional[float] = None + gain_less_than_5pct: Optional[float] = None + gain_more_than_5pct: Optional[float] = None def run(self): """Calculate intra-decile proportions for this specific decile.""" @@ -117,7 +119,7 @@ def compute_intra_decile_impacts( baseline_simulation: Simulation, reform_simulation: Simulation, income_variable: str = "household_net_income", - decile_variable: str | None = None, + decile_variable: Optional[str] = None, entity: str = "household", quantiles: int = 10, ) -> OutputCollection[IntraDecileImpact]: diff --git a/src/policyengine/outputs/local_authority_impact.py b/src/policyengine/outputs/local_authority_impact.py index fc91f3ec..20b17efe 100644 --- a/src/policyengine/outputs/local_authority_impact.py +++ b/src/policyengine/outputs/local_authority_impact.py @@ -5,7 +5,7 @@ that reweights all households to represent that local authority's demographics. """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import h5py import numpy as np @@ -35,7 +35,7 @@ class LocalAuthorityImpact(Output): year: str = "2025" # Results populated by run() - local_authority_results: list[dict] | None = None + local_authority_results: Optional[list[dict]] = None def run(self) -> None: """Load weight matrix and compute per-local-authority metrics.""" diff --git a/src/policyengine/outputs/poverty.py b/src/policyengine/outputs/poverty.py index 10db4682..85a761a5 100644 --- a/src/policyengine/outputs/poverty.py +++ b/src/policyengine/outputs/poverty.py @@ -1,7 +1,7 @@ """Poverty analysis output types.""" -from enum import StrEnum -from typing import Any +from enum import Enum +from typing import Any, Optional import pandas as pd from pydantic import ConfigDict @@ -9,7 +9,7 @@ from policyengine.core import Output, OutputCollection, Simulation -class UKPovertyType(StrEnum): +class UKPovertyType(str, Enum): """UK poverty measure types.""" ABSOLUTE_BHC = "absolute_bhc" @@ -18,7 +18,7 @@ class UKPovertyType(StrEnum): RELATIVE_AHC = "relative_ahc" -class USPovertyType(StrEnum): +class USPovertyType(str, Enum): """US poverty measure types.""" SPM = "spm" @@ -51,22 +51,22 @@ class Poverty(Output): simulation: Simulation poverty_variable: str - poverty_type: str | None = None + poverty_type: Optional[str] = None entity: str = "person" # Optional demographic filters - filter_variable: str | None = None - filter_variable_eq: Any | None = None - filter_variable_leq: Any | None = None - filter_variable_geq: Any | None = None + filter_variable: Optional[str] = None + filter_variable_eq: Optional[Any] = None + filter_variable_leq: Optional[Any] = None + filter_variable_geq: Optional[Any] = None # Convenience group label (set by by_age/by_gender/by_race wrappers) - filter_group: str | None = None + filter_group: Optional[str] = None # Results populated by run() - headcount: float | None = None - total_population: float | None = None - rate: float | None = None + headcount: Optional[float] = None + total_population: Optional[float] = None + rate: Optional[float] = None def run(self): """Calculate poverty headcount and rate.""" @@ -128,10 +128,10 @@ def run(self): def calculate_uk_poverty_rates( simulation: Simulation, - filter_variable: str | None = None, - filter_variable_eq: Any | None = None, - filter_variable_leq: Any | None = None, - filter_variable_geq: Any | None = None, + filter_variable: Optional[str] = None, + filter_variable_eq: Optional[Any] = None, + filter_variable_leq: Optional[Any] = None, + filter_variable_geq: Optional[Any] = None, ) -> OutputCollection[Poverty]: """Calculate all UK poverty rates for a simulation. @@ -151,7 +151,7 @@ def calculate_uk_poverty_rates( poverty = Poverty( simulation=simulation, poverty_variable=poverty_variable, - poverty_type=str(poverty_type), + poverty_type=poverty_type.value, entity="person", filter_variable=filter_variable, filter_variable_eq=filter_variable_eq, @@ -184,10 +184,10 @@ def calculate_uk_poverty_rates( def calculate_us_poverty_rates( simulation: Simulation, - filter_variable: str | None = None, - filter_variable_eq: Any | None = None, - filter_variable_leq: Any | None = None, - filter_variable_geq: Any | None = None, + filter_variable: Optional[str] = None, + filter_variable_eq: Optional[Any] = None, + filter_variable_leq: Optional[Any] = None, + filter_variable_geq: Optional[Any] = None, ) -> OutputCollection[Poverty]: """Calculate all US poverty rates for a simulation. @@ -207,7 +207,7 @@ def calculate_us_poverty_rates( poverty = Poverty( simulation=simulation, poverty_variable=poverty_variable, - poverty_type=str(poverty_type), + poverty_type=poverty_type.value, entity="person", filter_variable=filter_variable, filter_variable_eq=filter_variable_eq, diff --git a/src/policyengine/tax_benefit_models/uk/analysis.py b/src/policyengine/tax_benefit_models/uk/analysis.py index f7f5af5b..0a545b52 100644 --- a/src/policyengine/tax_benefit_models/uk/analysis.py +++ b/src/policyengine/tax_benefit_models/uk/analysis.py @@ -2,7 +2,7 @@ import tempfile from pathlib import Path -from typing import Any +from typing import Any, Optional import pandas as pd from microdf import MicroDataFrame @@ -65,7 +65,7 @@ class UKHouseholdInput(BaseModel): def calculate_household_impact( household_input: UKHouseholdInput, - policy: Policy | None = None, + policy: Optional[Policy] = None, ) -> UKHouseholdOutput: """Calculate tax and benefit impacts for a single UK household.""" n_people = len(household_input.people) diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index ec0f579b..47f78403 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -1,4 +1,5 @@ from pathlib import Path +from typing import Optional import pandas as pd from microdf import MicroDataFrame @@ -33,7 +34,7 @@ def entity_data(self) -> dict[str, MicroDataFrame]: class PolicyEngineUKDataset(Dataset): """UK dataset with multi-year entity-level data.""" - data: UKYearData | None = None + data: Optional[UKYearData] = None def model_post_init(self, __context): """Called after Pydantic initialization.""" diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index ff65be1b..edd5c069 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -1,7 +1,7 @@ import datetime from importlib import metadata from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pandas as pd from microdf import MicroDataFrame @@ -45,7 +45,7 @@ class PolicyEngineUK(TaxBenefitModel): uk_model = PolicyEngineUK() -def _get_runtime_data_build_metadata() -> dict[str, str | None]: +def _get_runtime_data_build_metadata() -> dict[str, Optional[str]]: try: from policyengine_uk.build_metadata import get_data_build_metadata except ModuleNotFoundError as exc: @@ -430,8 +430,8 @@ def load(self, simulation: "Simulation"): def _managed_release_bundle( dataset_uri: str, - dataset_source: str | None = None, -) -> dict[str, str | None]: + dataset_source: Optional[str] = None, +) -> dict[str, Optional[str]]: bundle = dict(uk_latest.release_bundle) bundle["runtime_dataset"] = dataset_logical_name(dataset_uri) bundle["runtime_dataset_uri"] = dataset_uri @@ -443,7 +443,7 @@ def _managed_release_bundle( def managed_microsimulation( *, - dataset: str | None = None, + dataset: Optional[str] = None, allow_unmanaged: bool = False, **kwargs, ): diff --git a/src/policyengine/tax_benefit_models/uk/outputs.py b/src/policyengine/tax_benefit_models/uk/outputs.py index 273a27c6..97032a9c 100644 --- a/src/policyengine/tax_benefit_models/uk/outputs.py +++ b/src/policyengine/tax_benefit_models/uk/outputs.py @@ -1,5 +1,7 @@ """UK-specific output templates.""" +from typing import Optional + from pydantic import ConfigDict from policyengine.core import Output, Simulation @@ -22,13 +24,13 @@ class ProgrammeStatistics(Output): is_tax: bool = False # Results populated by run() - baseline_total: float | None = None - reform_total: float | None = None - change: float | None = None - baseline_count: float | None = None - reform_count: float | None = None - winners: float | None = None - losers: float | None = None + baseline_total: Optional[float] = None + reform_total: Optional[float] = None + change: Optional[float] = None + baseline_count: Optional[float] = None + reform_count: Optional[float] = None + winners: Optional[float] = None + losers: Optional[float] = None def run(self): """Calculate statistics for this programme.""" diff --git a/src/policyengine/tax_benefit_models/us/analysis.py b/src/policyengine/tax_benefit_models/us/analysis.py index 375a4e5f..122ae2af 100644 --- a/src/policyengine/tax_benefit_models/us/analysis.py +++ b/src/policyengine/tax_benefit_models/us/analysis.py @@ -2,7 +2,7 @@ import tempfile from pathlib import Path -from typing import Any +from typing import Any, Optional, Union import pandas as pd from microdf import MicroDataFrame @@ -54,7 +54,7 @@ class USHouseholdInput(BaseModel): def calculate_household_impact( household_input: USHouseholdInput, - policy: Policy | None = None, + policy: Optional[Policy] = None, ) -> USHouseholdOutput: """Calculate tax and benefit impacts for a single US household.""" n_people = len(household_input.people) @@ -201,7 +201,7 @@ class PolicyReformAnalysis(BaseModel): def economic_impact_analysis( baseline_simulation: Simulation, reform_simulation: Simulation, - inequality_preset: USInequalityPreset | str = USInequalityPreset.STANDARD, + inequality_preset: Union[USInequalityPreset, str] = USInequalityPreset.STANDARD, ) -> PolicyReformAnalysis: """Perform comprehensive analysis of a policy reform. diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 7ea12f8e..da10733b 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -1,5 +1,6 @@ import warnings from pathlib import Path +from typing import Optional import pandas as pd from microdf import MicroDataFrame @@ -40,7 +41,7 @@ def entity_data(self) -> dict[str, MicroDataFrame]: class PolicyEngineUSDataset(Dataset): """US dataset with multi-year entity-level data.""" - data: USYearData | None = None + data: Optional[USYearData] = None def model_post_init(self, __context) -> None: """Called after Pydantic initialization.""" diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 2c560e3a..a896f5c4 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -1,7 +1,7 @@ import datetime from importlib import metadata from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pandas as pd from microdf import MicroDataFrame @@ -51,7 +51,7 @@ class PolicyEngineUS(TaxBenefitModel): us_model = PolicyEngineUS() -def _get_runtime_data_build_metadata() -> dict[str, str | None]: +def _get_runtime_data_build_metadata() -> dict[str, Optional[str]]: try: from policyengine_us.build_metadata import get_data_build_metadata except ModuleNotFoundError as exc: @@ -595,8 +595,8 @@ def _build_simulation_from_dataset(self, microsim, dataset, system): def _managed_release_bundle( dataset_uri: str, - dataset_source: str | None = None, -) -> dict[str, str | None]: + dataset_source: Optional[str] = None, +) -> dict[str, Optional[str]]: bundle = dict(us_latest.release_bundle) bundle["runtime_dataset"] = dataset_logical_name(dataset_uri) bundle["runtime_dataset_uri"] = dataset_uri @@ -608,7 +608,7 @@ def _managed_release_bundle( def managed_microsimulation( *, - dataset: str | None = None, + dataset: Optional[str] = None, allow_unmanaged: bool = False, **kwargs, ): diff --git a/src/policyengine/tax_benefit_models/us/outputs.py b/src/policyengine/tax_benefit_models/us/outputs.py index 63fd1a36..1dd6f001 100644 --- a/src/policyengine/tax_benefit_models/us/outputs.py +++ b/src/policyengine/tax_benefit_models/us/outputs.py @@ -1,5 +1,7 @@ """US-specific output templates.""" +from typing import Optional + from pydantic import ConfigDict from policyengine.core import Output, Simulation @@ -22,13 +24,13 @@ class ProgramStatistics(Output): is_tax: bool = False # Results populated by run() - baseline_total: float | None = None - reform_total: float | None = None - change: float | None = None - baseline_count: float | None = None - reform_count: float | None = None - winners: float | None = None - losers: float | None = None + baseline_total: Optional[float] = None + reform_total: Optional[float] = None + change: Optional[float] = None + baseline_count: Optional[float] = None + reform_count: Optional[float] = None + winners: Optional[float] = None + losers: Optional[float] = None def run(self): """Calculate statistics for this program.""" diff --git a/src/policyengine/utils/parametric_reforms.py b/src/policyengine/utils/parametric_reforms.py index 71476afa..025df22e 100644 --- a/src/policyengine/utils/parametric_reforms.py +++ b/src/policyengine/utils/parametric_reforms.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional, Union from policyengine_core.periods import period @@ -13,8 +13,8 @@ def reform_dict_from_parameter_values( - parameter_values: list[ParameterValue] | None, -) -> dict | None: + parameter_values: Optional[list[ParameterValue]], +) -> Optional[dict]: """ Convert a list of ParameterValue objects to a reform dict format. @@ -83,8 +83,8 @@ def modifier(simulation): def build_reform_dict( - policy_or_dynamic: Policy | Dynamic | None, -) -> dict | None: + policy_or_dynamic: Optional[Union[Policy, Dynamic]], +) -> Optional[dict]: """Extract a reform dict from a Policy or Dynamic object. If the object has parameter_values, converts them to reform dict format. @@ -103,7 +103,9 @@ def build_reform_dict( return None -def merge_reform_dicts(base: dict | None, override: dict | None) -> dict | None: +def merge_reform_dicts( + base: Optional[dict], override: Optional[dict] +) -> Optional[dict]: """Merge two reform dicts, with override values taking precedence. Either or both dicts can be None. When both have entries for the same diff --git a/src/policyengine/utils/plotting.py b/src/policyengine/utils/plotting.py index 4478a02d..2ca8e48c 100644 --- a/src/policyengine/utils/plotting.py +++ b/src/policyengine/utils/plotting.py @@ -1,5 +1,7 @@ """Plotting utilities for PolicyEngine visualisations.""" +from typing import Optional + import plotly.graph_objects as go # PolicyEngine brand colours @@ -26,12 +28,12 @@ def format_fig( fig: go.Figure, - title: str | None = None, - xaxis_title: str | None = None, - yaxis_title: str | None = None, + title: Optional[str] = None, + xaxis_title: Optional[str] = None, + yaxis_title: Optional[str] = None, show_legend: bool = True, - height: int | None = None, - width: int | None = None, + height: Optional[int] = None, + width: Optional[int] = None, ) -> go.Figure: """Apply PolicyEngine visual style to a plotly figure. diff --git a/tests/fixtures/parameter_labels_fixtures.py b/tests/fixtures/parameter_labels_fixtures.py index 0e22424c..d86c900f 100644 --- a/tests/fixtures/parameter_labels_fixtures.py +++ b/tests/fixtures/parameter_labels_fixtures.py @@ -1,7 +1,7 @@ """Fixtures for parameter_labels utility tests.""" from enum import Enum -from typing import Any +from typing import Any, Optional from unittest.mock import MagicMock @@ -24,7 +24,7 @@ class MockStateCode(Enum): def create_mock_parameter( name: str, - label: str | None = None, + label: Optional[str] = None, parent: Any = None, ) -> MagicMock: """Create a mock CoreParameter object.""" @@ -37,9 +37,9 @@ def create_mock_parameter( def create_mock_parent_node( name: str, - label: str | None = None, - breakdown: list[str] | None = None, - breakdown_labels: list[str] | None = None, + label: Optional[str] = None, + breakdown: Optional[list[str]] = None, + breakdown_labels: Optional[list[str]] = None, parent: Any = None, ) -> MagicMock: """Create a mock parent ParameterNode with optional breakdown metadata.""" @@ -58,8 +58,8 @@ def create_mock_parent_node( def create_mock_scale( name: str, - label: str | None = None, - scale_type: str | None = None, + label: Optional[str] = None, + scale_type: Optional[str] = None, ) -> MagicMock: """Create a mock ParameterScale object.""" scale = MagicMock() @@ -74,7 +74,7 @@ def create_mock_scale( def create_mock_variable( name: str, - possible_values: type[Enum] | None = None, + possible_values: Optional[type[Enum]] = None, ) -> MagicMock: """Create a mock Variable object with optional enum values.""" var = MagicMock() @@ -87,8 +87,8 @@ def create_mock_variable( def create_mock_system( - variables: dict[str, MagicMock] | None = None, - scales: list[MagicMock] | None = None, + variables: Optional[dict[str, MagicMock]] = None, + scales: Optional[list[MagicMock]] = None, ) -> MagicMock: """Create a mock tax-benefit system.""" system = MagicMock() diff --git a/tests/fixtures/parametric_reforms_fixtures.py b/tests/fixtures/parametric_reforms_fixtures.py index 98bc7aa2..6fcbd991 100644 --- a/tests/fixtures/parametric_reforms_fixtures.py +++ b/tests/fixtures/parametric_reforms_fixtures.py @@ -1,6 +1,7 @@ """Fixtures for parametric reforms tests.""" from datetime import date +from typing import Optional from unittest.mock import MagicMock import pytest @@ -23,7 +24,7 @@ def create_parameter_value( parameter: Parameter, value: float, start_date: date, - end_date: date | None = None, + end_date: Optional[date] = None, ) -> ParameterValue: """Create a ParameterValue for testing.""" return ParameterValue( diff --git a/tests/fixtures/variable_label_fixtures.py b/tests/fixtures/variable_label_fixtures.py index 1ce3572e..a4c01177 100644 --- a/tests/fixtures/variable_label_fixtures.py +++ b/tests/fixtures/variable_label_fixtures.py @@ -1,13 +1,14 @@ """Fixtures for variable label tests.""" +from typing import Optional from unittest.mock import MagicMock def create_mock_openfisca_variable( name: str, - label: str | None = None, + label: Optional[str] = None, entity_key: str = "person", - documentation: str | None = None, + documentation: Optional[str] = None, value_type: type = float, default_value=0, ) -> MagicMock: diff --git a/tests/test_intra_decile_impact.py b/tests/test_intra_decile_impact.py index dc4e6a96..04ae5412 100644 --- a/tests/test_intra_decile_impact.py +++ b/tests/test_intra_decile_impact.py @@ -1,5 +1,6 @@ """Unit tests for IntraDecileImpact and DecileImpact with decile_variable.""" +from typing import Optional from unittest.mock import MagicMock import numpy as np @@ -40,7 +41,7 @@ def _make_version(variable_name: str, entity: str) -> TaxBenefitModelVersion: return version -def _make_sim(household_data: dict, variables: list | None = None) -> MagicMock: +def _make_sim(household_data: dict, variables: Optional[list] = None) -> MagicMock: """Create a mock Simulation with household-level data.""" hh_df = MicroDataFrame( pd.DataFrame(household_data),