node: Add 'as_enum' on ScalarNode and 'get_enum' helper on MappingNode

This adds a method to ensure that a value is from a set of valid values
and raises an error message accordingly.

- Define Enum types for each of the relevant cases
- Adapt call places that were doing such things manually

6 files changed, 111 insertions, 13 deletions

diff --git a/src/buildstream/_gitsourcebase.py b/src/buildstream/_gitsourcebase.py
index fb6010b1e..68eb4c2fb 100644
--- a/src/buildstream/_gitsourcebase.py
+++ b/src/buildstream/_gitsourcebase.py
@@ -32,6 +32,7 @@ from configparser import RawConfigParser
 from .source import Source, SourceError, SourceFetcher
 from .types import Consistency, CoreWarnings
 from . import utils
+from .types import FastEnum
 from .utils import move_atomic, DirectoryExistsError
 
 GIT_MODULES = '.gitmodules'
@@ -42,6 +43,11 @@ WARN_UNLISTED_SUBMODULE = "unlisted-submodule"
 WARN_INVALID_SUBMODULE = "invalid-submodule"
 
 
+class _RefFormat(FastEnum):
+    SHA1 = "sha1"
+    GIT_DESCRIBE = "git-describe"
+
+
 # Because of handling of submodules, we maintain a _GitMirror
 # for the primary git source and also for each submodule it
 # might have at a given time
@@ -156,7 +162,7 @@ class _GitMirror(SourceFetcher):
             cwd=self.mirror)
         ref = output.rstrip('\n')
 
-        if self.source.ref_format == 'git-describe':
+        if self.source.ref_format == _RefFormat.GIT_DESCRIBE:
             # Prefix the ref with the closest tag, if available,
             # to make the ref human readable
             exit_code, output = self.source.check_output(
@@ -394,10 +400,7 @@ class _GitSourceBase(Source):
         self.mirror = self.BST_MIRROR_CLASS(self, '', self.original_url, ref, tags=tags, primary=True)
         self.tracking = node.get_str('track', None)
 
-        self.ref_format = node.get_str('ref-format', 'sha1')
-        if self.ref_format not in ['sha1', 'git-describe']:
-            provenance = node.get_scalar('ref-format').get_provenance()
-            raise SourceError("{}: Unexpected value for ref-format: {}".format(provenance, self.ref_format))
+        self.ref_format = node.get_enum('ref-format', _RefFormat, _RefFormat.SHA1)
 
         # At this point we now know if the source has a ref and/or a track.
         # If it is missing both then we will be unable to track or build.
diff --git a/src/buildstream/_project.py b/src/buildstream/_project.py
index f4a7466de..9428ab4f6 100644
--- a/src/buildstream/_project.py
+++ b/src/buildstream/_project.py
@@ -41,6 +41,7 @@ from ._projectrefs import ProjectRefs, ProjectRefStorage
 from ._versions import BST_FORMAT_VERSION
 from ._loader import Loader
 from .element import Element
+from .types import FastEnum
 from ._message import Message, MessageType
 from ._includes import Includes
 from ._workspaces import WORKSPACE_PROJECT_FILE
@@ -50,6 +51,13 @@ from ._workspaces import WORKSPACE_PROJECT_FILE
 _PROJECT_CONF_FILE = 'project.conf'
 
 
+# List of all places plugins can come from
+class PluginOrigins(FastEnum):
+    CORE = "core"
+    LOCAL = "local"
+    PIP = "pip"
+
+
 # HostMount()
 #
 # A simple object describing the behavior of
@@ -862,14 +870,8 @@ class Project():
                 'origin', 'sources', 'elements',
                 'package-name', 'path',
             ]
-            allowed_origins = ['core', 'local', 'pip']
             origin.validate_keys(allowed_origin_fields)
 
-            origin_value = origin.get_str('origin')
-            if origin_value not in allowed_origins:
-                raise LoadError("Origin '{}' is not one of the allowed types"
-                                .format(origin_value), LoadErrorReason.INVALID_YAML)
-
             # Store source versions for checking later
             source_versions = origin.get_mapping('sources', default={})
             for key in source_versions.keys():
@@ -888,7 +890,9 @@ class Project():
 
             # Store the origins if they're not 'core'.
             # core elements are loaded by default, so storing is unnecessary.
-            if origin.get_str('origin') != 'core':
+            origin_value = origin.get_enum('origin', PluginOrigins)
+
+            if origin_value != PluginOrigins.CORE:
                 self._store_origin(origin, 'sources', plugin_source_origins)
                 self._store_origin(origin, 'elements', plugin_element_origins)
 
@@ -928,7 +932,7 @@ class Project():
                 if group in origin_node:
                     del origin_node[group]
 
-            if origin_node.get_str('origin') == 'local':
+            if origin_node.get_enum('origin', PluginOrigins) == PluginOrigins.LOCAL:
                 path = self.get_path_from_node(origin.get_scalar('path'),
                                                check_is_dir=True)
                 # paths are passed in relative to the project, but must be absolute
diff --git a/src/buildstream/_types.pyx b/src/buildstream/_types.pyx
index 671f0bd2f..0f8d3690c 100644
--- a/src/buildstream/_types.pyx
+++ b/src/buildstream/_types.pyx
@@ -77,3 +77,6 @@ class MetaFastEnum(type):
 
     def __setattr__(self, key, value):
         raise ValueError("Adding new values dynamically is not supported")
