Skip to content

Commit

Permalink
Fix false positive with inference of type-annotated Enum classes.
Browse files Browse the repository at this point in the history
  • Loading branch information
mbyrnepr2 committed Aug 7, 2022
1 parent 74ad5b7 commit e263425
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 1 deletion.
4 changes: 4 additions & 0 deletions ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ Release date: TBA

Closes #1717

* Fix false positive with inference of type-annotated Enum classes.

Refs PyCQA/pylint#7265

What's New in astroid 2.12.2?
=============================
Release date: 2022-07-12
Expand Down
2 changes: 1 addition & 1 deletion astroid/brain/brain_namedtuple_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -379,7 +379,7 @@ def infer_enum_class(node: nodes.ClassDef) -> nodes.ClassDef:
continue

inferred_return_value = None
if isinstance(stmt, nodes.Assign):
if isinstance(stmt, (nodes.AnnAssign, nodes.Assign)):
if isinstance(stmt.value, nodes.Const):
if isinstance(stmt.value.value, str):
inferred_return_value = repr(stmt.value.value)
Expand Down
72 changes: 72 additions & 0 deletions tests/unittest_scoped_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -2495,6 +2495,78 @@ class Sounds(Enum):
assert sorted(actual) == ["bee", "cat"]


def test_enums_type_annotation_str_member() -> None:
"""A type-annotated member of an Enum class where:
- `member.value` is of type `nodes.Const` &
- `member.value.value` is of type `str`
is inferred as: `repr(member.value.value)`
"""
node = builder.extract_node(
"""
from enum import Enum
class Veg(Enum):
TOMATO: str = "sweet"
Veg.TOMATO.value
"""
)
inferred_member_value = node.inferred()[0]
assert isinstance(inferred_member_value, nodes.Const)
assert inferred_member_value.value == "sweet"


@pytest.mark.parametrize("annotation, value", [("int", 42), ("bytes", b"")])
def test_enums_type_annotation_non_str_member(annotation, value) -> None:
"""A type-annotated member of an Enum class where:
- `member.value` is of type `nodes.Const` &
- `member.value.value` is not of type `str`
is inferred as: `member.value.value`
"""

node = builder.extract_node(
f"""
from enum import Enum
class Veg(Enum):
TOMATO: {annotation} = {value}
Veg.TOMATO.value
"""
)
inferred_member_value = node.inferred()[0]
assert isinstance(inferred_member_value, nodes.Const)
assert inferred_member_value.value == value


@pytest.mark.parametrize(
"annotation, value",
[
("dict", {"variety": "beefeater"}),
("list", ["beefeater", "moneymaker"]),
("TypedDict", {"variety": "moneymaker"}),
],
)
def test_enums_type_annotations_non_const_member(annotation, value) -> None:
"""A type-annotated member of an Enum class where:
- `member.value` is not of type `nodes.Const`
is inferred as: `member.value.as_string()`.
"""

member = builder.extract_node(
f"""
from enum import Enum
class Veg(Enum):
TOMATO: {annotation} = {value}
Veg.TOMATO.value
"""
)

inferred_member_value = member.inferred()[0]
assert not isinstance(inferred_member_value, nodes.Const)
assert inferred_member_value.as_string() == repr(value)


def test_metaclass_cannot_infer_call_yields_an_instance() -> None:
node = builder.extract_node(
"""
Expand Down

0 comments on commit e263425

Please sign in to comment.