Ensure reconstructed Enums are idompotent

Previously, mypy would infer that `Foo(Foo.x)` is of type `Foo`. This is problematic as it means Enums can not be reconstructed for runtime saftey.

Hence, this commit extends the enum plugin and refines the return type of `Enum.__new__`, ensuring reconstructed Enums are idempotent

Fixes #19669

Author: alfiedouglas0

mypy/plugins/default.py

@@ -5,7 +5,7 @@
 
 import mypy.errorcodes as codes
 from mypy import message_registry
-from mypy.nodes import DictExpr, IntExpr, StrExpr, UnaryExpr
+from mypy.nodes import DictExpr, IntExpr, StrExpr, TypeInfo, UnaryExpr
 from mypy.plugin import (
     AttributeContext,
     ClassDefContext,
@@ -47,7 +47,12 @@
     dataclass_tag_callback,
     replace_function_sig_callback,
 )
-from mypy.plugins.enums import enum_member_callback, enum_name_callback, enum_value_callback
+from mypy.plugins.enums import (
+    enum_member_callback,
+    enum_name_callback,
+    enum_new_callback,
+    enum_value_callback,
+)
 from mypy.plugins.functools import (
     functools_total_ordering_maker_callback,
     functools_total_ordering_makers,
@@ -104,6 +109,12 @@ def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type]
             return partial_new_callback
         elif fullname == "enum.member":
             return enum_member_callback
+        elif (
+            (st := self.lookup_fully_qualified(fullname))
+            and isinstance(st.node, TypeInfo)
+            and getattr(st.node, "is_enum", False)
+        ):
+            return enum_new_callback
         return None
 
     def get_function_signature_hook(

mypy/plugins/enums.py

@@ -28,6 +28,8 @@
     LiteralType,
     ProperType,
     Type,
+    TypeVarType,
+    UnionType,
     get_proper_type,
     is_named_instance,
 )
@@ -297,3 +299,62 @@ def _extract_underlying_field_name(typ: Type) -> str | None:
     # as a string.
     assert isinstance(underlying_literal.value, str)
     return underlying_literal.value
+
+
+def enum_new_callback(ctx: mypy.plugin.FunctionContext) -> Type:
+    """This plugin refines the return type of `__new__`, ensuring reconstructed
+    Enums are idempotent.
+
+    By default, mypy will infer that `Foo(Foo.x)` is of type `Foo`. This plugin
+    ensures types are not loosened, meaning with this plugin enabled
+    `Foo(Foo.x)` is of type `Literal[Foo.x]?`.
+
+    This means with this plugin:
+    ```
+        reveal_type(Foo(Foo.x))  # mypy reveals Literal[Foo.x]?
+    ```
+
+    This plugin works by adjusting the return type of `__new__` to be the given
+    argument type, if and only if `__new__` comes from `enum.Enum`.
+
+    This plugin supports arguments that are Final, Literial, Union of Literials
+    and generic TypeVars.
+    """
+    base_ret = ctx.default_return_type
+    enum_inst = get_proper_type(base_ret)
+    if not isinstance(enum_inst, Instance):
+        return base_ret
+
+    info: TypeInfo = enum_inst.type
+    if not info.is_enum:
+        return base_ret
+
+    if _implements_new(info):
+        return base_ret
+
+    if not ctx.args or not ctx.args[0] or not ctx.arg_types or not ctx.arg_types[0]:
+        return base_ret
+
+    arg0_t = get_proper_type(ctx.arg_types[0][0])
+
+    if isinstance(arg0_t, Instance) and arg0_t.type is info:
+        return arg0_t
+    elif isinstance(arg0_t, LiteralType) and arg0_t.fallback.type is info:
+        return arg0_t
+    elif isinstance(arg0_t, UnionType):
+
+        def is_memeber(given_t: ProperType) -> bool:
+            return (isinstance(given_t, Instance) and given_t.type is info) or (
+                isinstance(given_t, LiteralType) and given_t.fallback.type is info
+            )
+
+        items = [get_proper_type(it) for it in arg0_t.items]
+        if items and all(is_memeber(item) for item in items):
+            return arg0_t
+    elif (isinstance(arg0_t, TypeVarType)) and isinstance(
+        get_proper_type(arg0_t.upper_bound), Instance
+    ):
+        if arg0_t.upper_bound.type is info:
+            return arg0_t
+
+    return base_ret