diff --git a/crates/ruff_benchmark/benches/ty.rs b/crates/ruff_benchmark/benches/ty.rs index 905f30b9bd..04c79bff8e 100644 --- a/crates/ruff_benchmark/benches/ty.rs +++ b/crates/ruff_benchmark/benches/ty.rs @@ -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); diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index 4e642b0218..3ec5acfa85 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -1542,6 +1542,96 @@ Quux. "); } + #[test] + fn enum_attributes() { + let test = cursor_test( + "\ +from enum import Enum + +class Answer(Enum): + NO = 0 + YES = 1 + +Answer. +", + ); + + 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 .mro() -> list[type] + name :: Any + value :: Any + __annotations__ :: dict[str, Any] + __base__ :: type | None + __bases__ :: tuple[type, ...] + __basicsize__ :: int + __bool__ :: bound method .__bool__() -> Literal[True] + __class__ :: + __contains__ :: bound method .__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 .__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 .__instancecheck__(instance: Any, /) -> bool + __itemsize__ :: int + __iter__ :: bound method .__iter__() -> Iterator[_EnumMemberT] + __len__ :: bound method .__len__() -> int + __members__ :: MappingProxyType[str, Unknown] + __module__ :: str + __mro__ :: tuple[, , ] + __name__ :: str + __ne__ :: def __ne__(self, value: object, /) -> bool + __new__ :: def __new__(cls, value: object) -> Self + __or__ :: bound method .__or__(value: Any, /) -> UnionType + __order__ :: str + __prepare__ :: bound method .__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 .__reversed__() -> Iterator[_EnumMemberT] + __ror__ :: bound method .__ror__(value: Any, /) -> UnionType + __setattr__ :: def __setattr__(self, name: str, value: Any, /) -> None + __signature__ :: bound method .__signature__() -> str + __sizeof__ :: def __sizeof__(self) -> int + __str__ :: def __str__(self) -> str + __subclasscheck__ :: bound method .__subclasscheck__(subclass: type, /) -> bool + __subclasses__ :: bound method .__subclasses__() -> list[Self] + __subclasshook__ :: bound method .__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 ._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() { diff --git a/crates/ty_python_semantic/resources/mdtest/annotations/literal.md b/crates/ty_python_semantic/resources/mdtest/annotations/literal.md index 865bf073f1..05b0868523 100644 --- a/crates/ty_python_semantic/resources/mdtest/annotations/literal.md +++ b/crates/ty_python_semantic/resources/mdtest/annotations/literal.md @@ -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] diff --git a/crates/ty_python_semantic/resources/mdtest/attributes.md b/crates/ty_python_semantic/resources/mdtest/attributes.md index bba2b4688d..260cee2f91 100644 --- a/crates/ty_python_semantic/resources/mdtest/attributes.md +++ b/crates/ty_python_semantic/resources/mdtest/attributes.md @@ -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 diff --git a/crates/ty_python_semantic/resources/mdtest/call/overloads.md b/crates/ty_python_semantic/resources/mdtest/call/overloads.md index a66e28ad7e..1f41f34087 100644 --- a/crates/ty_python_semantic/resources/mdtest/call/overloads.md +++ b/crates/ty_python_semantic/resources/mdtest/call/overloads.md @@ -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 diff --git a/crates/ty_python_semantic/resources/mdtest/comparison/enums.md b/crates/ty_python_semantic/resources/mdtest/comparison/enums.md new file mode 100644 index 0000000000..e4034adf59 --- /dev/null +++ b/crates/ty_python_semantic/resources/mdtest/comparison/enums.md @@ -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] +``` diff --git a/crates/ty_python_semantic/resources/mdtest/conditional/match.md b/crates/ty_python_semantic/resources/mdtest/conditional/match.md index 532c8c46cf..ec387e56be 100644 --- a/crates/ty_python_semantic/resources/mdtest/conditional/match.md +++ b/crates/ty_python_semantic/resources/mdtest/conditional/match.md @@ -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. diff --git a/crates/ty_python_semantic/resources/mdtest/dataclasses/fields.md b/crates/ty_python_semantic/resources/mdtest/dataclasses/fields.md index cafed930ec..1c9f27ccb2 100644 --- a/crates/ty_python_semantic/resources/mdtest/dataclasses/fields.md +++ b/crates/ty_python_semantic/resources/mdtest/dataclasses/fields.md @@ -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] ``` diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 5e1931f374..e6fbf723d5 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -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 NoMembers(Enum): ... + +def _(answer: Answer, no_members: NoMembers): + reveal_type(type(answer)) # revealed: + 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: diff --git a/crates/ty_python_semantic/resources/mdtest/ide_support/all_members.md b/crates/ty_python_semantic/resources/mdtest/ide_support/all_members.md index 05b4ceb75a..79629c971c 100644 --- a/crates/ty_python_semantic/resources/mdtest/ide_support/all_members.md +++ b/crates/ty_python_semantic/resources/mdtest/ide_support/all_members.md @@ -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. diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md index 4ba04803c6..fb1943c0fe 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md @@ -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 diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/is.md b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/is.md index c7d99c48b2..8736ee40fd 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/is.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/is.md @@ -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 diff --git a/crates/ty_python_semantic/resources/mdtest/type_of/basic.md b/crates/ty_python_semantic/resources/mdtest/type_of/basic.md index e733df9b94..b95e5154aa 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_of/basic.md +++ b/crates/ty_python_semantic/resources/mdtest/type_of/basic.md @@ -154,7 +154,8 @@ reveal_type(Foo.__mro__) # revealed: tuple[, @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: reveal_type(y) # revealed: + reveal_type(z) # revealed: ``` diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/is_assignable_to.md b/crates/ty_python_semantic/resources/mdtest/type_properties/is_assignable_to.md index ab730adb68..c499773536 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/is_assignable_to.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/is_assignable_to.md @@ -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`. diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/is_disjoint_from.md b/crates/ty_python_semantic/resources/mdtest/type_properties/is_disjoint_from.md index 81231b4ad5..6ec192fb8c 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/is_disjoint_from.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/is_disjoint_from.md @@ -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)) +``` diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/is_equivalent_to.md b/crates/ty_python_semantic/resources/mdtest/type_properties/is_equivalent_to.md index b309c266be..d2815b3cdb 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/is_equivalent_to.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/is_equivalent_to.md @@ -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)) diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/is_single_valued.md b/crates/ty_python_semantic/resources/mdtest/type_properties/is_single_valued.md index e947c24883..5c97c2524c 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/is_single_valued.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/is_single_valued.md @@ -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])) +``` diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/is_singleton.md b/crates/ty_python_semantic/resources/mdtest/type_properties/is_singleton.md index a3c41449bc..c50a968169 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/is_singleton.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/is_singleton.md @@ -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])) diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/is_subtype_of.md b/crates/ty_python_semantic/resources/mdtest/type_properties/is_subtype_of.md index 5726ce83f7..fa9dfc1eee 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/is_subtype_of.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/is_subtype_of.md @@ -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 diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/materialization.md b/crates/ty_python_semantic/resources/mdtest/type_properties/materialization.md index ebfb45801c..e6009f2b5b 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/materialization.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/materialization.md @@ -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 ``` diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/str_repr.md b/crates/ty_python_semantic/resources/mdtest/type_properties/str_repr.md index 1539daabe1..e267601929 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/str_repr.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/str_repr.md @@ -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 `` + reveal_type(repr(e)) # revealed: str + reveal_type(repr(f)) # revealed: LiteralString + reveal_type(repr(g)) # revealed: str ``` diff --git a/crates/ty_python_semantic/resources/mdtest/type_properties/truthiness.md b/crates/ty_python_semantic/resources/mdtest/type_properties/truthiness.md index 4cb8bdbc45..5e5febf3a2 100644 --- a/crates/ty_python_semantic/resources/mdtest/type_properties/truthiness.md +++ b/crates/ty_python_semantic/resources/mdtest/type_properties/truthiness.md @@ -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 +``` diff --git a/crates/ty_python_semantic/src/module_resolver/module.rs b/crates/ty_python_semantic/src/module_resolver/module.rs index c1d322f90e..4311679665 100644 --- a/crates/ty_python_semantic/src/module_resolver/module.rs +++ b/crates/ty_python_semantic/src/module_resolver/module.rs @@ -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", } } diff --git a/crates/ty_python_semantic/src/semantic_model.rs b/crates/ty_python_semantic/src/semantic_model.rs index d15eaff9d9..faa07b5f23 100644 --- a/crates/ty_python_semantic/src/semantic_model.rs +++ b/crates/ty_python_semantic/src/semantic_model.rs @@ -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))?, diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index e8fce04ec0..0f69f6e582 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -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 { diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index b0157c2e81..79a79ddb88 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/crates/ty_python_semantic/src/types/call/bind.rs @@ -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(), }; diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index fa429c4dce..86bed2e1f8 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -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; diff --git a/crates/ty_python_semantic/src/types/class_base.rs b/crates/ty_python_semantic/src/types/class_base.rs index f682f5d9f4..23f2aca396 100644 --- a/crates/ty_python_semantic/src/types/class_base.rs +++ b/crates/ty_python_semantic/src/types/class_base.rs @@ -148,6 +148,7 @@ impl<'db> ClassBase<'db> { | Type::DataclassTransformer(_) | Type::BytesLiteral(_) | Type::IntLiteral(_) + | Type::EnumLiteral(_) | Type::StringLiteral(_) | Type::LiteralString | Type::Tuple(_) diff --git a/crates/ty_python_semantic/src/types/display.rs b/crates/ty_python_semantic/src/types/display.rs index c778f7bf63..b99589dee9 100644 --- a/crates/ty_python_semantic/src/types/display.rs +++ b/crates/ty_python_semantic/src/types/display.rs @@ -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(_) ) } diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 53e534a9c3..85b3c56094 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -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, +} + +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, + _count: u32, + _class: ClassLiteral<'_>, +) -> salsa::CycleRecoveryAction> { + salsa::CycleRecoveryAction::Iterate +} + +#[allow(clippy::unnecessary_wraps)] +fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option { + Some(EnumMetadata::empty()) +} + /// List all members of an enum. -pub(crate) fn enum_members<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> Vec { +#[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 { + // 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> = FxHashSet::default(); + let mut enum_values: FxHashMap, 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::>(); + + if members.is_empty() { + // Enum subclasses without members are not considered enums. + return None; + } + + Some(EnumMetadata { members, aliases }) } diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index de075fefe9..9a98286afe 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/crates/ty_python_semantic/src/types/ide_support.rs @@ -119,6 +119,7 @@ impl<'db> AllMembers<'db> { | Type::BooleanLiteral(_) | Type::StringLiteral(_) | Type::BytesLiteral(_) + | Type::EnumLiteral(_) | Type::LiteralString | Type::Tuple(_) | Type::PropertyInstance(_) diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index c32cd62c97..72f9a01684 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -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) => { diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index 20a912c72d..1838255223 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -236,6 +236,7 @@ impl ClassInfoConstraintFunction { | Type::BoundMethod(_) | Type::BoundSuper(_) | Type::BytesLiteral(_) + | Type::EnumLiteral(_) | Type::Callable(_) | Type::DataclassDecorator(_) | Type::Never diff --git a/crates/ty_python_semantic/src/types/property_tests/type_generation.rs b/crates/ty_python_semantic/src/types/property_tests/type_generation.rs index 1d0a651ada..d0d7097af3 100644 --- a/crates/ty_python_semantic/src/types/property_tests/type_generation.rs +++ b/crates/ty_python_semantic/src/types/property_tests/type_generation.rs @@ -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), diff --git a/crates/ty_python_semantic/src/types/type_ordering.rs b/crates/ty_python_semantic/src/types/type_ordering.rs index 4054f58e17..58b62af789 100644 --- a/crates/ty_python_semantic/src/types/type_ordering.rs +++ b/crates/ty_python_semantic/src/types/type_ordering.rs @@ -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, diff --git a/crates/ty_python_semantic/src/types/visitor.rs b/crates/ty_python_semantic/src/types/visitor.rs index 1deaef62fc..3c9bb00a14 100644 --- a/crates/ty_python_semantic/src/types/visitor.rs +++ b/crates/ty_python_semantic/src/types/visitor.rs @@ -146,6 +146,7 @@ impl<'db> From> for TypeKind<'db> { | Type::BooleanLiteral(_) | Type::StringLiteral(_) | Type::BytesLiteral(_) + | Type::EnumLiteral(_) | Type::DataclassDecorator(_) | Type::DataclassTransformer(_) | Type::WrapperDescriptor(_)