Add type hint for a generic enum definition

jasujm · May 12, 2020 · c5deda2 · c5deda2
1 parent db90e7e
commit c5deda2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/python/enumecg/__init__.py b/python/enumecg/__init__.py
@@ -45,7 +45,7 @@ def generator(
 
 
 def generate(
-    enum,
+    enum: definitions.Enum,
     *,
     documentation: typing.Union[generators.DocumentationStyle, str, None] = None,
     primary_type: typing.Union[definitions.PrimaryType, str, None] = None,

diff --git a/python/enumecg/definitions.py b/python/enumecg/definitions.py
@@ -60,6 +60,14 @@ class EnumDefinition:
     enhanced_enum_documentation: typing.Optional[EnumDocumentation] = None
 
 
+Enum = typing.Union[EnumDefinition, typing.Mapping, py_enum.EnumMeta]
+"""Generic enum definition
+
+Types accepted by :func:`make_definition()` and other functions that
+are used to generate enhanced enum definition.
+"""
+
+
 def _make_definition_from_dict(enum_dict, *, primary_type, value_type):
     typename = enum_dict["typename"]
     members = enum_dict["members"]
@@ -122,7 +130,7 @@ def _extract_python_enum_attrs(enum):
 
 
 def make_definition(
-    enum,
+    enum: Enum,
     *,
     primary_type: typing.Optional[PrimaryType] = None,
     value_type: typing.Optional[str] = None,