Enums: advanced features

[sharkdp]

We have basic support for enums, but there is a long tail of things that could be improved:

• [x] Infer type for member names and member values.
• [x] Special-case handling of IntEnum, StrEnum, or other custom classes that derive from both Enum and some other class (infer appropriate types for member values)
• [x] auto() values
• [ ] Handling of enum.Flag (enum expansion may not be performed for subclasses of enum.Flag)
• [ ] Custom __new__ or __init__ methods in enums.
• [ ] Support for _generate_next_value_
• [ ] Special-case handling for call to enum class to retrieve member by value (e.g. Color("red"))
• [ ] Function syntax to create Enums

[Andre-Medina]

Noticed enums weren't typed quite right in the ty VSC extension, glad to see it being worked on.

should also consider EnumClass.ITEM.name being type str

from enum import StrEnum, IntEnum

class TempTest(StrEnum):
    ITEM = "data"
    ITEM_2 = "more_data"


should_be_enum_type = TempTest.ITEM
should_be_str = TempTest.ITEM.value
should_be_str_also = TempTest.ITEM.name

class TestInt(IntEnum):

    ITEM = 1
    ITEM_2 = 2

should_be_enum_type_also = TestInt.ITEM
should_be_int = TestInt.ITEM.value
should_be_str_still = TestInt.ITEM.name

With the extension in its current state, ty confusing some type hints:

[sharkdp]

@Andre-Medina It looks like StrEnum can't be imported? It's only available on 3.11 and later. How did you configure your Python version?

Once you have fixed that, should_be_enum_type should indeed be of type Literal[TempTest.ITEM]. Accessing .value and .name with more precise types is not yet supported, but it's listed as the first item in the enumeration above.

[Andre-Medina]

It looks like StrEnum can't be imported? It's only available on 3.11 and later. How did you configure your Python version?

Unsure, probably an issue with my local env, but good to hear .value and .name will be supported :)

[aidandj]

There is a backported package backports.strenum that I've been using for legacy code.

I was testing out the latest main, and there is a regression. Using that StrEnum package, I used to get a str for the revealed type on alpha20. But on the latest main (in ruff) I get Literal[<index of item>]

if sys.version_info >= (3, 11):
    from enum import StrEnum
else:
    from backports.strenum import StrEnum
from enum import Enum, auto


class SomeEnum(StrEnum):
    A = auto()
    B = auto()
    C = auto()
    D = auto()


reveal_type(SomeEnum.A.value)

With alpha 20, I get type string. With ruff commit: 706be0a6e7e09936511198f2ff8982915520d138 I get Literal[1]

[thejchap]

ah - this is being exposed now

it does look like the type could actually be Literal['a'] instead of builtins.str

from the docs:

Note Using auto with StrEnum results in the lower-cased member name as the value.

https://docs.python.org/3/library/enum.html#enum.StrEnum

[aidandj]

I didn't actually test it with 3.11 I'm realizing. I'll do that later.

Confirmed same behavior on 3.12 and commit 706be0a6e7e09936511198f2ff8982915520d138

[sharkdp]

Thank you for reporting this. We should fix this. That's why the "Special-case handling of IntEnum, StrEnum" and "auto() values" items are still unchecked in the description above. We could either patch this by not inferring a precise type for .value if we derive from StrEnum. Or we could implement the actual auto behavior for StrEnum, like @thejchap suggested, if it's not too hard.

[thejchap]

i can take a look at this next week

[spaceone]

I guess another example for this issue is:

$ ty check t.py 
t.py:10:22: error[invalid-type-form] Type arguments for `Literal` must be `None`, a literal value (int, bool, str, or bytes), or an enum member
Found 1 diagnostic
$ cat t.py

import enum
from typing import Literal


class Foo(enum.IntEnum):
    foo = 1
    bar = 2


def foo(bar: Literal[Foo.foo | Foo.bar]) -> None:
    pass

[AlexWaygood]

@spaceone that looks like a true positive to me. Foo.foo | Foo.bar is not an enum member. | unions are not allowed inside Literal[] slices. You probably meant to write Literal[Foo.foo, Foo.bar] (a comma instead of a |).

We could probably improve our diagnostic for that case.

[spaceone]

@AlexWaygood ah thank you. Then it's a undetected case of mypy.

→https://github.com/python/mypy/issues/20438