feat: add ability to provide a type name to enum when using from_enum

This commit is contained in:
Thomas Leonard 2022-06-24 18:00:55 +02:00
parent efe4b89015
commit 8f6a8f9c4a
2 changed files with 51 additions and 2 deletions

View File

@ -52,7 +52,10 @@ class EnumMeta(SubclassWithMeta_Meta):
return super(EnumMeta, cls).__call__(*args, **kwargs) return super(EnumMeta, cls).__call__(*args, **kwargs)
# return cls._meta.enum(*args, **kwargs) # return cls._meta.enum(*args, **kwargs)
def from_enum(cls, enum, description=None, deprecation_reason=None): # noqa: N805 def from_enum(
cls, enum, name=None, description=None, deprecation_reason=None
): # noqa: N805
name = name or enum.__name__
description = description or enum.__doc__ description = description or enum.__doc__
meta_dict = { meta_dict = {
"enum": enum, "enum": enum,
@ -60,7 +63,7 @@ class EnumMeta(SubclassWithMeta_Meta):
"deprecation_reason": deprecation_reason, "deprecation_reason": deprecation_reason,
} }
meta_class = type("Meta", (object,), meta_dict) meta_class = type("Meta", (object,), meta_dict)
return type(meta_class.enum.__name__, (Enum,), {"Meta": meta_class}) return type(name, (Enum,), {"Meta": meta_class})
class Enum(UnmountedType, BaseType, metaclass=EnumMeta): class Enum(UnmountedType, BaseType, metaclass=EnumMeta):

View File

@ -328,6 +328,52 @@ def test_enum_resolver_compat():
assert results.data["colorByName"] == Color.RED.name assert results.data["colorByName"] == Color.RED.name
def test_enum_with_name():
from enum import Enum as PyEnum
class Color(PyEnum):
RED = 1
YELLOW = 2
BLUE = 3
GColor = Enum.from_enum(Color, description="original colors")
UniqueGColor = Enum.from_enum(
Color, name="UniqueColor", description="unique colors"
)
class Query(ObjectType):
color = GColor(required=True)
unique_color = UniqueGColor(required=True)
schema = Schema(query=Query)
assert (
str(schema).strip()
== dedent(
'''
type Query {
color: Color!
uniqueColor: UniqueColor!
}
"""original colors"""
enum Color {
RED
YELLOW
BLUE
}
"""unique colors"""
enum UniqueColor {
RED
YELLOW
BLUE
}
'''
).strip()
)
def test_enum_resolver_invalid(): def test_enum_resolver_invalid():
from enum import Enum as PyEnum from enum import Enum as PyEnum