PyCharm 类型提示枚举迭代

Question

Python's enum class supports iteration, but PyCharm has trouble figuring this out.

from enum import Enum

class Color(Enum):
    RED = 0
    BLUE = 1

for color in Color:
    # Warning: Expected 'collections.Iterable', got 'Type[Color]' instead
    print(color)

Though the method EnumMeta.__iter__ exists, PyCharm has trouble figuring this out.

I don't mind manually adding type hinting to work around the problem, I'm just not sure what and where.

Answer 1

Maybe it's not cleanest solution, but following works for me:

from enum import Enum
import typing

class Color(Enum):
    RED = 0
    BLUE = 1

Color = Color  # type: typing.Union[typing.Type[Color], typing.Iterable]

PyCharm supports type hinting using format defined in PEP 484 (for Python versions lower than 3.5 in the form of comments, for 3.5 and higher in form of annotations).

Important note here is that on Python version lower than 3.5, importing typing module should be guarded in some way (PyCharm recognizes this import as valid, even without having typing module installed in site packages, however when code is run ImportError occurs).

Answer 2

A workaround that I used in PyCharm was through use of __members__ with .items()

from enum import Enum

class Color(Enum):
    RED = 0
    BLUE = 1

for name, color in Color.__members__.items():
    print(name, color)

Returning:

('RED', <Color.RED: 0>)
('BLUE', <Color.BLUE: 1>)

See the following link for reference:

https://cpython-test-docs.readthedocs.io/en/latest/library/enum.html#iteration