Infer members of Enums as instances of the Enum they belong to

Author: DanielNoord

ChangeLog

@@ -44,6 +44,11 @@ Release date: TBA
 
   Closes PyCQA/pylint#5776
 
+* Members of ``Enums`` are now correctly inferred as instances of the ``Enum`` they belong
+  to instead of instances of their own class.
+
+  Closes #744
+
 * Rename ``ModuleSpec`` -> ``module_type`` constructor parameter to match attribute
   name and improve typing. Use ``type`` instead.
 

astroid/bases.py

@@ -181,10 +181,26 @@ class BaseInstance(Proxy):
 
     special_attributes = None
 
+    def __init__(self, proxied=None):
+        self._explicit_instance_attrs: dict[str, list[nodes.NodeNG]] = {}
+        """Attributes that have been explicitly set during initialization
+        of the specific instance.
+
+        This dictionary can be used to differentiate between attributes assosciated to
+        the proxy and attributes that are specific to the instantiated instance.
+        """
+        super().__init__(proxied)
+
     def display_type(self):
         return "Instance of"
 
     def getattr(self, name, context=None, lookupclass=True):
+        # See if the attribute is set explicitly for this instance
+        try:
+            return self._explicit_instance_attrs[name]
+        except KeyError:
+            pass
+
         try:
             values = self._proxied.instance_attr(name, context)
         except AttributeInferenceError as exc:

astroid/brain/brain_namedtuple_enum.py

@@ -12,7 +12,7 @@
 from textwrap import dedent
 
 import astroid
