ORM: add EnumData plugin to to_aiida_type dispatch (#5314)

aiidateam · Jan 19, 2022 · 44577af · 44577af
1 parent ad28e9b
commit 44577af
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/aiida/orm/nodes/data/enum.py b/aiida/orm/nodes/data/enum.py
@@ -22,13 +22,19 @@ class Color(Enum):
 
 from aiida.common.lang import type_check
 
+from .base import to_aiida_type
 from .data import Data
 
 __all__ = ('EnumData',)
 
 EnumType = t.TypeVar('EnumType', bound=Enum)
 
 
+@to_aiida_type.register(Enum)
+def _(value):
+    return EnumData(member=value)
+
+
 class EnumData(Data):
     """Data plugin that allows to easily wrap an :class:`enum.Enum` member.
 

diff --git a/tests/orm/nodes/data/test_to_aiida_type.py b/tests/orm/nodes/data/test_to_aiida_type.py
@@ -11,6 +11,7 @@
 import pytest
 
 from aiida import orm
+from aiida.common.links import LinkType
 
 
 #yapf: disable
@@ -23,6 +24,7 @@
         (orm.Int, 5),
         (orm.List, [0, 1, 2]),
         (orm.Str, 'test-string'),
+        (orm.EnumData, LinkType.RETURN),
     )
 )
 # yapf: enable