mypy Oddity With Python Enum Lookups

user_3350050 0
Viewed 148

I'm trying to write a custom enum.IntEnum class that overrides the starting value and adds a string attribute (discussed here). mypy doesn't like when I look up enums by value. See the two numbered lines near the bottom.

from __future__ import annotations
import enum
import itertools

gooCount = itertools.count(42)

class Goo(enum.IntEnum):
    FOO = "fzz",
    MOO = "mrp",

    label: str

    def __new__(cls, label: str) -> Goo:
        value = next(gooCount)

        mbr = int.__new__(cls, value)

        mbr._value_ = value
        mbr.label = label

        return mbr

foo = Goo(42)                 # line 23
assert foo == 42
assert foo is Goo.FOO
assert foo.label == "fzz"

moo = Goo(43)                 # line 28
assert moo == 43
assert moo is Goo.MOO
assert moo.label == "mrp"

The code runs fine, but mypy complains about those two lines, which look up particular enums by their integer value.

% mypy goo.py
goo.py:23: error: Argument 1 to "Goo" has incompatible type "int"; expected "str"
goo.py:28: error: Argument 1 to "Goo" has incompatible type "int"; expected "str"

I expect it's somehow confused about the metaclass magic in enum. The __init__() method should probably be overloaded to accept an int for these lookups, or whatever additional types I'm adding, a str in this case. Is there a way to fix these errors without # type: ignore or typing.cast()?

Update 1: I wasn't able to make those overloads explicit by adding this,

<same as above>

from typing import Union, overload

<same as above>

class Goo(enum.IntEnum):

    <same as above>

    @overload
    def __init__(self, arg: int) -> None: ...

    @overload
    def __init__(self, arg: str) -> None: ...

    def __init__(self, arg: Union[int, str]) -> None:
        if isinstance(arg, int):
            super().__init__(arg)                       # line 32
        else:
            self.label = arg

<same as above>

That gives me this.

% mypy goo.py
goo.py:32: error: Too many arguments for "__init__" of "object"

Update 2: Calling enum.IntEnum.__init__() instead of super().__init__() on line 32 still runs and appeases mypy, but that was arrived at by trial and error, not understanding. Also, I can't seem to even reach that line, so this "works" too. That makes it appear as though the possibility of an int is entirely artificial, there just for mypy.

<same as above>

from typing import Union

<same as above>

class Goo(enum.IntEnum):

    <same as above>

    def __init__(self, arg: Union[int, str]) -> None:
        assert not isinstance(arg, int)
        self.label = arg

<same as above>

Is this my answer? Why?

mypy enums python
1 Answers

This looks like a limitation of mypy. The best answer I can find is the __init__() hack I mentioned in Update 2, repeated here but in full.

from __future__ import annotations
import enum
import itertools

goo_count = itertools.count(42)

class Goo(enum.IntEnum):
    FOO = "fzz",
    MOO = "mrp",

    label: str

    def __new__(cls, label: str) -> Goo:
        value = next(goo_count)

        mbr = int.__new__(cls, value)

        mbr._value_ = value
        mbr.label = label

        return mbr

    def __init__(self, arg: int) -> None:
        assert isinstance(arg, str)  # enum by-value lookups do not call __init__()

foo = Goo(42)
assert foo == 42
assert foo is Goo.FOO
assert foo.label == "fzz"

moo = Goo(43)
assert moo == 43
assert moo is Goo.MOO
assert moo.label == "mrp"

I'm not accepting this as the answer since it's just a workaround. Hopefully, mypy will eventually learn this documented use case of enums.

user_33500500
Related