-from astroid import arguments, inference_tip, nodes, util
+from astroid import arguments, bases, inference_tip, nodes, util
 from astroid.builder import AstroidBuilder, extract_node
 from astroid.context import InferenceContext
 from astroid.exceptions import (
@@ -351,106 +351,55 @@ def __mul__(self, other):
 
 def infer_enum_class(node: nodes.ClassDef) -> nodes.ClassDef:
     """Specific inference for enums."""
-    for basename in (b for cls in node.mro() for b in cls.basenames):
-        if node.root().name == "enum":
-            # Skip if the class is directly from enum module.
-            break
-        dunder_members = {}
-        target_names = set()
-        for local, values in node.locals.items():
-            if any(not isinstance(value, nodes.AssignName) for value in values):
-                continue
-
-            stmt = values[0].statement(future=True)
-            if isinstance(stmt, nodes.Assign):
-                if isinstance(stmt.targets[0], nodes.Tuple):
-                    targets = stmt.targets[0].itered()
-                else:
-                    targets = stmt.targets
-            elif isinstance(stmt, nodes.AnnAssign):
-                targets = [stmt.target]
+    if node.root().name == "enum":
+        # Skip if the class is directly from enum module.
+        return node
+    dunder_members: dict[str, bases.Instance] = {}
+    for local, values in node.locals.items():
+        if any(not isinstance(value, nodes.AssignName) for value in values):
+            continue
+
+        stmt = values[0].statement(future=True)
+        if isinstance(stmt, nodes.Assign):
+            if isinstance(stmt.targets[0], nodes.Tuple):
+                targets: list[nodes.NodeNG] = stmt.targets[0].itered()
             else:
+                targets = stmt.targets
+            value_node = stmt.value
+        elif isinstance(stmt, nodes.AnnAssign):
+            targets = [stmt.target]  # type: ignore[list-item] # .target shouldn't be None
+            value_node = stmt.value
+        else:
+            continue
+
+        new_targets: list[bases.Instance] = []
+        for target in targets:
+            if isinstance(target, nodes.Starred):
                 continue
 
-            inferred_return_value = None
-            if isinstance(stmt, nodes.Assign):
-                if isinstance(stmt.value, nodes.Const):
-                    if isinstance(stmt.value.value, str):
-                        inferred_return_value = repr(stmt.value.value)
-                    else:
-                        inferred_return_value = stmt.value.value
-                else:
-                    inferred_return_value = stmt.value.as_string()
-
-            new_targets = []
-            for target in targets:
-                if isinstance(target, nodes.Starred):
-                    continue
-                target_names.add(target.name)
-                # Replace all the assignments with our mocked class.
-                classdef = dedent(
-                    """
-                class {name}({types}):
-                    @property
-                    def value(self):
-                        return {return_value}
-                    @property
-                    def name(self):
-                        return "{name}"
-                """.format(
-                        name=target.name,
-                        types=", ".join(node.basenames),
-                        return_value=inferred_return_value,
-                    )
-                )
-                if "IntFlag" in basename:
-                    # Alright, we need to add some additional methods.
-                    # Unfortunately we still can't infer the resulting objects as
-                    # Enum members, but once we'll be able to do that, the following
-                    # should result in some nice symbolic execution
-                    classdef += INT_FLAG_ADDITION_METHODS.format(name=target.name)
-
-                fake = AstroidBuilder(
-                    AstroidManager(), apply_transforms=False
-                ).string_build(classdef)[target.name]
-                fake.parent = target.parent
-                for method in node.mymethods():
-                    fake.locals[method.name] = [method]
-                new_targets.append(fake.instantiate_class())
-                dunder_members[local] = fake
-            node.locals[local] = new_targets
-        members = nodes.Dict(parent=node)
-        members.postinit(
-            [
-                (nodes.Const(k, parent=members), nodes.Name(v.name, parent=members))
-                for k, v in dunder_members.items()
-            ]
-        )
-        node.locals["__members__"] = [members]
-        # The enum.Enum class itself defines two @DynamicClassAttribute data-descriptors
-        # "name" and "value" (which we override in the mocked class for each enum member
-        # above). When dealing with inference of an arbitrary instance of the enum
-        # class, e.g. in a method defined in the class body like:
-        #     class SomeEnum(enum.Enum):
-        #         def method(self):
-        #             self.name  # <- here
-        # In the absence of an enum member called "name" or "value", these attributes
-        # should resolve to the descriptor on that particular instance, i.e. enum member.
-        # For "value", we have no idea what that should be, but for "name", we at least
-        # know that it should be a string, so infer that as a guess.
-        if "name" not in target_names:
-            code = dedent(
-                """
-            @property
-            def name(self):
-                return ''
-            """
-            )
-            name_dynamicclassattr = AstroidBuilder(AstroidManager()).string_build(code)[
-                "name"
+            # Instantiate a class of the Enum with the value and name
+            # attributes set to the values of the assignment
+            # See: https://docs.python.org/3/library/enum.html#creating-an-enum
+            target_node = node.instantiate_class()
+            target_node._explicit_instance_attrs["value"] = [value_node]
+            target_node._explicit_instance_attrs["name"] = [
+                nodes.const_factory(target.name)
             ]
-            node.locals["name"] = [name_dynamicclassattr]
-        break
+
+            new_targets.append(target_node)
+            dunder_members[local] = target_node
+
+        node.locals[local] = new_targets
+
+    # Creation of the __members__ attribute of the Enum node
+    members = nodes.Dict(parent=node)
+    members.postinit(
+        [
+            (nodes.Const(k, parent=members), nodes.Name(v.name, parent=members))
+            for k, v in dunder_members.items()
+        ]
+    )
+    node.locals["__members__"] = [members]
     return node
 
 

astroid/nodes/node_classes.py

@@ -304,6 +304,8 @@ def __init__(
             parent=parent,
         )
 
+        Instance.__init__(self)
+
     def postinit(self, elts: list[NodeNG]) -> None:
         """Do some setup after initialisation.
 
@@ -1928,6 +1930,8 @@ def __init__(
             parent=parent,
         )
 
+        Instance.__init__(self)
+
     def __getattr__(self, name):
         # This is needed because of Proxy's __getattr__ method.
         # Calling object.__new__ on this class without calling
@@ -2277,6 +2281,7 @@ def __init__(
             end_col_offset=end_col_offset,
             parent=parent,
         )
+        Instance.__init__(self)
 
     def postinit(self, items: list[tuple[NodeNG, NodeNG]]) -> None:
         """Do some setup after initialisation.

tests/test_brain_ssl.py

@@ -40,4 +40,8 @@ def test_ssl_brain() -> None:
     # TLSVersion is inferred from the main module, not from the brain
     inferred_cert_required = next(module.body[4].value.infer())
     assert isinstance(inferred_cert_required, bases.Instance)
-    assert inferred_cert_required._proxied.name == "CERT_REQUIRED"
+    assert inferred_cert_required._proxied.name == "VerifyMode"
+
+    value_node = inferred_cert_required.getattr("value")[0]
+    assert isinstance(value_node, nodes.Const)
+    assert value_node.value == 2

tests/unittest_brain.py

@@ -751,11 +751,13 @@ def mymethod(self, x):
 
         enumeration = next(module["MyEnum"].infer())
         one = enumeration["one"]
-        self.assertEqual(one.pytype(), ".MyEnum.one")
+        self.assertEqual(one.pytype(), ".MyEnum")
 
         for propname in ("name", "value"):
-            prop = next(iter(one.getattr(propname)))
-            self.assertIn("builtins.property", prop.decoratornames())
+            # On the base Enum class 'name' and 'value' are properties
+            # decorated by DynamicClassAttribute.
+            prop = next(iter(one._proxied.getattr(propname)))
+            self.assertIn("types.DynamicClassAttribute", prop.decoratornames())
 
         meth = one.getattr("mymethod")[0]
         self.assertIsInstance(meth, astroid.FunctionDef)
@@ -1010,18 +1012,17 @@ def func(self):
         """
         i_name, i_value, c_name, c_value = astroid.extract_node(code)
 
-        # <instance>.name should be a string, <class>.name should be a property (that
+        # <instance>.name should be Uninferable, <class>.name should be a property (that
         # forwards the lookup to __getattr__)
         inferred = next(i_name.infer())
-        assert isinstance(inferred, nodes.Const)
-        assert inferred.pytype() == "builtins.str"
+        assert inferred is util.Uninferable
         inferred = next(c_name.infer())
         assert isinstance(inferred, objects.Property)
 
-        # Inferring .value should not raise InferenceError. It is probably Uninferable
-        # but we don't particularly care
-        next(i_value.infer())
-        next(c_value.infer())
+        inferred = next(i_value.infer())
+        assert inferred is util.Uninferable
+        inferred = next(c_value.infer())
+        assert isinstance(inferred, objects.Property)
 
     def test_enum_name_and_value_members_override_dynamicclassattr(self) -> None:
         code = """
@@ -1038,19 +1039,23 @@ def func(self):
         """
         i_name, i_value, c_name, c_value = astroid.extract_node(code)
 
-        # All of these cases should be inferred as enum members
-        inferred = next(i_name.infer())
-        assert isinstance(inferred, bases.Instance)
-        assert inferred.pytype() == ".TrickyEnum.name"
-        inferred = next(c_name.infer())
-        assert isinstance(inferred, bases.Instance)
-        assert inferred.pytype() == ".TrickyEnum.name"
-        inferred = next(i_value.infer())
-        assert isinstance(inferred, bases.Instance)
-        assert inferred.pytype() == ".TrickyEnum.value"
-        inferred = next(c_value.infer())
-        assert isinstance(inferred, bases.Instance)
-        assert inferred.pytype() == ".TrickyEnum.value"
+        # All of these cases should be inferred as enum instances
+        # and refer to the same instance
+        name_inner = next(i_name.infer())
+        assert isinstance(name_inner, bases.Instance)
+        assert name_inner.pytype() == ".TrickyEnum"
+        name_outer = next(c_name.infer())
+        assert isinstance(name_outer, bases.Instance)
+        assert name_outer.pytype() == ".TrickyEnum"
+        assert name_inner == name_outer
+
+        value_inner = next(i_value.infer())
+        assert isinstance(value_inner, bases.Instance)
+        assert value_inner.pytype() == ".TrickyEnum"
+        value_outer = next(c_value.infer())
+        assert isinstance(value_outer, bases.Instance)
+        assert value_outer.pytype() == ".TrickyEnum"
+        assert value_inner == value_outer
 
     def test_enum_subclass_member_name(self) -> None:
         ast_node = astroid.extract_node(
@@ -1168,7 +1173,15 @@ class MyEnum(PyEnum):
         )
         inferred = next(ast_node.infer())
         assert isinstance(inferred, bases.Instance)
-        assert inferred._proxied.name == "ENUM_KEY"
+        assert inferred._proxied.name == "MyEnum"
+
+        name_node = inferred.getattr("name")[0]
+        assert isinstance(name_node, nodes.Const)
+        assert name_node.value == "ENUM_KEY"
+
+        value_node = inferred.getattr("value")[0]
+        assert isinstance(value_node, nodes.Const)
+        assert value_node.value == "enum_value"
 
 
 @unittest.skipUnless(HAS_DATEUTIL, "This test requires the dateutil library.")