Skip to content

Commit 87b32ba

Browse files
committed
add GenericContext.specialize_constrained mdtest method
1 parent be5ad64 commit 87b32ba

File tree

4 files changed

+79
-8
lines changed

4 files changed

+79
-8
lines changed

crates/ty_python_semantic/src/types.rs

Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4337,6 +4337,14 @@ impl<'db> Type<'db> {
43374337
))
43384338
.into()
43394339
}
4340+
Type::KnownInstance(KnownInstanceType::GenericContext(tracked))
4341+
if name == "specialize_constrained" =>
4342+
{
4343+
Place::bound(Type::KnownBoundMethod(
4344+
KnownBoundMethodType::GenericContextSpecializeConstrained(tracked),
4345+
))
4346+
.into()
4347+
}
43404348

43414349
Type::ClassLiteral(class)
43424350
if name == "__get__" && class.is_known(db, KnownClass::FunctionType) =>
@@ -7169,6 +7177,7 @@ impl<'db> Type<'db> {
71697177
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
71707178
| KnownBoundMethodType::ConstraintSetSatisfies(_)
71717179
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
7180+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_)
71727181
)
71737182
| Type::DataclassDecorator(_)
71747183
| Type::DataclassTransformer(_)
@@ -7328,7 +7337,8 @@ impl<'db> Type<'db> {
73287337
| KnownBoundMethodType::ConstraintSetNever
73297338
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
73307339
| KnownBoundMethodType::ConstraintSetSatisfies(_)
7331-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_),
7340+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
7341+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_),
73327342
)
73337343
| Type::DataclassDecorator(_)
73347344
| Type::DataclassTransformer(_)
@@ -10823,6 +10833,9 @@ pub enum KnownBoundMethodType<'db> {
1082310833
ConstraintSetImpliesSubtypeOf(TrackedConstraintSet<'db>),
1082410834
ConstraintSetSatisfies(TrackedConstraintSet<'db>),
1082510835
ConstraintSetSatisfiedByAllTypeVars(TrackedConstraintSet<'db>),
10836+
10837+
// GenericContext methods
10838+
GenericContextSpecializeConstrained(GenericContext<'db>),
1082610839
}
1082710840

1082810841
pub(super) fn walk_method_wrapper_type<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
@@ -10852,7 +10865,8 @@ pub(super) fn walk_method_wrapper_type<'db, V: visitor::TypeVisitor<'db> + ?Size
1085210865
| KnownBoundMethodType::ConstraintSetNever
1085310866
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
1085410867
| KnownBoundMethodType::ConstraintSetSatisfies(_)
10855-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_) => {}
10868+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
10869+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_) => {}
1085610870
}
1085710871
}
1085810872

@@ -10928,6 +10942,10 @@ impl<'db> KnownBoundMethodType<'db> {
1092810942
| (
1092910943
KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_),
1093010944
KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_),
10945+
)
10946+
| (
10947+
KnownBoundMethodType::GenericContextSpecializeConstrained(_),
10948+
KnownBoundMethodType::GenericContextSpecializeConstrained(_),
1093110949
) => ConstraintSet::from(true),
1093210950

1093310951
(
@@ -10942,7 +10960,8 @@ impl<'db> KnownBoundMethodType<'db> {
1094210960
| KnownBoundMethodType::ConstraintSetNever
1094310961
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
1094410962
| KnownBoundMethodType::ConstraintSetSatisfies(_)
10945-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_),
10963+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
10964+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_),
1094610965
KnownBoundMethodType::FunctionTypeDunderGet(_)
1094710966
| KnownBoundMethodType::FunctionTypeDunderCall(_)
1094810967
| KnownBoundMethodType::PropertyDunderGet(_)
@@ -10954,7 +10973,8 @@ impl<'db> KnownBoundMethodType<'db> {
1095410973
| KnownBoundMethodType::ConstraintSetNever
1095510974
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
1095610975
| KnownBoundMethodType::ConstraintSetSatisfies(_)
10957-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_),
10976+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
10977+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_),
1095810978
) => ConstraintSet::from(false),
1095910979
}
1096010980
}
@@ -11019,6 +11039,11 @@ impl<'db> KnownBoundMethodType<'db> {
1101911039
.constraints(db)
1102011040
.iff(db, right_constraints.constraints(db)),
1102111041

11042+
(
11043+
KnownBoundMethodType::GenericContextSpecializeConstrained(left_generic_context),
11044+
KnownBoundMethodType::GenericContextSpecializeConstrained(right_generic_context),
11045+
) => ConstraintSet::from(left_generic_context == right_generic_context),
11046+
1102211047
(
1102311048
KnownBoundMethodType::FunctionTypeDunderGet(_)
1102411049
| KnownBoundMethodType::FunctionTypeDunderCall(_)
@@ -11031,7 +11056,8 @@ impl<'db> KnownBoundMethodType<'db> {
1103111056
| KnownBoundMethodType::ConstraintSetNever
1103211057
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
1103311058
| KnownBoundMethodType::ConstraintSetSatisfies(_)
11034-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_),
11059+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
11060+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_),
1103511061
KnownBoundMethodType::FunctionTypeDunderGet(_)
1103611062
| KnownBoundMethodType::FunctionTypeDunderCall(_)
1103711063
| KnownBoundMethodType::PropertyDunderGet(_)
@@ -11043,7 +11069,8 @@ impl<'db> KnownBoundMethodType<'db> {
1104311069
| KnownBoundMethodType::ConstraintSetNever
1104411070
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
1104511071
| KnownBoundMethodType::ConstraintSetSatisfies(_)
11046-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_),
11072+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
11073+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_),
1104711074
) => ConstraintSet::from(false),
1104811075
}
1104911076
}
@@ -11069,7 +11096,8 @@ impl<'db> KnownBoundMethodType<'db> {
1106911096
| KnownBoundMethodType::ConstraintSetNever
1107011097
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
1107111098
| KnownBoundMethodType::ConstraintSetSatisfies(_)
11072-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_) => self,
11099+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
11100+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_) => self,
1107311101
}
1107411102
}
1107511103

