[ty] Improve sys.version_info special casing (#19894)

This commit is contained in:
Alex Waygood 2025-08-13 14:39:13 +01:00 committed by GitHub
parent 79c949f0f7
commit 2f3c7ad1fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 76 additions and 51 deletions

View file

@ -142,6 +142,15 @@ reveal_type(os.stat_result.__mro__)
reveal_type(os.stat_result.__getitem__)
```
But perhaps the most commonly used tuple subclass instance is the singleton `sys.version_info`:
```py
import sys
# revealed: Overload[(self, index: Literal[-5, 0], /) -> Literal[3], (self, index: Literal[-4, 1], /) -> Literal[11], (self, index: Literal[-3, -1, 2, 4], /) -> int, (self, index: Literal[-2, 3], /) -> Literal["alpha", "beta", "candidate", "final"], (self, index: SupportsIndex, /) -> int | Literal["alpha", "beta", "candidate", "final"], (self, index: slice[Any, Any, Any], /) -> tuple[int | Literal["alpha", "beta", "candidate", "final"], ...]]
reveal_type(type(sys.version_info).__getitem__)
```
Because of the synthesized `__getitem__` overloads we synthesize for tuples and tuple subclasses,
tuples are naturally understood as being subtypes of protocols that have precise return types from
`__getitem__` method members:

View file

@ -20,7 +20,7 @@ use crate::types::function::{DataclassTransformerParams, KnownFunction};
use crate::types::generics::{GenericContext, Specialization, walk_specialization};
use crate::types::infer::nearest_enclosing_class;
use crate::types::signatures::{CallableSignature, Parameter, Parameters, Signature};
use crate::types::tuple::TupleSpec;
use crate::types::tuple::{TupleSpec, TupleType};
use crate::types::{
BareTypeAliasType, Binding, BoundSuperError, BoundSuperType, CallableType, DataclassParams,
DeprecatedInstance, KnownInstanceType, StringLiteralType, TypeAliasType, TypeMapping,
@ -1348,12 +1348,22 @@ impl<'db> ClassLiteral<'db> {
let class_definition =
semantic_index(db, self.file(db)).expect_single_definition(class_stmt);
if self.is_known(db, KnownClass::VersionInfo) {
let tuple_type = TupleType::new(db, TupleSpec::version_info_spec(db))
.expect("sys.version_info tuple spec should always be a valid tuple");
Box::new([
definition_expression_type(db, class_definition, &class_stmt.bases()[0]),
Type::from(tuple_type.to_class_type(db)),
])
} else {
class_stmt
.bases()
.iter()
.map(|base_node| definition_expression_type(db, class_definition, base_node))
.collect()
}
}
/// Return `Some()` if this class is known to be a [`SolidBase`], or `None` if it is not.
pub(super) fn as_solid_base(self, db: &'db dyn Db) -> Option<SolidBase<'db>> {

View file

@ -11,8 +11,8 @@ use crate::types::cyclic::PairVisitor;
use crate::types::enums::is_single_member_enum;
use crate::types::protocol_class::walk_protocol_interface;
use crate::types::tuple::{TupleSpec, TupleType};
use crate::types::{ClassBase, DynamicType, TypeMapping, TypeRelation, TypeTransformer, UnionType};
use crate::{Db, FxOrderSet, Program};
use crate::types::{ClassBase, DynamicType, TypeMapping, TypeRelation, TypeTransformer};
use crate::{Db, FxOrderSet};
pub(super) use synthesized_protocol::SynthesizedProtocolType;
@ -115,47 +115,6 @@ impl<'db> NominalInstanceType<'db> {
/// I.e., for the type `tuple[int, str]`, this will return the tuple spec `[int, str]`.
/// For a subclass of `tuple[int, str]`, it will return the same tuple spec.
pub(super) fn tuple_spec(&self, db: &'db dyn Db) -> Option<Cow<'db, TupleSpec<'db>>> {
fn own_tuple_spec_of_class<'db>(
db: &'db dyn Db,
class: ClassType<'db>,
) -> Option<Cow<'db, TupleSpec<'db>>> {
let (class_literal, specialization) = class.class_literal(db);
match class_literal.known(db)? {
KnownClass::Tuple => Some(
specialization
.and_then(|spec| Some(Cow::Borrowed(spec.tuple(db)?)))
.unwrap_or_else(|| Cow::Owned(TupleSpec::homogeneous(Type::unknown()))),
),
KnownClass::VersionInfo => {
let python_version = Program::get(db).python_version(db);
let int_instance_ty = KnownClass::Int.to_instance(db);
// TODO: just grab this type from typeshed (it's a `sys._ReleaseLevel` type alias there)
let release_level_ty = {
let elements: Box<[Type<'db>]> = ["alpha", "beta", "candidate", "final"]
.iter()
.map(|level| Type::string_literal(db, level))
.collect();
// For most unions, it's better to go via `UnionType::from_elements` or use `UnionBuilder`;
// those techniques ensure that union elements are deduplicated and unions are eagerly simplified
// into other types where necessary. Here, however, we know that there are no duplicates
// in this union, so it's probably more efficient to use `UnionType::new()` directly.
Type::Union(UnionType::new(db, elements))
};
Some(Cow::Owned(TupleSpec::from_elements([
Type::IntLiteral(python_version.major.into()),
Type::IntLiteral(python_version.minor.into()),
int_instance_ty,
release_level_ty,
int_instance_ty,
])))
}
_ => None,
}
}
match self.0 {
NominalInstanceInner::ExactTuple(tuple) => Some(Cow::Borrowed(tuple.tuple(db))),
NominalInstanceInner::NonTuple(class) => {
@ -169,7 +128,26 @@ impl<'db> NominalInstanceType<'db> {
class
.iter_mro(db)
.filter_map(ClassBase::into_class)
.find_map(|class| own_tuple_spec_of_class(db, class))
.find_map(|class| match class.known(db)? {
// N.B. this is a pure optimisation: iterating through the MRO would give us
// the correct tuple spec for `sys._version_info`, since we special-case the class
// in `ClassLiteral::explicit_bases()` so that it is inferred as inheriting from
// a tuple type with the correct spec for the user's configured Python version and platform.
KnownClass::VersionInfo => {
Some(Cow::Owned(TupleSpec::version_info_spec(db)))
}
KnownClass::Tuple => Some(
class
.into_generic_alias()
.and_then(|alias| {
Some(Cow::Borrowed(alias.specialization(db).tuple(db)?))
})
.unwrap_or_else(|| {
Cow::Owned(TupleSpec::homogeneous(Type::unknown()))
}),
),
_ => None,
})
}
}
}

View file

@ -30,7 +30,7 @@ use crate::types::{
UnionBuilder, UnionType, cyclic::PairVisitor,
};
use crate::util::subscript::{Nth, OutOfBoundsError, PyIndex, PySlice, StepSizeZeroError};
use crate::{Db, FxOrderSet};
use crate::{Db, FxOrderSet, Program};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum TupleLength {
@ -1162,6 +1162,34 @@ impl<'db> Tuple<Type<'db>> {
Tuple::Variable(_) => false,
}
}
/// Return the `TupleSpec` for the singleton `sys.version_info`
pub(crate) fn version_info_spec(db: &'db dyn Db) -> TupleSpec<'db> {
let python_version = Program::get(db).python_version(db);
let int_instance_ty = KnownClass::Int.to_instance(db);
// TODO: just grab this type from typeshed (it's a `sys._ReleaseLevel` type alias there)
let release_level_ty = {
let elements: Box<[Type<'db>]> = ["alpha", "beta", "candidate", "final"]
.iter()
.map(|level| Type::string_literal(db, level))
.collect();
// For most unions, it's better to go via `UnionType::from_elements` or use `UnionBuilder`;
// those techniques ensure that union elements are deduplicated and unions are eagerly simplified
// into other types where necessary. Here, however, we know that there are no duplicates
// in this union, so it's probably more efficient to use `UnionType::new()` directly.
Type::Union(UnionType::new(db, elements))
};
TupleSpec::from_elements([
Type::IntLiteral(python_version.major.into()),
Type::IntLiteral(python_version.minor.into()),
int_instance_ty,
release_level_ty,
int_instance_ty,
])
}
}
impl<T> From<FixedLengthTuple<T>> for Tuple<T> {