Skip to content

Commit

Permalink
Add fake "name" property to enum.Enum subclasses
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
nelfin committed Jun 10, 2021
1 parent 0245bd9 commit 1321913
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 1 deletion.
23 changes: 23 additions & 0 deletions astroid/brain/brain_namedtuple_enum.py
Expand Up @@ -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
Expand Down Expand Up @@ -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(
"""
Expand Down Expand Up @@ -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

Expand Down
57 changes: 56 additions & 1 deletion tests/unittest_brain.py
Expand Up @@ -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
Expand Down Expand Up @@ -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)

# <instance>.name should be a string, <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"
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):
Expand Down

0 comments on commit 1321913

Please sign in to comment.