python: make yaml generator a bit more generic

- pass a dict - add type information for parameters Signed-off-by: Axel Heider <[email protected]>
axel-h · Feb 3, 2022 · e5e32ad · e5e32ad
1 parent 1c3470b
commit e5e32ad
Showing 1 changed file with 16 additions and 10 deletions.
diff --git a/tools/hardware/outputs/yaml.py b/tools/hardware/outputs/yaml.py
@@ -8,14 +8,15 @@
 
 import argparse
 import yaml
-from typing import List
+from typing import List, Dict
 import hardware
 from hardware.config import Config
 from hardware.fdt import FdtParser
+from hardware.memory import Region
 from hardware.utils.rule import HardwareYaml
 
 
-def make_yaml_list_of_regions(regions) -> List:
+def make_yaml_list_of_regions(regions: List[Region]) -> List:
     return [
         {
             'start': r.base,
@@ -25,19 +26,19 @@ def make_yaml_list_of_regions(regions) -> List:
     ]
 
 
-def create_yaml_file(dev_mem, phys_mem, outputStream):
+def create_yaml_file(regions_dict: Dict[str, List[Region]], outputStream):
 
     yaml.add_representer(
         int,
         lambda dumper, data: yaml.ScalarNode('tag:yaml.org,2002:int', hex(data)))
 
-    yaml_obj = {
-        'devices': make_yaml_list_of_regions(dev_mem),
-        'memory':  make_yaml_list_of_regions(phys_mem)
-    }
-
     with outputStream:
-        yaml.dump(yaml_obj, outputStream)
+        yaml.dump(
+            {
+                key: make_yaml_list_of_regions(val)
+                for key, val in regions_dict.items()
+            },
+            outputStream)
 
 
 def get_kernel_devices(tree: FdtParser, hw_yaml: HardwareYaml):
@@ -61,7 +62,12 @@ def run(tree: FdtParser, hw_yaml: HardwareYaml, config: Config,
     dev_mem = hardware.utils.memory.get_addrspace_exclude(
         list(reserved) + phys_mem + kernel_devs, config)
 
-    create_yaml_file(dev_mem, phys_mem, args.yaml_out)
+    create_yaml_file(
+        {
+            'devices':  dev_mem,
+            'memory':   phys_mem
+        },
+        args.yaml_out)
 
 
 def add_args(parser):