From fd72577c731b6346a976122412c7d829e4163766 Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 17 Nov 2025 11:36:00 +0000 Subject: [PATCH] Add addition operator support for Policy and Dynamic classes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Implement __add__ method for Policy and Dynamic classes - Parameter values are appended when combining instances - Simulation modifiers are chained in sequence (first then second) - Handles edge cases where one or both modifiers are None - Enables intuitive composition: policy1 + policy2 + policy3 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- changelog_entry.yaml | 15 +++++++++++++-- src/policyengine/core/dynamic.py | 26 ++++++++++++++++++++++++++ src/policyengine/core/policy.py | 26 ++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 4c132743..aef6eaf9 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -1,4 +1,15 @@ - bump: minor changes: - - Just basemodels, no sqlmodels. - - Clean, working analysis at both household and macro level for uk and us. + added: + - Policy and Dynamic classes now support addition operator (__add__) to combine policies and dynamics. + - Parameter values are appended when combining policies/dynamics. + - Simulation modifiers are chained in sequence when combining policies/dynamics. + changed: + - Cleaned up unused dependencies (removed sqlalchemy, sqlmodel, alembic, psycopg2, pymysql, google-cloud-storage, getpass4, rich, ipywidgets, tqdm, blosc). + - Added missing dependencies (plotly, requests). + - Removed unused pe-migrate script entry point. + fixed: + - Fixed all ruff linting errors (148 errors resolved). + - Replaced star imports with explicit imports for better code clarity. + - Added explicit re-exports in __init__.py files. + - Silenced PyTables PerformanceWarning in US datasets. diff --git a/src/policyengine/core/dynamic.py b/src/policyengine/core/dynamic.py index 9b312952..3b6ba553 100644 --- a/src/policyengine/core/dynamic.py +++ b/src/policyengine/core/dynamic.py @@ -15,3 +15,29 @@ class Dynamic(BaseModel): simulation_modifier: Callable | None = None created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) + + def __add__(self, other: "Dynamic") -> "Dynamic": + """Combine two dynamics by appending parameter values and chaining simulation modifiers.""" + if not isinstance(other, Dynamic): + return NotImplemented + + # Combine simulation modifiers + combined_modifier = None + if self.simulation_modifier is not None and other.simulation_modifier is not None: + + def combined_modifier(sim): + sim = self.simulation_modifier(sim) + sim = other.simulation_modifier(sim) + return sim + + elif self.simulation_modifier is not None: + combined_modifier = self.simulation_modifier + elif other.simulation_modifier is not None: + combined_modifier = other.simulation_modifier + + return Dynamic( + name=f"{self.name} + {other.name}", + description=f"Combined dynamic: {self.name} and {other.name}", + parameter_values=self.parameter_values + other.parameter_values, + simulation_modifier=combined_modifier, + ) diff --git a/src/policyengine/core/policy.py b/src/policyengine/core/policy.py index 20587d85..3aeb19b9 100644 --- a/src/policyengine/core/policy.py +++ b/src/policyengine/core/policy.py @@ -15,3 +15,29 @@ class Policy(BaseModel): simulation_modifier: Callable | None = None created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) + + def __add__(self, other: "Policy") -> "Policy": + """Combine two policies by appending parameter values and chaining simulation modifiers.""" + if not isinstance(other, Policy): + return NotImplemented + + # Combine simulation modifiers + combined_modifier = None + if self.simulation_modifier is not None and other.simulation_modifier is not None: + + def combined_modifier(sim): + sim = self.simulation_modifier(sim) + sim = other.simulation_modifier(sim) + return sim + + elif self.simulation_modifier is not None: + combined_modifier = self.simulation_modifier + elif other.simulation_modifier is not None: + combined_modifier = other.simulation_modifier + + return Policy( + name=f"{self.name} + {other.name}", + description=f"Combined policy: {self.name} and {other.name}", + parameter_values=self.parameter_values + other.parameter_values, + simulation_modifier=combined_modifier, + )