From 19392206f69f1ed4aced8c3789126c9773a4cec5 Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Tue, 15 Oct 2024 03:17:21 -0700 Subject: [PATCH 01/10] feat: add `write_yaml` method in `parameter_node.py` to output the node data as a YAML file WIP fixes: #274 --- .../parameters/parameter_node.py | 22 +++++++++++++++---- tests/core/test_parameters.py | 10 +++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 5f05e33c2..7ef1b7bd2 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -1,16 +1,16 @@ import copy import os import typing -from typing import Iterable, List, Type, Union +from pathlib import Path +from typing import Iterable, Union + +import yaml from policyengine_core import commons, parameters, tools from policyengine_core.data_structures import Reference from policyengine_core.periods.instant_ import Instant from policyengine_core.tracers import TracingParameterNodeAtInstant - from .at_instant_like import AtInstantLike -from .parameter import Parameter -from .parameter_node_at_instant import ParameterNodeAtInstant from .config import COMMON_KEYS, FILE_EXTENSIONS from .helpers import ( load_parameter_file, @@ -19,6 +19,8 @@ _parse_child, _load_yaml_file, ) +from .parameter import Parameter +from .parameter_node_at_instant import ParameterNodeAtInstant EXCLUDED_PARAMETER_CHILD_NAMES = ["reference", "__pycache__"] @@ -274,3 +276,15 @@ def get_child(self, path: str) -> "ParameterNode": f"Could not find the parameter (failed at {name})." ) return node + + def write_yaml(self, file_path: Path) -> yaml: + data = {"description": self.description} + for key, value in self.children.items(): + name = key + value_at_instant = value.values_list[0] + data[name] = {"values": {value_at_instant.instant_str: value_at_instant.value}} + try: + with open(file_path, "w") as f: + yaml.dump(data, f) + except Exception as e: + print(f"Error when writing YAML file: {e}") diff --git a/tests/core/test_parameters.py b/tests/core/test_parameters.py index 2cb71e2d9..9afd7ca64 100644 --- a/tests/core/test_parameters.py +++ b/tests/core/test_parameters.py @@ -1,4 +1,5 @@ import tempfile +from pathlib import Path import pytest @@ -141,3 +142,12 @@ def test_name(): } parameter = ParameterNode("root", data=parameter_data) assert parameter.children["2010"].name == "root.2010" + + +def test_write_yaml(): + parameter_data = { + "description": "Parameter indexed by a numeric key", + "2010": {"values": {"2006-01-01": 0.0075}}, + } + parameter = ParameterNode("root", data=parameter_data) + parameter.write_yaml(Path("output.yaml")) From d9a37d5ee7f26abd6b38831ed977f5f041e09ed4 Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Wed, 16 Oct 2024 12:26:08 -0700 Subject: [PATCH 02/10] fix: modified `write_yaml` method in `parameter_node.py` to output the node data as a YAML file WIP fixes: #274 --- changelog_entry.yaml | 5 +++++ policyengine_core/parameters/parameter_node.py | 10 ++++++---- tests/core/test_parameters.py | 16 +++++++++++++--- 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29bb..23a8388f7 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,5 @@ +- bump: minor + changes: + added: + - `write_yaml` function to output ParameterNode data to a YAML file + - `test_write_yaml` test to produce a sample output diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 7ef1b7bd2..c262cda90 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -278,13 +278,15 @@ def get_child(self, path: str) -> "ParameterNode": return node def write_yaml(self, file_path: Path) -> yaml: - data = {"description": self.description} + data = {"name": self.name, "description": self.description} for key, value in self.children.items(): name = key - value_at_instant = value.values_list[0] - data[name] = {"values": {value_at_instant.instant_str: value_at_instant.value}} + value_dict = {} + data[name] = {"values": value_dict} + for value_at_instant in value.values_list: + value_dict[value_at_instant.instant_str] = value_at_instant.value try: with open(file_path, "w") as f: - yaml.dump(data, f) + yaml.dump(data, f, sort_keys=False) except Exception as e: print(f"Error when writing YAML file: {e}") diff --git a/tests/core/test_parameters.py b/tests/core/test_parameters.py index 9afd7ca64..b14945612 100644 --- a/tests/core/test_parameters.py +++ b/tests/core/test_parameters.py @@ -145,9 +145,19 @@ def test_name(): def test_write_yaml(): - parameter_data = { - "description": "Parameter indexed by a numeric key", - "2010": {"values": {"2006-01-01": 0.0075}}, + parameter_data={ + 'amount': { + 'values': { + "2015-01-01": {'value': 550}, + "2016-01-01": {'value': 600} + } + }, + 'min_age': { + 'values': { + "2015-01-01": {'value': 25}, + "2016-01-01": {'value': 18} + } + }, } parameter = ParameterNode("root", data=parameter_data) parameter.write_yaml(Path("output.yaml")) From 56582af4809dd4fc481d894bbb4533b4976968af Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Wed, 16 Oct 2024 12:36:01 -0700 Subject: [PATCH 03/10] chore: fix format WIP fixes: #274 --- policyengine_core/parameters/parameter_node.py | 4 +++- tests/core/test_parameters.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index c262cda90..e8f6d034c 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -284,7 +284,9 @@ def write_yaml(self, file_path: Path) -> yaml: value_dict = {} data[name] = {"values": value_dict} for value_at_instant in value.values_list: - value_dict[value_at_instant.instant_str] = value_at_instant.value + value_dict[value_at_instant.instant_str] = ( + value_at_instant.value + ) try: with open(file_path, "w") as f: yaml.dump(data, f, sort_keys=False) diff --git a/tests/core/test_parameters.py b/tests/core/test_parameters.py index b14945612..0a7694c04 100644 --- a/tests/core/test_parameters.py +++ b/tests/core/test_parameters.py @@ -145,17 +145,17 @@ def test_name(): def test_write_yaml(): - parameter_data={ - 'amount': { - 'values': { - "2015-01-01": {'value': 550}, - "2016-01-01": {'value': 600} + parameter_data = { + "amount": { + "values": { + "2015-01-01": {"value": 550}, + "2016-01-01": {"value": 600}, } }, - 'min_age': { - 'values': { - "2015-01-01": {'value': 25}, - "2016-01-01": {'value': 18} + "min_age": { + "values": { + "2015-01-01": {"value": 25}, + "2016-01-01": {"value": 18}, } }, } From 5e8bc8b546e55fa8c28a22495e3886010f31ab38 Mon Sep 17 00:00:00 2001 From: Jingyun Du <61254755+SylviaDu99@users.noreply.github.com> Date: Wed, 16 Oct 2024 12:37:24 -0700 Subject: [PATCH 04/10] Update changelog_entry.yaml --- changelog_entry.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 23a8388f7..369eef8c0 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -1,5 +1,5 @@ - bump: minor changes: added: - - `write_yaml` function to output ParameterNode data to a YAML file - - `test_write_yaml` test to produce a sample output + - write_yaml function to output ParameterNode data to a YAML file + - test_write_yaml test to produce a sample output From a0a6bfa5d317e5685621932586356dc575563aed Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 21 Oct 2024 03:33:15 +0200 Subject: [PATCH 05/10] test: Add sample data --- tests/core/test_parameters.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/core/test_parameters.py b/tests/core/test_parameters.py index 0a7694c04..339284b8d 100644 --- a/tests/core/test_parameters.py +++ b/tests/core/test_parameters.py @@ -150,13 +150,29 @@ def test_write_yaml(): "values": { "2015-01-01": {"value": 550}, "2016-01-01": {"value": 600}, - } + }, + "branch_name": "default", + "description": "The amount of the basic income", + "documentation": None, + "file_path": "test_path/to/file/1", + # "metadata": Unclear yet what form this dict takes + "modified": False, + "trace": True, + "tracer": None, }, "min_age": { "values": { "2015-01-01": {"value": 25}, "2016-01-01": {"value": 18}, - } + }, + "branch_name": "labor_supply_1", + "description": "The minimum age to receive the basic income", + "documentation": None, + "file_path": "test_path/to/file/2", + # "metadata": Unclear yet what form this dict takes + "modified": True, + "trace": False, + "tracer": None, }, } parameter = ParameterNode("root", data=parameter_data) From 4fb27f8ce8a57c3821ff874267746221b6531b29 Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Mon, 28 Oct 2024 14:21:39 -0700 Subject: [PATCH 06/10] feat: add function in AtInstantLike to return desired attributes in dictionary form WIP fixes: #274 --- .../parameters/at_instant_like.py | 29 +++++++++++++++++++ .../parameters/parameter_node.py | 12 ++------ tests/core/test_parameters.py | 10 ------- 3 files changed, 31 insertions(+), 20 deletions(-) diff --git a/policyengine_core/parameters/at_instant_like.py b/policyengine_core/parameters/at_instant_like.py index ecde86143..5718af48f 100644 --- a/policyengine_core/parameters/at_instant_like.py +++ b/policyengine_core/parameters/at_instant_like.py @@ -19,3 +19,32 @@ def get_at_instant(self, instant: Instant) -> Any: @abc.abstractmethod def _get_at_instant(self, instant): ... + + def get_attr_dict(self) -> dict: + attr_dict = {} + attr_list = [ + "name", + "description", + "documentation", + "file_path", + "metadata", + "trace", + "tracer", + "branch_name", + "modified", + "values_list", + ] + for attr in attr_list: + if hasattr(self, attr): + attr_dict[attr] = getattr(self, attr) + if hasattr(self, "children"): + for child_name, child in self.children.items(): + attr_dict[child_name] = child.get_attr_dict() + if hasattr(self, "values_list"): + value_dict = {} + attr_dict["values_list"] = value_dict + for value_at_instant in self.values_list: + value_dict[value_at_instant.instant_str] = ( + value_at_instant.value + ) + return attr_dict diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index e8f6d034c..6bb45717f 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -278,17 +278,9 @@ def get_child(self, path: str) -> "ParameterNode": return node def write_yaml(self, file_path: Path) -> yaml: - data = {"name": self.name, "description": self.description} - for key, value in self.children.items(): - name = key - value_dict = {} - data[name] = {"values": value_dict} - for value_at_instant in value.values_list: - value_dict[value_at_instant.instant_str] = ( - value_at_instant.value - ) + data = self.get_attr_dict() try: with open(file_path, "w") as f: - yaml.dump(data, f, sort_keys=False) + yaml.dump(data, f, sort_keys=True) except Exception as e: print(f"Error when writing YAML file: {e}") diff --git a/tests/core/test_parameters.py b/tests/core/test_parameters.py index 339284b8d..15e2d687d 100644 --- a/tests/core/test_parameters.py +++ b/tests/core/test_parameters.py @@ -151,28 +151,18 @@ def test_write_yaml(): "2015-01-01": {"value": 550}, "2016-01-01": {"value": 600}, }, - "branch_name": "default", "description": "The amount of the basic income", "documentation": None, - "file_path": "test_path/to/file/1", - # "metadata": Unclear yet what form this dict takes "modified": False, - "trace": True, - "tracer": None, }, "min_age": { "values": { "2015-01-01": {"value": 25}, "2016-01-01": {"value": 18}, }, - "branch_name": "labor_supply_1", "description": "The minimum age to receive the basic income", "documentation": None, - "file_path": "test_path/to/file/2", - # "metadata": Unclear yet what form this dict takes "modified": True, - "trace": False, - "tracer": None, }, } parameter = ParameterNode("root", data=parameter_data) From b4731ec2fea68a819c1aa1f416d1848670c949b0 Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Tue, 29 Oct 2024 21:20:52 -0700 Subject: [PATCH 07/10] feat: simplified function in AtInstantLike to return desired attributes in dictionary form; write_yaml in ParameterNode class now do most of the job WIP fixes: #274 --- .../parameters/at_instant_like.py | 32 +++---------------- .../parameters/parameter_node.py | 20 +++++++++++- 2 files changed, 24 insertions(+), 28 deletions(-) diff --git a/policyengine_core/parameters/at_instant_like.py b/policyengine_core/parameters/at_instant_like.py index 5718af48f..5f74de696 100644 --- a/policyengine_core/parameters/at_instant_like.py +++ b/policyengine_core/parameters/at_instant_like.py @@ -20,31 +20,9 @@ def get_at_instant(self, instant: Instant) -> Any: @abc.abstractmethod def _get_at_instant(self, instant): ... - def get_attr_dict(self) -> dict: - attr_dict = {} - attr_list = [ - "name", - "description", - "documentation", - "file_path", - "metadata", - "trace", - "tracer", - "branch_name", - "modified", - "values_list", - ] - for attr in attr_list: - if hasattr(self, attr): - attr_dict[attr] = getattr(self, attr) - if hasattr(self, "children"): - for child_name, child in self.children.items(): - attr_dict[child_name] = child.get_attr_dict() - if hasattr(self, "values_list"): - value_dict = {} - attr_dict["values_list"] = value_dict - for value_at_instant in self.values_list: - value_dict[value_at_instant.instant_str] = ( - value_at_instant.value - ) + def get_attr_dict(self, exclusion_list) -> dict: + attr_dict = self.__dict__ + for attr_name in exclusion_list: + if attr_name in attr_dict.keys(): + del attr_dict[attr_name] return attr_dict diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 6bb45717f..d645e9f3d 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -21,6 +21,7 @@ ) from .parameter import Parameter from .parameter_node_at_instant import ParameterNodeAtInstant +from .parameter_at_instant import ParameterAtInstant EXCLUDED_PARAMETER_CHILD_NAMES = ["reference", "__pycache__"] @@ -278,7 +279,24 @@ def get_child(self, path: str) -> "ParameterNode": return node def write_yaml(self, file_path: Path) -> yaml: - data = self.get_attr_dict() + exclusion_list = ["parent", "children", "_at_instant_cache"] + data = self.get_attr_dict(exclusion_list) + for attr_name in data.keys(): + attr_value = data.get(attr_name) + if type(attr_value) in [ + ParameterNode, + parameters.ParameterScale, + Parameter, + ]: + child_data = attr_value.get_attr_dict(exclusion_list) + data[attr_name] = child_data + if "values_list" in child_data.keys(): + value_dict = {} + for value_at_instant in child_data["values_list"]: + value_dict[value_at_instant.instant_str] = ( + value_at_instant.value + ) + child_data["values_list"] = value_dict try: with open(file_path, "w") as f: yaml.dump(data, f, sort_keys=True) From eaceba11041da2b4713cf80bbdfd9b74ad270ce0 Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Mon, 11 Nov 2024 17:28:53 -0800 Subject: [PATCH 08/10] feat: add get_attr_dict function in both super class and all subclasses WIP fixes: #274 --- .../parameters/at_instant_like.py | 9 ++-- policyengine_core/parameters/parameter.py | 28 ++++++++++--- .../parameters/parameter_node.py | 41 +++++++++---------- .../parameters/parameter_scale.py | 16 ++++++++ tests/core/test_parameters.py | 17 +++++++- 5 files changed, 78 insertions(+), 33 deletions(-) diff --git a/policyengine_core/parameters/at_instant_like.py b/policyengine_core/parameters/at_instant_like.py index 5f74de696..844fa2b66 100644 --- a/policyengine_core/parameters/at_instant_like.py +++ b/policyengine_core/parameters/at_instant_like.py @@ -20,9 +20,6 @@ def get_at_instant(self, instant: Instant) -> Any: @abc.abstractmethod def _get_at_instant(self, instant): ... - def get_attr_dict(self, exclusion_list) -> dict: - attr_dict = self.__dict__ - for attr_name in exclusion_list: - if attr_name in attr_dict.keys(): - del attr_dict[attr_name] - return attr_dict + @abc.abstractmethod + def get_attr_dict(self) -> dict: + raise NotImplementedError diff --git a/policyengine_core/parameters/parameter.py b/policyengine_core/parameters/parameter.py index fd792030d..09c032ae1 100644 --- a/policyengine_core/parameters/parameter.py +++ b/policyengine_core/parameters/parameter.py @@ -2,14 +2,15 @@ import os from typing import Dict, List, Optional -from policyengine_core.errors import ParameterParsingError -from .at_instant_like import AtInstantLike -from .parameter_at_instant import ParameterAtInstant +import numpy -from .helpers import _validate_parameter, _compose_name -from .config import COMMON_KEYS from policyengine_core.commons.misc import empty_clone +from policyengine_core.errors import ParameterParsingError from policyengine_core.periods import INSTANT_PATTERN, period as get_period +from .at_instant_like import AtInstantLike +from .config import COMMON_KEYS +from .helpers import _validate_parameter, _compose_name +from .parameter_at_instant import ParameterAtInstant class Parameter(AtInstantLike): @@ -45,6 +46,8 @@ class Parameter(AtInstantLike): """ + _exclusion_list = ["parent", "_at_instant_cache"] + def __init__( self, name: str, data: dict, file_path: Optional[str] = None ) -> None: @@ -233,3 +236,18 @@ def relative_change(self, start_instant, end_instant): if end_value is None or start_value is None: return None return end_value / start_value - 1 + + def get_attr_dict(self) -> dict: + data = self.__dict__.copy() + for attr in self._exclusion_list: + if attr in data.keys(): + del data[attr] + if "values_list" in data.keys(): + value_dict = {} + for value_at_instant in data["values_list"]: + value = value_at_instant.value + if type(value) is numpy.float64: + value = float(value) + value_dict[value_at_instant.instant_str] = value + data["values_list"] = value_dict + return data diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index d645e9f3d..917fd1d9d 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -7,7 +7,6 @@ import yaml from policyengine_core import commons, parameters, tools -from policyengine_core.data_structures import Reference from policyengine_core.periods.instant_ import Instant from policyengine_core.tracers import TracingParameterNodeAtInstant from .at_instant_like import AtInstantLike @@ -21,7 +20,6 @@ ) from .parameter import Parameter from .parameter_node_at_instant import ParameterNodeAtInstant -from .parameter_at_instant import ParameterAtInstant EXCLUDED_PARAMETER_CHILD_NAMES = ["reference", "__pycache__"] @@ -35,6 +33,8 @@ class ParameterNode(AtInstantLike): None # By default, no restriction on the keys ) + _exclusion_list = ["parent", "_at_instant_cache"] + parent: "ParameterNode" = None """The parent of the node, or None if the node is the root of the tree.""" @@ -278,27 +278,26 @@ def get_child(self, path: str) -> "ParameterNode": ) return node + def get_attr_dict(self) -> dict: + data = self.__dict__.copy() + for attr in self._exclusion_list: + if attr in data.keys(): + del data[attr] + if "children" in data.keys(): + child_dict = data.get("children") + for child_name, child in child_dict.items(): + data[child_name] = child.get_attr_dict() + del data["children"] + return data + + class NoAliasDumper(yaml.SafeDumper): + def ignore_aliases(self, data): + return True + def write_yaml(self, file_path: Path) -> yaml: - exclusion_list = ["parent", "children", "_at_instant_cache"] - data = self.get_attr_dict(exclusion_list) - for attr_name in data.keys(): - attr_value = data.get(attr_name) - if type(attr_value) in [ - ParameterNode, - parameters.ParameterScale, - Parameter, - ]: - child_data = attr_value.get_attr_dict(exclusion_list) - data[attr_name] = child_data - if "values_list" in child_data.keys(): - value_dict = {} - for value_at_instant in child_data["values_list"]: - value_dict[value_at_instant.instant_str] = ( - value_at_instant.value - ) - child_data["values_list"] = value_dict + data = self.get_attr_dict() try: with open(file_path, "w") as f: - yaml.dump(data, f, sort_keys=True) + yaml.dump(data, f, sort_keys=True, Dumper=self.NoAliasDumper) except Exception as e: print(f"Error when writing YAML file: {e}") diff --git a/policyengine_core/parameters/parameter_scale.py b/policyengine_core/parameters/parameter_scale.py index 39dcd2297..349afdeb3 100644 --- a/policyengine_core/parameters/parameter_scale.py +++ b/policyengine_core/parameters/parameter_scale.py @@ -24,6 +24,8 @@ class ParameterScale(AtInstantLike): # 'unit' and 'reference' are only listed here for backward compatibility _allowed_keys = config.COMMON_KEYS.union({"brackets"}) + _exclusion_list = ["parent", "_at_instant_cache"] + def __init__(self, name: str, data: dict, file_path: str): """ :param name: name of the scale, eg "taxes.some_scale" @@ -169,3 +171,17 @@ def _get_at_instant(self, instant: Instant) -> TaxScaleLike: threshold = bracket.threshold scale.add_bracket(threshold, rate * base) return scale + + def get_attr_dict(self) -> dict: + data = self.__dict__.copy() + for attr in self._exclusion_list: + if attr in data.keys(): + del data[attr] + if "brackets" in data.keys(): + node_list = data["brackets"] + i = 0 + for node in node_list: + node_list[i] = node.get_attr_dict() + i += 1 + data["brackets"] = node_list + return data diff --git a/tests/core/test_parameters.py b/tests/core/test_parameters.py index 15e2d687d..68369e099 100644 --- a/tests/core/test_parameters.py +++ b/tests/core/test_parameters.py @@ -1,6 +1,6 @@ import tempfile from pathlib import Path - +import os import pytest from policyengine_core.parameters import ( @@ -9,6 +9,7 @@ ParameterNotFoundError, load_parameter_file, ) +from policyengine_core.tools.test_runner import yaml def test_get_at_instant(tax_benefit_system): @@ -167,3 +168,17 @@ def test_write_yaml(): } parameter = ParameterNode("root", data=parameter_data) parameter.write_yaml(Path("output.yaml")) + + try: + with open("output.yaml", "r") as file: + data = yaml.safe_load(file) + os.remove("output.yaml") + except yaml.YAMLError as e: + pytest.fail(f"Output is not valid YAML: {e}") + + +# from policyengine_us import Microsimulation +# def test_yaml_us(): +# baseline = Microsimulation() +# tbs = baseline.tax_benefit_system +# tbs.parameters.gov.write_yaml("test_output.yaml") From 731c92120a4462af07635bab53706ed44bc937f3 Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Sat, 23 Nov 2024 15:28:56 -0800 Subject: [PATCH 09/10] minor: set sort_keys to false when output to yaml file to make it more readable; also added some comments WIP fixes: #274 --- policyengine_core/parameters/parameter.py | 1 + policyengine_core/parameters/parameter_node.py | 3 ++- policyengine_core/parameters/parameter_scale.py | 1 + 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/policyengine_core/parameters/parameter.py b/policyengine_core/parameters/parameter.py index 09c032ae1..38dbb090e 100644 --- a/policyengine_core/parameters/parameter.py +++ b/policyengine_core/parameters/parameter.py @@ -47,6 +47,7 @@ class Parameter(AtInstantLike): """ _exclusion_list = ["parent", "_at_instant_cache"] + """The keys to be excluded from the node when output to a yaml file.""" def __init__( self, name: str, data: dict, file_path: Optional[str] = None diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 917fd1d9d..4f44cc0ae 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -34,6 +34,7 @@ class ParameterNode(AtInstantLike): ) _exclusion_list = ["parent", "_at_instant_cache"] + """The keys to be excluded from the node when output to a yaml file.""" parent: "ParameterNode" = None """The parent of the node, or None if the node is the root of the tree.""" @@ -298,6 +299,6 @@ def write_yaml(self, file_path: Path) -> yaml: data = self.get_attr_dict() try: with open(file_path, "w") as f: - yaml.dump(data, f, sort_keys=True, Dumper=self.NoAliasDumper) + yaml.dump(data, f, sort_keys=False, Dumper=self.NoAliasDumper) except Exception as e: print(f"Error when writing YAML file: {e}") diff --git a/policyengine_core/parameters/parameter_scale.py b/policyengine_core/parameters/parameter_scale.py index 349afdeb3..851462de0 100644 --- a/policyengine_core/parameters/parameter_scale.py +++ b/policyengine_core/parameters/parameter_scale.py @@ -25,6 +25,7 @@ class ParameterScale(AtInstantLike): _allowed_keys = config.COMMON_KEYS.union({"brackets"}) _exclusion_list = ["parent", "_at_instant_cache"] + """The keys to be excluded from the node when output to a yaml file.""" def __init__(self, name: str, data: dict, file_path: str): """ From ffb84a7b8f1538e20ae189f96dd7dde2ddbdd1e9 Mon Sep 17 00:00:00 2001 From: SylviaDu99 Date: Sat, 23 Nov 2024 19:40:50 -0800 Subject: [PATCH 10/10] minor: use OrderedDict in `get_attr_dict` functions to provide more readable output WIP fixes: #274 --- policyengine_core/parameters/parameter.py | 7 ++++--- policyengine_core/parameters/parameter_node.py | 8 +++++--- policyengine_core/parameters/parameter_scale.py | 7 ++++--- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/policyengine_core/parameters/parameter.py b/policyengine_core/parameters/parameter.py index 38dbb090e..5127b3e32 100644 --- a/policyengine_core/parameters/parameter.py +++ b/policyengine_core/parameters/parameter.py @@ -3,7 +3,7 @@ from typing import Dict, List, Optional import numpy - +from collections import OrderedDict from policyengine_core.commons.misc import empty_clone from policyengine_core.errors import ParameterParsingError from policyengine_core.periods import INSTANT_PATTERN, period as get_period @@ -239,7 +239,7 @@ def relative_change(self, start_instant, end_instant): return end_value / start_value - 1 def get_attr_dict(self) -> dict: - data = self.__dict__.copy() + data = OrderedDict(self.__dict__.copy()) for attr in self._exclusion_list: if attr in data.keys(): del data[attr] @@ -251,4 +251,5 @@ def get_attr_dict(self) -> dict: value = float(value) value_dict[value_at_instant.instant_str] = value data["values_list"] = value_dict - return data + data.move_to_end("values_list") + return dict(data) diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 4f44cc0ae..2dd0fd8cf 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -5,6 +5,7 @@ from typing import Iterable, Union import yaml +from collections import OrderedDict from policyengine_core import commons, parameters, tools from policyengine_core.periods.instant_ import Instant @@ -280,7 +281,7 @@ def get_child(self, path: str) -> "ParameterNode": return node def get_attr_dict(self) -> dict: - data = self.__dict__.copy() + data = OrderedDict(self.__dict__.copy()) for attr in self._exclusion_list: if attr in data.keys(): del data[attr] @@ -288,8 +289,9 @@ def get_attr_dict(self) -> dict: child_dict = data.get("children") for child_name, child in child_dict.items(): data[child_name] = child.get_attr_dict() - del data["children"] - return data + data.move_to_end(child_name) + del data["children"] + return dict(data) class NoAliasDumper(yaml.SafeDumper): def ignore_aliases(self, data): diff --git a/policyengine_core/parameters/parameter_scale.py b/policyengine_core/parameters/parameter_scale.py index 851462de0..2d9a1c14d 100644 --- a/policyengine_core/parameters/parameter_scale.py +++ b/policyengine_core/parameters/parameter_scale.py @@ -2,7 +2,7 @@ import os import typing from typing import Any, Iterable - +from collections import OrderedDict from policyengine_core import commons, parameters, tools from policyengine_core.errors import ParameterParsingError from policyengine_core.parameters import AtInstantLike, config, helpers @@ -174,7 +174,7 @@ def _get_at_instant(self, instant: Instant) -> TaxScaleLike: return scale def get_attr_dict(self) -> dict: - data = self.__dict__.copy() + data = OrderedDict(self.__dict__.copy()) for attr in self._exclusion_list: if attr in data.keys(): del data[attr] @@ -185,4 +185,5 @@ def get_attr_dict(self) -> dict: node_list[i] = node.get_attr_dict() i += 1 data["brackets"] = node_list - return data + data.move_to_end("brackets") + return dict(data)