Type narrowing of indexing a StrEnum created in a function

[pchar2]

Bug Report

When indexing a StrEnum created in a function, the behaviour is not the same as created in an executed module.

To Reproduce

from enum import StrEnum

def create_str_enum(value: str, names: list[str]) -> StrEnum:
    return StrEnum(value=value, names=names)
    
WithoutBug: StrEnum = StrEnum(value="WithoutBug", names=["foo"])
WithBug: StrEnum = create_str_enum("WithBug", ["foo"])
    
WithoutBug["foo"]
WithBug["foo"]

playground

Expected Behavior

No error.

Actual Behavior

main.py:10: error: Invalid index type "str" for "StrEnum"; expected type "SupportsIndex | slice[Any, Any, Any]" [index]

Your Environment

• Mypy version used: all version
• Mypy command-line flags: /
• Mypy configuration options from mypy.ini (and other config files): /
• Python version used: 3.11-3.12

[sterliakov]

We special-case literal StrEnum calls, and that logic does not run for custom wrappers. The annotation is ignored, that's a bug - your annotation on WithoutBug should be considered and rejected, so that both cases result in an error.

The typical case we aim to support is

WithoutBug = StrEnum(value="WithoutBug", names=["foo"])

(without annotation). The magic happens here:

mypy/mypy/semanal_enum.py

Lines 83 to 106 in 58f7628

	def check_enum_call(
	self, node: Expression, var_name: str, is_func_scope: bool
	) -> TypeInfo | None:
	"""Check if a call defines an Enum.
	Example:
	A = enum.Enum('A', 'foo bar')
	is equivalent to:
	class A(enum.Enum):
	foo = 1
	bar = 2
	"""
	if not isinstance(node, CallExpr):
	return None
	call = node
	callee = call.callee
	if not isinstance(callee, RefExpr):
	return None
	fullname = callee.fullname
	if fullname not in ENUM_BASES:
	return None

But there's another problem: enums are really special. Expression StrEnum("Name", ["val"]) does not create a enum instance (even though stubs says so), it creates a class (enum.EnumType instance). It is magic: derived classes do not work that way, only Enum and its standard subclasses. So correct definition of "creates some StrEnum" is

def create_str_enum(value: str, names: list[str]) -> type[StrEnum]:
    # __new__ returns an instance in stubs, but this literal constructs a class.
    return StrEnum(value=value, names=names)  # type: ignore[return-value]

The error you show is correct for StrEnum instances: WithBug.foo['foo'] raises while trying to index string by another string.

I suggest to keep this open to track the remaining issue: the assignment below should be flagged.

from enum import StrEnum

Bad: StrEnum = StrEnum(value="Bad", names=["foo"])

Noticed: StrEnum = Bad  # E: Incompatible types in assignment (expression has type "type[Bad]", variable has type "StrEnum")  [assignment]

https://mypy-play.net/?mypy=latest&python=3.12&gist=2feee172745e3df4b87e41acdb0f8a6a

[pchar2]

Hello @sterliakov

So the correct way to have it work would be:

from enum import StrEnum

def create_str_enum(value: str, names: list[str]) -> type[StrEnum]:
    return StrEnum(value=value, names=names)  # type: ignore[return-value]
    
MyStrEnum: type[StrEnum] = create_str_enum("StrEnum", ["foo"])
MyStrEnum["foo"]

The error you show is correct for StrEnum instances: WithBug.foo['foo'] raises while trying to index string by another string.

I did not indexed a value, but the enum, WithBug['foo'] .

I suggest to keep this open to track the remaining issue: the assignment below should be flagged.

Okay.

Thank you for the detailed answer.

[pchar2]

There is an issue with returning a type of StrEnum. Accessing the attribute directly now raises an error.

from enum import StrEnum

MyStrEnum: StrEnum = StrEnum("MyStrEnum", ["foo"])
MyStrEnum("foo")
MyStrEnum["foo"]
MyStrEnum.foo

def create_str_enum_2(value: str, names: list[str]) -> StrEnum:
    return StrEnum(value=value, names=names)  # type: ignore[return-value]

MyStrEnum2: StrEnum = create_str_enum_2("MyStrEnum2", ["foo"])
MyStrEnum2("foo")  # error: "StrEnum" not callable  [operator]
MyStrEnum2["foo"]  # error: Invalid index type "str" for "StrEnum"; expected type "SupportsIndex | slice[Any, Any, Any]"  [index]
MyStrEnum2.foo  # error: "StrEnum" has no attribute "foo"  [attr-defined]


def create_str_enum_3(value: str, names: list[str]) -> type[StrEnum]:
    return StrEnum(value=value, names=names)  # type: ignore[return-value]

MyStrEnum3: type[StrEnum] = create_str_enum_3("MyStrEnum3", ["foo"])
MyStrEnum3("foo")
MyStrEnum3["foo"]
MyStrEnum3.foo  # error: "type[StrEnum]" has no attribute "foo"  [attr-defined]

playground

[sterliakov]

Accessing the attribute directly now raises an error.

Yeah, that's right. Since the magic does not happen without explicit enum.StrEnum call, you're left with what typeshed defines for that class.

def create_str_enum_3(value: str, names: list[str]) -> type[StrEnum]:
    return StrEnum(value=value, names=names)  # type: ignore[return-value]

MyStrEnum3: type[StrEnum] = create_str_enum_3("MyStrEnum3", ["bar"])
MyStrEnum3("foo")
MyStrEnum3["foo"]
MyStrEnum3.foo  # error: "type[StrEnum]" has no attribute "foo"  [attr-defined]

Now there's no foo attribute, and the only way to tell that is by inspecting values, not types, of the arguments provided. Again, this only works with explicit StrEnum because mypy makes extra effort to support such calls (even the spec does not mandate that, listing functional style support as optional: https://typing.python.org/en/latest/spec/enums.html#enum-definition). It does not cross function boundary.

(side note: there's no need to annotate every variable, hinting variables that are assigned from non-generic, non-overloaded typed calls is redundant unless you need to widen the type explicitly, MyStrEnum3 = create_str_enum_3("MyStrEnum3", ["foo"]) is completely equivalent)