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,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

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,11 @@ 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 = self.get_attr_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"))