Fix handling of boolean attributes in codelists (#106)

nomenclature/code.py

@@ -1,13 +1,14 @@
-from re import compile, match
 from typing import Union, List, Dict
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, StrictStr, StrictInt, StrictFloat, StrictBool
 
 
 class Code(BaseModel):
     """A simple class for a mapping of a "code" to its attributes"""
 
     name: str
-    attributes: Dict[str, Union[str, int, float, bool, List, None]]
+    attributes: Dict[
+        str, Union[StrictStr, StrictInt, StrictFloat, StrictBool, List, None]
+    ]
 
     @classmethod
     def from_dict(cls, mapping):

nomenclature/codelist.py

@@ -249,7 +249,9 @@ def to_yaml(self, path=None):
 
         # translate to list of nested dicts, replace None by empty field, write to file
         stream = (
-            yaml.dump([{code: attrs} for code, attrs in self.mapping.items()])
+            yaml.dump(
+                [{code: attrs} for code, attrs in self.mapping.items()], sort_keys=False
+            )
             .replace(": null\n", ":\n")
             .replace(": nan\n", ":\n")
         )

tests/data/simple_codelist/foo.yaml

@@ -1,3 +1,4 @@
 - Some Variable:
     definition: Some basic variable
     unit:  # testing that importing an empty unit (ie dimensionless variable) works
+    bool: true  # testing that bool is not cast to string by pydantic

tests/test_codelist.py

@@ -14,6 +14,7 @@ def test_simple_codelist():
 
     assert "Some Variable" in code
     assert code["Some Variable"]["unit"] is None  # this is a dimensionless variable
+    assert type(code["Some Variable"]["bool"]) == bool  # this is a boolean
 
 
 def test_codelist_to_yaml():
@@ -23,8 +24,9 @@ def test_codelist_to_yaml():
     assert code.to_yaml() == (
         "- Some Variable:\n"
         "    definition: Some basic variable\n"
-        "    file: simple_codelist/foo.yaml\n"
         "    unit:\n"
+        "    bool: true\n"
+        "    file: simple_codelist/foo.yaml\n"
     )