Skip to content

Commit

Permalink
fix: modified write_yaml method in parameter_node.py to output th…
Browse files Browse the repository at this point in the history
…e node data as a YAML file

WIP

fixes: PolicyEngine#274
  • Loading branch information
SylviaDu99 committed Oct 16, 2024
1 parent 1939220 commit d9a37d5
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 7 deletions.
5 changes: 5 additions & 0 deletions changelog_entry.yaml
Original file line number Diff line number Diff line change
@@ -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
10 changes: 6 additions & 4 deletions policyengine_core/parameters/parameter_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
16 changes: 13 additions & 3 deletions tests/core/test_parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"))

0 comments on commit d9a37d5

Please sign in to comment.