diff --git a/crates/ty_python_semantic/resources/mdtest/annotations/literal.md b/crates/ty_python_semantic/resources/mdtest/annotations/literal.md index 3bd9e54c85..897be97e77 100644 --- a/crates/ty_python_semantic/resources/mdtest/annotations/literal.md +++ b/crates/ty_python_semantic/resources/mdtest/annotations/literal.md @@ -39,6 +39,8 @@ def f(): reveal_type(a7) # revealed: None reveal_type(a8) # revealed: Literal[1] reveal_type(b1) # revealed: Literal[Color.RED] + # TODO should be `Literal[MissingT.MISSING]` + reveal_type(b2) # revealed: @Todo(functional `Enum` syntax) # error: [invalid-type-form] invalid1: Literal[3 + 4] @@ -66,6 +68,208 @@ a_list: list[int] = [1, 2, 3] invalid6: Literal[a_list[0]] ``` +## Parameterizing with a type alias + +`typing.Literal` can also be parameterized with a type alias for any literal type or union of +literal types. + +### PEP 695 type alias + +```toml +[environment] +python-version = "3.12" +``` + +```py +from typing import Literal +from enum import Enum + +import mod + +class E(Enum): + A = 1 + B = 2 + +type SingleInt = Literal[1] +type SingleStr = Literal["foo"] +type SingleBytes = Literal[b"bar"] +type SingleBool = Literal[True] +type SingleNone = Literal[None] +type SingleEnum = Literal[E.A] +type UnionLiterals = Literal[1, "foo", b"bar", True, None, E.A] +# We support this because it is an equivalent type to the following union of literals, but maybe +# we should not, because it doesn't use `Literal` form? Other type checkers do not. +type AnEnum1 = E +type AnEnum2 = Literal[E.A, E.B] +# Similarly, we support this because it is equivalent to `Literal[True, False]`. +type Bool1 = bool +type Bool2 = Literal[True, False] + +def _( + single_int: Literal[SingleInt], + single_str: Literal[SingleStr], + single_bytes: Literal[SingleBytes], + single_bool: Literal[SingleBool], + single_none: Literal[SingleNone], + single_enum: Literal[SingleEnum], + union_literals: Literal[UnionLiterals], + an_enum1: Literal[AnEnum1], + an_enum2: Literal[AnEnum2], + bool1: Literal[Bool1], + bool2: Literal[Bool2], + multiple: Literal[SingleInt, SingleStr, SingleEnum], + single_int_other_module: Literal[mod.SingleInt], +): + reveal_type(single_int) # revealed: Literal[1] + reveal_type(single_str) # revealed: Literal["foo"] + reveal_type(single_bytes) # revealed: Literal[b"bar"] + reveal_type(single_bool) # revealed: Literal[True] + reveal_type(single_none) # revealed: None + reveal_type(single_enum) # revealed: Literal[E.A] + reveal_type(union_literals) # revealed: Literal[1, "foo", b"bar", True, E.A] | None + reveal_type(an_enum1) # revealed: E + reveal_type(an_enum2) # revealed: E + reveal_type(bool1) # revealed: bool + reveal_type(bool2) # revealed: bool + reveal_type(multiple) # revealed: Literal[1, "foo", E.A] + reveal_type(single_int_other_module) # revealed: Literal[2] +``` + +`mod.py`: + +```py +from typing import Literal + +type SingleInt = Literal[2] +``` + +### PEP 613 type alias + +```py +from typing import Literal, TypeAlias +from enum import Enum + +class E(Enum): + A = 1 + B = 2 + +SingleInt: TypeAlias = Literal[1] +SingleStr: TypeAlias = Literal["foo"] +SingleBytes: TypeAlias = Literal[b"bar"] +SingleBool: TypeAlias = Literal[True] +SingleNone: TypeAlias = Literal[None] +SingleEnum: TypeAlias = Literal[E.A] +UnionLiterals: TypeAlias = Literal[1, "foo", b"bar", True, None, E.A] +AnEnum1: TypeAlias = E +AnEnum2: TypeAlias = Literal[E.A, E.B] +Bool1: TypeAlias = bool +Bool2: TypeAlias = Literal[True, False] + +def _( + single_int: Literal[SingleInt], + single_str: Literal[SingleStr], + single_bytes: Literal[SingleBytes], + single_bool: Literal[SingleBool], + single_none: Literal[SingleNone], + single_enum: Literal[SingleEnum], + union_literals: Literal[UnionLiterals], + # Could also not error + an_enum1: Literal[AnEnum1], # error: [invalid-type-form] + an_enum2: Literal[AnEnum2], + # Could also not error + bool1: Literal[Bool1], # error: [invalid-type-form] + bool2: Literal[Bool2], + multiple: Literal[SingleInt, SingleStr, SingleEnum], +): + # TODO should be `Literal[1]` + reveal_type(single_int) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal["foo"]` + reveal_type(single_str) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[b"bar"]` + reveal_type(single_bytes) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[True]` + reveal_type(single_bool) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `None` + reveal_type(single_none) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[E.A]` + reveal_type(single_enum) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[1, "foo", b"bar", True, E.A] | None` + reveal_type(union_literals) # revealed: @Todo(Inference of subscript on special form) + # Could also be `E` + reveal_type(an_enum1) # revealed: Unknown + # TODO should be `E` + reveal_type(an_enum2) # revealed: @Todo(Inference of subscript on special form) + # Could also be `bool` + reveal_type(bool1) # revealed: Unknown + # TODO should be `bool` + reveal_type(bool2) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[1, "foo", E.A]` + reveal_type(multiple) # revealed: @Todo(Inference of subscript on special form) +``` + +### Implicit type alias + +```py +from typing import Literal +from enum import Enum + +class E(Enum): + A = 1 + B = 2 + +SingleInt = Literal[1] +SingleStr = Literal["foo"] +SingleBytes = Literal[b"bar"] +SingleBool = Literal[True] +SingleNone = Literal[None] +SingleEnum = Literal[E.A] +UnionLiterals = Literal[1, "foo", b"bar", True, None, E.A] +# For implicit type aliases, we may not want to support this. It's simpler not to, and no other +# type checker does. +AnEnum1 = E +AnEnum2 = Literal[E.A, E.B] +# For implicit type aliases, we may not want to support this. +Bool1 = bool +Bool2 = Literal[True, False] + +def _( + single_int: Literal[SingleInt], + single_str: Literal[SingleStr], + single_bytes: Literal[SingleBytes], + single_bool: Literal[SingleBool], + single_none: Literal[SingleNone], + single_enum: Literal[SingleEnum], + union_literals: Literal[UnionLiterals], + an_enum1: Literal[AnEnum1], # error: [invalid-type-form] + an_enum2: Literal[AnEnum2], + bool1: Literal[Bool1], # error: [invalid-type-form] + bool2: Literal[Bool2], + multiple: Literal[SingleInt, SingleStr, SingleEnum], +): + # TODO should be `Literal[1]` + reveal_type(single_int) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal["foo"]` + reveal_type(single_str) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[b"bar"]` + reveal_type(single_bytes) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[True]` + reveal_type(single_bool) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `None` + reveal_type(single_none) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[E.A]` + reveal_type(single_enum) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[1, "foo", b"bar", True, E.A] | None` + reveal_type(union_literals) # revealed: @Todo(Inference of subscript on special form) + reveal_type(an_enum1) # revealed: Unknown + # TODO should be `E` + reveal_type(an_enum2) # revealed: @Todo(Inference of subscript on special form) + reveal_type(bool1) # revealed: Unknown + # TODO should be `bool` + reveal_type(bool2) # revealed: @Todo(Inference of subscript on special form) + # TODO should be `Literal[1, "foo", E.A]` + reveal_type(multiple) # revealed: @Todo(Inference of subscript on special form) +``` + ## Shortening unions of literals When a Literal is parameterized with more than one value, it’s treated as exactly to equivalent to diff --git a/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md b/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md index adf0de358d..ceb588d7fe 100644 --- a/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md +++ b/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md @@ -259,7 +259,7 @@ class Color(Enum): RED = "red" f: dict[list[Literal[1]], list[Literal[Color.RED]]] = {[1]: [Color.RED, Color.RED]} -reveal_type(f) # revealed: dict[list[Literal[1]], list[Literal[Color.RED]]] +reveal_type(f) # revealed: dict[list[Literal[1]], list[Color]] class X[T]: def __init__(self, value: T): ... diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index be2fb264d8..b20f332999 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -1153,6 +1153,23 @@ impl<'db> Type<'db> { matches!(self, Type::FunctionLiteral(..)) } + /// Detects types which are valid to appear inside a `Literal[…]` type annotation. + pub(crate) fn is_literal_or_union_of_literals(&self, db: &'db dyn Db) -> bool { + match self { + Type::Union(union) => union + .elements(db) + .iter() + .all(|ty| ty.is_literal_or_union_of_literals(db)), + Type::StringLiteral(_) + | Type::BytesLiteral(_) + | Type::IntLiteral(_) + | Type::BooleanLiteral(_) + | Type::EnumLiteral(_) => true, + Type::NominalInstance(_) => self.is_none(db) || self.is_bool(db) || self.is_enum(db), + _ => false, + } + } + pub(crate) fn is_union_of_single_valued(&self, db: &'db dyn Db) -> bool { self.as_union().is_some_and(|union| { union.elements(db).iter().all(|ty| { diff --git a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs index 0d72548e49..50d22fac10 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs @@ -6,7 +6,6 @@ use crate::types::diagnostic::{ self, INVALID_TYPE_FORM, NON_SUBSCRIPTABLE, report_invalid_argument_number_to_special_form, report_invalid_arguments_to_annotated, report_invalid_arguments_to_callable, }; -use crate::types::enums::is_enum_class; use crate::types::signatures::Signature; use crate::types::string_annotation::parse_string_annotation; use crate::types::tuple::{TupleSpecBuilder, TupleType}; @@ -1369,7 +1368,6 @@ impl<'db> TypeInferenceBuilder<'db, '_> { parameters: &'param ast::Expr, ) -> Result, Vec<&'param ast::Expr>> { Ok(match parameters { - // TODO handle type aliases ast::Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => { let value_ty = self.infer_expression(value, TypeContext::default()); if matches!(value_ty, Type::SpecialForm(SpecialFormType::Literal)) { @@ -1421,27 +1419,6 @@ impl<'db> TypeInferenceBuilder<'db, '_> { literal @ ast::Expr::NumberLiteral(number) if number.value.is_int() => { self.infer_expression(literal, TypeContext::default()) } - // For enum values - ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => { - let value_ty = self.infer_expression(value, TypeContext::default()); - - if is_enum_class(self.db(), value_ty) { - let ty = value_ty - .member(self.db(), &attr.id) - .place - .ignore_possibly_undefined() - .unwrap_or(Type::unknown()); - self.store_expression_type(parameters, ty); - ty - } else { - self.store_expression_type(parameters, Type::unknown()); - if value_ty.is_todo() { - value_ty - } else { - return Err(vec![parameters]); - } - } - } // for negative and positive numbers ast::Expr::UnaryOp(u) if matches!(u.op, ast::UnaryOp::USub | ast::UnaryOp::UAdd) @@ -1451,6 +1428,35 @@ impl<'db> TypeInferenceBuilder<'db, '_> { self.store_expression_type(parameters, ty); ty } + // enum members and aliases to literal types + ast::Expr::Name(_) | ast::Expr::Attribute(_) => { + let subscript_ty = self.infer_expression(parameters, TypeContext::default()); + // TODO handle implicit type aliases also + match subscript_ty { + // type aliases to literal types + Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) => { + let value_ty = type_alias.value_type(self.db()); + if value_ty.is_literal_or_union_of_literals(self.db()) { + return Ok(value_ty); + } + } + // `Literal[SomeEnum.Member]` + Type::EnumLiteral(_) => { + return Ok(subscript_ty); + } + // `Literal[SingletonEnum.Member]`, where `SingletonEnum.Member` simplifies to + // just `SingletonEnum`. + Type::NominalInstance(_) if subscript_ty.is_enum(self.db()) => { + return Ok(subscript_ty); + } + // suppress false positives for e.g. members of functional-syntax enums + Type::Dynamic(DynamicType::Todo(_)) => { + return Ok(subscript_ty); + } + _ => {} + } + return Err(vec![parameters]); + } _ => { self.infer_expression(parameters, TypeContext::default()); return Err(vec![parameters]);