Inference rules (#109)

* use custom loader class, disable inference
* additional custom loader tests
* add inference tests to cmd tests

Author: rsgalloway

lib/envstack/node.py

@@ -250,30 +250,35 @@ def resolve(self, env: dict = os.environ):
 
 
 class CustomLoader(yaml.SafeLoader):
+    """Custom Loader class to preserve order of keys and ensure required
+    keys are first in the mapping."""
+
     required_keys = {"include", "all", "darwin", "linux", "windows"}
 
     def construct_mapping(self, node: yaml.Node, deep: bool = False):
-        mapping = super().construct_mapping(node, deep=deep)
-        for key, value in mapping.items():
+        """Construct a mapping from a YAML node, preserving the order of keys
+        and ensuring"""
+        # keep YAML merge keys (<<) working
+        self.flatten_mapping(node)
+
+        mapping = {}
+
+        for key_node, value_node in node.value:
+            # never implicit-resolve scalar KEYS (YES/NO/ON/OFF/null/date/etc)
+            if isinstance(key_node, yaml.ScalarNode):
+                key = key_node.value
+            else:
+                key = self.construct_object(key_node, deep=deep)
+
+            value = self.construct_object(value_node, deep=deep)
+            mapping[key] = value
+
+        # preserve order of keys and ensure required keys are first in the mapping
+        for key, value in list(mapping.items()):
             if key in self.required_keys:
                 continue
-            try:
-                if node.tag == Base64Node.yaml_tag:
-                    mapping[key] = Base64Node(value)
-                elif node.tag == EncryptedNode.yaml_tag:
-                    mapping[key] = EncryptedNode(value)
-                elif node.tag == AESGCMNode.yaml_tag:
-                    mapping[key] = AESGCMNode(value)
-                elif node.tag == FernetNode.yaml_tag:
-                    mapping[key] = FernetNode(value)
-                elif node.tag == MD5Node.yaml_tag:
-                    mapping[key] = MD5Node(value)
-                else:
-                    mapping[key] = Template(value)
-            except Exception as e:
-                raise yaml.constructor.ConstructorError(
-                    None, None, f"Error parsing template: {e}", node.start_mark
-                )
+            mapping[key] = value
+
         return mapping
 
 

lib/envstack/util.py

@@ -45,7 +45,13 @@
 import yaml
 
 from envstack import config
-from envstack.node import AESGCMNode, Base64Node, EncryptedNode, FernetNode
+from envstack.node import (
+    AESGCMNode,
+    Base64Node,
+    EncryptedNode,
+    FernetNode,
+    CustomLoader,
+)
 
 # default memoization cache timeout in seconds
 CACHE_TIMEOUT = 5
@@ -643,8 +649,7 @@ def validate_yaml(file_path: str):
     """
     try:
         with open(file_path, "r") as stream:
-            data = yaml.safe_load(stream.read())
-            # data = yaml.load(stream.read(), Loader=CustomLoader)
+            data = yaml.load(stream.read(), Loader=CustomLoader)
         if not isinstance(data, dict):
             raise yaml.YAMLError("invalid data structure")
         return data

tests/fixtures/env/thing.env

@@ -5,6 +5,8 @@ all: &all
   LOG_LEVEL: ${LOG_LEVEL:=INFO}
   FLOAT: 1.0
   INT: 5
+  YES: 1
+  TRUE: 1
   NUMBER_LIST: [1, 2, 3]
   CHAR_LIST: ['a', 'b', 'c', "${HELLO}"]
   DICT: {a: 1, b: 2, c: "${INT}"}

tests/test_cmds.py

@@ -159,6 +159,8 @@ def test_thing(self):
 PYTHONPATH=${DEPLOY_ROOT}/lib/python:${PYTHONPATH}
 ROOT=${HOME}/.local/pipe
 STACK=thing
+TRUE=1
+YES=1
 """
         command = "%s thing" % self.envstack_bin
         output = subprocess.check_output(command, shell=True, universal_newlines=True)
@@ -437,6 +439,8 @@ def test_thing(self):
   INT: 5
   LOG_LEVEL: ${LOG_LEVEL:=INFO}
   NUMBER_LIST: [1, 2, 3]
+  TRUE: 1
+  YES: 1
 darwin:
   <<: *all
   ROOT: ${HOME}/Library/Application Support/pipe
@@ -502,6 +506,8 @@ def test_thing_encrypted(self):
   INT: !encrypt NQ==
   LOG_LEVEL: !encrypt JHtMT0dfTEVWRUw6PUlORk99
   NUMBER_LIST: !encrypt WzEsIDIsIDNd
+  TRUE: !encrypt MQ==
+  YES: !encrypt MQ==
 darwin:
   <<: *all
   ROOT: !encrypt JHtIT01FfS9MaWJyYXJ5L0FwcGxpY2F0aW9uIFN1cHBvcnQvcGlwZQ==

tests/test_node.py

@@ -170,6 +170,11 @@ def test_str(self):
 
 
 class TestCustomLoader(unittest.TestCase):
+    """Test the CustomLoader class"""
+
+    def load(self, s: str):
+        return yaml.load(s, Loader=CustomLoader)
+
     def test_construct_mapping(self):
         """test the CustomLoader construct_mapping method"""
         envfile = os.path.join(envpath, "project.env")
@@ -178,8 +183,56 @@ def test_construct_mapping(self):
         mapping = loader.construct_mapping(node)
         self.assertIsInstance(mapping, dict)
 
+    def test_keys_stay_strings_booleanish(self):
+        """test that keys that look like booleans are not converted to bools
+        by the resolver"""
+        data = self.load("YES: 1\nNO: 2\nON: 3\nOFF: 4\ntrue: 5\nFalse: 6\n")
+        assert "YES" in data and data["YES"] == 1
+        assert "NO" in data and data["NO"] == 2
+        assert "ON" in data and data["ON"] == 3
+        assert "OFF" in data and data["OFF"] == 4
+        assert "true" in data and data["true"] == 5
+        assert "False" in data and data["False"] == 6
+        # ensure the resolver didn't produce bool keys
+        assert True not in data
+        assert False not in data
+
+    def test_keys_stay_strings_nullish(self):
+        """test that keys that look like nulls are not converted to None by
+        the resolver"""
+        data = self.load("NULL: 1\nNull: 2\nnull: 3\n~: 4\n")
+        assert data["NULL"] == 1
+        assert data["Null"] == 2
+        assert data["null"] == 3
+        assert "~" in data and data["~"] == 4
+        assert None not in data
+
+    def test_keys_stay_strings_timestampish(self):
+        """test that keys that look like timestamps are not converted to
+        datetime objects"""
+        # YAML 1.1 can infer timestamps
+        data = self.load("2024-01-02: 1\n2024-01-02T03:04:05Z: 2\n")
+        assert data["2024-01-02"] == 1
+        assert data["2024-01-02T03:04:05Z"] == 2
+
+    def test_yaml_merge_key_works(self):
+        """test that the YAML merge key (<<) works correctly with anchors and
+        aliases"""
+        s = """
+all: &all
+  A: 1
+dev:
+  <<: *all
+  B: 2
+"""
+        data = self.load(s)
+        assert data["dev"]["A"] == 1
+        assert data["dev"]["B"] == 2
+
 
 class TestCustomDumper(unittest.TestCase):
+    """Test the CustomDumper class"""
+
     def test_init(self):
         """test the CustomDumper __init__ method"""
         dumper = CustomDumper(None)

tests/test_util.py

@@ -131,9 +131,7 @@ def test_tokenized_value_single(self):
             "ROOT": "/var/tmp",
         }
         result = evaluate_modifiers(expression, environ)
