# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html # For details: https://github.com/pylint-dev/astroid/blob/main/LICENSE # Copyright (c) https://github.com/pylint-dev/astroid/blob/main/CONTRIBUTORS.txt from __future__ import annotations import unittest import pytest import astroid from astroid import bases, builder, nodes, objects from astroid.exceptions import InferenceError class EnumBrainTest(unittest.TestCase): def test_simple_enum(self) -> None: module = builder.parse( """ import enum class MyEnum(enum.Enum): one = "one" two = "two" def mymethod(self, x): return 5 """ ) enumeration = next(module["MyEnum"].infer()) one = enumeration["one"] self.assertEqual(one.pytype(), ".MyEnum.one") for propname in ("name", "value"): prop = next(iter(one.getattr(propname))) self.assertIn("builtins.property", prop.decoratornames()) meth = one.getattr("mymethod")[0] self.assertIsInstance(meth, astroid.FunctionDef) def test_looks_like_enum_false_positive(self) -> None: # Test that a class named Enumeration is not considered a builtin enum. module = builder.parse( """ class Enumeration(object): def __init__(self, name, enum_list): pass test = 42 """ ) enumeration = module["Enumeration"] test = next(enumeration.igetattr("test")) self.assertEqual(test.value, 42) def test_user_enum_false_positive(self) -> None: # Test that a user-defined class named Enum is not considered a builtin enum. ast_node = astroid.extract_node( """ class Enum: pass class Color(Enum): red = 1 Color.red #@ """ ) assert isinstance(ast_node, nodes.NodeNG) inferred = ast_node.inferred() self.assertEqual(len(inferred), 1) self.assertIsInstance(inferred[0], astroid.Const) self.assertEqual(inferred[0].value, 1) def test_ignores_with_nodes_from_body_of_enum(self) -> None: code = """ import enum class Error(enum.Enum): Foo = "foo" Bar = "bar" with "error" as err: pass """ node = builder.extract_node(code) inferred = next(node.infer()) assert "err" in inferred.locals assert len(inferred.locals["err"]) == 1 def test_enum_multiple_base_classes(self) -> None: module = builder.parse( """ import enum class Mixin: pass class MyEnum(Mixin, enum.Enum): one = 1 """ ) enumeration = next(module["MyEnum"].infer()) one = enumeration["one"] clazz = one.getattr("__class__")[0] self.assertTrue( clazz.is_subtype_of(".Mixin"), "Enum instance should share base classes with generating class", ) def test_int_enum(self) -> None: module = builder.parse( """ import enum class MyEnum(enum.IntEnum): one = 1 """ ) enumeration = next(module["MyEnum"].infer()) one = enumeration["one"] clazz = one.getattr("__class__")[0] self.assertTrue( clazz.is_subtype_of("builtins.int"), "IntEnum based enums should be a subtype of int", ) def test_enum_func_form_is_class_not_instance(self) -> None: cls, instance = builder.extract_node( """ from enum import Enum f = Enum('Audience', ['a', 'b', 'c']) f #@ f(1) #@ """ ) inferred_cls = next(cls.infer()) self.assertIsInstance(inferred_cls, bases.Instance) inferred_instance = next(instance.infer()) self.assertIsInstance(inferred_instance, bases.Instance) self.assertIsInstance(next(inferred_instance.igetattr("name")), nodes.Const) self.assertIsInstance(next(inferred_instance.igetattr("value")), nodes.Const) def test_enum_func_form_iterable(self) -> None: instance = builder.extract_node( """ from enum import Enum Animal = Enum('Animal', 'ant bee cat dog') Animal """ ) inferred = next(instance.infer()) self.assertIsInstance(inferred, astroid.Instance) self.assertTrue(inferred.getattr("__iter__")) def test_enum_func_form_subscriptable(self) -> None: instance, name = builder.extract_node( """ from enum import Enum Animal = Enum('Animal', 'ant bee cat dog') Animal['ant'] #@ Animal['ant'].name #@ """ ) instance = next(instance.infer()) self.assertIsInstance(instance, astroid.Instance) inferred = next(name.infer()) self.assertIsInstance(inferred, astroid.Const) def test_enum_func_form_has_dunder_members(self) -> None: instance = builder.extract_node( """ from enum import Enum Animal = Enum('Animal', 'ant bee cat dog') for i in Animal.__members__: i #@ """ ) instance = next(instance.infer()) self.assertIsInstance(instance, astroid.Const) self.assertIsInstance(instance.value, str) def test_infer_enum_value_as_the_right_type(self) -> None: string_value, int_value = builder.extract_node( """ from enum import Enum class A(Enum): a = 'a' b = 1 A.a.value #@ A.b.value #@ """ ) inferred_string = string_value.inferred() assert any( isinstance(elem, astroid.Const) and elem.value == "a" for elem in inferred_string ) inferred_int = int_value.inferred() assert any( isinstance(elem, astroid.Const) and elem.value == 1 for elem in inferred_int ) def test_mingled_single_and_double_quotes_does_not_crash(self) -> None: node = builder.extract_node( """ from enum import Enum class A(Enum): a = 'x"y"' A.a.value #@ """ ) inferred_string = next(node.infer()) assert inferred_string.value == 'x"y"' def test_special_characters_does_not_crash(self) -> None: node = builder.extract_node( """ import enum class Example(enum.Enum): NULL = '\\N{NULL}' Example.NULL.value """ ) inferred_string = next(node.infer()) assert inferred_string.value == "\N{NULL}" def test_dont_crash_on_for_loops_in_body(self) -> None: node = builder.extract_node( """ class Commands(IntEnum): _ignore_ = 'Commands index' _init_ = 'value string' BEL = 0x07, 'Bell' Commands = vars() for index in range(4): Commands[f'DC{index + 1}'] = 0x11 + index, f'Device Control {index + 1}' Commands """ ) inferred = next(node.infer()) assert isinstance(inferred, astroid.ClassDef) def test_enum_tuple_list_values(self) -> None: tuple_node, list_node = builder.extract_node( """ import enum class MyEnum(enum.Enum): a = (1, 2) b = [2, 4] MyEnum.a.value #@ MyEnum.b.value #@ """ ) inferred_tuple_node = next(tuple_node.infer()) inferred_list_node = next(list_node.infer()) assert isinstance(inferred_tuple_node, astroid.Tuple) assert isinstance(inferred_list_node, astroid.List) assert inferred_tuple_node.as_string() == "(1, 2)" assert inferred_list_node.as_string() == "[2, 4]" def test_enum_starred_is_skipped(self) -> None: code = """ from enum import Enum class ContentType(Enum): TEXT, PHOTO, VIDEO, GIF, YOUTUBE, *_ = [1, 2, 3, 4, 5, 6] ContentType.TEXT #@ """ node = astroid.extract_node(code) next(node.infer()) def test_enum_name_is_str_on_self(self) -> None: 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()) next(c_value.infer()) def test_enum_name_and_value_members_override_dynamicclassattr(self) -> None: 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" def test_enum_subclass_member_name(self) -> None: ast_node = astroid.extract_node( """ from enum import Enum class EnumSubclass(Enum): pass class Color(EnumSubclass): red = 1 Color.red.name #@ """ ) assert isinstance(ast_node, nodes.NodeNG) inferred = ast_node.inferred() self.assertEqual(len(inferred), 1) self.assertIsInstance(inferred[0], astroid.Const) self.assertEqual(inferred[0].value, "red") def test_enum_subclass_member_value(self) -> None: ast_node = astroid.extract_node( """ from enum import Enum class EnumSubclass(Enum): pass class Color(EnumSubclass): red = 1 Color.red.value #@ """ ) assert isinstance(ast_node, nodes.NodeNG) inferred = ast_node.inferred() self.assertEqual(len(inferred), 1) self.assertIsInstance(inferred[0], astroid.Const) self.assertEqual(inferred[0].value, 1) def test_enum_subclass_member_method(self) -> None: # See Pylint issue #2626 ast_node = astroid.extract_node( """ from enum import Enum class EnumSubclass(Enum): def hello_pylint(self) -> str: return self.name class Color(EnumSubclass): red = 1 Color.red.hello_pylint() #@ """ ) assert isinstance(ast_node, nodes.NodeNG) inferred = ast_node.inferred() self.assertEqual(len(inferred), 1) self.assertIsInstance(inferred[0], astroid.Const) self.assertEqual(inferred[0].value, "red") def test_enum_subclass_different_modules(self) -> None: # See Pylint issue #2626 astroid.extract_node( """ from enum import Enum class EnumSubclass(Enum): pass """, "a", ) ast_node = astroid.extract_node( """ from a import EnumSubclass class Color(EnumSubclass): red = 1 Color.red.value #@ """ ) assert isinstance(ast_node, nodes.NodeNG) inferred = ast_node.inferred() self.assertEqual(len(inferred), 1) self.assertIsInstance(inferred[0], astroid.Const) self.assertEqual(inferred[0].value, 1) def test_members_member_ignored(self) -> None: ast_node = builder.extract_node( """ from enum import Enum class Animal(Enum): a = 1 __members__ = {} Animal.__members__ #@ """ ) inferred = next(ast_node.infer()) self.assertIsInstance(inferred, astroid.Dict) self.assertTrue(inferred.locals) def test_enum_as_renamed_import(self) -> None: """Originally reported in https://github.com/pylint-dev/pylint/issues/5776.""" ast_node: nodes.Attribute = builder.extract_node( """ from enum import Enum as PyEnum class MyEnum(PyEnum): ENUM_KEY = "enum_value" MyEnum.ENUM_KEY """ ) inferred = next(ast_node.infer()) assert isinstance(inferred, bases.Instance) assert inferred._proxied.name == "ENUM_KEY" def test_class_named_enum(self) -> None: """Test that the user-defined class named `Enum` is not inferred as `enum.Enum`""" astroid.extract_node( """ class Enum: def __init__(self, one, two): self.one = one self.two = two def pear(self): ... """, "module_with_class_named_enum", ) attribute_nodes = astroid.extract_node( """ import module_with_class_named_enum module_with_class_named_enum.Enum("apple", "orange") #@ typo_module_with_class_named_enum.Enum("apple", "orange") #@ """ ) name_nodes = astroid.extract_node( """ from module_with_class_named_enum import Enum Enum("apple", "orange") #@ TypoEnum("apple", "orange") #@ """ ) # Test that both of the successfully inferred `Name` & `Attribute` # nodes refer to the user-defined Enum class. for inferred in (attribute_nodes[0].inferred()[0], name_nodes[0].inferred()[0]): assert isinstance(inferred, astroid.Instance) assert inferred.name == "Enum" assert inferred.qname() == "module_with_class_named_enum.Enum" assert "pear" in inferred.locals # Test that an `InferenceError` is raised when an attempt is made to # infer a `Name` or `Attribute` node & they cannot be found. for node in (attribute_nodes[1], name_nodes[1]): with pytest.raises(InferenceError): node.inferred()