@@ -11087,7 +11115,8 @@ impl<'db> KnownBoundMethodType<'db> {
1108711115
| KnownBoundMethodType::ConstraintSetNever
1108811116
| KnownBoundMethodType::ConstraintSetImpliesSubtypeOf(_)
1108911117
| KnownBoundMethodType::ConstraintSetSatisfies(_)
11090-
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_) => {
11118+
| KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(_)
11119+
| KnownBoundMethodType::GenericContextSpecializeConstrained(_) => {
1109111120
KnownClass::ConstraintSet
1109211121
}
1109311122
}
@@ -11249,6 +11278,19 @@ impl<'db> KnownBoundMethodType<'db> {
1124911278
Some(KnownClass::Bool.to_instance(db)),
1125011279
)))
1125111280
}
11281+
11282+
KnownBoundMethodType::GenericContextSpecializeConstrained(_) => {
11283+
Either::Right(std::iter::once(Signature::new(
11284+
Parameters::new([Parameter::positional_only(Some(Name::new_static(
11285+
"constraints",
11286+
)))
11287+
.with_annotated_type(KnownClass::ConstraintSet.to_instance(db))]),
11288+
Some(UnionType::from_elements(
11289+
db,
11290+
[KnownClass::Specialization.to_instance(db), Type::none(db)],
11291+
)),
11292+
)))
11293+
}
1125211294
}
1125311295
}
1125411296
}

crates/ty_python_semantic/src/types/call/bind.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1253,6 +1253,26 @@ impl<'db> Bindings<'db> {
12531253
overload.set_return_type(Type::BooleanLiteral(result));
12541254
}
12551255

1256+
Type::KnownBoundMethod(
1257+
KnownBoundMethodType::GenericContextSpecializeConstrained(generic_context),
1258+
) => {
1259+
let [Some(constraints)] = overload.parameter_types() else {
1260+
continue;
1261+
};
1262+
let Type::KnownInstance(KnownInstanceType::ConstraintSet(constraints)) =
1263+
constraints
1264+
else {
1265+
continue;
1266+
};
1267+
if let Ok(specialization) =
1268+
generic_context.specialize_constrained(db, constraints.constraints(db))
1269+
{
1270+
overload.set_return_type(Type::KnownInstance(
1271+
KnownInstanceType::Specialization(specialization),
1272+
));
1273+
}
1274+
}
1275+
12561276
Type::ClassLiteral(class) => match class.known(db) {
12571277
Some(KnownClass::Bool) => match overload.parameter_types() {
12581278
[Some(arg)] => overload.set_return_type(arg.bool(db).into_type(db)),

crates/ty_python_semantic/src/types/display.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -541,6 +541,9 @@ impl Display for DisplayRepresentation<'_> {
541541
Type::KnownBoundMethod(KnownBoundMethodType::ConstraintSetSatisfiedByAllTypeVars(
542542
_,
543543
)) => f.write_str("bound method `ConstraintSet.satisfied_by_all_typevars`"),
544+
Type::KnownBoundMethod(KnownBoundMethodType::GenericContextSpecializeConstrained(
545+
_,
546+
)) => f.write_str("bound method `GenericContext.specialize_constrained`"),
544547
Type::WrapperDescriptor(kind) => {
545548
let (method, object) = match kind {
546549
WrapperDescriptorKind::FunctionTypeDunderGet => ("__get__", "function"),

crates/ty_vendored/ty_extensions/ty_extensions.pyi

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,12 @@ class GenericContext:
9797
alias.
9898
"""
9999

100+
def specialize_constrained(self, constraints: ConstraintSet) -> Specialization | None:
101+
"""
102+
Returns a specialization of this generic context that satisfies the
103+
given constraints, or None if the constraints cannot be satisfied.
104+
"""
105+
100106
class Specialization:
101107
"""A mapping of typevars to specific types"""
102108

0 commit comments

Comments
 (0)