diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29b..23a8388f 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 7ef1b7bd..c262cda9 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 9afd7ca6..b1494561 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"))