feat: add function in AtInstantLike to return desired attributes in d…

…ictionary form WIP fixes: PolicyEngine#274
SylviaDu99 · Oct 28, 2024 · 4fb27f8 · 4fb27f8
1 parent ab9a573
commit 4fb27f8
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 20 deletions.
diff --git 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
@@ -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
@@ -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)