diff options
| -rw-r--r-- | ChangeLog | 3 | ||||
| -rw-r--r-- | astroid/brain/py2stdlib.py | 4 | ||||
| -rw-r--r-- | astroid/tests/unittest_brain.py | 33 |
3 files changed, 38 insertions, 2 deletions
@@ -2,6 +2,9 @@ Change log for the astroid package (used to be astng) ===================================================== -- + * Generated enum member stubs now support IntEnum and multiple + base classes. + * astroid.builder.AstroidBuilder.string_build and astroid.builder.AstroidBuilder.file_build are now raising AstroidBuildingException when the parsing of the string raises diff --git a/astroid/brain/py2stdlib.py b/astroid/brain/py2stdlib.py index e63ea86..4231fe7 100644 --- a/astroid/brain/py2stdlib.py +++ b/astroid/brain/py2stdlib.py @@ -304,7 +304,7 @@ def infer_enum_class(node): for target in targets: # Replace all the assignments with our mocked class. classdef = dedent(''' - class %(name)s(object): + class %(name)s(%(types)s): @property def value(self): # Not the best return. @@ -312,7 +312,7 @@ def infer_enum_class(node): @property def name(self): return %(name)r - ''' % {'name': target.name}) + ''' % {'name': target.name, 'types': ', '.join(node.basenames)}) fake = AstroidBuilder(MANAGER).string_build(classdef)[target.name] fake.parent = target.parent for method in node.mymethods(): diff --git a/astroid/tests/unittest_brain.py b/astroid/tests/unittest_brain.py index 91a9a79..b0e0608 100644 --- a/astroid/tests/unittest_brain.py +++ b/astroid/tests/unittest_brain.py @@ -252,5 +252,38 @@ class EnumBrainTest(unittest.TestCase): meth = one.getattr('mymethod')[0] self.assertIsInstance(meth, astroid.Function) + def test_enum_multiple_base_classes(self): + module = AstroidBuilder().string_build(dedent(""" + import enum + + class Mixin: + pass + + class MyEnum(Mixin, enum.Enum): + one = 1 + """)) + enum = next(module['MyEnum'].infer()) + one = enum['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): + module = AstroidBuilder().string_build(dedent(""" + import enum + + class MyEnum(enum.IntEnum): + one = 1 + """)) + + enum = next(module['MyEnum'].infer()) + one = enum['one'] + + clazz = one.getattr('__class__')[0] + int_type = '{}.{}'.format(bases.BUILTINS, 'int') + self.assertTrue(clazz.is_subtype_of(int_type), + 'IntEnum based enums should be a subtype of int') + if __name__ == '__main__': unittest.main() |