Repository: strawberry-graphql/strawberry

Issue:
Add support for enums



Discussion:

Resolving patch:
diff --git a/strawberry/__init__.py b/strawberry/__init__.py
--- a/strawberry/__init__.py
+++ b/strawberry/__init__.py
@@ -1,8 +1,9 @@
 __version__ = "0.1.0"
 
 
+from .enum import enum  # noqa
 from .field import field  # noqa
 from .mutation import mutation, subscription  # noqa
-from .schema import Schema  # noqa
-from .type import type, input  # noqa
 from .scalars import ID  # noqa
+from .schema import Schema  # noqa
+from .type import input, type  # noqa
diff --git a/strawberry/enum.py b/strawberry/enum.py
new file mode 100644
--- /dev/null
+++ b/strawberry/enum.py
@@ -0,0 +1,46 @@
+from enum import EnumMeta
+
+from graphql import GraphQLEnumType, GraphQLEnumValue
+from graphql.utilities.schema_printer import print_type
+
+from .exceptions import NotAnEnum
+from .type_converter import REGISTRY
+
+
+def _process_enum(cls, name="", description=""):
+    if not isinstance(cls, EnumMeta):
+        raise NotAnEnum()
+
+    if not name:
+        name = cls.__name__
+
+    REGISTRY[name] = cls
+
+    def repr_(self):
+        return print_type(self.field)
+
+    setattr(cls, "__repr__", repr_)
+
+    cls.field = GraphQLEnumType(
+        name=name,
+        values=[(item.name, GraphQLEnumValue(item.value)) for item in cls],
+        description=description,
+    )
+
+    return cls
+
+
+def enum(_cls=None, *, name="", description=""):
+    """Registers the enum in the GraphQL type system.
+
+    If name is passed, the name of the GraphQL type will be
+    the value passed of name instead of the Enum class name.
+    """
+
+    def wrap(cls):
+        return _process_enum(cls, name, description)
+
+    if not _cls:
+        return wrap
+
+    return wrap(_cls)
diff --git a/strawberry/exceptions.py b/strawberry/exceptions.py
--- a/strawberry/exceptions.py
+++ b/strawberry/exceptions.py
@@ -3,6 +3,13 @@
 from typing import List, Set
 
 
+class NotAnEnum(Exception):
+    def __init__(self):
+        message = "strawberry.enum can only be used with subclasses of Enum"
+
+        super().__init__(message)
+
+
 class MissingReturnAnnotationError(Exception):
     """The field is missing the return annotation"""
 


Test patch:
diff --git a/tests/test_enum.py b/tests/test_enum.py
new file mode 100644
--- /dev/null
+++ b/tests/test_enum.py
@@ -0,0 +1,118 @@
+from enum import Enum
+
+import pytest
+
+import strawberry
+from graphql import GraphQLEnumType, GraphQLEnumValue
+from strawberry.exceptions import NotAnEnum
+
+
+def test_create_enum():
+    @strawberry.enum
+    class StringTest(Enum):
+        A = "c"
+        B = "i"
+        C = "a"
+        D = "o"
+
+    assert StringTest.field
+
+    assert type(StringTest.field) == GraphQLEnumType
+    assert StringTest.field.name == "StringTest"
+
+    assert StringTest.field.values == {
+        "A": GraphQLEnumValue("c"),
+        "B": GraphQLEnumValue("i"),
+        "C": GraphQLEnumValue("a"),
+        "D": GraphQLEnumValue("o"),
+    }
+
+    @strawberry.enum
+    class IntTest(Enum):
+        A = 1
+        B = 2
+        C = 3
+
+    assert IntTest.field
+
+    assert type(IntTest.field) == GraphQLEnumType
+    assert IntTest.field.name == "IntTest"
+
+    assert IntTest.field.values == {
+        "A": GraphQLEnumValue(1),
+        "B": GraphQLEnumValue(2),
+        "C": GraphQLEnumValue(3),
+    }
+
+    @strawberry.enum
+    class ComplexTest(Enum):
+        MERCURY = (3.303e23, 2.4397e6)
+        VENUS = (4.869e24, 6.0518e6)
+        EARTH = (5.976e24, 6.37814e6)
+        MARS = (6.421e23, 3.3972e6)
+        JUPITER = (1.9e27, 7.1492e7)
+        SATURN = (5.688e26, 6.0268e7)
+        URANUS = (8.686e25, 2.5559e7)
+        NEPTUNE = (1.024e26, 2.4746e7)
+
+        def __init__(self, mass, radius):
+            self.mass = mass
+            self.radius = radius
+
+        @property
+        def surface_gravity(self):
+            G = 6.67300e-11
+            return G * self.mass / (self.radius * self.radius)
+
+    assert ComplexTest.field
+
+    assert type(ComplexTest.field) == GraphQLEnumType
+    assert ComplexTest.field.name == "ComplexTest"
+
+    assert ComplexTest.field.values == {
+        "MERCURY": GraphQLEnumValue((3.303e23, 2.4397e6)),
+        "VENUS": GraphQLEnumValue((4.869e24, 6.0518e6)),
+        "EARTH": GraphQLEnumValue((5.976e24, 6.37814e6)),
+        "MARS": GraphQLEnumValue((6.421e23, 3.3972e6)),
+        "JUPITER": GraphQLEnumValue((1.9e27, 7.1492e7)),
+        "SATURN": GraphQLEnumValue((5.688e26, 6.0268e7)),
+        "URANUS": GraphQLEnumValue((8.686e25, 2.5559e7)),
+        "NEPTUNE": GraphQLEnumValue((1.024e26, 2.4746e7)),
+    }
+
+
+def test_create_enum_with_custom_name():
+    @strawberry.enum(name="NewName")
+    class Test(Enum):
+        A = 1
+        B = 2
+        C = 3
+
+    assert Test.field
+
+    assert type(Test.field) == GraphQLEnumType
+    assert Test.field.name == "NewName"
+
+
+def test_create_enum_with_arguments():
+    @strawberry.enum(name="NewName", description="This is the description")
+    class Test(Enum):
+        A = 1
+        B = 2
+        C = 3
+
+    assert Test.field
+
+    assert type(Test.field) == GraphQLEnumType
+    assert Test.field.name == "NewName"
+    assert Test.field.description == "This is the description"
+
+
+def test_raises_error_when_using_enum_with_a_not_enum_class():
+    with pytest.raises(NotAnEnum) as e:
+
+        @strawberry.enum
+        class NormalClass:
+            hello = "world"
+
+    assert ("strawberry.enum can only be used with subclasses of Enum",) == e.value.args