_yaml: Add a 'get_node' on Mapping, when return type can be of multiple typesbschubert/node-api-noget

Sometimes, we might want to have a list or a string, in which case, we
can now use 'get_node', which can return multiple types.

This method doesn't contain a 'default' value, as building a 'Node' from
it would be hard. We therefore just have a 'allow_none' flag that allows
returning 'None' when the value is not found.

- includes: use the new 'get_node' instead of trying multiple versions

3 files changed, 27 insertions, 12 deletions

diff --git a/src/buildstream/_includes.py b/src/buildstream/_includes.py
index 9bed8189c..2831c9765 100644
--- a/src/buildstream/_includes.py
+++ b/src/buildstream/_includes.py
@@ -35,19 +35,14 @@ class Includes:
         if current_loader is None:
             current_loader = self._loader
 
-        try:
-            includes = node.get_str('(@)', default=None)
-            if includes is not None:
-                includes = [includes]
-        except LoadError:
-            try:
-                includes = node.get_sequence('(@)').as_str_list()
-            except LoadError:
-                provenance = _yaml.node_get_provenance(node, key='(@)')
-                raise LoadError(LoadErrorReason.INVALID_DATA,
-                                "{}: {} must either be list or str".format(provenance, _yaml.node_sanitize(node)))
+        includes_node = node.get_node('(@)', allowed_types=[_yaml.ScalarNode, _yaml.SequenceNode], allow_none=True)
+
+        if includes_node:
+            if type(includes_node) is _yaml.ScalarNode:  # pylint: disable=unidiomatic-typecheck
+                includes = [includes_node.as_str()]
+            else:
+                includes = includes_node.as_str_list()
 
-        if includes:
             include_provenance = _yaml.node_get_provenance(node, key='(@)')
             _yaml.node_del(node, '(@)')
 
diff --git a/src/buildstream/_yaml.pxd b/src/buildstream/_yaml.pxd
index 3fffda426..8deb17a78 100644
--- a/src/buildstream/_yaml.pxd
+++ b/src/buildstream/_yaml.pxd
@@ -31,6 +31,7 @@ cdef class Node:
 cdef class MappingNode(Node):
     cdef Node get(self, str key, default, default_constructor)
     cpdef MappingNode get_mapping(self, str key, default=*)
+    cpdef Node get_node(self, str key, list allowed_types, bint allow_none=*)
     cpdef ScalarNode get_scalar(self, str key, default=*)
     cpdef SequenceNode get_sequence(self, str key, object default=*)
     cpdef bint get_bool(self, str key, default=*) except *
diff --git a/src/buildstream/_yaml.pyx b/src/buildstream/_yaml.pyx
index 8142a1820..4a49a9832 100644
--- a/src/buildstream/_yaml.pyx
+++ b/src/buildstream/_yaml.pyx
@@ -156,6 +156,25 @@ cdef class MappingNode(Node):
 
         return value
 
+    cpdef Node get_node(self, str key, list allowed_types, bint allow_none = False):
+        cdef value = self.value.get(key, _sentinel)
+
+        if value is _sentinel:
+            if allow_none:
+                return None
+
+            provenance = node_get_provenance(self)
+            raise LoadError(LoadErrorReason.INVALID_DATA,
+                            "{}: Dictionary did not contain expected key '{}'".format(provenance, key))
+
+        if type(value) not in allowed_types:
+            provenance = node_get_provenance(self)
+            raise LoadError(LoadErrorReason.INVALID_DATA,
+                            "{}: Value of '{}' is not one of the following: {}.".format(
+                                provenance, key, ", ".join(allowed_types)))
+
+        return value
+
     cpdef ScalarNode get_scalar(self, str key, object default=_sentinel):
         value = self.get(key, default, ScalarNode)