Python
/
ruff
mirror of https://github.com/astral-sh/ruff
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

simple subtyping for bidirectional inference

This commit is contained in:
Ibraheem Ahmed 2025-12-02 00:09:39 -05:00
parent 392a8e4e50
commit 623210232e
10 changed files with 190 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
crates/ty_python_semantic/resources/mdtest/assignment/annotations.md
View File

@ -542,8 +542,7 @@ e: list[Any] | None = [1]
reveal_type(e) # revealed: list[Any]
f: list[Any] | None = f2(1)
# TODO: Better constraint solver.
reveal_type(f) # revealed: list[int] | None
reveal_type(f) # revealed: list[Any] | None
g: list[Any] | dict[Any, Any] = f3(1)
# TODO: Better constraint solver.
@ -600,6 +599,48 @@ reveal_type(x7) # revealed: Contravariant[Any]
reveal_type(x8) # revealed: Invariant[Any]
```
## Declared type preference sees through subtyping
```toml
[environment]
python-version = "3.12"
```
```py
from typing import Any, Iterable, Literal, MutableSequence, Sequence
x1: Sequence[Any] = [1, 2, 3]
reveal_type(x1) # revealed: list[Any]
x2: MutableSequence[Any] = [1, 2, 3]
reveal_type(x2) # revealed: list[Any]
x3: Iterable[Any] = [1, 2, 3]
reveal_type(x3) # revealed: list[Any]
class X[T]:
value: T
def __init__(self, value: T): ...
class A[T](X[T]): ...
def a[T](value: T) -> A[T]:
return A(value)
x4: A[object] = A(1)
reveal_type(x4) # revealed: A[object]
x5: X[object] = A(1)
reveal_type(x5) # revealed: A[object]
x6: X[object] | None = A(1)
reveal_type(x6) # revealed: A[object]
x7: X[object] | None = a(1)
reveal_type(x7) # revealed: A[object]
```
## Narrow generic unions
```toml

2
crates/ty_python_semantic/resources/mdtest/dataclasses/fields.md
View File

@ -37,7 +37,7 @@ class Data:
content: list[int] = field(default_factory=list)
timestamp: datetime = field(default_factory=datetime.now, init=False)
# revealed: (self: Data, content: list[int] = list[int]) -> None
# revealed: (self: Data, content: list[int] = Unknown) -> None
reveal_type(Data.__init__)
data = Data([1, 2, 3])

37
crates/ty_python_semantic/resources/mdtest/literal_promotion.md
View File

