Skip to content

Commit 347ea0c

Browse files
authored
Fix inference of Enums when they are imported under an alias (#1587)
1 parent 3abd45e commit 347ea0c

File tree

3 files changed

+19
-3
lines changed

3 files changed

+19
-3
lines changed

ChangeLog

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ Release date: TBA
3636
* On Python versions >= 3.9, ``astroid`` now understands subscripting
3737
builtin classes such as ``enumerate`` or ``staticmethod``.
3838

39+
* Fixed inference of ``Enums`` when they are imported under an alias.
40+
41+
Closes PyCQA/pylint#5776
42+
3943
* Rename ``ModuleSpec`` -> ``module_type`` constructor parameter to match attribute
4044
name and improve typing. Use ``type`` instead.
4145

astroid/brain/brain_namedtuple_enum.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -349,11 +349,9 @@ def __mul__(self, other):
349349
"""
350350

351351

352-
def infer_enum_class(node):
352+
def infer_enum_class(node: nodes.ClassDef) -> nodes.ClassDef:
353353
"""Specific inference for enums."""
354354
for basename in (b for cls in node.mro() for b in cls.basenames):
355-
if basename not in ENUM_BASE_NAMES:
356-
continue
357355
if node.root().name == "enum":
358356
# Skip if the class is directly from enum module.
359357
break

tests/unittest_brain.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1156,6 +1156,20 @@ class Animal(Enum):
11561156
self.assertIsInstance(inferred, astroid.Dict)
11571157
self.assertTrue(inferred.locals)
11581158

1159+
def test_enum_as_renamed_import(self) -> None:
1160+
"""Originally reported in https://github.com/PyCQA/pylint/issues/5776."""
1161+
ast_node: nodes.Attribute = builder.extract_node(
1162+
"""
1163+
from enum import Enum as PyEnum
1164+
class MyEnum(PyEnum):
1165+
ENUM_KEY = "enum_value"
1166+
MyEnum.ENUM_KEY
1167+
"""
1168+
)
1169+
inferred = next(ast_node.infer())
1170+
assert isinstance(inferred, bases.Instance)
1171+
assert inferred._proxied.name == "ENUM_KEY"
1172+
11591173

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

0 commit comments

Comments
 (0)