From daf4197be12aaf8faf87931e0b9722a6f150d7b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20van=20Noord?= <13665637+DanielNoord@users.noreply.github.com> Date: Wed, 6 Jul 2022 17:01:19 +0200 Subject: [PATCH] Infer members of ``Enums`` as instances of the ``Enum`` they belong to --- ChangeLog | 5 + astroid/bases.py | 18 ++++ astroid/brain/brain_namedtuple_enum.py | 139 ++++++++----------------- astroid/nodes/node_classes.py | 3 + tests/test_brain_ssl.py | 6 +- tests/unittest_brain.py | 67 +++++++----- 6 files changed, 117 insertions(+), 121 deletions(-) diff --git a/ChangeLog b/ChangeLog index 4d3ac2a4d4..9355fc7e77 100644 --- a/ChangeLog +++ b/ChangeLog @@ -67,6 +67,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. diff --git a/astroid/bases.py b/astroid/bases.py index a154838c9c..4794d44b25 100644 --- a/astroid/bases.py +++ b/astroid/bases.py @@ -200,10 +200,28 @@ class BaseInstance(Proxy): special_attributes = None + _proxied: nodes.ClassDef + + def __init__(self, proxied: nodes.ClassDef | None = None) -> 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: diff --git a/astroid/brain/brain_namedtuple_enum.py b/astroid/brain/brain_namedtuple_enum.py index 955f0c1667..24cf15cfc8 100644 --- a/astroid/brain/brain_namedtuple_enum.py +++ b/astroid/brain/brain_namedtuple_enum.py @@ -357,106 +357,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 + 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 = stmt.targets[0].itered() - else: - targets = stmt.targets - elif isinstance(stmt, nodes.AnnAssign): - targets = [stmt.target] + 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() + # 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["__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" - ] - 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 diff --git a/astroid/nodes/node_classes.py b/astroid/nodes/node_classes.py index 8511020b63..d45be1879a 100644 --- a/astroid/nodes/node_classes.py +++ b/astroid/nodes/node_classes.py @@ -287,6 +287,8 @@ def __init__( parent=parent, ) + Instance.__init__(self, self._proxied) + def postinit(self, elts: list[NodeNG]) -> None: """Do some setup after initialisation. @@ -2269,6 +2271,7 @@ def __init__( end_col_offset=end_col_offset, parent=parent, ) + Instance.__init__(self, self._proxied) def postinit( self, items: list[tuple[SuccessfulInferenceResult, SuccessfulInferenceResult]] diff --git a/tests/test_brain_ssl.py b/tests/test_brain_ssl.py index f14efade2b..de0b17cb2b 100644 --- a/tests/test_brain_ssl.py +++ b/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 diff --git a/tests/unittest_brain.py b/tests/unittest_brain.py index 1eab27639f..67c2289066 100644 --- a/tests/unittest_brain.py +++ b/tests/unittest_brain.py @@ -18,7 +18,7 @@ import astroid from astroid import MANAGER, bases, builder, nodes, objects, test_utils, util from astroid.bases import Instance -from astroid.const import PY39_PLUS +from astroid.const import PY39_PLUS, PY311_PLUS from astroid.exceptions import ( AttributeInferenceError, InferenceError, @@ -771,11 +771,17 @@ 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 in < 3.11. + prop = next(iter(one._proxied.getattr(propname))) + if PY311_PLUS: + expected_name = "enum.property" + else: + expected_name = "types.DynamicClassAttribute" + self.assertIn(expected_name, prop.decoratornames()) meth = one.getattr("mymethod")[0] self.assertIsInstance(meth, astroid.FunctionDef) @@ -1030,18 +1036,17 @@ def func(self): """ i_name, i_value, c_name, c_value = astroid.extract_node(code) - # .name should be a string, .name should be a property (that + # .name should be Uninferable, .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 = """ @@ -1058,19 +1063,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( @@ -1188,7 +1197,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.")