feat: support Annotated pattern for enum registration (#4293)

* feat: support Annotated pattern for enum registration

* fix: address review feedback on enum docs and tests

* refactor: replace EnumDefinition with functools.partial

* refactor: replace functools.partial with EnumAnnotation class

Author: bellini666

RELEASE.md

@@ -0,0 +1,30 @@
+Release type: minor
+
+Enums can now be registered via `Annotated`. The preferred way is still using
+`@strawberry.enum` as a decorator, but when you need to expose an existing enum
+under a different name or alias, `Annotated` works as a proper type alias in all
+type checkers:
+
+```python
+from typing import Annotated
+from enum import Enum
+import strawberry
+
+
+class IceCreamFlavour(Enum):
+    VANILLA = "vanilla"
+    STRAWBERRY = "strawberry"
+    CHOCOLATE = "chocolate"
+
+
+MyIceCreamFlavour = Annotated[
+    IceCreamFlavour, strawberry.enum(description="Ice cream flavours")
+]
+
+
+@strawberry.type
+class Query:
+    @strawberry.field
+    def flavour(self) -> MyIceCreamFlavour:
+        return IceCreamFlavour.VANILLA
+```

docs/editors/mypy.md

@@ -29,3 +29,33 @@ plugins = ["pydantic.mypy", "strawberry.ext.mypy_plugin"]
 
 The strawberry plugin synthesises `__init__`, `to_pydantic()` and
 `from_pydantic()` on pydantic-decorated classes so that mypy can see them.
+
+## Enums
+
+The preferred way to register an enum is with the decorator:
+
+```python
+from enum import Enum
+import strawberry
+
+
+@strawberry.enum
+class IceCreamFlavour(Enum):
+    VANILLA = "vanilla"
+    STRAWBERRY = "strawberry"
+```
+
+If you need to expose an existing enum under a different name or alias, use
+`Annotated` instead of assigning `strawberry.enum(IceCreamFlavour)` to a
+variable — mypy treats the latter as a value, not a type, so it cannot be used
+in annotations.
+
+```python
+from typing import Annotated
+import strawberry
+
+MyIceCreamFlavour = Annotated[IceCreamFlavour, strawberry.enum(description="...")]
+```
+
+`MyIceCreamFlavour` is a proper type alias that mypy and Pyright accept in
+annotations.

strawberry/annotation.py

@@ -29,7 +29,7 @@
     get_object_definition,
     has_object_definition,
 )
-from strawberry.types.enum import StrawberryEnumDefinition
+from strawberry.types.enum import EnumAnnotation, StrawberryEnumDefinition
 from strawberry.types.enum import enum as strawberry_enum
 from strawberry.types.lazy_type import LazyType
 from strawberry.types.maybe import _annotation_is_maybe
@@ -189,7 +189,7 @@ def _resolve_evaled_type(self, evaled_type: Any) -> StrawberryType | type:
         # Everything remaining should be a raw annotation that needs to be turned into
         # a StrawberryType
         if self._is_enum(evaled_type):
-            return self.create_enum(evaled_type)
+            return self.create_enum(evaled_type, args)
         if self._is_optional(evaled_type, args):
             return self.create_optional(evaled_type)
         if self._is_union(evaled_type, args):
@@ -215,10 +215,20 @@ def create_concrete_type(self, evaled_type: type) -> type:
             return evaled_type.__strawberry_definition__.resolve_generic(evaled_type)
         raise ValueError(f"Not supported {evaled_type}")
 
-    def create_enum(self, evaled_type: Any) -> StrawberryEnumDefinition:
+    def create_enum(
+        self, evaled_type: Any, args: list[Any] | None = None
+    ) -> StrawberryEnumDefinition:
+        enum_annotation: EnumAnnotation | None = None
+        if args:
+            enum_annotation = next(
+                (a for a in args if isinstance(a, EnumAnnotation)), None
+            )
+
         try:
             return evaled_type.__strawberry_definition__
         except AttributeError:
+            if enum_annotation is not None:
+                return enum_annotation(evaled_type).__strawberry_definition__
             return strawberry_enum(evaled_type).__strawberry_definition__
 
     def create_list(self, evaled_type: Any) -> StrawberryList:

strawberry/types/enum.py

@@ -105,6 +105,23 @@ class MyEnum(Enum):
 GraphqlEnumNameFrom = Literal["key", "value"]
 
 
+@dataclasses.dataclass
+class EnumAnnotation:
+    name: str | None = None
+    description: str | None = None
+    directives: Iterable[object] = ()
+    graphql_name_from: GraphqlEnumNameFrom = "key"
+
+    def __call__(self, cls: EnumType) -> EnumType:
+        return _process_enum(
+            cls,
+            self.name,
+            self.description,
+            directives=self.directives,
+            graphql_name_from=self.graphql_name_from,
+        )
+
+
 def _process_enum(
     cls: EnumType,
     name: str | None = None,
@@ -238,20 +255,17 @@ class MyEnum(Enum):
     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: EnumType) -> EnumType:
