From f5b5a14f3f35eb215b9592e456cb5018044ea39a Mon Sep 17 00:00:00 2001 From: sgrekov Date: Mon, 13 Feb 2023 17:22:21 +0200 Subject: [PATCH 1/3] Corrected enum metaclass to fix pickle.dumps() --- graphene/tests/issues/test_881.py | 27 +++++++++++++++++++++++++++ graphene/types/enum.py | 4 +++- 2 files changed, 30 insertions(+), 1 deletion(-) create mode 100644 graphene/tests/issues/test_881.py diff --git a/graphene/tests/issues/test_881.py b/graphene/tests/issues/test_881.py new file mode 100644 index 000000000..960c58392 --- /dev/null +++ b/graphene/tests/issues/test_881.py @@ -0,0 +1,27 @@ +import pickle + +from ...types.enum import Enum + + +class MyEnum(Enum): + # is defined outside of test because pickle unable to dump class inside ot pytest function + A = 'a' + B = 1 + + +def test_enums_pickling(): + a = MyEnum.A + pickled = pickle.dumps(a) + restored = pickle.loads(pickled) + assert type(a) is type(restored) + assert a == restored + assert a.value == restored.value + assert a.name == restored.name + + b = MyEnum.B + pickled = pickle.dumps(b) + restored = pickle.loads(pickled) + assert type(a) is type(restored) + assert b == restored + assert b.value == restored.value + assert b.name == restored.name diff --git a/graphene/types/enum.py b/graphene/types/enum.py index 58e65c69e..79bafe00b 100644 --- a/graphene/types/enum.py +++ b/graphene/types/enum.py @@ -31,9 +31,11 @@ def __new__(cls, name_, bases, classdict, **options): # with the enum values. enum_members.pop("Meta", None) enum = PyEnum(cls.__name__, enum_members) - return SubclassWithMeta_Meta.__new__( + obj = SubclassWithMeta_Meta.__new__( cls, name_, bases, dict(classdict, __enum__=enum), **options ) + globals()[name_] = obj.__enum__ + return obj def get(cls, value): return cls._meta.enum(value) From dac4ba4e1261526fb012955db7be0f5a4eb6b839 Mon Sep 17 00:00:00 2001 From: sgrekov Date: Tue, 14 Feb 2023 09:01:38 +0200 Subject: [PATCH 2/3] considered case with colliding class names (try to distinguish by file name) --- graphene/tests/issues/test_881.py | 2 +- graphene/types/enum.py | 40 +++++++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/graphene/tests/issues/test_881.py b/graphene/tests/issues/test_881.py index 960c58392..a08f79b67 100644 --- a/graphene/tests/issues/test_881.py +++ b/graphene/tests/issues/test_881.py @@ -5,7 +5,7 @@ class MyEnum(Enum): # is defined outside of test because pickle unable to dump class inside ot pytest function - A = 'a' + A = "a" B = 1 diff --git a/graphene/types/enum.py b/graphene/types/enum.py index 79bafe00b..b16a63b89 100644 --- a/graphene/types/enum.py +++ b/graphene/types/enum.py @@ -1,11 +1,40 @@ +import inspect +import sys from enum import Enum as PyEnum - +from typing import Any, Dict from graphene.utils.subclass_with_meta import SubclassWithMeta_Meta from .base import BaseOptions, BaseType from .unmountedtype import UnmountedType +class _ModuleItemHelper: + _enum_metas: Dict[str, Any] = {} + + def __getattr__(self, name: str) -> Any: + try: + return globals()[name] + except KeyError: + if name == "__path__": + return None + if len(_ModuleItemHelper._enum_metas[name]) == 1: + return next(iter(_ModuleItemHelper._enum_metas[name].values())) + else: + # if there are more than 1 class with the name - take first by stack + for fr in inspect.stack(): + f_name = inspect.getmodulename(fr.filename) + if f_name in _ModuleItemHelper._enum_metas[name]: + return _ModuleItemHelper._enum_metas[name][f_name] + raise + + def __setattr__(self, name: str, value: Any) -> None: + cls, path = value + if not _ModuleItemHelper._enum_metas.get(name): + _ModuleItemHelper._enum_metas[name] = {} + + _ModuleItemHelper._enum_metas[name][path.split(".")[-1]] = cls + + def eq_enum(self, other): if isinstance(other, self.__class__): return self is other @@ -34,7 +63,11 @@ def __new__(cls, name_, bases, classdict, **options): obj = SubclassWithMeta_Meta.__new__( cls, name_, bases, dict(classdict, __enum__=enum), **options ) - globals()[name_] = obj.__enum__ + # globals()[name_] = obj.__enum__ + if enum_members.get("__module__"): + setattr( + sys.modules[__name__], name_, (obj.__enum__, enum_members["__module__"]) + ) return obj def get(cls, value): @@ -118,3 +151,6 @@ def get_type(cls): is mounted (as a Field, InputField or Argument) """ return cls + + +sys.modules[__name__] = _ModuleItemHelper() # type: ignore From fa20db005422b3ba852457c8165d15aed3f8221c Mon Sep 17 00:00:00 2001 From: sgrekov Date: Tue, 14 Feb 2023 21:43:38 +0200 Subject: [PATCH 3/3] reverted simple solution back (without attempt to support duplicate Enum class names) --- graphene/tests/issues/test_881.py | 6 ++--- graphene/types/enum.py | 40 ++----------------------------- 2 files changed, 5 insertions(+), 41 deletions(-) diff --git a/graphene/tests/issues/test_881.py b/graphene/tests/issues/test_881.py index a08f79b67..f97b59176 100644 --- a/graphene/tests/issues/test_881.py +++ b/graphene/tests/issues/test_881.py @@ -3,14 +3,14 @@ from ...types.enum import Enum -class MyEnum(Enum): +class PickleEnum(Enum): # is defined outside of test because pickle unable to dump class inside ot pytest function A = "a" B = 1 def test_enums_pickling(): - a = MyEnum.A + a = PickleEnum.A pickled = pickle.dumps(a) restored = pickle.loads(pickled) assert type(a) is type(restored) @@ -18,7 +18,7 @@ def test_enums_pickling(): assert a.value == restored.value assert a.name == restored.name - b = MyEnum.B + b = PickleEnum.B pickled = pickle.dumps(b) restored = pickle.loads(pickled) assert type(a) is type(restored) diff --git a/graphene/types/enum.py b/graphene/types/enum.py index b16a63b89..79bafe00b 100644 --- a/graphene/types/enum.py +++ b/graphene/types/enum.py @@ -1,40 +1,11 @@ -import inspect -import sys from enum import Enum as PyEnum -from typing import Any, Dict + from graphene.utils.subclass_with_meta import SubclassWithMeta_Meta from .base import BaseOptions, BaseType from .unmountedtype import UnmountedType -class _ModuleItemHelper: - _enum_metas: Dict[str, Any] = {} - - def __getattr__(self, name: str) -> Any: - try: - return globals()[name] - except KeyError: - if name == "__path__": - return None - if len(_ModuleItemHelper._enum_metas[name]) == 1: - return next(iter(_ModuleItemHelper._enum_metas[name].values())) - else: - # if there are more than 1 class with the name - take first by stack - for fr in inspect.stack(): - f_name = inspect.getmodulename(fr.filename) - if f_name in _ModuleItemHelper._enum_metas[name]: - return _ModuleItemHelper._enum_metas[name][f_name] - raise - - def __setattr__(self, name: str, value: Any) -> None: - cls, path = value - if not _ModuleItemHelper._enum_metas.get(name): - _ModuleItemHelper._enum_metas[name] = {} - - _ModuleItemHelper._enum_metas[name][path.split(".")[-1]] = cls - - def eq_enum(self, other): if isinstance(other, self.__class__): return self is other @@ -63,11 +34,7 @@ def __new__(cls, name_, bases, classdict, **options): obj = SubclassWithMeta_Meta.__new__( cls, name_, bases, dict(classdict, __enum__=enum), **options ) - # globals()[name_] = obj.__enum__ - if enum_members.get("__module__"): - setattr( - sys.modules[__name__], name_, (obj.__enum__, enum_members["__module__"]) - ) + globals()[name_] = obj.__enum__ return obj def get(cls, value): @@ -151,6 +118,3 @@ def get_type(cls): is mounted (as a Field, InputField or Argument) """ return cls - - -sys.modules[__name__] = _ModuleItemHelper() # type: ignore