diff --git a/astroid/brain/brain_namedtuple_enum.py b/astroid/brain/brain_namedtuple_enum.py index 0410c7195f..7ffd724a0e 100644 --- a/astroid/brain/brain_namedtuple_enum.py +++ b/astroid/brain/brain_namedtuple_enum.py @@ -362,6 +362,7 @@ def infer_enum_class(node): # 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 @@ -391,6 +392,7 @@ def infer_enum_class(node): 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( """ @@ -429,6 +431,27 @@ def name(self): ] ) 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(MANAGER).string_build(code)["name"] + node.locals["name"] = [name_dynamicclassattr] break return node diff --git a/tests/unittest_brain.py b/tests/unittest_brain.py index 427df6fba7..bae5136a0d 100644 --- a/tests/unittest_brain.py +++ b/tests/unittest_brain.py @@ -43,7 +43,7 @@ import pytest import astroid -from astroid import MANAGER, bases, builder, nodes, test_utils, util +from astroid import MANAGER, bases, builder, nodes, objects, test_utils, util try: import multiprocessing # pylint: disable=unused-import @@ -993,6 +993,61 @@ class ContentType(Enum): node = astroid.extract_node(code) next(node.infer()) + def test_enum_name_is_str_on_self(self): + code = """ + from enum import Enum + class TestEnum(Enum): + def func(self): + self.name #@ + self.value #@ + TestEnum.name #@ + TestEnum.value #@ + """ + i_name, i_value, c_name, c_value = astroid.extract_node(code) + + # .name should be a string, .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" + 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()) + inferred = next(c_value.infer()) + assert isinstance(inferred, objects.Property) + + def test_enum_name_and_value_members_override_dynamicclassattr(self): + code = """ + from enum import Enum + class TrickyEnum(Enum): + name = 1 + value = 2 + + def func(self): + self.name #@ + self.value #@ + TrickyEnum.name #@ + TrickyEnum.value #@ + """ + 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" + @unittest.skipUnless(HAS_DATEUTIL, "This test requires the dateutil library.") class DateutilBrainTest(unittest.TestCase):