Add function to output ParameterNode as YAML

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

policyengine_core/parameters/at_instant_like.py

@@ -19,3 +19,10 @@ 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

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,9 @@
     _parse_child,
     _load_yaml_file,
 )
+from .parameter import Parameter
+from .parameter_node_at_instant import ParameterNodeAtInstant
+from .parameter_at_instant import ParameterAtInstant
 
 EXCLUDED_PARAMETER_CHILD_NAMES = ["reference", "__pycache__"]
 
@@ -274,3 +277,28 @@ 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:
+        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)
+        except Exception as e:
+            print(f"Error when writing YAML file: {e}")

tests/core/test_parameters.py

@@ -1,4 +1,5 @@
 import tempfile
+from pathlib import Path
 
 import pytest
 
@@ -141,3 +142,28 @@ def test_name():
     }
     parameter = ParameterNode("root", data=parameter_data)
     assert parameter.children["2010"].name == "root.2010"
+
+
+def test_write_yaml():
+    parameter_data = {
+        "amount": {
+            "values": {
+                "2015-01-01": {"value": 550},
+                "2016-01-01": {"value": 600},
+            },
+            "description": "The amount of the basic income",
+            "documentation": None,
+            "modified": False,
+        },
+        "min_age": {
+            "values": {
+                "2015-01-01": {"value": 25},
+                "2016-01-01": {"value": 18},
+            },
+            "description": "The minimum age to receive the basic income",
+            "documentation": None,
+            "modified": True,
+        },
+    }
+    parameter = ParameterNode("root", data=parameter_data)
+    parameter.write_yaml(Path("output.yaml"))