policyengine 3.1.6__py3-none-any.whl → 3.1.8__py3-none-any.whl

policyengine/core/parameter.py

@@ -1,3 +1,4 @@
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
 from pydantic import BaseModel, Field
@@ -5,6 +6,9 @@ from pydantic import BaseModel, Field
 from .parameter_value import ParameterValue
 from .tax_benefit_model_version import TaxBenefitModelVersion
 
+if TYPE_CHECKING:
+    from .parameter_value import ParameterValue
+
 
 class Parameter(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4()))

policyengine/core/parameter_value.py

@@ -1,14 +1,16 @@
 from datetime import datetime
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
 from pydantic import BaseModel, Field
 
-from .parameter import Parameter
+if TYPE_CHECKING:
+    from .parameter import Parameter
 
 
 class ParameterValue(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
-    parameter: Parameter | None = None
+    parameter: "Parameter | None" = None
     value: float | int | str | bool | list | None = None
     start_date: datetime
     end_date: datetime | None = None

policyengine/tax_benefit_models/uk/datasets.py

@@ -40,14 +40,37 @@ class PolicyEngineUKDataset(Dataset):
             self.load()
 
     def save(self) -> None:
-        """Save dataset to HDF5 file."""
+        """Save dataset to HDF5 file.
+
+        Converts object columns to categorical dtype to avoid slow pickle serialization.
+        """
         filepath = Path(self.filepath)
         if not filepath.parent.exists():
             filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        # Convert DataFrames and optimize object columns to categorical
+        person_df = pd.DataFrame(self.data.person)
+        benunit_df = pd.DataFrame(self.data.benunit)
+        household_df = pd.DataFrame(self.data.household)
+
+        # Convert object columns to categorical to avoid pickle serialization
+        for col in person_df.columns:
+            if person_df[col].dtype == "object":
+                person_df[col] = person_df[col].astype("category")
+
+        for col in benunit_df.columns:
+            if benunit_df[col].dtype == "object":
+                benunit_df[col] = benunit_df[col].astype("category")
+
+        for col in household_df.columns:
+            if household_df[col].dtype == "object":
+                household_df[col] = household_df[col].astype("category")
+
         with pd.HDFStore(filepath, mode="w") as store:
-            store["person"] = pd.DataFrame(self.data.person)
-            store["benunit"] = pd.DataFrame(self.data.benunit)
-            store["household"] = pd.DataFrame(self.data.household)
+            # Use format='table' to support categorical dtypes
+            store.put("person", person_df, format="table")
+            store.put("benunit", benunit_df, format="table")
+            store.put("household", household_df, format="table")
 
     def load(self) -> None:
         """Load dataset from HDF5 file into this instance."""

{policyengine-3.1.6.dist-info → policyengine-3.1.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: policyengine
-Version: 3.1.6
+Version: 3.1.8
 Summary: A package to conduct policy analysis using PolicyEngine tax-benefit models.
 Author-email: PolicyEngine <hello@policyengine.org>
 License:                     GNU AFFERO GENERAL PUBLIC LICENSE

{policyengine-3.1.6.dist-info → policyengine-3.1.8.dist-info}/RECORD

@@ -1,12 +1,12 @@
 policyengine/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-policyengine/__pycache__/__init__.cpython-313.pyc,sha256=uYQXIlDDDMSRk17xgnGViCXX-MefhsWgf4kL0hHrVaQ,175
+policyengine/__pycache__/__init__.cpython-313.pyc,sha256=LQBTMmPPnJo-w-7_mvuE542bVf07bu1ONsDcLBRHLH8,175
 policyengine/core/__init__.py,sha256=KBVhkqzkvjWLDDwk96vquQKL63ZFuLen5AzBOBnO9pg,912
 policyengine/core/dataset.py,sha256=iJr9-J6w11uMRYy3EEJO9Gveku1m71AA1yzeo-0SiCs,16094
 policyengine/core/dataset_version.py,sha256=6KeFCRGQto_Yyl4QY4Vo2JFythjaXrNAOHQiwRGESyM,378
 policyengine/core/dynamic.py,sha256=ng9BjDzxdwjJ0e7zoqXFmq33E1SRbaaPYfW7pjRSSzI,1641
 policyengine/core/output.py,sha256=cCW4vbzkLdQaT_nJTyDJBl7Hubm7nZeRuR7aVG1dKvg,643
-policyengine/core/parameter.py,sha256=wePNpoeSEwI2NLE1gENKvS4pqWd-GpN8f8-iZzxNHlo,526
-policyengine/core/parameter_value.py,sha256=b0ts1kbWcwjPSYnZm2rlCylmTLPJRLxDL8z3RmxM5OI,377
+policyengine/core/parameter.py,sha256=8RKKuGCDW1OoHPoXI9vF3JY6COc1qhUtMolXTUxPoEs,626
+policyengine/core/parameter_value.py,sha256=ZRBZWFYtaY9TqdgjrCymzOZNmuKOBZsrWBET24DIJ_Q,434
 policyengine/core/policy.py,sha256=ExMrUDMvNk_uuOL0cSm0UCzDyGka0t_yk6x4U0Kp6Ww,1635
 policyengine/core/simulation.py,sha256=yvvved75XMcGP3Bj9E2tmKRxvI-DQVZv7k4uTETwBm0,1134
 policyengine/core/tax_benefit_model.py,sha256=2Yc1RlQrUG7djDMZbJOQH4Ns86_lOnLeISCGR4-9zMo,176
@@ -20,7 +20,7 @@ policyengine/tax_benefit_models/uk.py,sha256=HzAG_dORmsj1NJ9pd9WrqwgZPe9DUDrZ1wV
 policyengine/tax_benefit_models/us.py,sha256=G51dAmHo8NJLb2mnbne6iO5eNaatCGUd_2unvawwF84,946
 policyengine/tax_benefit_models/uk/__init__.py,sha256=AiA74iED5FEryvUCMfVZi6pYDYuTfQcj9B01h8J5xFA,1105
 policyengine/tax_benefit_models/uk/analysis.py,sha256=O4eYJYF7tsgiuLuiWMU0OXq7ss6U8-vzlg6nC2U8sgU,3175
-policyengine/tax_benefit_models/uk/datasets.py,sha256=-lmj4eG2my2GGmMMkxI1iXobGQW5irBgylEwyV0xU6c,8039
+policyengine/tax_benefit_models/uk/datasets.py,sha256=_puLrvNo700HFv17TLGXufXRFcrLS9-pBVjYwp1VFac,8989
 policyengine/tax_benefit_models/uk/model.py,sha256=djCLLBeik9ZGFqe5WXUNVcu1JXjh0U8MXakvncasUj0,9754
 policyengine/tax_benefit_models/uk/outputs.py,sha256=2mYLwQW4QNvrOHtHfm_ACqE9gbmuLxvcCyldRU46s0o,3543
 policyengine/tax_benefit_models/us/__init__.py,sha256=zP-UUQqOc9g0ymyHkweJdi4RVXQDKSR6SUxavUKvV0s,1101
@@ -32,8 +32,8 @@ policyengine/utils/__init__.py,sha256=1X-VYAWLyB9A0YRHwsGWrqQHns1WfeZ7ISC6DMU5my
 policyengine/utils/dates.py,sha256=HnAqyl8S8EOYp8ibsnMTmECYoDWCSqwL-7A2_qKgxSc,1510
 policyengine/utils/parametric_reforms.py,sha256=4P3U39-4pYTU4BN6JjgmVLUkCkBhRfZJ6UIWTlsjyQE,1155
 policyengine/utils/plotting.py,sha256=ZAzTWz38vIaW0c3Nt4Un1kfrNoXLyHCDd1pEJIlsRg4,5335
-policyengine-3.1.6.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-policyengine-3.1.6.dist-info/METADATA,sha256=Sdl_MCwOgD_eSUaXhrD7J2BmUK0Pgt2pSKXdEVuDak4,45889
-policyengine-3.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-policyengine-3.1.6.dist-info/top_level.txt,sha256=_23UPobfkneHQkpJ0e0OmDJfhCUfoXj_F2sTckCGOH4,13
-policyengine-3.1.6.dist-info/RECORD,,
+policyengine-3.1.8.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+policyengine-3.1.8.dist-info/METADATA,sha256=h0r-m8cflbRqVtwajLnbgxmrKmowdeWUpXHVwLRLIGw,45889
+policyengine-3.1.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+policyengine-3.1.8.dist-info/top_level.txt,sha256=_23UPobfkneHQkpJ0e0OmDJfhCUfoXj_F2sTckCGOH4,13
+policyengine-3.1.8.dist-info/RECORD,,