+
+    def __iter__(self):
+        return iter(self._value_to_entry.values())
diff --git a/src/buildstream/node.pxd b/src/buildstream/node.pxd
index 18520146d..02e95d06f 100644
--- a/src/buildstream/node.pxd
+++ b/src/buildstream/node.pxd
@@ -46,6 +46,7 @@ cdef class MappingNode(Node):
 
     # Public Methods
     cpdef bint get_bool(self, str key, default=*) except *
+    cpdef object get_enum(self, str key, object constraint, object default=*)
     cpdef int get_int(self, str key, default=*) except *
     cpdef MappingNode get_mapping(self, str key, default=*)
     cpdef Node get_node(self, str key, list allowed_types=*, bint allow_none=*)
@@ -78,6 +79,7 @@ cdef class ScalarNode(Node):
 
     # Public Methods
     cpdef bint as_bool(self) except *
+    cpdef object as_enum(self, object constraint)
     cpdef int as_int(self) except *
     cpdef str as_str(self)
     cpdef bint is_none(self)
diff --git a/src/buildstream/node.pyx b/src/buildstream/node.pyx
index ea63151b9..98a785868 100644
--- a/src/buildstream/node.pyx
+++ b/src/buildstream/node.pyx
@@ -329,6 +329,45 @@ cdef class ScalarNode(Node):
                             .format(provenance, path, bool.__name__, self.value),
                             LoadErrorReason.INVALID_DATA)
 
+    cpdef object as_enum(self, object constraint):
+        """Get the value of the node as an enum member from `constraint`
+
+        The constraint must be a :class:`buildstream.types.FastEnum` or a plain python Enum.
+
+        For example you could do:
+
+        .. code-block:: python
+
+            from buildstream.types import FastEnum
+
+            class SupportedCompressions(FastEnum):
+              NONE = "none"
+              GZIP = "gzip"
+              XZ = "xz"
+
+
+            x = config.get_scalar('compress').as_enum(SupportedCompressions)
+
+            if x == SupportedCompressions.GZIP:
+                print("Using GZIP")
+
+        Args:
+            constraint (:class:`buildstream.types.FastEnum` or :class:`Enum`): an enum from which to extract the value
+                                                                               for the current node.
+
+        Returns:
+            :class:`FastEnum` or :class:`Enum`: the value contained in the node, as a member of `constraint`
+        """
+        try:
+            return constraint(self.value)
+        except ValueError:
+            provenance = self.get_provenance()
+            path = provenance._toplevel._find(self)[-1]
+            valid_values = [str(v.value) for v in constraint]
+            raise LoadError("{}: Value of '{}' should be one of '{}'".format(
+                                provenance, path, ", ".join(valid_values)),
+                            LoadErrorReason.INVALID_DATA)
+
     cpdef int as_int(self) except *:
         """Get the value of the node as an integer.
 
@@ -511,6 +550,44 @@ cdef class MappingNode(Node):
         cdef ScalarNode scalar = self.get_scalar(key, default)
         return scalar.as_bool()
 
+    cpdef object get_enum(self, str key, object constraint, object default=_sentinel):
+        """Get the value of the node as an enum member from `constraint`
+
+        Args:
+            key (str): key for which to get the value
+            constraint (:class:`buildstream.types.FastEnum` or :class:`Enum`): an enum from which to extract the value
+                                                                               for the current node.
+            default (object): default value to return if `key` is not in the mapping
+
+        Raises:
+            :class:`buildstream._exceptions.LoadError`: if the value is not is not found or not part of the
+                                                        provided enum.
+
+        Returns:
+            :class:`buildstream.types.Enum` or :class:`Enum`: the value contained in the node, as a member of
+                                                              `constraint`
+        """
+        cdef object value = self.value.get(key, _sentinel)
+
+        if value is _sentinel:
+            if default is _sentinel:
+                provenance = self.get_provenance()
+                raise LoadError("{}: Dictionary did not contain expected key '{}'".format(provenance, key),
+                                LoadErrorReason.INVALID_DATA)
+
+            if default is None:
+                return None
+            else:
+                return constraint(default)
+
+        if type(value) is not ScalarNode:
+            provenance = value.get_provenance()
+            raise LoadError("{}: Value of '{}' is not of the expected type 'scalar'"
+                                .format(provenance, key),
+                            LoadErrorReason.INVALID_DATA)
+
+        return (<ScalarNode> value).as_enum(constraint)
+
     cpdef int get_int(self, str key, object default=_sentinel) except *:
         """get_int(key, default=sentinel)
 
diff --git a/src/buildstream/types.py b/src/buildstream/types.py
index 56f5b99f0..ff4a3cc57 100644
--- a/src/buildstream/types.py
+++ b/src/buildstream/types.py
@@ -46,6 +46,15 @@ class FastEnum(metaclass=MetaFastEnum):
 
     _value_to_entry = dict()  # A dict of all values mapping to the entries in the enum
 
+    @classmethod
+    def values(cls):
+        """Get all the possible values for the enum.
+
+        Returns:
+            list: the list of all possible values for the enum
+        """
+        return cls._value_to_entry.keys()
+
     def __new__(cls, value):
         try:
             return cls._value_to_entry[value]