Merge pull request #192 from CitrineInformatics/feature/robust-enum

Augment base enumeration class to behave well in string contexts

Author: kroenlein

gemd/entity/attribute/base_attribute.py

@@ -49,7 +49,7 @@ def __init__(self,
         self._value = None
         self._template = None
         self._origin = None
-        self._file_links = None,
+        self._file_links = None
 
         self.value = value
         self.template = template
@@ -106,15 +106,15 @@ def _template_type() -> Type:
         """Get the expected type of template for this object (property of child)."""
 
     @property
-    def origin(self) -> str:
+    def origin(self) -> Origin:
         """Get origin."""
         return self._origin
 
     @origin.setter
     def origin(self, origin: Union[Origin, str]):
         if origin is None:
             raise ValueError("origin must be specified (but may be `unknown`)")
-        self._origin = Origin.get_value(origin)
+        self._origin = Origin.from_str(origin, exception=True)
 
     @property
     def file_links(self) -> List[FileLink]:

gemd/entity/object/material_run.py

@@ -96,13 +96,13 @@ def measurements(self) -> List["MeasurementRun"]:
         return self._measurements
 
     @property
-    def sample_type(self) -> str:
+    def sample_type(self) -> SampleType:
         """Get the sample type."""
         return self._sample_type
 
     @sample_type.setter
     def sample_type(self, sample_type: Union[SampleType, str]):
-        self._sample_type = SampleType.get_value(sample_type)
+        self._sample_type = SampleType.from_str(sample_type, exception=True)
 
     @staticmethod
     def _spec_type() -> Type:

gemd/entity/setters.py

@@ -6,7 +6,7 @@
 T = TypeVar('T')
 
 
-def validate_list(obj: Union[Iterable[T], T],
+def validate_list(obj: Optional[Union[Iterable[T], T]],
                   typ: Union[Iterable[Type], Type],
                   *,
                   trigger: Callable[[T], Optional[T]] = None) -> ValidList:

gemd/enumeration/base_enumeration.py

@@ -1,20 +1,79 @@
 """Base class for all enumerations."""
+from deprecation import deprecated
 from enum import Enum
+from typing import Optional
 
 
-class BaseEnumeration(Enum):
-    """Enumeration class that can convert between enumerations and associated values."""
+class BaseEnumeration(str, Enum):
+    """
+    Enumeration class that can convert between enumerations and associated values.
 
-    def __init__(self, *_):
-        """Ensure that there are no duplicates in the enumeration."""
-        cls = self.__class__
-        if any(self.value == e.value for e in cls):
-            raise ValueError("Duplicates not allowed in enumerated set of values {}".format(cls))
-        if not isinstance(self.value, str):
+    BaseEnumeration is a powerful support class for string enumerations.  It inherits
+    from both str and Enum to enable a class with str capabilities but still a
+    restricted data space.  All constructors are case-insensitive on input and a given
+    enumeration can recognize multiple synonyms for input, though only one value will
+    correspond to the value itsself.  For example:
+
+    ```
+    Fruits(BaseEnumeration):
+        APPLE = "Apple"
+        AVOCADO = "Avocado", "Alligator Pear"
+    ```
+
+    will recognize "apple", "APPLE" and "  aPpLe  " as referring to Fruits.APPLE,
+    and "avocado" and "alligator pear" as referring to Fruits.AVOCADO.  However,
+    since str(Fruits.AVOCADO) is "Avocado", Fruits.AVOCADO != "Alligator Pear".
+
+    """
+
+    def __new__(cls, value: str, *args):
+        """Overloaded to allow for synonyms."""
+        if any(not isinstance(x, str) for x in (value,) + args):
             raise ValueError("All values of enum {} must be strings".format(cls))
+        if cls.from_str(value, exception=False) is not None:
+            raise ValueError("Duplicates not allowed in enumerated set of values {}".format(cls))
+        obj = str.__new__(cls, value)
+        obj._value_ = value
+        obj.synonyms = frozenset(args)
+        obj.matches = frozenset([obj.lower()]).union(x.lower() for x in obj.synonyms)
+        return obj
 
     @classmethod
-    def get_value(cls, name):
+    def from_str(cls, val: str, *, exception: bool = False) -> Optional["BaseEnumeration"]:
+        """
+        Given a string value, return the Enumeration object that matches.
+
+        Parameters
+        ----------
+        val: str
+            The string to match against.  Leading and trailing whitespace is ignored.
+            Case is ignored.
+        exception: bool
+            Whether to raise an error if the string doesn't match anything.  Default: False.
+
+        Returns
+        -------
+        BaseEnumeration
+            The matching enumerated element, or None
+
+        :param val:
+        :param exception:
+        :return:
+
+        """
+        if val is None:
+            result = None
+        else:
+            result = next((x for x in cls if str.lower(val).strip() in x.matches), None)
+        if exception and result is None:
+            raise ValueError(f"{val} is not a valid {cls}; valid choices are {[x for x in cls]}")
+        return result
+
+    @classmethod
+    @deprecated(deprecated_in="1.15.0",
+                removed_in="2.0.0",
+                details="Enumerations autocast to values now.")
+    def get_value(cls, name: str) -> str:
         """
         Return a valid value associated with name.
 
@@ -23,14 +82,13 @@ def get_value(cls, name):
         """
         if name is None:
             return None
-        if any(name == e.value for e in cls):
-            return name
-        if any(name == e for e in cls):
-            return name.value
-        raise ValueError("'{}' is not a valid choice for enumeration {}".format(name, cls))
+        return cls.from_str(name, exception=True).value
 
     @classmethod
-    def get_enum(cls, name):
+    @deprecated(deprecated_in="1.15.0",
+                removed_in="2.0.0",
+                details="Use from_str for retreiving the correct Enum object.")
+    def get_enum(cls, name: str) -> "BaseEnumeration":
         """
         Return the enumeration associated with name.
 
@@ -39,8 +97,8 @@ def get_enum(cls, name):
         """
         if name is None:
             return None
-        if any(name == e.value for e in cls):
-            return next(e for e in cls if e.value == name)
-        if any(name == e for e in cls):
-            return name
-        raise ValueError("'{}' is not a valid choice for enumeration {}".format(name, cls))
+        return cls.from_str(name, exception=True)
+
+    def __str__(self):
+        """Return the value of the enumeration object."""
+        return self.value

gemd/json/gemd_encoder.py

@@ -2,7 +2,6 @@
 from uuid import UUID
 
 from gemd.entity.dict_serializable import DictSerializable
-from gemd.enumeration.base_enumeration import BaseEnumeration
 
 
 class GEMDEncoder(JSONEncoder):
@@ -12,8 +11,6 @@ def default(self, o):
         """Default encoder implementation."""
         if isinstance(o, DictSerializable):
             return o.as_dict()
-        elif isinstance(o, BaseEnumeration):
-            return o.value
         elif isinstance(o, UUID):
             return str(o)
         else:

setup.py

@@ -4,7 +4,7 @@
 packages.append("")
 
 setup(name='gemd',
-      version='1.14.1',
+      version='1.15.0',
       python_requires='>=3.7',
       url='http://github.com/CitrineInformatics/gemd-python',
       description="Python binding for Citrine's GEMD data model",

tests/enumeration/test_enumeration.py

@@ -13,16 +13,24 @@ class GoodClass(BaseEnumeration):
         RED = "Red"
         BLUE = "Blue"
 
-    assert GoodClass.get_value("Red") == "Red"
-    assert GoodClass.get_value(GoodClass.BLUE) == "Blue"
-    assert GoodClass.get_value(None) is None
-    assert GoodClass.get_enum("Red") == GoodClass.RED
-    assert GoodClass.get_enum(GoodClass.BLUE) == GoodClass.BLUE
-    assert GoodClass.get_enum(None) is None
-    with pytest.raises(ValueError):
-        GoodClass.get_value("Green")
-    with pytest.raises(ValueError):
-        GoodClass.get_enum("Green")
+    with pytest.deprecated_call():
+        assert GoodClass.get_value("Red") == "Red"
+    with pytest.deprecated_call():
+        assert GoodClass.get_value(GoodClass.BLUE) == "Blue"
+    with pytest.deprecated_call():
+        assert GoodClass.get_value(None) is None
+    with pytest.deprecated_call():
+        assert GoodClass.get_enum("Red") == GoodClass.RED
+    with pytest.deprecated_call():
+        assert GoodClass.get_enum(GoodClass.BLUE) == GoodClass.BLUE
+    with pytest.deprecated_call():
+        assert GoodClass.get_enum(None) is None
+    with pytest.deprecated_call():
+        with pytest.raises(ValueError):
+            GoodClass.get_value("Green")
+    with pytest.deprecated_call():
+        with pytest.raises(ValueError):
+            GoodClass.get_enum("Green")
 
 
 def test_json_serde():
@@ -47,3 +55,29 @@ class BadClass1(BaseEnumeration):
         class BadClass2(BaseEnumeration):
             FIRST = "one"
             SECOND = 2
+
+
+def test_string_enum():
+    """Test that the synonym mechanism works."""
+
+    class TestEnum(BaseEnumeration):
+        ONE = "One", "1"
+        TWO = "Two", "2"
+
+    assert TestEnum.ONE == "One", "Equality failed"
+    assert str(TestEnum.ONE) == "One", "Equality failed, cast"
+    assert TestEnum.ONE != "1", "Equality worked for synonym"
+    assert TestEnum.from_str("One") == TestEnum.ONE, "from_str worked"
+    assert TestEnum.from_str("ONE") == TestEnum.ONE, "from_str, caps worked"
+    assert TestEnum.from_str("one") == TestEnum.ONE, "from_str, lower worked"
+    assert TestEnum.from_str("1") == TestEnum.ONE, "from_str, synonym worked"
+    assert TestEnum.from_str(None) is None, "from_str, bad returned None"
+    assert TestEnum.from_str("1.0") is None, "from_str, bad returned None"
+    with pytest.raises(ValueError, match="valid"):
+        TestEnum.from_str("1.0", exception=True)
+    for key in TestEnum.TWO.synonyms:
+        assert key != TestEnum.TWO, f"Synonym {key} was equal?"
+        assert TestEnum.from_str(key) == TestEnum.TWO, f"from_str didn't resolve {key}"
+        assert (
+                TestEnum.from_str(key.upper()) == TestEnum.TWO
+        ), f"from_str didn't resolve {key.upper()}"

tests/json/test_json.py

@@ -112,7 +112,8 @@ def test_enumeration_serde():
     """An enumeration should get serialized as a string."""
     condition = Condition(name="A condition", notes=Origin.UNKNOWN)
     copy_condition = GEMDJson().copy(condition)
-    assert copy_condition.notes == Origin.get_value(condition.notes)
+    assert copy_condition.notes == Origin.UNKNOWN.value
+    assert not isinstance(copy_condition.notes, Origin)
 
 
 def test_attribute_serde():