Use Python enums to enumerate the possible primary types in EnumECG

CHANGELOG.rst

@@ -12,6 +12,7 @@ Changed
   - Change the format of ``dict`` representation of enumerators to
     have more explicit ordering of members
   - Improvements to the documentation
+  - Use Python enums to enumerate the possible primary types in EnumECG
 
 Version 0.3
 -----------

python/enumecg/__init__.py

@@ -11,7 +11,7 @@
 
 import typing
 
-from . import generators
+from . import definitions, generators
 
 
 def generator(
@@ -34,7 +34,7 @@ def generate(
     enum,
     *,
     documentation: typing.Optional[str] = None,
-    primary_type: typing.Optional[str] = None,
+    primary_type: typing.Union[definitions.PrimaryType, str, None] = None,
     value_type: typing.Optional[str] = None
 ) -> str:
     """Generate code for an enhanced enum
@@ -61,12 +61,16 @@ def generate(
     Parameters:
         enum: The enum definition
         documentation: See :ref:`enumecg-documentation-generation`.
-        primary_type: See :ref:`enumecg-primary-enum`.
+        primary_type: A string or an enumerator indicating the
+                      primary type. See :ref:`enumecg-primary-enum`.
         value_type: See :ref:`enumerator-value-type`.
 
     Returns:
         The enhanced enum definition created from the ``enum`` description.
+
     """
+    if primary_type is not None:
+        primary_type = definitions.PrimaryType(primary_type)
     return str(
         generator(documentation=documentation).generate_enum_definitions(
             enum, primary_type=primary_type, value_type=value_type

python/enumecg/cli.py

@@ -13,7 +13,7 @@
 
 from . import generate
 from .generators import CodeGenerator
-from .definitions import PRIMARY_TYPE_CHOICES
+from .definitions import PrimaryType
 
 
 @click.command()
@@ -24,7 +24,7 @@
 )
 @click.option(
     "--primary-type",
-    type=click.Choice(PRIMARY_TYPE_CHOICES),
+    type=click.Choice([e.value for e in PrimaryType.__members__.values()]),
     help="Primary enumeration type",
 )
 @click.option("--value-type", help="Enumerator value type")

python/enumecg/definitions.py

@@ -15,12 +15,16 @@
 
 from . import utils
 
-PRIMARY_TYPE_CHOICES = ["label", "enhanced"]
-"""Possible primary types when generating enum definitions
 
-These are the accepted choices for the ``primary_type`` argument in
-:func:`make_definition()`.
-"""
+class PrimaryType(py_enum.Enum):
+    """Possible primary types when generating enum definitions
+
+    These are the accepted choices for the ``primary_type`` argument
+    in :func:`make_definition()`.
+    """
+
+    label = "label"
+    enhanced = "enhanced"
 
 
 @dataclasses.dataclass
@@ -60,12 +64,12 @@ def _make_definition_from_dict(enum_dict, *, primary_type, value_type):
     member_formatter = utils.NameFormatter(*(member["name"] for member in members))
     label_enum_typename = (
         formatter.join(formatter.parts[0] + ["label"])
-        if primary_type != "label"
+        if primary_type != PrimaryType.label
         else typename
     )
     enhanced_enum_typename = (
         formatter.join(["enhanced"] + formatter.parts[0])
-        if primary_type != "enhanced"
+        if primary_type != PrimaryType.enhanced
         else typename
     )
     type_deducer = utils.CppTypeDeducer(
@@ -97,9 +101,9 @@ def _make_definition_from_dict(enum_dict, *, primary_type, value_type):
             for (n, member) in enumerate(members)
         ],
         associate_namespace_name=formatter.join(formatter.parts[0], pluralize=True),
-        label_enum_documentation=documentation if primary_type == "label" else None,
+        label_enum_documentation=documentation if primary_type == PrimaryType.label else None,
         enhanced_enum_documentation=documentation
-        if primary_type == "enhanced"
+        if primary_type == PrimaryType.enhanced
         else None,
     )
 
@@ -115,7 +119,7 @@ def _extract_python_enum_attrs(enum):
 def make_definition(
     enum,
     *,
-    primary_type: typing.Optional[str] = None,
+    primary_type: typing.Optional[PrimaryType] = None,
     value_type: typing.Optional[str] = None,
 ) -> EnumDefinition:
     """Make :class:`EnumDefinition` instance from various types
@@ -133,8 +137,10 @@ def make_definition(
 
     Parameters:
         enum: The enum definition.
-        primary_type: See :ref:`enumecg-primary-enum`.
+        primary_type: A :class:`PrimaryType` enumerator indicating the
+                      primary type. See :ref:`enumecg-primary-enum`.
         value_type: See :ref:`enumerator-value-type`.
+
     """
     if isinstance(enum, EnumDefinition):
         return enum

python/tests/test_cli.py

@@ -5,6 +5,7 @@
 
 from enumecg import generate
 from enumecg.cli import cli
+from enumecg.definitions import PrimaryType
 
 
 @pytest.fixture
@@ -34,11 +35,11 @@ def test_cli_should_have_documentation_option(cli_runner, enum_file, status_defi
     assert result.output == generate(status_definition, documentation="doxygen") + "\n"
 
 
-@pytest.mark.parametrize("primary_type", ["label", "enhanced"])
+@pytest.mark.parametrize("primary_type", PrimaryType)
 def test_cli_should_have_primeray_type_option(
     cli_runner, enum_file, status_definition_dict, primary_type
 ):
-    result = cli_runner.invoke(cli, ["--primary-type", primary_type, str(enum_file)])
+    result = cli_runner.invoke(cli, ["--primary-type", primary_type.value, str(enum_file)])
     assert (
         result.output
         == generate(status_definition_dict, primary_type=primary_type) + "\n"

python/tests/test_definitions.py

@@ -1,6 +1,6 @@
 import pytest
 
-from enumecg.definitions import EnumDefinition, EnumMemberDefinition, make_definition
+from enumecg.definitions import EnumDefinition, EnumMemberDefinition, make_definition, PrimaryType
 
 from .conftest import Status
 
@@ -30,7 +30,7 @@ def test_make_definition_with_label_type_as_primary(
     status_definition.label_enum_typename = "Status"
     status_definition.label_enum_documentation = status_documentation
     assert (
-        make_definition(status_definition_dict, primary_type="label")
+        make_definition(status_definition_dict, primary_type=PrimaryType.label)
         == status_definition
     )
 
@@ -41,7 +41,7 @@ def test_make_definition_with_enhanced_type_as_primary(
     status_definition.enhanced_enum_typename = "Status"
     status_definition.enhanced_enum_documentation = status_documentation
     make_definition(
-        status_definition_dict, primary_type="enhanced"
+        status_definition_dict, primary_type=PrimaryType.enhanced
     ) == status_definition