bpo-45535: Improve output of Enum `dir()` (GH-29316)

Modify the ``EnumType.__dir__()`` and ``Enum.__dir__()`` to ensure
that user-defined methods and methods inherited from mixin classes always
show up in the output of `help()`. This change also makes it easier for
IDEs to provide auto-completion.
This commit is contained in:
Alex Waygood 2021-12-02 16:49:52 +00:00 committed by GitHub
parent cb8f491f46
commit b2afdc95cc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 386 additions and 53 deletions

View file

@ -997,11 +997,12 @@ Plain :class:`Enum` classes always evaluate as :data:`True`.
"""""""""""""""""""""""""""""
If you give your enum subclass extra methods, like the `Planet`_
class below, those methods will show up in a :func:`dir` of the member,
but not of the class::
class below, those methods will show up in a :func:`dir` of the member and the
class. Attributes defined in an :func:`__init__` method will only show up in a
:func:`dir` of the member::
>>> dir(Planet)
['EARTH', 'JUPITER', 'MARS', 'MERCURY', 'NEPTUNE', 'SATURN', 'URANUS', 'VENUS', '__class__', '__doc__', '__members__', '__module__']
['EARTH', 'JUPITER', 'MARS', 'MERCURY', 'NEPTUNE', 'SATURN', 'URANUS', 'VENUS', '__class__', '__doc__', '__init__', '__members__', '__module__', 'surface_gravity']
>>> dir(Planet.EARTH)
['__class__', '__doc__', '__module__', 'mass', 'name', 'radius', 'surface_gravity', 'value']

View file

@ -162,7 +162,8 @@ Data Types
.. method:: EnumType.__dir__(cls)
Returns ``['__class__', '__doc__', '__members__', '__module__']`` and the
names of the members in *cls*::
names of the members in ``cls``. User-defined methods and methods from
mixin classes will also be included::
>>> dir(Color)
['BLUE', 'GREEN', 'RED', '__class__', '__doc__', '__members__', '__module__']
@ -260,7 +261,7 @@ Data Types
.. method:: Enum.__dir__(self)
Returns ``['__class__', '__doc__', '__module__', 'name', 'value']`` and
any public methods defined on *self.__class__*::
any public methods defined on ``self.__class__`` or a mixin class::
>>> from datetime import date
>>> class Weekday(Enum):

View file

@ -635,10 +635,60 @@ class EnumType(type):
super().__delattr__(attr)
def __dir__(self):
return (
['__class__', '__doc__', '__members__', '__module__']
+ self._member_names_
)
# Start off with the desired result for dir(Enum)
cls_dir = {'__class__', '__doc__', '__members__', '__module__'}
add_to_dir = cls_dir.add
mro = self.__mro__
this_module = globals().values()
is_from_this_module = lambda cls: any(cls is thing for thing in this_module)
first_enum_base = next(cls for cls in mro if is_from_this_module(cls))
enum_dict = Enum.__dict__
sentinel = object()
# special-case __new__
ignored = {'__new__', *filter(_is_sunder, enum_dict)}
add_to_ignored = ignored.add
# We want these added to __dir__
# if and only if they have been user-overridden
enum_dunders = set(filter(_is_dunder, enum_dict))
# special-case __new__
if self.__new__ is not first_enum_base.__new__:
add_to_dir('__new__')
for cls in mro:
# Ignore any classes defined in this module
if cls is object or is_from_this_module(cls):
continue
cls_lookup = cls.__dict__
# If not an instance of EnumType,
# ensure all attributes excluded from that class's `dir()` are ignored here.
if not isinstance(cls, EnumType):
cls_lookup = set(cls_lookup).intersection(dir(cls))
for attr_name in cls_lookup:
# Already seen it? Carry on
if attr_name in cls_dir or attr_name in ignored:
continue
# Sunders defined in Enum.__dict__ are already in `ignored`,
# But sunders defined in a subclass won't be (we want all sunders excluded).
elif _is_sunder(attr_name):
add_to_ignored(attr_name)
# Not an "enum dunder"? Add it to dir() output.
elif attr_name not in enum_dunders:
add_to_dir(attr_name)
# Is an "enum dunder", and is defined by a class from enum.py? Ignore it.
elif getattr(self, attr_name) is getattr(first_enum_base, attr_name, sentinel):
add_to_ignored(attr_name)
# Is an "enum dunder", and is either user-defined or defined by a mixin class?
# Add it to dir() output.
else:
add_to_dir(attr_name)
# sort the output before returning it, so that the result is deterministic.
return sorted(cls_dir)
def __getattr__(cls, name):
"""
@ -985,13 +1035,10 @@ class Enum(metaclass=EnumType):
"""
Returns all members and all public methods
"""
added_behavior = [
m
for cls in self.__class__.mro()
for m in cls.__dict__
if m[0] != '_' and m not in self._member_map_
] + [m for m in self.__dict__ if m[0] != '_']
return (['__class__', '__doc__', '__module__'] + added_behavior)
cls = type(self)
to_exclude = {'__members__', '__init__', '__new__', *cls._member_names_}
filtered_self_dict = (name for name in self.__dict__ if not name.startswith('_'))
return sorted({'name', 'value', *dir(cls), *filtered_self_dict} - to_exclude)
def __format__(self, format_spec):
"""

View file

@ -203,59 +203,340 @@ class TestEnum(unittest.TestCase):
IDES_OF_MARCH = 2013, 3, 15
self.Holiday = Holiday
def test_dir_on_class(self):
Season = self.Season
self.assertEqual(
set(dir(Season)),
set(['__class__', '__doc__', '__members__', '__module__',
'SPRING', 'SUMMER', 'AUTUMN', 'WINTER']),
)
class DateEnum(date, Enum): pass
self.DateEnum = DateEnum
def test_dir_on_item(self):
Season = self.Season
self.assertEqual(
set(dir(Season.WINTER)),
set(['__class__', '__doc__', '__module__', 'name', 'value']),
)
class FloatEnum(float, Enum): pass
self.FloatEnum = FloatEnum
def test_dir_with_added_behavior(self):
class Test(Enum):
class Wowser(Enum):
this = 'that'
these = 'those'
def wowser(self):
"""Wowser docstring"""
return ("Wowser! I'm %s!" % self.name)
self.assertEqual(
set(dir(Test)),
set(['__class__', '__doc__', '__members__', '__module__', 'this', 'these']),
)
self.assertEqual(
set(dir(Test.this)),
set(['__class__', '__doc__', '__module__', 'name', 'value', 'wowser']),
)
@classmethod
def classmethod_wowser(cls): pass
@staticmethod
def staticmethod_wowser(): pass
self.Wowser = Wowser
class IntWowser(IntEnum):
this = 1
these = 2
def wowser(self):
"""Wowser docstring"""
return ("Wowser! I'm %s!" % self.name)
@classmethod
def classmethod_wowser(cls): pass
@staticmethod
def staticmethod_wowser(): pass
self.IntWowser = IntWowser
class FloatWowser(float, Enum):
this = 3.14
these = 4.2
def wowser(self):
"""Wowser docstring"""
return ("Wowser! I'm %s!" % self.name)
@classmethod
def classmethod_wowser(cls): pass
@staticmethod
def staticmethod_wowser(): pass
self.FloatWowser = FloatWowser
class WowserNoMembers(Enum):
def wowser(self): pass
@classmethod
def classmethod_wowser(cls): pass
@staticmethod
def staticmethod_wowser(): pass
class SubclassOfWowserNoMembers(WowserNoMembers): pass
self.WowserNoMembers = WowserNoMembers
self.SubclassOfWowserNoMembers = SubclassOfWowserNoMembers
class IntWowserNoMembers(IntEnum):
def wowser(self): pass
@classmethod
def classmethod_wowser(cls): pass
@staticmethod
def staticmethod_wowser(): pass
self.IntWowserNoMembers = IntWowserNoMembers
class FloatWowserNoMembers(float, Enum):
def wowser(self): pass
@classmethod
def classmethod_wowser(cls): pass
@staticmethod
def staticmethod_wowser(): pass
self.FloatWowserNoMembers = FloatWowserNoMembers
class EnumWithInit(Enum):
def __init__(self, greeting, farewell):
self.greeting = greeting
self.farewell = farewell
ENGLISH = 'hello', 'goodbye'
GERMAN = 'Guten Morgen', 'Auf Wiedersehen'
def some_method(self): pass
self.EnumWithInit = EnumWithInit
def test_dir_on_sub_with_behavior_on_super(self):
# see issue22506
class SuperEnum(Enum):
class SuperEnum1(Enum):
def invisible(self):
return "did you see me?"
class SubEnum(SuperEnum):
class SubEnum1(SuperEnum1):
sample = 5
self.assertEqual(
set(dir(SubEnum.sample)),
set(['__class__', '__doc__', '__module__', 'name', 'value', 'invisible']),
)
self.SubEnum1 = SubEnum1
def test_dir_on_sub_with_behavior_including_instance_dict_on_super(self):
# see issue40084
class SuperEnum(IntEnum):
class SuperEnum2(IntEnum):
def __new__(cls, value, description=""):
obj = int.__new__(cls, value)
obj._value_ = value
obj.description = description
return obj
class SubEnum(SuperEnum):
class SubEnum2(SuperEnum2):
sample = 5
self.assertTrue({'description'} <= set(dir(SubEnum.sample)))
self.SubEnum2 = SubEnum2
def test_dir_basics_for_all_enums(self):
enums_for_tests = (
# Generic enums in enum.py
Enum,
IntEnum,
StrEnum,
# Generic enums defined outside of enum.py
self.DateEnum,
self.FloatEnum,
# Concrete enums derived from enum.py generics
self.Grades,
self.Season,
# Concrete enums derived from generics defined outside of enum.py
self.Konstants,
self.Holiday,
# Standard enum with added behaviour & members
self.Wowser,
# Mixin-enum-from-enum.py with added behaviour & members
self.IntWowser,
# Mixin-enum-from-oustide-enum.py with added behaviour & members
self.FloatWowser,
# Equivalents of the three immediately above, but with no members
self.WowserNoMembers,
self.IntWowserNoMembers,
self.FloatWowserNoMembers,
# Enum with members and an __init__ method
self.EnumWithInit,
# Special cases to test
self.SubEnum1,
self.SubEnum2
)
for cls in enums_for_tests:
with self.subTest(cls=cls):
cls_dir = dir(cls)
# test that dir is deterministic
self.assertEqual(cls_dir, dir(cls))
# test that dir is sorted
self.assertEqual(list(cls_dir), sorted(cls_dir))
# test that there are no dupes in dir
self.assertEqual(len(cls_dir), len(set(cls_dir)))
# test that there are no sunders in dir
self.assertFalse(any(enum._is_sunder(attr) for attr in cls_dir))
self.assertNotIn('__new__', cls_dir)
for attr in ('__class__', '__doc__', '__members__', '__module__'):
with self.subTest(attr=attr):
self.assertIn(attr, cls_dir)
def test_dir_for_enum_with_members(self):
enums_for_test = (
# Enum with members
self.Season,
# IntEnum with members
self.Grades,
# Two custom-mixin enums with members
self.Konstants,
self.Holiday,
# several enums-with-added-behaviour and members
self.Wowser,
self.IntWowser,
self.FloatWowser,
# An enum with an __init__ method and members
self.EnumWithInit,
# Special cases to test
self.SubEnum1,
self.SubEnum2
)
for cls in enums_for_test:
cls_dir = dir(cls)
member_names = cls._member_names_
with self.subTest(cls=cls):
self.assertTrue(all(member_name in cls_dir for member_name in member_names))
for member in cls:
member_dir = dir(member)
# test that dir is deterministic
self.assertEqual(member_dir, dir(member))
# test that dir is sorted
self.assertEqual(list(member_dir), sorted(member_dir))
# test that there are no dupes in dir
self.assertEqual(len(member_dir), len(set(member_dir)))
for attr_name in cls_dir:
with self.subTest(attr_name=attr_name):
if attr_name in {'__members__', '__init__', '__new__', *member_names}:
self.assertNotIn(attr_name, member_dir)
else:
self.assertIn(attr_name, member_dir)
self.assertFalse(any(enum._is_sunder(attr) for attr in member_dir))
def test_dir_for_enums_with_added_behaviour(self):
enums_for_test = (
self.Wowser,
self.IntWowser,
self.FloatWowser,
self.WowserNoMembers,
self.SubclassOfWowserNoMembers,
self.IntWowserNoMembers,
self.FloatWowserNoMembers
)
for cls in enums_for_test:
with self.subTest(cls=cls):
self.assertIn('wowser', dir(cls))
self.assertIn('classmethod_wowser', dir(cls))
self.assertIn('staticmethod_wowser', dir(cls))
self.assertTrue(all(
all(attr in dir(member) for attr in ('wowser', 'classmethod_wowser', 'staticmethod_wowser'))
for member in cls
))
self.assertEqual(dir(self.WowserNoMembers), dir(self.SubclassOfWowserNoMembers))
# Check classmethods are present
self.assertIn('from_bytes', dir(self.IntWowser))
self.assertIn('from_bytes', dir(self.IntWowserNoMembers))
def test_help_output_on_enum_members(self):
added_behaviour_enums = (
self.Wowser,
self.IntWowser,
self.FloatWowser
)
for cls in added_behaviour_enums:
with self.subTest(cls=cls):
rendered_doc = pydoc.render_doc(cls.this)
self.assertIn('Wowser docstring', rendered_doc)
if cls in {self.IntWowser, self.FloatWowser}:
self.assertIn('float(self)', rendered_doc)
def test_dir_for_enum_with_init(self):
EnumWithInit = self.EnumWithInit
cls_dir = dir(EnumWithInit)
self.assertIn('__init__', cls_dir)
self.assertIn('some_method', cls_dir)
self.assertNotIn('greeting', cls_dir)
self.assertNotIn('farewell', cls_dir)
member_dir = dir(EnumWithInit.ENGLISH)
self.assertNotIn('__init__', member_dir)
self.assertIn('some_method', member_dir)
self.assertIn('greeting', member_dir)
self.assertIn('farewell', member_dir)
def test_mixin_dirs(self):
from datetime import date
enums_for_test = (
# generic mixins from enum.py
(IntEnum, int),
(StrEnum, str),
# generic mixins from outside enum.py
(self.FloatEnum, float),
(self.DateEnum, date),
# concrete mixin from enum.py
(self.Grades, int),
# concrete mixin from outside enum.py
(self.Holiday, date),
# concrete mixin from enum.py with added behaviour
(self.IntWowser, int),
# concrete mixin from outside enum.py with added behaviour
(self.FloatWowser, float)
)
enum_dict = Enum.__dict__
enum_dir = dir(Enum)
enum_module_names = enum.__all__
is_from_enum_module = lambda cls: cls.__name__ in enum_module_names
is_enum_dunder = lambda attr: enum._is_dunder(attr) and attr in enum_dict
def attr_is_inherited_from_object(cls, attr_name):
for base in cls.__mro__:
if attr_name in base.__dict__:
return base is object
return False
# General tests
for enum_cls, mixin_cls in enums_for_test:
with self.subTest(enum_cls=enum_cls):
cls_dir = dir(enum_cls)
cls_dict = enum_cls.__dict__
mixin_attrs = [
x for x in dir(mixin_cls)
if not attr_is_inherited_from_object(cls=mixin_cls, attr_name=x)
]
first_enum_base = next(
base for base in enum_cls.__mro__
if is_from_enum_module(base)
)
for attr in mixin_attrs:
with self.subTest(attr=attr):
if enum._is_sunder(attr):
# Unlikely, but no harm in testing
self.assertNotIn(attr, cls_dir)
elif attr in {'__class__', '__doc__', '__members__', '__module__'}:
self.assertIn(attr, cls_dir)
elif is_enum_dunder(attr):
if is_from_enum_module(enum_cls):
self.assertNotIn(attr, cls_dir)
elif getattr(enum_cls, attr) is getattr(first_enum_base, attr):
self.assertNotIn(attr, cls_dir)
else:
self.assertIn(attr, cls_dir)
else:
self.assertIn(attr, cls_dir)
# Some specific examples
int_enum_dir = dir(IntEnum)
self.assertIn('imag', int_enum_dir)
self.assertIn('__rfloordiv__', int_enum_dir)
self.assertNotIn('__format__', int_enum_dir)
self.assertNotIn('__hash__', int_enum_dir)
self.assertNotIn('__init_subclass__', int_enum_dir)
self.assertNotIn('__subclasshook__', int_enum_dir)
class OverridesFormatOutsideEnumModule(Enum):
def __format__(self, *args, **kwargs):
return super().__format__(*args, **kwargs)
SOME_MEMBER = 1
self.assertIn('__format__', dir(OverridesFormatOutsideEnumModule))
self.assertIn('__format__', dir(OverridesFormatOutsideEnumModule.SOME_MEMBER))
def test_dir_on_sub_with_behavior_on_super(self):
# see issue22506
self.assertEqual(
set(dir(self.SubEnum1.sample)),
set(['__class__', '__doc__', '__module__', 'name', 'value', 'invisible']),
)
def test_dir_on_sub_with_behavior_including_instance_dict_on_super(self):
# see issue40084
self.assertTrue({'description'} <= set(dir(self.SubEnum2.sample)))
def test_enum_in_enum_out(self):
Season = self.Season
@ -4156,7 +4437,8 @@ class TestIntEnumConvert(unittest.TestCase):
self.assertEqual(test_type.CONVERT_TEST_NAME_E, 5)
# Ensure that test_type only picked up names matching the filter.
self.assertEqual([name for name in dir(test_type)
if name[0:2] not in ('CO', '__')],
if name[0:2] not in ('CO', '__')
and name not in dir(IntEnum)],
[], msg='Names other than CONVERT_TEST_* found.')
@unittest.skipUnless(python_version == (3, 8),
@ -4207,7 +4489,8 @@ class TestStrEnumConvert(unittest.TestCase):
self.assertEqual(test_type.CONVERT_STR_TEST_2, 'goodbye')
# Ensure that test_type only picked up names matching the filter.
self.assertEqual([name for name in dir(test_type)
if name[0:2] not in ('CO', '__')],
if name[0:2] not in ('CO', '__')
and name not in dir(StrEnum)],
[], msg='Names other than CONVERT_STR_* found.')
def test_convert_repr_and_str(self):

View file

@ -0,0 +1 @@
Improve output of ``dir()`` with Enums.