Support iteration on enums

[bbc2]

Iterating on enums makes mypy (0.4.5) raise errors even though it's OK:

import enum


class Colors(enum.Enum):
    RED = 0
    BLUE = 1


print([color.name for color in Colors])

> python3 color.py 
['RED', 'BLUE']

> mypy color.py 
color.py:9: error: Iterable expected
color.py:9: error: "Colors" has no attribute "__iter__"

(follow-up issue created as suggested in #529)

[elazarg]

At first sight, we only need #2193

class Enum:
    ...
    @classmethod
    def __iter__(cls: Type[T]) -> Iterator[T]: ...

And somehow making Enum itself iterable?

But it's not so simple, since Enum is a class appearing in the context of an expression, and therefore "coerce" to a callable somewhere in the code. Thus the above code still results in error: Iterable expected. At the point where the error is emitted, the information that it's enum does not seem to exist anymore - it is a simple callable.

[dokai]

In terms of making a class itself iterable it seems you have to use a metaclass instead of using just a @classmethod in the class definition.

In [7]: class A:
   ...:     @classmethod
   ...:     def __iter__(cls):
   ...:         return iter(range(3))
   ...:

In [8]: list(A)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-8-ed10f13c6f39> in <module>()
----> 1 list(A)

TypeError: 'type' object is not iterable

In [9]: class MetaIter(type):
   ...:     def __iter__(self):
   ...:         return iter(range(3))
   ...:

In [10]: class B(metaclass=MetaIter): pass

In [11]: list(B)
Out[11]: [0, 1, 2]

[gvanrossum]

Yes, that's what I meant. But that's runtime. How to do it in mypy is probably quite different because it doesn't interpret the metaclass.

[gvanrossum]

Did you check the is_type_obj() method? It should still know... Also, I doubt that defining a classmethod is the same as having it on the metaclass. IIRC that doesn't even work in CPython.

[elazarg]

I didn't check yet.

As you said, it's not how mypy work. I think it should, but I guess it will require a big change.

(I wonder whether mypy should mimic Smalltalk's object model. It's very elegant, and perhaps we can think of Python as "hiding the theory" (e.g. that type(type) != type) for the convenience of the user. Similarly to the -type-in-type flag in COQ, which collapses the universe of types into a single one.
Yes, I know, I get carried away sometimes. Sorry)

[gvanrossum]

Fixed by #1136.

[levsa]

Should be python/typeshed#1136

[levsa]

I'm seeing this problem when using a subclass of IntEnum on mypy v0520. Should it be fixed in 0520?

test.py:

from enum import IntEnum

class MyEnum(IntEnum):
	A = 1
	B = 2
	C = 3

for x in MyEnum:
	print(x)

> mypy test.py
test.py:8: error: Iterable expected
test.py:8: error: Type[MyEnum] has no attribute "__iter__"

[ilevkivskyi]

@levsa
I just tried and IntEnum still fails on master, but we have a separate issue for it #3622

[ngaya-ll]

@gvanrossum Doesn't seem to be working currently (see #4400). Can this issue be re-opened?

[ethanhs]

@ngaya-ll Hm, that is a good catch. I will reopen this.

[elazarg]

@ngaya-ll this seems to be a problem in the unit test, not in enum implementation. @ethanhs - unless reproduced in actual run, IMO this issue is still closed.

[elazarg]

Correction: it should be opened, since it is fixed by python/typeshed#1755 which is not merged yet. Sorry @ethanhs.

[ngaya-ll]

Interesting. The original test case proposed by @bbc2 passes, but mypy still doesn't accept enum types as Iterable. Should I file a separate ticket?

[elazarg]

The original test passes because this kind of iteration uses structural typing already. The compatibility checks use structural subtyping (protocols) only if there is no explicit (nominal) protocol declared. That's what python/typeshed#1755 fixes - the problem is in typeshed. However, said PR uses a feature that reveals some other, seemingly unrelated problem in mypy: #4272.

[srittau]

For reference, using mypy 0.560, the following:

from enum import Enum

class Foo(Enum):
    X = "x"

reveal_type(list(Foo))
reveal_type(set(Foo))

will print

foo.py:6: error: Revealed type is 'builtins.list[enum.Enum*]'
foo.py:7: error: Revealed type is 'builtins.set[enum.Enum*]'

I'd expect this to be builtins.list[Foo] and builtins.set[Foo].

[elazarg]

This example seems to work on master

[srittau]

I will report back when 0.570 is released.

[ilevkivskyi]

All examples I have found in this issue now work on mypy master.

[otchita]

I'm seeing this issue when using a subclass of MultiValueEnum on mypy==0.782

test.py:

from aenum import MultiValueEnum

class MyEnum(MultiValueEnum):
    A = 1, 2
    B = 3
    C = 4

for x in MyEnum:
    print(x)

> mypy test.py

test.py:8: error: "Type[MyEnum]" has no attribute "__iter__" (not iterable)

[JelleZijlstra]

@otchita that might be an issue with the stubs for aenum.

[harahu]

@ilevkivskyi I'm seeing this issue when trying to express a function that should work for enums in general. Simplified example:

from enum import Enum
from typing import Type


class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3


def enum_getter(x: Type[Enum]) -> Enum:
    return next(iter(x))


x = enum_getter(Color)

This code runs in practice, but i get the following type checking error:
error: Argument 1 to "iter" has incompatible type "Type[Color]"; expected "SupportsIter[<nothing>]"

[colatkinson]

Confirming that I'm running into the same behavior as @harahu on 0.991. Possibly related: #12553.

Interestingly, if that example is modified slightly, it does work:

from enum import Enum
from typing import Type


class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3


def enum_getter(x: Type[Enum]) -> Enum:
    for v in x:
        return v
    assert False, "unreachable"


x = enum_getter(Color)

So it seems that as far as mypy is concerned Enums support iteration, but not iter().

This also makes the error go away:

from enum import Enum
from typing import Type, Iterable


class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3


def enum_getter(x: Type[Enum]) -> Enum:
    y: Iterable[Enum] = x
    return next(iter(y))


x = enum_getter(Color)

Not sure if this is exactly the same problem as the original -- perhaps a maintainer could advise as to whether a new issue is warranted.

[ilevkivskyi]

Yes, this is most likely #12553 (which seems easy to fix, so I may submit PR there soon).