-        return _process_enum(
-            cls,
-            name,
-            description,
-            directives=directives,
-            graphql_name_from=graphql_name_from,
-        )
+    wrapper = EnumAnnotation(
+        name=name,
+        description=description,
+        directives=directives,
+        graphql_name_from=graphql_name_from,
+    )
 
     if not cls:
-        return wrap
+        return wrapper
 
-    return wrap(cls)
+    return wrapper(cls)
 
 
 WithStrawberryEnumDefinition = WithStrawberryDefinition["StrawberryEnumDefinition"]
@@ -265,6 +279,7 @@ def has_enum_definition(obj: Any) -> TypeGuard[type[WithStrawberryEnumDefinition
 
 
 __all__ = [
+    "EnumAnnotation",
     "EnumValue",
     "EnumValueDefinition",
     "StrawberryEnumDefinition",

tests/enums/test_enum.py

@@ -1,4 +1,5 @@
 from enum import Enum, IntEnum
+from typing import Annotated
 
 import pytest
 
@@ -236,3 +237,60 @@ class Query:
 
     with pytest.raises(TypeError, match=r"Expected name to be a string."):
         strawberry.Schema(Query)
+
+
+def test_annotated_enum_with_description():
+    class IceCreamFlavour(Enum):
+        VANILLA = "vanilla"
+        STRAWBERRY = "strawberry"
+        CHOCOLATE = "chocolate"
+
+    MyIceCreamFlavour = Annotated[
+        IceCreamFlavour, strawberry.enum(description="Ice cream flavours")
+    ]
+
+    @strawberry.type
+    class Query:
+        @strawberry.field
+        def flavour(self) -> MyIceCreamFlavour:
+            return IceCreamFlavour.VANILLA
+
+    schema = strawberry.Schema(Query)
+    schema_str = str(schema)
+    assert '"""Ice cream flavours"""' in schema_str
+    assert "enum IceCreamFlavour" in schema_str
+
+
+def test_annotated_enum_with_name():
+    class Flavour(Enum):
+        VANILLA = "vanilla"
+        STRAWBERRY = "strawberry"
+
+    MyIceCreamFlavour = Annotated[Flavour, strawberry.enum(name="IceCreamFlavour")]
+
+    @strawberry.type
+    class Query:
+        @strawberry.field
+        def flavour(self) -> MyIceCreamFlavour:
+            return Flavour.VANILLA
+
+    schema = strawberry.Schema(Query)
+    assert "IceCreamFlavour" in str(schema)
+
+
+def test_annotated_enum_on_already_decorated():
+    @strawberry.enum
+    class Flavour(Enum):
+        VANILLA = "vanilla"
+        STRAWBERRY = "strawberry"
+
+    MyFlavour = Annotated[Flavour, strawberry.enum(description="Flavours")]
+
+    @strawberry.type
+    class Query:
+        @strawberry.field
+        def flavour(self) -> MyFlavour:
+            return Flavour.VANILLA
+
+    schema = strawberry.Schema(Query)
+    assert "enum Flavour" in str(schema)

tests/typecheckers/test_enum.py

@@ -279,6 +279,59 @@ def test_enum_with_manual_decorator_and_name():
     )
 
 
+CODE_WITH_ANNOTATED = """
+from enum import Enum
+from typing import Annotated
+
+import strawberry
+
+class IceCreamFlavour(Enum):
+    VANILLA = "vanilla"
+    STRAWBERRY = "strawberry"
+    CHOCOLATE = "chocolate"
+
+MyIceCreamFlavour = Annotated[IceCreamFlavour, strawberry.enum(description="Flavours")]
+
+x: MyIceCreamFlavour = IceCreamFlavour.VANILLA
+reveal_type(x)
+"""
+
+
+def test_enum_with_annotated():
+    results = typecheck(CODE_WITH_ANNOTATED)
+
+    assert results.pyright == snapshot(
+        [
+            Result(
+                type="information",
+                message='Type of "x" is "Literal[IceCreamFlavour.VANILLA]"',
+                line=15,
+                column=13,
+            ),
+        ]
+    )
+    assert results.mypy == snapshot(
+        [
+            Result(
+                type="note",
+                message='Revealed type is "mypy_test.IceCreamFlavour"',
+                line=15,
+                column=13,
+            ),
+        ]
+    )
+    assert results.ty == snapshot(
+        [
+            Result(
+                type="information",
+                message="Revealed type: `Literal[IceCreamFlavour.VANILLA]`",
+                line=15,
+                column=13,
+            ),
+        ]
+    )
+
+
 CODE_WITH_DEPRECATION_REASON = """
 from enum import Enum