@ -341,3 +341,40 @@ reveal_type(x21) # revealed: X[Literal[1]]
x22: X[Literal[1]] | None = x(1)
reveal_type(x22) # revealed: X[Literal[1]]
```
## Literal annotations see through subtyping
```py
from typing import Iterable, Literal, MutableSequence, Sequence
x1: Sequence[Literal[1, 2, 3]] = [1, 2, 3]
reveal_type(x1) # revealed: list[Literal[1, 2, 3]]
x2: MutableSequence[Literal[1, 2, 3]] = [1, 2, 3]
reveal_type(x2) # revealed: list[Literal[1, 2, 3]]
x3: Iterable[Literal[1, 2, 3]] = [1, 2, 3]
reveal_type(x3) # revealed: list[Literal[1, 2, 3]]
class X[T]:
value: T
def __init__(self, value: T): ...
class A[T](X[T]): ...
def a[T](value: T) -> A[T]:
return A(value)
x4: A[Literal[1]] = A(1)
reveal_type(x4) # revealed: A[Literal[1]]
x5: X[Literal[1]] = A(1)
reveal_type(x5) # revealed: A[Literal[1]]
x6: X[Literal[1]] | None = A(1)
reveal_type(x6) # revealed: A[Literal[1]]
x7: X[Literal[1]] | None = a(1)
reveal_type(x7) # revealed: A[Literal[1]]
```

2
crates/ty_python_semantic/resources/mdtest/type_compendium/tuple.md
View File

@ -55,7 +55,7 @@ def f(x: Iterable[int], y: list[str], z: Never, aa: list[Never], bb: LiskovUncom
reveal_type(tuple((1, 2))) # revealed: tuple[Literal[1], Literal[2]]
reveal_type(tuple([1])) # revealed: tuple[Unknown | int, ...]
reveal_type(tuple([1])) # revealed: tuple[object, ...]
# error: [invalid-argument-type]
reveal_type(tuple[int]([1])) # revealed: tuple[int]

70
crates/ty_python_semantic/src/types.rs
View File

@ -1000,7 +1000,10 @@ impl<'db> Type<'db> {
}
/// If this type is a class instance, returns its specialization.
pub(crate) fn class_specialization(self, db: &'db dyn Db) -> Option<Specialization<'db>> {
pub(crate) fn class_specialization(
self,
db: &'db dyn Db,
) -> Option<(ClassType<'db>, Specialization<'db>)> {
self.specialization_of_optional(db, None)
}
@ -1011,15 +1014,17 @@ impl<'db> Type<'db> {
expected_class: ClassLiteral<'_>,
) -> Option<Specialization<'db>> {
self.specialization_of_optional(db, Some(expected_class))
.map(|(_, specialization)| specialization)
}
fn specialization_of_optional(
self,
db: &'db dyn Db,
expected_class: Option<ClassLiteral<'_>>,
) -> Option<Specialization<'db>> {
) -> Option<(ClassType<'db>, Specialization<'db>)> {
let class_type = match self {
Type::NominalInstance(instance) => instance,
Type::ProtocolInstance(instance) => instance.to_nominal_instance()?,
Type::TypeAlias(alias) => alias.value_type(db).as_nominal_instance()?,
_ => return None,
}
@ -1030,7 +1035,48 @@ impl<'db> Type<'db> {
return None;
}
specialization
Some((class_type, specialization?))
}
/// Given a type variable `T` from the generic context of a class `C`:
/// - If `self` is a specialized instance of `C`, returns the type assigned to `T` on `self`.
/// - If `self` is a specialized instance of some class `A`, and `C` is a subclass of `A`
/// such that the type variable `U` on `A` is specialized to `T`, returns the type
/// assigned to `U` on `self`.
pub(crate) fn find_type_var_from(
self,
db: &'db dyn Db,
bound_typevar: BoundTypeVarInstance<'db>,
class: ClassLiteral<'db>,
) -> Option<Type<'db>> {
if let Some(specialization) = self.specialization_of(db, class) {
return specialization.get(db, bound_typevar);
}
// TODO: We should use the constraint solver here to determine the type mappings for more
// complex subtyping relationships, e.g., `type[C[T]]` to `Callable[..., T]`, or unions
// containing multiple generic elements.
for base in class.iter_mro(db, None) {
let Some(ClassType::Generic(class)) = base.into_class() else {
continue;
};
for (base_typevar, base_ty) in class
.specialization(db)
.generic_context(db)
.variables(db)
.zip(class.specialization(db).types(db))
{
if *base_ty == Type::TypeVar(bound_typevar) {
// TODO: This is potentially quadratic.
if let Some(ty) = self.find_type_var_from(db, base_typevar, class.origin(db)) {
return Some(ty);
}
}
}
}
None
}
/// Returns the top materialization (or upper bound materialization) of this type, which is the
@ -3852,20 +3898,20 @@ impl<'db> Type<'db> {
return;
};
let tcx_specialization = tcx.annotation.and_then(|tcx| {
tcx.filter_union(db, |ty| ty.specialization_of(db, class_literal).is_some())
.specialization_of(db, class_literal)
});
for (typevar, ty) in specialization
for (type_var, ty) in specialization
.generic_context(db)
.variables(db)
.zip(specialization.types(db))
{
let variance = typevar.variance_with_polarity(db, polarity);
let tcx = TypeContext::new(tcx_specialization.and_then(|spec| spec.get(db, typevar)));
let variance = type_var.variance_with_polarity(db, polarity);
let tcx = tcx.and_then(|tcx| {
tcx.filter_union(db, |ty| {
ty.find_type_var_from(db, type_var, class_literal).is_some()
})
.find_type_var_from(db, type_var, class_literal)
});
f(typevar, *ty, variance, tcx);
f(type_var, *ty, variance, tcx);
visitor.visit(*ty, || {
ty.visit_specialization_impl(db, tcx, variance, f, visitor);

2
crates/ty_python_semantic/src/types/bound_super.rs
View File

@ -321,7 +321,7 @@ impl<'db> BoundSuperType<'db> {
Type::NominalInstance(instance) => SuperOwnerKind::Instance(instance),
Type::ProtocolInstance(protocol) => {
if let Some(nominal_instance) = protocol.as_nominal_type() {
if let Some(nominal_instance) = protocol.to_nominal_instance() {
SuperOwnerKind::Instance(nominal_instance)
} else {
return Err(BoundSuperError::AbstractOwnerType {

28
crates/ty_python_semantic/src/types/call/bind.rs
View File

@ -2818,10 +2818,32 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
// Prefer the declared type of generic classes.
let preferred_type_mappings = return_with_tcx.and_then(|(return_ty, tcx)| {
tcx.filter_union(self.db, |ty| ty.class_specialization(self.db).is_some())
.class_specialization(self.db)?;
let tcx = tcx.filter_union(self.db, |ty| ty.class_specialization(self.db).is_some());
let return_ty =
return_ty.filter_union(self.db, |ty| ty.class_specialization(self.db).is_some());
// TODO: We should use the constraint solver here to determine the type mappings for more
// complex subtyping relationships, e.g., `type[C[T]]` to `Callable[..., T]`, or unions
// containing multiple generic elements.
if let Some((class_literal, _)) = return_ty.class_specialization(self.db)
&& let Some(generic_alias) = class_literal.into_generic_alias()
{
let specialization = generic_alias.specialization(self.db);
for (class_type_var, return_ty) in specialization
.generic_context(self.db)
.variables(self.db)
.zip(specialization.types(self.db))
{
if let Some(ty) = tcx.find_type_var_from(
self.db,
class_type_var,
generic_alias.origin(self.db),
) {
builder.infer(*return_ty, ty).ok()?;
}
}
}
builder.infer(return_ty, tcx).ok()?;
Some(builder.type_mappings().clone())
});

6
crates/ty_python_semantic/src/types/infer.rs
View File

@ -381,6 +381,12 @@ impl<'db> TypeContext<'db> {
}
}
pub(crate) fn and_then(self, f: impl FnOnce(Type<'db>) -> Option<Type<'db>>) -> Self {
Self {
annotation: self.annotation.and_then(f),
}
}
pub(crate) fn is_typealias(&self) -> bool {
self.annotation
.is_some_and(|ty| ty.is_typealias_special_form())

46
crates/ty_python_semantic/src/types/infer/builder.rs
View File

@ -7497,41 +7497,24 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
annotation.filter_disjoint_elements(self.db(), collection_ty, inferable)
});
// Extract the annotated type of `T`, if provided.
let annotated_elt_tys = tcx
.known_specialization(self.db(), collection_class)
.map(|specialization| specialization.types(self.db()));
// Create a set of constraints to infer a precise type for `T`.
let mut builder = SpecializationBuilder::new(self.db(), inferable);
match annotated_elt_tys {
// The annotated type acts as a constraint for `T`.
//
// Note that we infer the annotated type _before_ the elements, to more closely match the
// order of any unions as written in the type annotation.
Some(annotated_elt_tys) => {
for (elt_ty, annotated_elt_ty) in iter::zip(elt_tys.clone(), annotated_elt_tys) {
builder
.infer(Type::TypeVar(elt_ty), *annotated_elt_ty)
.ok()?;
}
}
for elt_ty in elt_tys.clone() {
let elt_tcx = tcx
.annotation
// The annotated type acts as a constraint for `T`.
//
// Note that we infer the annotated type _before_ the elements, to more closely match the
// order of any unions as written in the type annotation.
.and_then(|tcx| tcx.find_type_var_from(self.db(), elt_ty, class_literal))
// If a valid type annotation was not provided, avoid restricting the type of the collection
// by unioning the inferred type with `Unknown`.
.unwrap_or(Type::unknown());
// If a valid type annotation was not provided, avoid restricting the type of the collection
// by unioning the inferred type with `Unknown`.
None => {
for elt_ty in elt_tys.clone() {
builder.infer(Type::TypeVar(elt_ty), Type::unknown()).ok()?;
}
}
builder.infer(Type::TypeVar(elt_ty), elt_tcx).ok()?;
}
let elt_tcxs = match annotated_elt_tys {
None => Either::Left(iter::repeat(TypeContext::default())),
Some(tys) => Either::Right(tys.iter().map(|ty| TypeContext::new(Some(*ty)))),
};
for elts in elts {
// An unpacking expression for a dictionary.
if let &[None, Some(value)] = elts.as_slice() {
@ -7554,10 +7537,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
// The inferred type of each element acts as an additional constraint on `T`.
for (elt, elt_ty, elt_tcx) in itertools::izip!(elts, elt_tys.clone(), elt_tcxs.clone())
{
for (elt, elt_ty) in iter::zip(elts, elt_tys.clone()) {
let Some(elt) = elt else { continue };
let elt_tcx =
tcx.and_then(|tcx| tcx.find_type_var_from(self.db(), elt_ty, class_literal));
let inferred_elt_ty = infer_elt_expression(self, elt, elt_tcx);
// Simplify the inference based on the declared type of the element.

6
crates/ty_python_semantic/src/types/instance.rs
View File

@ -165,7 +165,7 @@ impl<'db> Type<'db> {
// This matches the behaviour of other type checkers, and is required for us to
// recognise `str` as a subtype of `Container[str]`.
structurally_satisfied.or(db, || {
let Some(nominal_instance) = protocol.as_nominal_type() else {
let Some(nominal_instance) = protocol.to_nominal_instance() else {
return ConstraintSet::from(false);
};
@ -175,7 +175,7 @@ impl<'db> Type<'db> {
// `Q`'s members in a Liskov-incompatible way.
let type_to_test = self
.as_protocol_instance()
.and_then(ProtocolInstanceType::as_nominal_type)
.and_then(ProtocolInstanceType::to_nominal_instance)
.map(Type::NominalInstance)
.unwrap_or(self);
@ -650,7 +650,7 @@ impl<'db> ProtocolInstanceType<'db> {
/// If this is a synthesized protocol that does not correspond to a class definition
/// in source code, return `None`. These are "pure" abstract types, that cannot be
/// treated in a nominal way.
pub(super) fn as_nominal_type(self) -> Option<NominalInstanceType<'db>> {
pub(super) fn to_nominal_instance(self) -> Option<NominalInstanceType<'db>> {
match self.inner {
Protocol::FromClass(class) => {
Some(NominalInstanceType(NominalInstanceInner::NonTuple(*class)))