Skip to content

Commit 1321913

Browse files
committed
Add fake "name" property to enum.Enum subclasses
Ref pylint-dev/pylint#1932. Ref pylint-dev/pylint#2062. The enum.Enum class itself defines two @DynamicClassAttribute data-descriptors "name" and "value" which behave differently when looked up on an instance or on the class. 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 we should assume that "self.name" is the string name of some enum member, unless the enum itself defines a "name" member.
1 parent 0245bd9 commit 1321913

File tree

2 files changed

+79
-1
lines changed

2 files changed

+79
-1
lines changed

astroid/brain/brain_namedtuple_enum.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,7 @@ def infer_enum_class(node):
362362
# Skip if the class is directly from enum module.
363363
break
364364
dunder_members = {}
365+
target_names = set()
365366
for local, values in node.locals.items():
366367
if any(not isinstance(value, nodes.AssignName) for value in values):
367368
continue
@@ -391,6 +392,7 @@ def infer_enum_class(node):
391392
for target in targets:
392393
if isinstance(target, nodes.Starred):
393394
continue
395+
target_names.add(target.name)
394396
# Replace all the assignments with our mocked class.
395397
classdef = dedent(
396398
"""
@@ -429,6 +431,27 @@ def name(self):
429431
]
430432
)
431433
node.locals["__members__"] = [members]
434+
# The enum.Enum class itself defines two @DynamicClassAttribute data-descriptors
435+
# "name" and "value" (which we override in the mocked class for each enum member
436+
# above). When dealing with inference of an arbitrary instance of the enum
437+
# class, e.g. in a method defined in the class body like:
438+
# class SomeEnum(enum.Enum):
439+
# def method(self):
440+
# self.name # <- here
441+
# In the absence of an enum member called "name" or "value", these attributes
442+
# should resolve to the descriptor on that particular instance, i.e. enum member.
443+
# For "value", we have no idea what that should be, but for "name", we at least
444+
# know that it should be a string, so infer that as a guess.
445+
if "name" not in target_names:
446+
code = dedent(
447+
"""
448+
@property
449+
def name(self):
450+
return ''
451+
"""
452+
)
453+
name_dynamicclassattr = AstroidBuilder(MANAGER).string_build(code)["name"]
454+
node.locals["name"] = [name_dynamicclassattr]
432455
break
433456
return node
434457

tests/unittest_brain.py

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
import pytest
4444

4545
import astroid
46-
from astroid import MANAGER, bases, builder, nodes, test_utils, util
46+
from astroid import MANAGER, bases, builder, nodes, objects, test_utils, util
4747

4848
try:
4949
import multiprocessing # pylint: disable=unused-import
@@ -993,6 +993,61 @@ class ContentType(Enum):
993993
node = astroid.extract_node(code)
994994
next(node.infer())
995995

996+
def test_enum_name_is_str_on_self(self):
997+
code = """
998+
from enum import Enum
999+
class TestEnum(Enum):
1000+
def func(self):
1001+
self.name #@
1002+
self.value #@
1003+
TestEnum.name #@
1004+
TestEnum.value #@
1005+
"""
1006+
i_name, i_value, c_name, c_value = astroid.extract_node(code)
1007+
1008+
# <instance>.name should be a string, <class>.name should be a property (that
1009+
# forwards the lookup to __getattr__)
1010+
inferred = next(i_name.infer())
1011+
assert isinstance(inferred, nodes.Const)
1012+
assert inferred.pytype() == "builtins.str"
1013+
inferred = next(c_name.infer())
1014+
assert isinstance(inferred, objects.Property)
1015+
1016+
# Inferring .value should not raise InferenceError. It is probably Uninferable
1017+
# but we don't particularly care
1018+
next(i_value.infer())
1019+
inferred = next(c_value.infer())
1020+
assert isinstance(inferred, objects.Property)
1021+
1022+
def test_enum_name_and_value_members_override_dynamicclassattr(self):
1023+
code = """
1024+
from enum import Enum
1025+
class TrickyEnum(Enum):
1026+
name = 1
1027+
value = 2
1028+
1029+
def func(self):
1030+
self.name #@
1031+
self.value #@
1032+
TrickyEnum.name #@
1033+
TrickyEnum.value #@
1034+
"""
1035+
i_name, i_value, c_name, c_value = astroid.extract_node(code)
1036+
1037+
# All of these cases should be inferred as enum members
1038+
inferred = next(i_name.infer())
1039+
assert isinstance(inferred, bases.Instance)
1040+
assert inferred.pytype() == ".TrickyEnum.name"
1041+
inferred = next(c_name.infer())
1042+
assert isinstance(inferred, bases.Instance)
1043+
assert inferred.pytype() == ".TrickyEnum.name"
1044+
inferred = next(i_value.infer())
1045+
assert isinstance(inferred, bases.Instance)
1046+
assert inferred.pytype() == ".TrickyEnum.value"
1047+
inferred = next(c_value.infer())
1048+
assert isinstance(inferred, bases.Instance)
1049+
assert inferred.pytype() == ".TrickyEnum.value"
1050+
9961051

9971052
@unittest.skipUnless(HAS_DATEUTIL, "This test requires the dateutil library.")
9981053
class DateutilBrainTest(unittest.TestCase):

0 commit comments

Comments
 (0)