-        self.assertEqual(
-            result, r"/var/tmp/{foo}"
-        )
+        self.assertEqual(result, r"/var/tmp/{foo}")
 
     def test_tokenized_value_trailing_slash(self):
         """Test tokenized value with a single token and trailing slash."""
@@ -142,9 +140,7 @@ def test_tokenized_value_trailing_slash(self):
             "ROOT": "/var/tmp",
         }
         result = evaluate_modifiers(expression, environ)
-        self.assertEqual(
-            result, r"/var/tmp/{foo}/"
-        )
+        self.assertEqual(result, r"/var/tmp/{foo}/")
 
     def test_tokenized_value_two_tokens(self):
         """Test tokenized value with two tokens."""
@@ -153,9 +149,7 @@ def test_tokenized_value_two_tokens(self):
             "ROOT": "/var/tmp",
         }
         result = evaluate_modifiers(expression, environ)
-        self.assertEqual(
-            result, r"/var/tmp/{foo}/{bar}"
-        )
+        self.assertEqual(result, r"/var/tmp/{foo}/{bar}")
 
     def test_tokenized_value_three_tokens(self):
         """Test tokenized value with three tokens."""
@@ -164,9 +158,7 @@ def test_tokenized_value_three_tokens(self):
             "ROOT": "/var/tmp",
         }
         result = evaluate_modifiers(expression, environ)
-        self.assertEqual(
-            result, r"/var/tmp/{foo}/{bar}/{baz}"
-        )
+        self.assertEqual(result, r"/var/tmp/{foo}/{bar}/{baz}")
 
     def test_tokenized_value_three_tokens_slash(self):
         """Test tokenized value with three tokens and a trailing slash."""
@@ -175,9 +167,7 @@ def test_tokenized_value_three_tokens_slash(self):
             "ROOT": "/var/tmp",
         }
         result = evaluate_modifiers(expression, environ)
-        self.assertEqual(
-            result, r"/var/tmp/{foo}/{bar}/{baz}/"
-        )
+        self.assertEqual(result, r"/var/tmp/{foo}/{bar}/{baz}/")
 
     def test_default_value_with_default_args(self):
         """Test default value with default args."""