[ty] Enum literal types (#19328)

## Summary

Add a new `Type::EnumLiteral(…)` variant and infer this type for member
accesses on enums.

**Example**: No more `@Todo` types here:
```py
from enum import Enum

class Answer(Enum):
    YES = 1
    NO = 2

    def is_yes(self) -> bool:
        return self == Answer.YES

reveal_type(Answer.YES)  # revealed: Literal[Answer.YES]
reveal_type(Answer.YES == Answer.NO)  # revealed: Literal[False]
reveal_type(Answer.YES.is_yes())  # revealed: bool
```

## Test Plan

* Many new Markdown tests for the new type variant
* Added enum literal types to property tests, ran property tests

## Ecosystem analysis

Summary:

Lots of false positives removed. All of the new diagnostics are
either new true positives (the majority) or known problems. Click for
detailed analysis</summary>

Details:

```diff
AutoSplit (https://github.com/Toufool/AutoSplit)
+ error[call-non-callable] src/capture_method/__init__.py:137:9: Method `__getitem__` of type `bound method CaptureMethodDict.__getitem__(key: Never, /) -> type[CaptureMethodBase]` is not callable on object of type `CaptureMethodDict`
+ error[call-non-callable] src/capture_method/__init__.py:147:9: Method `__getitem__` of type `bound method CaptureMethodDict.__getitem__(key: Never, /) -> type[CaptureMethodBase]` is not callable on object of type `CaptureMethodDict`
+ error[call-non-callable] src/capture_method/__init__.py:148:1: Method `__getitem__` of type `bound method CaptureMethodDict.__getitem__(key: Never, /) -> type[CaptureMethodBase]` is not callable on object of type `CaptureMethodDict`
```

New true positives. That `__getitem__` method is apparently annotated
with `Never` to prevent developers from using it.


```diff
dd-trace-py (https://github.com/DataDog/dd-trace-py)
+ error[invalid-assignment] ddtrace/vendor/psutil/_common.py:29:5: Object of type `None` is not assignable to `Literal[AddressFamily.AF_INET6]`
+ error[invalid-assignment] ddtrace/vendor/psutil/_common.py:33:5: Object of type `None` is not assignable to `Literal[AddressFamily.AF_UNIX]`
```

Arguably true positives:
e0a772c28b/ddtrace/vendor/psutil/_common.py (L29)

```diff
ignite (https://github.com/pytorch/ignite)
+ error[invalid-argument-type] tests/ignite/engine/test_custom_events.py:190:34: Argument to bound method `__call__` is incorrect: Expected `((...) -> Unknown) | None`, found `Literal["123"]`
+ error[invalid-argument-type] tests/ignite/engine/test_custom_events.py:220:37: Argument to function `default_event_filter` is incorrect: Expected `Engine`, found `None`
+ error[invalid-argument-type] tests/ignite/engine/test_custom_events.py:220:43: Argument to function `default_event_filter` is incorrect: Expected `int`, found `None`
+ error[call-non-callable] tests/ignite/engine/test_custom_events.py:561:9: Object of type `CustomEvents` is not callable
+ error[invalid-argument-type] tests/ignite/metrics/test_frequency.py:50:38: Argument to bound method `attach` is incorrect: Expected `Events`, found `CallableEventWithFilter`
```

All true positives. Some of them are inside `pytest.raises(TypeError,
…)` blocks 🙃

```diff
meson (https://github.com/mesonbuild/meson)
+ error[invalid-argument-type] unittests/internaltests.py:243:51: Argument to bound method `__init__` is incorrect: Expected `bool`, found `Literal[MachineChoice.HOST]`
+ error[invalid-argument-type] unittests/internaltests.py:271:51: Argument to bound method `__init__` is incorrect: Expected `bool`, found `Literal[MachineChoice.HOST]`
```

New true positives. Enum literals can not be assigned to `bool`, even if
their value types are `0` and `1`.

```diff
poetry (https://github.com/python-poetry/poetry)
+ error[invalid-assignment] src/poetry/console/exceptions.py:101:5: Object of type `Literal[""]` is not assignable to `InitVar[str]`
```

New false positive, missing support for `InitVar`.

```diff
prefect (https://github.com/PrefectHQ/prefect)
+ error[invalid-argument-type] src/integrations/prefect-dask/tests/test_task_runners.py:193:17: Argument is incorrect: Expected `StateType`, found `Literal[StateType.COMPLETED]`
```

This is confusing. There are two definitions
([one](74d8cd93ee/src/prefect/client/schemas/objects.py (L89-L100)),
[two](https://github.com/PrefectHQ/prefect/blob/main/src/prefect/server/schemas/states.py#L40))
of the `StateType` enum. Here, we're trying to assign one to the other.
I don't think that should be allowed, so this is a true positive (?).

```diff
python-htmlgen (https://github.com/srittau/python-htmlgen)
+ error[invalid-assignment] test_htmlgen/form.py:51:9: Object of type `str` is not assignable to attribute `autocomplete` of type `Autocomplete | None`
+ error[invalid-assignment] test_htmlgen/video.py:38:9: Object of type `str` is not assignable to attribute `preload` of type `Preload | None`
```

True positives. [The stubs are
wrong](01e3b911ac/htmlgen/form.pyi (L8-L10)).
These should not contain type annotations, but rather just `OFF = ...`.

```diff
rotki (https://github.com/rotki/rotki)
+ error[invalid-argument-type] rotkehlchen/tests/unit/test_serialization.py:62:30: Argument to bound method `deserialize` is incorrect: Expected `str`, found `Literal[15]`
```

New true positive.

```diff
vision (https://github.com/pytorch/vision)
+ error[unresolved-attribute] test/test_extended_models.py:302:17: Type `type[WeightsEnum]` has no attribute `DEFAULT`
+ error[unresolved-attribute] test/test_extended_models.py:302:58: Type `type[WeightsEnum]` has no attribute `DEFAULT`
```

Also new true positives. No `DEFAULT` member exists on `WeightsEnum`.
This commit is contained in:
David Peter 2025-07-15 21:31:53 +02:00 committed by GitHub
parent a0d4e1f854
commit a1edb69ea5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 1016 additions and 75 deletions

View file

@ -2,6 +2,7 @@
use ruff_benchmark::criterion;
use ruff_benchmark::real_world_projects::{InstalledProject, RealWorldProject};
use std::fmt::Write;
use std::ops::Range;
use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
@ -441,6 +442,37 @@ fn benchmark_complex_constrained_attributes_2(criterion: &mut Criterion) {
});
}
fn benchmark_many_enum_members(criterion: &mut Criterion) {
const NUM_ENUM_MEMBERS: usize = 512;
setup_rayon();
let mut code = String::new();
writeln!(&mut code, "from enum import Enum").ok();
writeln!(&mut code, "class E(Enum):").ok();
for i in 0..NUM_ENUM_MEMBERS {
writeln!(&mut code, " m{i} = {i}").ok();
}
writeln!(&mut code).ok();
for i in 0..NUM_ENUM_MEMBERS {
writeln!(&mut code, "print(E.m{i})").ok();
}
criterion.bench_function("ty_micro[many_enum_members]", |b| {
b.iter_batched_ref(
|| setup_micro_case(&code),
|case| {
let Case { db, .. } = case;
let result = db.check();
assert_eq!(result.len(), 0);
},
BatchSize::SmallInput,
);
});
}
struct ProjectBenchmark<'a> {
project: InstalledProject<'a>,
fs: MemoryFileSystem,
@ -591,6 +623,7 @@ criterion_group!(
benchmark_many_tuple_assignments,
benchmark_complex_constrained_attributes_1,
benchmark_complex_constrained_attributes_2,
benchmark_many_enum_members,
);
criterion_group!(project, anyio, attrs, hydra, datetype);
criterion_main!(check_file, micro, project);

View file

@ -1542,6 +1542,96 @@ Quux.<CURSOR>
");
}
#[test]
fn enum_attributes() {
let test = cursor_test(
"\
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
Answer.<CURSOR>
",
);
insta::with_settings!({
// See above: filter out some members which contain @Todo types that are
// rendered differently in release mode.
filters => [(r"(?m)\s*__(call|reduce_ex)__.+$", "")]},
{
assert_snapshot!(test.completions_without_builtins_with_types(), @r"
NO :: Literal[Answer.NO]
YES :: Literal[Answer.YES]
mro :: bound method <class 'Answer'>.mro() -> list[type]
name :: Any
value :: Any
__annotations__ :: dict[str, Any]
__base__ :: type | None
__bases__ :: tuple[type, ...]
__basicsize__ :: int
__bool__ :: bound method <class 'Answer'>.__bool__() -> Literal[True]
__class__ :: <class 'EnumMeta'>
__contains__ :: bound method <class 'Answer'>.__contains__(value: object) -> bool
__copy__ :: def __copy__(self) -> Self
__deepcopy__ :: def __deepcopy__(self, memo: Any) -> Self
__delattr__ :: def __delattr__(self, name: str, /) -> None
__dict__ :: MappingProxyType[str, Any]
__dictoffset__ :: int
__dir__ :: def __dir__(self) -> list[str]
__doc__ :: str | None
__eq__ :: def __eq__(self, value: object, /) -> bool
__flags__ :: int
__format__ :: def __format__(self, format_spec: str) -> str
__getattribute__ :: def __getattribute__(self, name: str, /) -> Any
__getitem__ :: bound method <class 'Answer'>.__getitem__(name: str) -> _EnumMemberT
__getstate__ :: def __getstate__(self) -> object
__hash__ :: def __hash__(self) -> int
__init__ :: def __init__(self) -> None
__init_subclass__ :: def __init_subclass__(cls) -> None
__instancecheck__ :: bound method <class 'Answer'>.__instancecheck__(instance: Any, /) -> bool
__itemsize__ :: int
__iter__ :: bound method <class 'Answer'>.__iter__() -> Iterator[_EnumMemberT]
__len__ :: bound method <class 'Answer'>.__len__() -> int
__members__ :: MappingProxyType[str, Unknown]
__module__ :: str
__mro__ :: tuple[<class 'Answer'>, <class 'Enum'>, <class 'object'>]
__name__ :: str
__ne__ :: def __ne__(self, value: object, /) -> bool
__new__ :: def __new__(cls, value: object) -> Self
__or__ :: bound method <class 'Answer'>.__or__(value: Any, /) -> UnionType
__order__ :: str
__prepare__ :: bound method <class 'EnumMeta'>.__prepare__(cls: str, bases: tuple[type, ...], **kwds: Any) -> _EnumDict
__qualname__ :: str
__reduce__ :: def __reduce__(self) -> str | tuple[Any, ...]
__repr__ :: def __repr__(self) -> str
__reversed__ :: bound method <class 'Answer'>.__reversed__() -> Iterator[_EnumMemberT]
__ror__ :: bound method <class 'Answer'>.__ror__(value: Any, /) -> UnionType
__setattr__ :: def __setattr__(self, name: str, value: Any, /) -> None
__signature__ :: bound method <class 'Answer'>.__signature__() -> str
__sizeof__ :: def __sizeof__(self) -> int
__str__ :: def __str__(self) -> str
__subclasscheck__ :: bound method <class 'Answer'>.__subclasscheck__(subclass: type, /) -> bool
__subclasses__ :: bound method <class 'Answer'>.__subclasses__() -> list[Self]
__subclasshook__ :: bound method <class 'Answer'>.__subclasshook__(subclass: type, /) -> bool
__text_signature__ :: str | None
__type_params__ :: tuple[TypeVar | ParamSpec | TypeVarTuple, ...]
__weakrefoffset__ :: int
_generate_next_value_ :: def _generate_next_value_(name: str, start: int, count: int, last_values: list[Any]) -> Any
_ignore_ :: str | list[str]
_member_map_ :: dict[str, Enum]
_member_names_ :: list[str]
_missing_ :: bound method <class 'Answer'>._missing_(value: object) -> Any
_name_ :: str
_order_ :: str
_value2member_map_ :: dict[Any, Enum]
_value_ :: Any
");
}
);
}
// We don't yet take function parameters into account.
#[test]
fn call_prefix1() {

View file

@ -35,8 +35,7 @@ def f():
reveal_type(a6) # revealed: Literal[True]
reveal_type(a7) # revealed: None
reveal_type(a8) # revealed: Literal[1]
# TODO: This should be Color.RED
reveal_type(b1) # revealed: @Todo(Attribute access on enum classes)
reveal_type(b1) # revealed: Literal[Color.RED]
# error: [invalid-type-form]
invalid1: Literal[3 + 4]

View file

@ -2350,19 +2350,17 @@ reveal_type(C().x) # revealed: int
## Enum classes
Enums are not supported yet; attribute access on an enum class is inferred as `Todo`.
```py
import enum
reveal_type(enum.Enum.__members__) # revealed: @Todo(Attribute access on enum classes)
reveal_type(enum.Enum.__members__) # revealed: MappingProxyType[str, Unknown]
class Foo(enum.Enum):
BAR = 1
reveal_type(Foo.BAR) # revealed: @Todo(Attribute access on enum classes)
reveal_type(Foo.BAR.value) # revealed: @Todo(Attribute access on enum classes)
reveal_type(Foo.__members__) # revealed: @Todo(Attribute access on enum classes)
reveal_type(Foo.BAR) # revealed: Literal[Foo.BAR]
reveal_type(Foo.BAR.value) # revealed: Any
reveal_type(Foo.__members__) # revealed: MappingProxyType[str, Unknown]
```
## References

View file

@ -398,12 +398,11 @@ from overloaded import SomeEnum, A, B, C, f
def _(x: SomeEnum):
reveal_type(f(SomeEnum.A)) # revealed: A
# TODO: This should be `B` once enums are supported and are expanded
reveal_type(f(SomeEnum.B)) # revealed: A
# TODO: This should be `C` once enums are supported and are expanded
reveal_type(f(SomeEnum.C)) # revealed: A
# TODO: This should be `A | B | C` once enums are supported and are expanded
reveal_type(f(x)) # revealed: A
reveal_type(f(SomeEnum.B)) # revealed: B
reveal_type(f(SomeEnum.C)) # revealed: C
# TODO: This should not be an error. The return type should be `A | B | C` once enums are expanded
# error: [no-matching-overload]
reveal_type(f(x)) # revealed: Unknown
```
### No matching overloads

View file

@ -0,0 +1,21 @@
# Comparison: Enums
```py
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
reveal_type(Answer.NO == Answer.NO) # revealed: Literal[True]
reveal_type(Answer.NO == Answer.YES) # revealed: Literal[False]
reveal_type(Answer.NO != Answer.NO) # revealed: Literal[False]
reveal_type(Answer.NO != Answer.YES) # revealed: Literal[True]
reveal_type(Answer.NO is Answer.NO) # revealed: Literal[True]
reveal_type(Answer.NO is Answer.YES) # revealed: Literal[False]
reveal_type(Answer.NO is not Answer.NO) # revealed: Literal[False]
reveal_type(Answer.NO is not Answer.YES) # revealed: Literal[True]
```

View file

@ -226,6 +226,28 @@ def _(target: str):
reveal_type(y) # revealed: Literal[1]
```
## Matching on enums
```py
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
def _(answer: Answer):
y = 0
match answer:
case Answer.YES:
reveal_type(answer) # revealed: Literal[Answer.YES]
y = 1
case Answer.NO:
reveal_type(answer) # revealed: Literal[Answer.NO]
y = 2
reveal_type(y) # revealed: Literal[0, 1, 2]
```
## Or match
A `|` pattern matches if any of the subpatterns match.

View file

@ -12,7 +12,7 @@ class Member:
tag: str | None = field(default=None, init=False)
# TODO: this should not include the `tag` parameter, since it has `init=False` set
# revealed: (self: Member, name: str, role: str = Unknown, tag: str | None = Unknown) -> None
# revealed: (self: Member, name: str, role: str = Literal["user"], tag: str | None = None) -> None
reveal_type(Member.__init__)
alice = Member(name="Alice", role="admin")
@ -28,8 +28,5 @@ bob = Member(name="Bob", tag="VIP")
```py
from dataclasses import field
# TODO: this should be `Literal[1]`. This is currently blocked on enum support, because
# the `dataclasses.field` overloads make use of a `_MISSING_TYPE` enum, for which we
# infer a @Todo type, and therefore pick the wrong overload.
reveal_type(field(default=1)) # revealed: Unknown
reveal_type(field(default=1)) # revealed: Literal[1]
```

View file

@ -10,8 +10,9 @@ class Color(Enum):
GREEN = 2
BLUE = 3
reveal_type(Color.RED) # revealed: @Todo(Attribute access on enum classes)
reveal_type(Color.RED.value) # revealed: @Todo(Attribute access on enum classes)
reveal_type(Color.RED) # revealed: Literal[Color.RED]
# TODO: This could be `Literal[1]`
reveal_type(Color.RED.value) # revealed: Any
# TODO: Should be `Color` or `Literal[Color.RED]`
reveal_type(Color["RED"]) # revealed: Unknown
@ -85,6 +86,21 @@ class Answer(Enum):
reveal_type(enum_members(Answer))
```
Enum members are allowed to be marked `Final` (without a type), even if unnecessary:
```py
from enum import Enum
from typing import Final
from ty_extensions import enum_members
class Answer(Enum):
YES: Final = 1
NO: Final = 2
# revealed: tuple[Literal["YES"], Literal["NO"]]
reveal_type(enum_members(Answer))
```
### Non-member attributes with disallowed type
Methods, callables, descriptors (including properties), and nested classes that are defined in the
@ -208,6 +224,27 @@ class Answer(Enum):
# revealed: tuple[Literal["YES"], Literal["NO"]]
reveal_type(enum_members(Answer))
reveal_type(Answer.DEFINITELY) # revealed: Literal[Answer.YES]
```
If a value is duplicated, we also treat that as an alias:
```py
from enum import Enum
class Color(Enum):
RED = 1
GREEN = 2
red = 1
green = 2
# revealed: tuple[Literal["RED"], Literal["GREEN"]]
reveal_type(enum_members(Color))
# revealed: Literal[Color.RED]
reveal_type(Color.red)
```
### Using `auto()`
@ -354,9 +391,8 @@ class Answer(Enum):
# revealed: tuple[Literal["name"], Literal["value"]]
reveal_type(enum_members(Answer))
# TODO: These should be `Answer` or `Literal[Answer.name]`/``Literal[Answer.value]`
reveal_type(Answer.name) # revealed: @Todo(Attribute access on enum classes)
reveal_type(Answer.value) # revealed: @Todo(Attribute access on enum classes)
reveal_type(Answer.name) # revealed: Literal[Answer.name]
reveal_type(Answer.value) # revealed: Literal[Answer.value]
```
## Iterating over enum members
@ -377,6 +413,75 @@ for color in Color:
reveal_type(list(Color)) # revealed: list[Unknown]
```
## Methods / non-member attributes
Methods and non-member attributes defined in the enum class can be accessed on enum members:
```py
from enum import Enum
class Answer(Enum):
YES = 1
NO = 2
def is_yes(self) -> bool:
return self == Answer.YES
constant: int = 1
reveal_type(Answer.YES.is_yes()) # revealed: bool
reveal_type(Answer.YES.constant) # revealed: int
class MyEnum(Enum):
def some_method(self) -> None:
pass
class MyAnswer(MyEnum):
YES = 1
NO = 2
reveal_type(MyAnswer.YES.some_method()) # revealed: None
```
## Accessing enum members from `type[…]`
```py
from enum import Enum
class Answer(Enum):
YES = 1
NO = 2
def _(answer: type[Answer]) -> None:
reveal_type(answer.YES) # revealed: Literal[Answer.YES]
reveal_type(answer.NO) # revealed: Literal[Answer.NO]
```
## Calling enum variants
```py
from enum import Enum
from typing import Callable
import sys
class Printer(Enum):
STDOUT = 1
STDERR = 2
def __call__(self, msg: str) -> None:
if self == Printer.STDOUT:
print(msg)
elif self == Printer.STDERR:
print(msg, file=sys.stderr)
Printer.STDOUT("Hello, world!")
Printer.STDERR("An error occurred!")
callable: Callable[[str], None] = Printer.STDOUT
callable("Hello again!")
callable = Printer.STDERR
callable("Another error!")
```
## Properties of enum types
### Implicitly final
@ -391,14 +496,13 @@ class Color(Enum):
GREEN = 2
BLUE = 3
# TODO: This should emit an error
# error: [subclass-of-final-class] "Class `ExtendedColor` cannot inherit from final class `Color`"
class ExtendedColor(Color):
YELLOW = 4
def f(color: Color):
if isinstance(color, int):
# TODO: This should be `Never`
reveal_type(color) # revealed: Color & int
reveal_type(color) # revealed: Never
```
An `Enum` subclass without any defined members can be subclassed:
@ -419,6 +523,43 @@ class Answer(MyEnum):
reveal_type(enum_members(Answer))
```
### Meta-type
```py
from enum import Enum
class Answer(Enum):
YES = 1
NO = 2
reveal_type(type(Answer.YES)) # revealed: <class 'Answer'>
class NoMembers(Enum): ...
def _(answer: Answer, no_members: NoMembers):
reveal_type(type(answer)) # revealed: <class 'Answer'>
reveal_type(type(no_members)) # revealed: type[NoMembers]
```
### Cyclic references
```py
from enum import Enum
from typing import Literal
from ty_extensions import enum_members
class Answer(Enum):
YES = 1
NO = 2
@classmethod
def yes(cls) -> "Literal[Answer.YES]":
return Answer.YES
# revealed: tuple[Literal["YES"], Literal["NO"]]
reveal_type(enum_members(Answer))
```
## Custom enum types
To do: <https://typing.python.org/en/latest/spec/enums.html#enum-definition>

View file

@ -192,6 +192,21 @@ static_assert("__doc__" in all_members(len))
static_assert("__doc__" in all_members("a".startswith))
```
### Enums
```py
from ty_extensions import all_members, static_assert
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
static_assert("NO" in all_members(Answer))
static_assert("YES" in all_members(Answer))
static_assert("__members__" in all_members(Answer))
```
### Unions
For unions, `all_members` will only return members that are available on all elements of the union.

View file

@ -14,16 +14,73 @@ def _(flag: bool):
## `!=` for other singleton types
```py
def _(flag: bool):
x = True if flag else False
### Bool
```py
def _(x: bool):
if x != False:
reveal_type(x) # revealed: Literal[True]
else:
reveal_type(x) # revealed: Literal[False]
```
### Enums
```py
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
def _(answer: Answer):
if answer != Answer.NO:
# TODO: This should be simplified to `Literal[Answer.YES]`
reveal_type(answer) # revealed: Answer & ~Literal[Answer.NO]
else:
# TODO: This should be `Literal[Answer.NO]`
reveal_type(answer) # revealed: Answer
```
This narrowing behavior is only safe if the enum has no custom `__eq__`/`__ne__` method:
```py
from enum import Enum
class AmbiguousEnum(Enum):
NO = 0
YES = 1
def __ne__(self, other) -> bool:
return True
def _(answer: AmbiguousEnum):
if answer != AmbiguousEnum.NO:
reveal_type(answer) # revealed: AmbiguousEnum
else:
reveal_type(answer) # revealed: AmbiguousEnum
```
Similar if that method is inherited from a base class:
```py
from enum import Enum
class Mixin:
def __eq__(self, other) -> bool:
return True
class AmbiguousEnum(Mixin, Enum):
NO = 0
YES = 1
def _(answer: AmbiguousEnum):
if answer == AmbiguousEnum.NO:
reveal_type(answer) # revealed: AmbiguousEnum
else:
reveal_type(answer) # revealed: AmbiguousEnum
```
## `x != y` where `y` is of literal type
```py

View file

@ -65,6 +65,23 @@ def _(flag1: bool, flag2: bool):
reveal_type(x) # revealed: Literal[1]
```
## `is` for enums
```py
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
def _(answer: Answer):
if answer is Answer.NO:
reveal_type(answer) # revealed: Literal[Answer.NO]
else:
# TODO: This should be `Literal[Answer.YES]`
reveal_type(answer) # revealed: Answer & ~Literal[Answer.NO]
```
## `is` for `EllipsisType` (Python 3.10+)
```toml

View file

@ -154,7 +154,8 @@ reveal_type(Foo.__mro__) # revealed: tuple[<class 'Foo'>, @Todo(GenericAlias in
## `@final` classes
`type[]` types are eagerly converted to class-literal types if a class decorated with `@final` is
used as the type argument. This applies to standard-library classes and user-defined classes:
used as the type argument. This applies to standard-library classes and user-defined classes. The
same also applies to enum classes with members, which are implicitly final:
```toml
[environment]
@ -164,11 +165,17 @@ python-version = "3.10"
```py
from types import EllipsisType
from typing import final
from enum import Enum
@final
class Foo: ...
def _(x: type[Foo], y: type[EllipsisType]):
class Answer(Enum):
NO = 0
YES = 1
def _(x: type[Foo], y: type[EllipsisType], z: type[Answer]):
reveal_type(x) # revealed: <class 'Foo'>
reveal_type(y) # revealed: <class 'EllipsisType'>
reveal_type(z) # revealed: <class 'Answer'>
```

View file

@ -124,6 +124,27 @@ static_assert(not is_assignable_to(Literal[b"foo"], Literal["foo"]))
static_assert(not is_assignable_to(Literal["foo"], Literal[b"foo"]))
```
### Enum literals
```py
from ty_extensions import static_assert, is_assignable_to
from typing_extensions import Literal
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
static_assert(is_assignable_to(Literal[Answer.YES], Literal[Answer.YES]))
static_assert(is_assignable_to(Literal[Answer.YES], Answer))
static_assert(is_assignable_to(Literal[Answer.YES, Answer.NO], Answer))
# TODO: this should not be an error
# error: [static-assert-error]
static_assert(is_assignable_to(Answer, Literal[Answer.YES, Answer.NO]))
static_assert(not is_assignable_to(Literal[Answer.YES], Literal[Answer.NO]))
```
### Slice literals
The type of a slice literal is currently inferred as a specialization of `slice`.

View file

@ -297,6 +297,11 @@ static_assert(is_disjoint_from(None, Intersection[int, Not[str]]))
```py
from typing_extensions import Literal, LiteralString
from ty_extensions import Intersection, Not, TypeOf, is_disjoint_from, static_assert, AlwaysFalsy, AlwaysTruthy
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
static_assert(is_disjoint_from(Literal[True], Literal[False]))
static_assert(is_disjoint_from(Literal[True], Literal[1]))
@ -310,6 +315,10 @@ static_assert(is_disjoint_from(Literal[b"a"], LiteralString))
static_assert(is_disjoint_from(Literal[b"a"], Literal[b"b"]))
static_assert(is_disjoint_from(Literal[b"a"], Literal["a"]))
static_assert(is_disjoint_from(Literal[Answer.YES], Literal[Answer.NO]))
static_assert(is_disjoint_from(Literal[Answer.YES], int))
static_assert(not is_disjoint_from(Literal[Answer.YES], Answer))
static_assert(is_disjoint_from(type[object], TypeOf[Literal]))
static_assert(is_disjoint_from(type[str], LiteralString))
@ -705,3 +714,28 @@ static_assert(not is_disjoint_from(TypeOf[Deque], Callable[..., Any]))
static_assert(not is_disjoint_from(Callable[..., Any], TypeOf[OrderedDict]))
static_assert(not is_disjoint_from(TypeOf[OrderedDict], Callable[..., Any]))
```
## Custom enum classes
```py
from enum import Enum
from ty_extensions import is_disjoint_from, static_assert
from typing_extensions import Literal
class MyEnum(Enum):
def special_method(self):
pass
class MyAnswer(MyEnum):
NO = 0
YES = 1
class UnrelatedClass:
pass
static_assert(is_disjoint_from(Literal[MyAnswer.NO], Literal[MyAnswer.YES]))
static_assert(is_disjoint_from(Literal[MyAnswer.NO], UnrelatedClass))
static_assert(not is_disjoint_from(Literal[MyAnswer.NO], MyAnswer))
static_assert(not is_disjoint_from(Literal[MyAnswer.NO], MyEnum))
```

View file

@ -15,6 +15,11 @@ materializations of `B`, and all materializations of `B` are also materializatio
```py
from typing_extensions import Literal, LiteralString, Never
from ty_extensions import Unknown, is_equivalent_to, static_assert, TypeOf, AlwaysTruthy, AlwaysFalsy
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
static_assert(is_equivalent_to(Literal[1, 2], Literal[1, 2]))
static_assert(is_equivalent_to(type[object], type))
@ -25,6 +30,13 @@ static_assert(not is_equivalent_to(Literal[1, 0], Literal[1, 2]))
static_assert(not is_equivalent_to(Literal[1, 2], Literal[1, 2, 3]))
static_assert(not is_equivalent_to(Literal[1, 2, 3], Literal[1, 2]))
static_assert(is_equivalent_to(Literal[Answer.YES], Literal[Answer.YES]))
# TODO: these should be equivalent
# error: [static-assert-error]
static_assert(is_equivalent_to(Literal[Answer.YES, Answer.NO], Answer))
static_assert(not is_equivalent_to(Literal[Answer.YES], Literal[Answer.NO]))
static_assert(not is_equivalent_to(Literal[Answer.YES], Answer))
static_assert(is_equivalent_to(Never, Never))
static_assert(is_equivalent_to(AlwaysTruthy, AlwaysTruthy))
static_assert(is_equivalent_to(AlwaysFalsy, AlwaysFalsy))

View file

@ -34,3 +34,49 @@ static_assert(is_single_valued(TypeOf[A().method]))
static_assert(is_single_valued(TypeOf[types.FunctionType.__get__]))
static_assert(is_single_valued(TypeOf[A.method.__get__]))
```
An enum literal is only considered single-valued if it has no custom `__eq__`/`__ne__` method, or if
these methods always return `True`/`False`, respectively. Otherwise, the single member of the enum
literal type might not compare equal to itself.
```py
from ty_extensions import is_single_valued, static_assert, TypeOf
from enum import Enum
class NormalEnum(Enum):
NO = 0
YES = 1
class ComparesEqualEnum(Enum):
NO = 0
YES = 1
def __eq__(self, other: object) -> Literal[True]:
return True
class CustomEqEnum(Enum):
NO = 0
YES = 1
def __eq__(self, other: object) -> bool:
return False
class CustomNeEnum(Enum):
NO = 0
YES = 1
def __ne__(self, other: object) -> bool:
return False
static_assert(is_single_valued(Literal[NormalEnum.NO]))
static_assert(is_single_valued(Literal[NormalEnum.YES]))
static_assert(is_single_valued(Literal[ComparesEqualEnum.NO]))
static_assert(is_single_valued(Literal[ComparesEqualEnum.YES]))
static_assert(not is_single_valued(Literal[CustomEqEnum.NO]))
static_assert(not is_single_valued(Literal[CustomEqEnum.YES]))
static_assert(not is_single_valued(Literal[CustomNeEnum.NO]))
static_assert(not is_single_valued(Literal[CustomNeEnum.YES]))
```

View file

@ -7,10 +7,17 @@ A type is a singleton type iff it has exactly one inhabitant.
```py
from typing_extensions import Literal, Never, Callable
from ty_extensions import is_singleton, static_assert
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
static_assert(is_singleton(None))
static_assert(is_singleton(Literal[True]))
static_assert(is_singleton(Literal[False]))
static_assert(is_singleton(Literal[Answer.YES]))
static_assert(is_singleton(Literal[Answer.NO]))
static_assert(is_singleton(type[bool]))

View file

@ -90,6 +90,11 @@ static_assert(is_subtype_of(C, object))
```py
from typing_extensions import Literal, LiteralString
from ty_extensions import is_subtype_of, static_assert, TypeOf, JustFloat
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
# Boolean literals
static_assert(is_subtype_of(Literal[True], bool))
@ -115,6 +120,16 @@ static_assert(is_subtype_of(LiteralString, object))
# Bytes literals
static_assert(is_subtype_of(Literal[b"foo"], bytes))
static_assert(is_subtype_of(Literal[b"foo"], object))
# Enum literals
static_assert(is_subtype_of(Literal[Answer.YES], Literal[Answer.YES]))
static_assert(is_subtype_of(Literal[Answer.YES], Answer))
static_assert(is_subtype_of(Literal[Answer.YES, Answer.NO], Answer))
# TODO: this should not be an error
# error: [static-assert-error]
static_assert(is_subtype_of(Answer, Literal[Answer.YES, Answer.NO]))
static_assert(not is_subtype_of(Literal[Answer.YES], Literal[Answer.NO]))
```
## Heterogeneous tuple types

View file

@ -80,6 +80,11 @@ The top / bottom (and only) materialization of any fully static type is just its
```py
from typing import Any, Literal
from ty_extensions import TypeOf, bottom_materialization, top_materialization
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
reveal_type(top_materialization(int)) # revealed: int
reveal_type(bottom_materialization(int)) # revealed: int
@ -93,6 +98,9 @@ reveal_type(bottom_materialization(Literal[True])) # revealed: Literal[True]
reveal_type(top_materialization(Literal["abc"])) # revealed: Literal["abc"]
reveal_type(bottom_materialization(Literal["abc"])) # revealed: Literal["abc"]
reveal_type(top_materialization(Literal[Answer.YES])) # revealed: Literal[Answer.YES]
reveal_type(bottom_materialization(Literal[Answer.YES])) # revealed: Literal[Answer.YES]
reveal_type(top_materialization(int | str)) # revealed: int | str
reveal_type(bottom_materialization(int | str)) # revealed: int | str
```

View file

@ -2,26 +2,35 @@
```py
from typing_extensions import Literal, LiteralString
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
def _(
a: Literal[1],
b: Literal[True],
c: Literal[False],
d: Literal["ab'cd"],
e: LiteralString,
f: int,
e: Literal[Answer.YES],
f: LiteralString,
g: int,
):
reveal_type(str(a)) # revealed: Literal["1"]
reveal_type(str(b)) # revealed: Literal["True"]
reveal_type(str(c)) # revealed: Literal["False"]
reveal_type(str(d)) # revealed: Literal["ab'cd"]
reveal_type(str(e)) # revealed: LiteralString
reveal_type(str(f)) # revealed: str
reveal_type(str(e)) # revealed: Literal["Answer.YES"]
reveal_type(str(f)) # revealed: LiteralString
reveal_type(str(g)) # revealed: str
reveal_type(repr(a)) # revealed: Literal["1"]
reveal_type(repr(b)) # revealed: Literal["True"]
reveal_type(repr(c)) # revealed: Literal["False"]
reveal_type(repr(d)) # revealed: Literal["'ab\\'cd'"]
reveal_type(repr(e)) # revealed: LiteralString
reveal_type(repr(f)) # revealed: str
# TODO: this could be `<Answer.YES: 1>`
reveal_type(repr(e)) # revealed: str
reveal_type(repr(f)) # revealed: LiteralString
reveal_type(repr(g)) # revealed: str
```

View file

@ -143,3 +143,65 @@ reveal_type(bool(A().method)) # revealed: Literal[True]
reveal_type(bool(f.__get__)) # revealed: Literal[True]
reveal_type(bool(FunctionType.__get__)) # revealed: Literal[True]
```
## Enums
```py
from enum import Enum
from typing import Literal
class NormalEnum(Enum):
NO = 0
YES = 1
class FalsyEnum(Enum):
NO = 0
YES = 1
def __bool__(self) -> Literal[False]:
return False
class AmbiguousEnum(Enum):
NO = 0
YES = 1
def __bool__(self) -> bool:
return self is AmbiguousEnum.YES
class AmbiguousBase(Enum):
def __bool__(self) -> bool:
return True
class AmbiguousEnum2(AmbiguousBase):
NO = 0
YES = 1
class CustomLenEnum(Enum):
NO = 0
YES = 1
def __len__(self):
return 0
# TODO: these could be `Literal[True]`
reveal_type(bool(NormalEnum.NO)) # revealed: bool
reveal_type(bool(NormalEnum.YES)) # revealed: bool
# TODO: these could be `Literal[False]`
reveal_type(bool(FalsyEnum.NO)) # revealed: bool
reveal_type(bool(FalsyEnum.YES)) # revealed: bool
# All of the following must be `bool`:
reveal_type(bool(AmbiguousEnum.NO)) # revealed: bool
reveal_type(bool(AmbiguousEnum.YES)) # revealed: bool
reveal_type(bool(AmbiguousEnum2.NO)) # revealed: bool
reveal_type(bool(AmbiguousEnum2.YES)) # revealed: bool
reveal_type(bool(AmbiguousEnum2.NO)) # revealed: bool
reveal_type(bool(AmbiguousEnum2.YES)) # revealed: bool
reveal_type(bool(CustomLenEnum.NO)) # revealed: bool
reveal_type(bool(CustomLenEnum.YES)) # revealed: bool
```

View file

@ -257,6 +257,8 @@ pub enum KnownModule {
#[cfg(test)]
#[strum(serialize = "unittest.mock")]
UnittestMock,
#[cfg(test)]
Uuid,
}
impl KnownModule {
@ -278,6 +280,8 @@ impl KnownModule {
Self::ImportLib => "importlib",
#[cfg(test)]
Self::UnittestMock => "unittest.mock",
#[cfg(test)]
Self::Uuid => "uuid",
}
}

View file

@ -230,6 +230,7 @@ impl<'db> Completion<'db> {
| Type::StringLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_) => CompletionKind::Value,
Type::EnumLiteral(_) => CompletionKind::Enum,
Type::ProtocolInstance(_) => CompletionKind::Interface,
Type::TypeVar(_) => CompletionKind::TypeParameter,
Type::Union(union) => union.elements(db).iter().find_map(|&ty| imp(db, ty))?,

View file

@ -39,6 +39,7 @@ use crate::types::call::{Binding, Bindings, CallArguments, CallableBinding};
pub(crate) use crate::types::class_base::ClassBase;
use crate::types::context::{LintDiagnosticGuard, LintDiagnosticGuardBuilder};
use crate::types::diagnostic::{INVALID_TYPE_FORM, UNSUPPORTED_BOOL_CONVERSION};
use crate::types::enums::enum_metadata;
use crate::types::function::{
DataclassTransformerParams, FunctionSpans, FunctionType, KnownFunction,
};
@ -568,6 +569,8 @@ pub enum Type<'db> {
BooleanLiteral(bool),
/// A string literal whose value is known
StringLiteral(StringLiteralType<'db>),
/// A singleton type that represents a specific enum member
EnumLiteral(EnumLiteralType<'db>),
/// A string known to originate only from literal values, but whose value is not known (unlike
/// `StringLiteral` above).
LiteralString,
@ -702,6 +705,7 @@ impl<'db> Type<'db> {
| Type::StringLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::SpecialForm(_)
| Type::KnownInstance(_)
| Type::AlwaysFalsy
@ -963,6 +967,7 @@ impl<'db> Type<'db> {
Type::IntLiteral(_) => Some(KnownClass::Int.to_instance(db)),
Type::BytesLiteral(_) => Some(KnownClass::Bytes.to_instance(db)),
Type::ModuleLiteral(_) => Some(KnownClass::ModuleType.to_instance(db)),
Type::EnumLiteral(literal) => Some(literal.enum_class_instance(db)),
_ => None,
}
}
@ -1044,6 +1049,7 @@ impl<'db> Type<'db> {
| Type::AlwaysTruthy
| Type::BooleanLiteral(_)
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::StringLiteral(_)
| Type::Never
| Type::WrapperDescriptor(_)
@ -1079,6 +1085,7 @@ impl<'db> Type<'db> {
| Type::StringLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::SpecialForm(_)
| Type::KnownInstance(_)
| Type::AlwaysFalsy
@ -1146,6 +1153,10 @@ impl<'db> Type<'db> {
Type::Union(union) => union.try_map(db, |element| element.into_callable(db)),
Type::EnumLiteral(enum_literal) => {
enum_literal.enum_class_instance(db).into_callable(db)
}
Type::Never
| Type::DataclassTransformer(_)
| Type::AlwaysTruthy
@ -1376,13 +1387,15 @@ impl<'db> Type<'db> {
| Type::BytesLiteral(_)
| Type::ClassLiteral(_)
| Type::FunctionLiteral(_)
| Type::ModuleLiteral(_),
| Type::ModuleLiteral(_)
| Type::EnumLiteral(_),
Type::StringLiteral(_)
| Type::IntLiteral(_)
| Type::BytesLiteral(_)
| Type::ClassLiteral(_)
| Type::FunctionLiteral(_)
| Type::ModuleLiteral(_),
| Type::ModuleLiteral(_)
| Type::EnumLiteral(_),
) => false,
(Type::Callable(self_callable), Type::Callable(other_callable)) => {
@ -1414,7 +1427,8 @@ impl<'db> Type<'db> {
| Type::BooleanLiteral(_)
| Type::IntLiteral(_)
| Type::BytesLiteral(_)
| Type::ModuleLiteral(_),
| Type::ModuleLiteral(_)
| Type::EnumLiteral(_),
_,
) => (self.literal_fallback_instance(db))
.is_some_and(|instance| instance.has_relation_to(db, target, relation)),
@ -1754,6 +1768,7 @@ impl<'db> Type<'db> {
| Type::IntLiteral(..)
| Type::StringLiteral(..)
| Type::BytesLiteral(..)
| Type::EnumLiteral(..)
| Type::FunctionLiteral(..)
| Type::BoundMethod(..)
| Type::MethodWrapper(..)
@ -1767,6 +1782,7 @@ impl<'db> Type<'db> {
| Type::IntLiteral(..)
| Type::StringLiteral(..)
| Type::BytesLiteral(..)
| Type::EnumLiteral(..)
| Type::FunctionLiteral(..)
| Type::BoundMethod(..)
| Type::MethodWrapper(..)
@ -1794,6 +1810,7 @@ impl<'db> Type<'db> {
| Type::DataclassDecorator(..)
| Type::DataclassTransformer(..)
| Type::IntLiteral(..)
| Type::EnumLiteral(..)
| Type::StringLiteral(..)
| Type::LiteralString,
)
@ -1810,6 +1827,7 @@ impl<'db> Type<'db> {
| Type::DataclassDecorator(..)
| Type::DataclassTransformer(..)
| Type::IntLiteral(..)
| Type::EnumLiteral(..)
| Type::StringLiteral(..)
| Type::LiteralString,
Type::Tuple(..),
@ -1822,6 +1840,7 @@ impl<'db> Type<'db> {
| Type::StringLiteral(..)
| Type::LiteralString
| Type::BytesLiteral(..)
| Type::EnumLiteral(..)
| Type::FunctionLiteral(..)
| Type::BoundMethod(..)
| Type::MethodWrapper(..)
@ -1834,6 +1853,7 @@ impl<'db> Type<'db> {
| Type::StringLiteral(..)
| Type::LiteralString
| Type::BytesLiteral(..)
| Type::EnumLiteral(..)
| Type::FunctionLiteral(..)
| Type::BoundMethod(..)
| Type::MethodWrapper(..)
@ -1902,7 +1922,9 @@ impl<'db> Type<'db> {
| Type::FunctionLiteral(..)
| Type::ModuleLiteral(..)
| Type::GenericAlias(..)
| Type::IntLiteral(..)),
| Type::IntLiteral(..)
| Type::EnumLiteral(..)
),
Type::ProtocolInstance(protocol),
)
| (
@ -1915,7 +1937,8 @@ impl<'db> Type<'db> {
| Type::FunctionLiteral(..)
| Type::ModuleLiteral(..)
| Type::GenericAlias(..)
| Type::IntLiteral(..)),
| Type::IntLiteral(..)
| Type::EnumLiteral(..)),
) => any_protocol_members_absent_or_disjoint(db, protocol, ty, visitor),
// This is the same as the branch above --
@ -2029,6 +2052,12 @@ impl<'db> Type<'db> {
!KnownClass::Bytes.is_subclass_of(db, instance.class)
}
(Type::EnumLiteral(enum_literal), instance@Type::NominalInstance(_))
| (instance@Type::NominalInstance(_), Type::EnumLiteral(enum_literal)) => {
!enum_literal.enum_class_instance(db).is_subtype_of(db, instance)
}
(Type::EnumLiteral(..), _) | (_, Type::EnumLiteral(..)) => true,
// A class-literal type `X` is always disjoint from an instance type `Y`,
// unless the type expressing "all instances of `Z`" is a subtype of of `Y`,
// where `Z` is `X`'s metaclass.
@ -2216,7 +2245,8 @@ impl<'db> Type<'db> {
| Type::WrapperDescriptor(..)
| Type::ClassLiteral(..)
| Type::GenericAlias(..)
| Type::ModuleLiteral(..) => true,
| Type::ModuleLiteral(..)
| Type::EnumLiteral(..) => true,
Type::SpecialForm(special_form) => {
// Nearly all `SpecialForm` types are singletons, but if a symbol could validly
// originate from either `typing` or `typing_extensions` then this is not guaranteed.
@ -2298,6 +2328,24 @@ impl<'db> Type<'db> {
| Type::SpecialForm(..)
| Type::KnownInstance(..) => true,
Type::EnumLiteral(_) => {
let check_dunder = |dunder_name, allowed_return_value| {
let call_result = self.try_call_dunder_with_policy(
db,
dunder_name,
&mut CallArguments::positional([Type::unknown()]),
MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK,
);
let call_result = call_result.as_ref();
call_result.is_ok_and(|bindings| {
bindings.return_type(db) == Type::BooleanLiteral(allowed_return_value)
}) || call_result
.is_err_and(|err| matches!(err, CallDunderError::MethodNotAvailable))
};
check_dunder("__eq__", true) && check_dunder("__ne__", false)
}
Type::ProtocolInstance(..) => {
// See comment in the `Type::ProtocolInstance` branch for `Type::is_singleton`.
false
@ -2457,6 +2505,7 @@ impl<'db> Type<'db> {
| Type::StringLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::Tuple(_)
| Type::TypeVar(_)
| Type::NominalInstance(_)
@ -2586,6 +2635,9 @@ impl<'db> Type<'db> {
KnownClass::Str.to_instance(db).instance_member(db, name)
}
Type::BytesLiteral(_) => KnownClass::Bytes.to_instance(db).instance_member(db, name),
Type::EnumLiteral(enum_literal) => enum_literal
.enum_class_instance(db)
.instance_member(db, name),
Type::Tuple(tuple) => tuple
.to_class_type(db)
.map(|class| class.instance_member(db, name))
@ -3095,6 +3147,7 @@ impl<'db> Type<'db> {
| Type::IntLiteral(..)
| Type::StringLiteral(..)
| Type::BytesLiteral(..)
| Type::EnumLiteral(..)
| Type::LiteralString
| Type::Tuple(..)
| Type::TypeVar(..)
@ -3189,8 +3242,27 @@ impl<'db> Type<'db> {
return class_attr_plain;
}
if self.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) {
return PlaceAndQualifiers::todo("Attribute access on enum classes");
if let Some(enum_class) = match self {
Type::ClassLiteral(literal) => Some(literal),
Type::SubclassOf(subclass_of) => subclass_of
.subclass_of()
.into_class()
.map(|class| class.class_literal(db).0),
_ => None,
} {
if let Some(metadata) = enum_metadata(db, enum_class) {
if let Some(resolved_name) = metadata.resolve_member(&name) {
return Place::Type(
Type::EnumLiteral(EnumLiteralType::new(
db,
enum_class,
resolved_name,
)),
Boundness::Bound,
)
.into();
}
}
}
let class_attr_fallback = Self::try_call_dunder_get_on_attribute(
@ -3424,6 +3496,12 @@ impl<'db> Type<'db> {
Truthiness::Ambiguous
}
Type::EnumLiteral(_) => {
// We currently make no attempt to infer the precise truthiness, but it's not impossible to do so.
// Note that custom `__bool__` or `__len__` methods on the class or superclasses affect the outcome.
Truthiness::Ambiguous
}
Type::IntLiteral(num) => Truthiness::from(*num != 0),
Type::BooleanLiteral(bool) => Truthiness::from(*bool),
Type::StringLiteral(str) => Truthiness::from(!str.value(db).is_empty()),
@ -4369,6 +4447,8 @@ impl<'db> Type<'db> {
// TODO: some `SpecialForm`s are callable (e.g. TypedDicts)
Type::SpecialForm(_) => CallableBinding::not_callable(self).into(),
Type::EnumLiteral(enum_literal) => enum_literal.enum_class_instance(db).bindings(db),
Type::PropertyInstance(_)
| Type::KnownInstance(_)
| Type::AlwaysFalsy
@ -4836,6 +4916,7 @@ impl<'db> Type<'db> {
Type::Intersection(_) => Some(todo_type!("Type::Intersection.to_instance")),
Type::BooleanLiteral(_)
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::FunctionLiteral(_)
| Type::Callable(..)
| Type::MethodWrapper(_)
@ -4904,6 +4985,7 @@ impl<'db> Type<'db> {
Type::SubclassOf(_)
| Type::BooleanLiteral(_)
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::AlwaysTruthy
| Type::AlwaysFalsy
| Type::IntLiteral(_)
@ -5165,6 +5247,7 @@ impl<'db> Type<'db> {
Type::BooleanLiteral(_) | Type::TypeIs(_) => KnownClass::Bool.to_class_literal(db),
Type::BytesLiteral(_) => KnownClass::Bytes.to_class_literal(db),
Type::IntLiteral(_) => KnownClass::Int.to_class_literal(db),
Type::EnumLiteral(enum_literal) => Type::ClassLiteral(enum_literal.enum_class(db)),
Type::FunctionLiteral(_) => KnownClass::FunctionType.to_class_literal(db),
Type::BoundMethod(_) => KnownClass::MethodType.to_class_literal(db),
Type::MethodWrapper(_) => KnownClass::MethodWrapperType.to_class_literal(db),
@ -5341,7 +5424,8 @@ impl<'db> Type<'db> {
| Type::BooleanLiteral(_)
| Type::LiteralString
| Type::StringLiteral(_)
| Type::BytesLiteral(_) => match type_mapping {
| Type::BytesLiteral(_)
| Type::EnumLiteral(_) => match type_mapping {
TypeMapping::Specialization(_) |
TypeMapping::PartialSpecialization(_) => self,
TypeMapping::PromoteLiterals => self.literal_fallback_instance(db)
@ -5463,6 +5547,7 @@ impl<'db> Type<'db> {
| Type::LiteralString
| Type::StringLiteral(_)
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::BoundSuper(_)
| Type::SpecialForm(_)
| Type::KnownInstance(_) => {}
@ -5479,6 +5564,14 @@ impl<'db> Type<'db> {
match self {
Type::IntLiteral(_) | Type::BooleanLiteral(_) => self.repr(db),
Type::StringLiteral(_) | Type::LiteralString => *self,
Type::EnumLiteral(enum_literal) => Type::string_literal(
db,
&format!(
"{enum_class}.{name}",
enum_class = enum_literal.enum_class(db).name(db),
name = enum_literal.name(db)
),
),
Type::SpecialForm(special_form) => Type::string_literal(db, special_form.repr()),
Type::KnownInstance(known_instance) => Type::StringLiteral(StringLiteralType::new(
db,
@ -5555,6 +5648,8 @@ impl<'db> Type<'db> {
| Self::LiteralString
| Self::IntLiteral(_)
| Self::BytesLiteral(_)
// TODO: For enum literals, it would be even better to jump to the definition of the specific member
| Self::EnumLiteral(_)
| Self::MethodWrapper(_)
| Self::WrapperDescriptor(_)
| Self::DataclassDecorator(_)
@ -8224,6 +8319,33 @@ impl<'db> BytesLiteralType<'db> {
}
}
/// A singleton type corresponding to a specific enum member.
///
/// For the enum variant `Answer.YES` of the enum below, this type would store
/// a reference to `Answer` in `enum_class` and the name "YES" in `name`.
/// ```py
/// class Answer(Enum):
/// NO = 0
/// YES = 1
/// ```
#[salsa::interned(debug)]
#[derive(PartialOrd, Ord)]
pub struct EnumLiteralType<'db> {
/// A reference to the enum class this literal belongs to
enum_class: ClassLiteral<'db>,
/// The name of the enum member
name: Name,
}
// The Salsa heap is tracked separately.
impl get_size2::GetSize for EnumLiteralType<'_> {}
impl<'db> EnumLiteralType<'db> {
pub fn enum_class_instance(self, db: &'db dyn Db) -> Type<'db> {
self.enum_class(db).to_non_generic_instance(db)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum BoundSuperError<'db> {
InvalidPivotClassType {

View file

@ -658,12 +658,19 @@ impl<'db> Bindings<'db> {
Some(KnownFunction::EnumMembers) => {
if let [Some(ty)] = overload.parameter_types() {
let return_ty = match ty {
Type::ClassLiteral(class) => TupleType::from_elements(
db,
enums::enum_members(db, *class)
.into_iter()
.map(|member| Type::string_literal(db, &member)),
),
Type::ClassLiteral(class) => {
if let Some(metadata) = enums::enum_metadata(db, *class) {
TupleType::from_elements(
db,
metadata
.members
.iter()
.map(|member| Type::string_literal(db, member)),
)
} else {
Type::unknown()
}
}
_ => Type::unknown(),
};

View file

@ -14,6 +14,7 @@ use crate::semantic_index::place::NodeWithScopeKind;
use crate::semantic_index::{DeclarationWithConstraint, SemanticIndex, attribute_declarations};
use crate::types::context::InferContext;
use crate::types::diagnostic::{INVALID_LEGACY_TYPE_VARIABLE, INVALID_TYPE_ALIAS_TYPE};
use crate::types::enums::enum_metadata;
use crate::types::function::{DataclassTransformerParams, KnownFunction};
use crate::types::generics::{GenericContext, Specialization, walk_specialization};
use crate::types::infer::nearest_enclosing_class;
@ -1074,6 +1075,7 @@ impl<'db> ClassLiteral<'db> {
pub(super) fn is_final(self, db: &'db dyn Db) -> bool {
self.known_function_decorators(db)
.contains(&KnownFunction::Final)
|| enum_metadata(db, self).is_some()
}
/// Attempt to resolve the [method resolution order] ("MRO") for this class.
@ -2178,6 +2180,10 @@ impl<'db> ClassLiteral<'db> {
}
}
pub(super) fn to_non_generic_instance(self, db: &'db dyn Db) -> Type<'db> {
Type::instance(db, ClassType::NonGeneric(self))
}
/// Return this class' involvement in an inheritance cycle, if any.
///
/// A class definition like this will fail at runtime,
@ -2597,6 +2603,76 @@ impl KnownClass {
}
}
/// Return `true` if this class is a subclass of `enum.Enum` *and* has enum members, i.e.
/// if it is an "actual" enum, not `enum.Enum` itself or a similar custom enum class.
pub(crate) const fn is_enum_subclass_with_members(self) -> bool {
match self {
KnownClass::Bool
| KnownClass::Object
| KnownClass::Bytes
| KnownClass::Bytearray
| KnownClass::Type
| KnownClass::Int
| KnownClass::Float
| KnownClass::Complex
| KnownClass::Str
| KnownClass::List
| KnownClass::Tuple
| KnownClass::Set
| KnownClass::FrozenSet
| KnownClass::Dict
| KnownClass::Slice
| KnownClass::Property
| KnownClass::BaseException
| KnownClass::Exception
| KnownClass::BaseExceptionGroup
| KnownClass::ExceptionGroup
| KnownClass::Staticmethod
| KnownClass::Classmethod
| KnownClass::Super
| KnownClass::Enum
| KnownClass::Auto
| KnownClass::Member
| KnownClass::Nonmember
| KnownClass::ABCMeta
| KnownClass::GenericAlias
| KnownClass::ModuleType
| KnownClass::FunctionType
| KnownClass::MethodType
| KnownClass::MethodWrapperType
| KnownClass::WrapperDescriptorType
| KnownClass::UnionType
| KnownClass::GeneratorType
| KnownClass::AsyncGeneratorType
| KnownClass::NoneType
| KnownClass::Any
| KnownClass::StdlibAlias
| KnownClass::SpecialForm
| KnownClass::TypeVar
| KnownClass::ParamSpec
| KnownClass::ParamSpecArgs
| KnownClass::ParamSpecKwargs
| KnownClass::TypeVarTuple
| KnownClass::TypeAliasType
| KnownClass::NoDefaultType
| KnownClass::NamedTuple
| KnownClass::NewType
| KnownClass::SupportsIndex
| KnownClass::Iterable
| KnownClass::ChainMap
| KnownClass::Counter
| KnownClass::DefaultDict
| KnownClass::Deque
| KnownClass::OrderedDict
| KnownClass::VersionInfo
| KnownClass::EllipsisType
| KnownClass::NotImplementedType
| KnownClass::Field
| KnownClass::KwOnly
| KnownClass::NamedTupleFallback => false,
}
}
/// Return `true` if this class is a protocol class.
///
/// In an ideal world, perhaps we wouldn't hardcode this knowledge here;

View file

@ -148,6 +148,7 @@ impl<'db> ClassBase<'db> {
| Type::DataclassTransformer(_)
| Type::BytesLiteral(_)
| Type::IntLiteral(_)
| Type::EnumLiteral(_)
| Type::StringLiteral(_)
| Type::LiteralString
| Type::Tuple(_)

View file

@ -44,7 +44,8 @@ impl Display for DisplayType<'_> {
Type::IntLiteral(_)
| Type::BooleanLiteral(_)
| Type::StringLiteral(_)
| Type::BytesLiteral(_) => {
| Type::BytesLiteral(_)
| Type::EnumLiteral(_) => {
write!(f, "Literal[{representation}]")
}
_ => representation.fmt(f),
@ -192,6 +193,14 @@ impl Display for DisplayRepresentation<'_> {
escape.bytes_repr(TripleQuotes::No).write(f)
}
Type::EnumLiteral(enum_literal) => {
write!(
f,
"{enum_class}.{name}",
enum_class = enum_literal.enum_class(self.db).name(self.db),
name = enum_literal.name(self.db),
)
}
Type::Tuple(specialization) => specialization.tuple(self.db).display(self.db).fmt(f),
Type::TypeVar(typevar) => f.write_str(typevar.name(self.db)),
Type::AlwaysTruthy => f.write_str("AlwaysTruthy"),
@ -800,6 +809,7 @@ impl Display for DisplayUnionType<'_> {
| Type::StringLiteral(_)
| Type::BytesLiteral(_)
| Type::BooleanLiteral(_)
| Type::EnumLiteral(_)
)
}

View file

@ -1,19 +1,76 @@
use rustc_hash::FxHashSet;
use ruff_python_ast::name::Name;
use rustc_hash::FxHashMap;
use crate::{
Db,
place::{Place, place_from_bindings, place_from_declarations},
place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
semantic_index::{place_table, use_def_map},
types::{ClassLiteral, KnownClass, MemberLookupPolicy, Type},
types::{ClassLiteral, DynamicType, KnownClass, MemberLookupPolicy, Type, TypeQualifiers},
};
#[derive(Debug, PartialEq, Eq, get_size2::GetSize)]
pub(crate) struct EnumMetadata {
pub(crate) members: Box<[Name]>,
pub(crate) aliases: FxHashMap<Name, Name>,
}
impl EnumMetadata {
fn empty() -> Self {
EnumMetadata {
members: Box::new([]),
aliases: FxHashMap::default(),
}
}
pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> {
if self.members.contains(name) {
Some(name)
} else {
self.aliases.get(name)
}
}
}
#[allow(clippy::ref_option)]
fn enum_metadata_cycle_recover(
_db: &dyn Db,
_value: &Option<EnumMetadata>,
_count: u32,
_class: ClassLiteral<'_>,
) -> salsa::CycleRecoveryAction<Option<EnumMetadata>> {
salsa::CycleRecoveryAction::Iterate
}
#[allow(clippy::unnecessary_wraps)]
fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option<EnumMetadata> {
Some(EnumMetadata::empty())
}
/// List all members of an enum.
pub(crate) fn enum_members<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> Vec<String> {
#[allow(clippy::ref_option, clippy::unnecessary_wraps)]
#[salsa::tracked(returns(ref), cycle_fn=enum_metadata_cycle_recover, cycle_initial=enum_metadata_cycle_initial, heap_size=get_size2::GetSize::get_heap_size)]
pub(crate) fn enum_metadata<'db>(
db: &'db dyn Db,
class: ClassLiteral<'db>,
) -> Option<EnumMetadata> {
// This is a fast path to avoid traversing the MRO of known classes
if class
.known(db)
.is_some_and(|known_class| !known_class.is_enum_subclass_with_members())
{
return None;
}
// TODO: This check needs to be extended (`EnumMeta`/`EnumType`)
if !Type::ClassLiteral(class).is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) {
return None;
}
let scope_id = class.body_scope(db);
let use_def_map = use_def_map(db, scope_id);
let table = place_table(db, scope_id);
let mut enum_values: FxHashSet<Type<'db>> = FxHashSet::default();
let mut enum_values: FxHashMap<Type<'db>, Name> = FxHashMap::default();
// TODO: handle `StrEnum` which uses lowercase names as values when using `auto()`.
let mut auto_counter = 0;
@ -33,13 +90,12 @@ pub(crate) fn enum_members<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> Ve
None
};
use_def_map
let mut aliases = FxHashMap::default();
let members = use_def_map
.all_end_of_scope_bindings()
.filter_map(|(place_id, bindings)| {
let name = table
.place_expr(place_id)
.as_name()
.map(ToString::to_string)?;
let name = table.place_expr(place_id).as_name()?;
if name.starts_with("__") && !name.ends_with("__") {
// Skip private attributes
@ -116,21 +172,31 @@ pub(crate) fn enum_members<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> Ve
if matches!(
value_ty,
Type::IntLiteral(_) | Type::StringLiteral(_) | Type::BytesLiteral(_)
) && !enum_values.insert(value_ty)
{
return None;
) {
if let Some(previous) = enum_values.insert(value_ty, name.clone()) {
aliases.insert(name.clone(), previous);
return None;
}
}
let declarations = use_def_map.end_of_scope_declarations(place_id);
let declared = place_from_declarations(db, declarations);
match declared.map(|d| d.place) {
Ok(Place::Unbound) => {
match declared {
Ok(PlaceAndQualifiers {
place: Place::Type(Type::Dynamic(DynamicType::Unknown), _),
qualifiers,
}) if qualifiers.contains(TypeQualifiers::FINAL) => {}
Ok(PlaceAndQualifiers {
place: Place::Unbound,
..
}) => {
// Undeclared attributes are considered members
}
Ok(Place::Type(Type::NominalInstance(instance), _))
if instance.class.is_known(db, KnownClass::Member) =>
{
Ok(PlaceAndQualifiers {
place: Place::Type(Type::NominalInstance(instance), _),
..
}) if instance.class.is_known(db, KnownClass::Member) => {
// If the attribute is specifically declared with `enum.member`, it is considered a member
}
_ => {
@ -141,5 +207,13 @@ pub(crate) fn enum_members<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> Ve
Some(name)
})
.collect()
.cloned()
.collect::<Box<_>>();
if members.is_empty() {
// Enum subclasses without members are not considered enums.
return None;
}
Some(EnumMetadata { members, aliases })
}

View file

@ -119,6 +119,7 @@ impl<'db> AllMembers<'db> {
| Type::BooleanLiteral(_)
| Type::StringLiteral(_)
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::LiteralString
| Type::Tuple(_)
| Type::PropertyInstance(_)

View file

@ -3425,6 +3425,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
| Type::IntLiteral(..)
| Type::StringLiteral(..)
| Type::BytesLiteral(..)
| Type::EnumLiteral(..)
| Type::LiteralString
| Type::Tuple(..)
| Type::SpecialForm(..)
@ -6382,6 +6383,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
| Type::StringLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::Tuple(_)
| Type::BoundSuper(_)
| Type::TypeVar(_)
@ -6710,6 +6712,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
| Type::StringLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::Tuple(_)
| Type::BoundSuper(_)
| Type::TypeVar(_)
@ -6738,6 +6741,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
| Type::StringLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::Tuple(_)
| Type::BoundSuper(_)
| Type::TypeVar(_)
@ -7390,6 +7394,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
KnownClass::Bytes.to_instance(self.db()),
range,
),
(Type::EnumLiteral(literal_1), Type::EnumLiteral(literal_2))
if op == ast::CmpOp::Eq =>
{
Ok(Type::BooleanLiteral(literal_1 == literal_2))
}
(Type::EnumLiteral(literal_1), Type::EnumLiteral(literal_2))
if op == ast::CmpOp::NotEq =>
{
Ok(Type::BooleanLiteral(literal_1 != literal_2))
}
(Type::Tuple(_), Type::NominalInstance(instance))
if instance.class.is_known(self.db(), KnownClass::VersionInfo) =>
{

View file

@ -236,6 +236,7 @@ impl ClassInfoConstraintFunction {
| Type::BoundMethod(_)
| Type::BoundSuper(_)
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::Callable(_)
| Type::DataclassDecorator(_)
| Type::Never

View file

@ -2,8 +2,8 @@ use crate::db::tests::TestDb;
use crate::place::{builtins_symbol, known_module_symbol};
use crate::types::tuple::TupleType;
use crate::types::{
BoundMethodType, CallableType, IntersectionBuilder, KnownClass, Parameter, Parameters,
Signature, SpecialFormType, SubclassOfType, Type, UnionType,
BoundMethodType, CallableType, EnumLiteralType, IntersectionBuilder, KnownClass, Parameter,
Parameters, Signature, SpecialFormType, SubclassOfType, Type, UnionType,
};
use crate::{Db, KnownModule};
use hashbrown::HashSet;
@ -25,6 +25,8 @@ pub(crate) enum Ty {
StringLiteral(&'static str),
LiteralString,
BytesLiteral(&'static str),
// An enum literal variant, using `uuid.SafeUUID` as base
EnumLiteral(&'static str),
// BuiltinInstance("str") corresponds to an instance of the builtin `str` class
BuiltinInstance(&'static str),
/// Members of the `abc` stdlib module
@ -135,6 +137,14 @@ impl Ty {
Ty::BooleanLiteral(b) => Type::BooleanLiteral(b),
Ty::LiteralString => Type::LiteralString,
Ty::BytesLiteral(s) => Type::bytes_literal(db, s.as_bytes()),
Ty::EnumLiteral(name) => Type::EnumLiteral(EnumLiteralType::new(
db,
known_module_symbol(db, KnownModule::Uuid, "SafeUUID")
.place
.expect_type()
.expect_class_literal(),
Name::new(name),
)),
Ty::BuiltinInstance(s) => builtins_symbol(db, s)
.place
.expect_type()
@ -252,6 +262,9 @@ fn arbitrary_core_type(g: &mut Gen, fully_static: bool) -> Ty {
Ty::LiteralString,
Ty::BytesLiteral(""),
Ty::BytesLiteral("\x00"),
Ty::EnumLiteral("safe"),
Ty::EnumLiteral("unsafe"),
Ty::EnumLiteral("unknown"),
Ty::KnownClassInstance(KnownClass::Object),
Ty::KnownClassInstance(KnownClass::Str),
Ty::KnownClassInstance(KnownClass::Int),

View file

@ -64,6 +64,10 @@ pub(super) fn union_or_intersection_elements_ordering<'db>(
(Type::BytesLiteral(_), _) => Ordering::Less,
(_, Type::BytesLiteral(_)) => Ordering::Greater,
(Type::EnumLiteral(left), Type::EnumLiteral(right)) => left.cmp(right),
(Type::EnumLiteral(_), _) => Ordering::Less,
(_, Type::EnumLiteral(_)) => Ordering::Greater,
(Type::FunctionLiteral(left), Type::FunctionLiteral(right)) => left.cmp(right),
(Type::FunctionLiteral(_), _) => Ordering::Less,
(_, Type::FunctionLiteral(_)) => Ordering::Greater,

View file

@ -146,6 +146,7 @@ impl<'db> From<Type<'db>> for TypeKind<'db> {
| Type::BooleanLiteral(_)
| Type::StringLiteral(_)
| Type::BytesLiteral(_)
| Type::EnumLiteral(_)
| Type::DataclassDecorator(_)
| Type::DataclassTransformer(_)
| Type::WrapperDescriptor(_)