From 87d1569d9e7f60c379c784dc0220ac9783470896 Mon Sep 17 00:00:00 2001 From: "E. Madison Bray" Date: Tue, 24 Oct 2023 12:21:51 +0200 Subject: [PATCH 1/2] Fixes enum value inference in cases where the value type is of a user-defined data type class (with __new__). This fixes a regression introduced by #10057 to fix #10000. The `not ti.fullname.startswith("builtins.") clause seemed to be intended to catch enums with a built-in data type like int or bytes, but this is overly broad. It should allow any type so long as it is not itself an enum.Enum subclass. --- mypy/plugins/enums.py | 8 ++++++-- test-data/unit/check-enum.test | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 7869a8b5cdfa..e7dd8fd86f7f 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -104,13 +104,17 @@ def _infer_value_type_with_auto_fallback( def _implements_new(info: TypeInfo) -> bool: """Check whether __new__ comes from enum.Enum or was implemented in a - subclass. In the latter case, we must infer Any as long as mypy can't infer + subclass of enum.Enum. In the latter case, we must infer Any as long as mypy can't infer the type of _value_ from assignments in __new__. + + If, however, __new__ comes from a user-defined class that is not an Enum subclass (i.e. + the data type) this is allowed, because we should in general infer that an enum entry's + value has that type. """ type_with_new = _first( ti for ti in info.mro - if ti.names.get("__new__") and not ti.fullname.startswith("builtins.") + if ti.is_enum and ti.names.get("__new__") ) if type_with_new is None: return False diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 6779ae266454..9563a7ec88c8 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -1366,6 +1366,20 @@ reveal_type(a._value_) # N: Revealed type is "Any" [builtins fixtures/primitives.pyi] [typing fixtures/typing-medium.pyi] +[case testValueTypeWithUserDataType] +from enum import Enum +from typing import Any + +class Data: + def __new__(cls, value: Any) -> Data: pass + +class DataEnum(Data, Enum): + A = Data(1) + +reveal_type(DataEnum.A) # N: Revealed type is "Literal[__main__.DataEnum.A]?" +reveal_type(DataEnum.A.value) # N: Revealed type is "__main__.Data" +reveal_type(DataEnum.A._value_) # N: Revealed type is "__main__.Data" + [case testEnumNarrowedToTwoLiterals] # Regression test: two literals of an enum would be joined # as the full type, regardless of the amount of elements From 05881ca442457d895085f063b5994bb076d472b0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 24 Oct 2023 10:37:01 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypy/plugins/enums.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index e7dd8fd86f7f..388d72f25001 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -111,11 +111,7 @@ def _implements_new(info: TypeInfo) -> bool: the data type) this is allowed, because we should in general infer that an enum entry's value has that type. """ - type_with_new = _first( - ti - for ti in info.mro - if ti.is_enum and ti.names.get("__new__") - ) + type_with_new = _first(ti for ti in info.mro if ti.is_enum and ti.names.get("__new__")) if type_with_new is None: return False return type_with_new.fullname not in ("enum.Enum", "enum.IntEnum", "enum.StrEnum")