Native support for Python's enum.Enum

[jamespfennell]

Hi! Would you consider adding native support for Python's enum.Enum types, from the standard library enum package? Something like this:

class MyEnum(enum.Enum):
    FIRST = 1
    SECOND = 2

# Assuming PyEnum is the correct ScalarValidator:
schema = strictyaml.Map({"test": strictyaml.PyEnum(MyEnum)})

parsed = strictyaml.load("test: FIRST", schema)
# Result is {"test": MyEnum.FIRST}

parsed = strictyaml.load("test: THIRD", schema)
# Raises a YAMLSerializationError exception

It's possible to workaround this case by using strictyaml.Enum, and then manually performing the type cast after:

schema = strictyaml.Map({"test": strictyaml.Enum(elem.name for name in MyEnum)})

parsed = strictyaml.load("test: FIRST", schema)
# Result is {"test": "FIRST"}
elem = MyEnum[parsed["test"]]

Having to do the type cast manually is a bit of a pain, and also if you want to output yaml again you need to convert the enum element to a string before doing so. On the other hand, the workaround handles the most important thing which is validating the input, so I understand if you don't want to bite.

I'll post an implementation of PyEnum next.

[jamespfennell]

Possible implementation:

import strictyaml
import enum

from strictyaml.exceptions import YAMLSerializationError
from strictyaml import Map, ScalarValidator


class PyEnum(ScalarValidator):

    def __init__(self, enum_):
        self._enum = enum_
        assert issubclass(
            self._enum,
            enum.Enum
        ), "argument must be a enum.Enum or subclass thereof"

    def validate_scalar(self, chunk):
        try:
            val = self._enum[chunk.contents]
        except KeyError:
            chunk.expecting_but_found(
                "when expecting one of: {0}".format(", ".join(elem.name for elem in self._enum))
            )
        else:
            return val

    def to_yaml(self, data):
        if data not in self._enum:
            raise YAMLSerializationError(
                "Got '{0}' when  expecting one of: {1}".format(
                    data, ", ".join(str(elem) for elem in self._enum)
                )
            )
        return data.name

    def __repr__(self):
        return u"PyEnum({0})".format(repr(self._enum))


class MyEnum(enum.Enum):
    FIRST = 1
    SECOND = 2


schema = Map({
    "test": PyEnum(MyEnum)
})

result = strictyaml.load("test: FIRST", schema)
# Outputs:
# YAML(OrderedDict([('test', <MyEnum.FIRST: 1>)]))

result['test'] = MyEnum.SECOND
print(result.as_yaml())
# Outputs:
# "test: SECOND"

result['test'] = "THIRD"
print(result.as_yaml())
# Raises:
# strictyaml.exceptions.YAMLSerializationError: Got 'THIRD' when  expecting one of: MyEnum.FIRST, MyEnum.SECOND

result = strictyaml.load("test: THIRD", schema)
# Raises:
# strictyaml.exceptions.YAMLValidationError: when expecting one of: FIRST, SECOND
# found arbitrary text
#   in "<unicode string>", line 1, column 1:
#     test: THIRD
#      ^ (line: 1)

PyEnum(Map)
# Raise:
# AssertionError: argument must be a enum.Enum or subclass thereof

[crdoconnor]

Interesting idea. Yeah, I'll have a think about how to do this :+1:

[jamespfennell]

Cool! Just to be clear, the second post above contains a full implementation of the idea. I could create a pull request if you like.

[crdoconnor]

I know. I was thinking of combining it with the existing Enum type though - and I also wanted to TDD this code.

On Sat, 14 Sep 2019, 16:13 James Fennell, [email protected] wrote:

Cool! Just to be clear, the second post above contains a full implementation of the idea. I could create a pull request if you like.

— You are receiving this because you commented. Reply to this email directly, view it on GitHub https://github.com/crdoconnor/strictyaml/issues/73?email_source=notifications&email_token=ABOJKNIJADC4K5XECOKMFL3QJT5QHA5CNFSM4IWR6FDKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOD6W5V3Y#issuecomment-531487471, or mute the thread https://github.com/notifications/unsubscribe-auth/ABOJKNO3PNINPSCK52A4KQLQJT5QHANCNFSM4IWR6FDA .

[jamespfennell]

Nice idea! I guess you could have the Enum constructor check if the input is a list versus enum.Enum and go from there.

Here are some tests I already wrote for PyEnum. I'm incorporating the code above into a project for which I generally write unit tests:

class TestPyEnum(unittest.TestCase):
    class MyEnum(enum.Enum):
        FIRST = 1
        SECOND = 2

    schema = Map({"test": systemconfigreader.PyEnum(MyEnum)})

    def test_read_valid(self):
        """[System config reader | PyEnum] Read valid"""
        result = strictyaml.load("test: FIRST", self.schema)

        self.assertDictEqual({"test": self.MyEnum.FIRST}, dict(result.data))

    def test_read_invalid(self):
        """[System config reader | PyEnum] Read invalid"""
        self.assertRaises(
            YAMLValidationError, lambda: strictyaml.load("test: THIRD", self.schema)
        )

    def test_write_valid(self):
        """[System config reader | PyEnum] Write valid"""
        yaml = strictyaml.as_document(
            {"test": self.MyEnum.FIRST}, schema=self.schema
        ).as_yaml()

        self.assertEqual("test: FIRST", yaml.strip())

    def test_write_invalid(self):
        """[System config reader | PyEnum] Write invalid"""
        self.assertRaises(
            YAMLSerializationError,
            lambda: strictyaml.as_document(
                {"test": "FIRST"}, schema=self.schema
            ).as_yaml(),
        )

    def test_instantiate_without_enum(self):
        """[System config reader | PyEnum] Instantiate without enum"""

        class DummyClass:
            pass

        self.assertRaises(AssertionError, lambda: systemconfigreader.PyEnum(DummyClass))