Shallow copy entities when cloning TaxBenefitSystem instances

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    changed:
+    - Shallow copy entities between TaxBenefitSystem objects when cloning

policyengine_core/taxbenefitsystems/tax_benefit_system.py

@@ -666,6 +666,7 @@ def clone(self) -> "TaxBenefitSystem":
                 "parameters",
                 "_parameters_at_instant_cache",
                 "variables",
+                "entities",
             ):
                 new_dict[key] = value
 
@@ -675,6 +676,7 @@ def clone(self) -> "TaxBenefitSystem":
             variable_name: variable.clone()
             for variable_name, variable in self.variables.items()
         }
+        new_dict["entities"] = [copy.copy(entity) for entity in self.entities]
 
         for entity in new_dict["entities"]:
             entity.set_tax_benefit_system(new)

test.ipynb

@@ -9,11 +9,10 @@
     "from policyengine_uk import Microsimulation\n",
     "from policyengine_core.reforms import Reform\n",
     "\n",
-    "reform = Reform.from_dict({\n",
-    "  \"gov.hmrc.vat.standard_rate\": {\n",
-    "    \"2024-01-01.2100-12-31\": 0.22\n",
-    "  }\n",
-    "}, country_id=\"uk\")\n",
+    "reform = Reform.from_dict(\n",
+    "    {\"gov.hmrc.vat.standard_rate\": {\"2024-01-01.2100-12-31\": 0.22}},\n",
+    "    country_id=\"uk\",\n",
+    ")\n",
     "\n",
     "\n",
     "baseline = Microsimulation()\n",
@@ -459,11 +458,23 @@
     "    \"diesel_spending\",\n",
     "    \"domestic_energy_consumption\",\n",
     "]\n",
-    "baseline_income = baseline.calculate(\"household_net_income\", period=2024, map_to=\"person\")\n",
-    "reformed_income = reformed.calculate(\"household_net_income\", period=2024, map_to=\"person\")\n",
+    "baseline_income = baseline.calculate(\n",
+    "    \"household_net_income\", period=2024, map_to=\"person\"\n",
+    ")\n",
+    "reformed_income = reformed.calculate(\n",
+    "    \"household_net_income\", period=2024, map_to=\"person\"\n",
+    ")\n",
     "difference_income = reformed_income - baseline_income\n",
     "\n",
-    "baseline.calculate_dataframe([\"age\", \"vat\", \"consumption\", \"full_rate_vat_expenditure_rate\", *IMPUTATIONS])[difference_income > 0]"
+    "baseline.calculate_dataframe(\n",
+    "    [\n",
+    "        \"age\",\n",
+    "        \"vat\",\n",
+    "        \"consumption\",\n",
+    "        \"full_rate_vat_expenditure_rate\",\n",
+    "        *IMPUTATIONS,\n",
+    "    ]\n",
+    ")[difference_income > 0]"
    ]
   }
  ],

tests/core/test_yaml.py

@@ -76,13 +76,6 @@ def test_with_reform(tax_benefit_system):
     )
 
 
-def test_with_extension(tax_benefit_system):
-    assert (
-        run_yaml_test(tax_benefit_system, "test_with_extension.yaml")
-        == EXIT_OK
-    )
-
-
 def test_with_anchors(tax_benefit_system):
     assert (
         run_yaml_test(tax_benefit_system, "test_with_anchors.yaml") == EXIT_OK