Compare commits
1 Commits
cjm/protos
...
david/sqla
| Author | SHA1 | Date | |
|---|---|---|---|
|
352628e986 |
@@ -106,45 +106,36 @@ reveal_type(admin_users) # revealed: Sequence[User]
|
||||
We can also specify particular columns to select:
|
||||
|
||||
```py
|
||||
reveal_type(User.id) # revealed: InstrumentedAttribute[int]
|
||||
stmt = select(User.id, User.name)
|
||||
# TODO: should be `Select[tuple[int, str]]`
|
||||
reveal_type(stmt) # revealed: Select[tuple[Unknown, Unknown]]
|
||||
reveal_type(stmt) # revealed: Select[tuple[int, str]]
|
||||
|
||||
ids_and_names = session.execute(stmt).all()
|
||||
# TODO: should be `Sequence[Row[tuple[int, str]]]`
|
||||
reveal_type(ids_and_names) # revealed: Sequence[Row[tuple[Unknown, Unknown]]]
|
||||
reveal_type(ids_and_names) # revealed: Sequence[Row[tuple[int, str]]]
|
||||
|
||||
for row in session.execute(stmt):
|
||||
# TODO: should be `Row[tuple[int, str]]`
|
||||
reveal_type(row) # revealed: Row[tuple[Unknown, Unknown]]
|
||||
reveal_type(row) # revealed: Row[tuple[int, str]]
|
||||
|
||||
for user_id, name in session.execute(stmt).tuples():
|
||||
# TODO: should be `int`
|
||||
reveal_type(user_id) # revealed: Unknown
|
||||
# TODO: should be `str`
|
||||
reveal_type(name) # revealed: Unknown
|
||||
reveal_type(user_id) # revealed: int
|
||||
reveal_type(name) # revealed: str
|
||||
|
||||
result = session.execute(stmt)
|
||||
row = result.one_or_none()
|
||||
assert row is not None
|
||||
(user_id, name) = row._tuple()
|
||||
# TODO: should be `int`
|
||||
reveal_type(user_id) # revealed: Unknown
|
||||
# TODO: should be `str`
|
||||
reveal_type(name) # revealed: Unknown
|
||||
reveal_type(user_id) # revealed: int
|
||||
reveal_type(name) # revealed: str
|
||||
|
||||
stmt = select(User.id).where(User.name == "Alice")
|
||||
|
||||
# TODO: should be `Select[tuple[int]]`
|
||||
reveal_type(stmt) # revealed: Select[tuple[Unknown]]
|
||||
reveal_type(stmt) # revealed: Select[tuple[int]]
|
||||
|
||||
alice_id = session.scalars(stmt).first()
|
||||
# TODO: should be `int | None`
|
||||
reveal_type(alice_id) # revealed: Unknown | None
|
||||
reveal_type(alice_id) # revealed: int | None
|
||||
|
||||
alice_id = session.scalar(stmt)
|
||||
# TODO: should be `int | None`
|
||||
reveal_type(alice_id) # revealed: Unknown | None
|
||||
reveal_type(alice_id) # revealed: int | None
|
||||
```
|
||||
|
||||
Using the legacy `query` API also works:
|
||||
@@ -203,8 +194,6 @@ async def test_async(session: AsyncSession):
|
||||
stmt = select(User.id, User.name)
|
||||
result = await session.execute(stmt)
|
||||
for user_id, name in result.tuples():
|
||||
# TODO: should be `int`
|
||||
reveal_type(user_id) # revealed: Unknown
|
||||
# TODO: should be `str`
|
||||
reveal_type(name) # revealed: Unknown
|
||||
reveal_type(user_id) # revealed: int
|
||||
reveal_type(name) # revealed: str
|
||||
```
|
||||
|
||||
@@ -3010,31 +3010,6 @@ class Bar(Protocol[S]):
|
||||
z: S | Bar[S]
|
||||
```
|
||||
|
||||
### Recursive generic protocols with growing specializations
|
||||
|
||||
This snippet caused a stack overflow in <https://github.com/astral-sh/ty/issues/1736> because the
|
||||
type parameter grows with each recursive call (`C[set[T]]` leads to `C[set[set[T]]]`, then
|
||||
`C[set[set[set[T]]]]`, etc.):
|
||||
|
||||
```toml
|
||||
[environment]
|
||||
python-version = "3.12"
|
||||
```
|
||||
|
||||
```py
|
||||
from typing import Protocol
|
||||
|
||||
class C[T](Protocol):
|
||||
a: "C[set[T]]"
|
||||
|
||||
def takes_c(c: C[set[int]]) -> None: ...
|
||||
def f(c: C[int]) -> None:
|
||||
# The key thing is that we don't stack overflow while checking this.
|
||||
# The cycle detection assumes compatibility when it detects potential
|
||||
# infinite recursion between protocol specializations.
|
||||
takes_c(c)
|
||||
```
|
||||
|
||||
### Recursive legacy generic protocol
|
||||
|
||||
```py
|
||||
|
||||
@@ -335,6 +335,12 @@ pub enum KnownModule {
|
||||
#[cfg(test)]
|
||||
Uuid,
|
||||
Warnings,
|
||||
#[strum(serialize = "sqlalchemy.sql.selectable")]
|
||||
SqlalchemySqlSelectable,
|
||||
#[strum(serialize = "sqlalchemy.sql._selectable_constructors")]
|
||||
SqlalchemySqlSelectableConstructors,
|
||||
#[strum(serialize = "sqlalchemy.orm.attributes")]
|
||||
SqlalchemyOrmAttributes,
|
||||
}
|
||||
|
||||
impl KnownModule {
|
||||
@@ -363,6 +369,9 @@ impl KnownModule {
|
||||
#[cfg(test)]
|
||||
Self::Uuid => "uuid",
|
||||
Self::Templatelib => "string.templatelib",
|
||||
Self::SqlalchemySqlSelectable => "sqlalchemy.sql.selectable",
|
||||
Self::SqlalchemySqlSelectableConstructors => "sqlalchemy.sql._selectable_constructors",
|
||||
Self::SqlalchemyOrmAttributes => "sqlalchemy.orm.attributes",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -378,7 +387,20 @@ impl KnownModule {
|
||||
if search_path.is_standard_library() {
|
||||
Self::from_str(name.as_str()).ok()
|
||||
} else {
|
||||
None
|
||||
// For non-stdlib search paths, check for known third-party modules
|
||||
Self::try_from_third_party_name(name)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a known module for third-party packages, if applicable.
|
||||
fn try_from_third_party_name(name: &ModuleName) -> Option<Self> {
|
||||
match name.as_str() {
|
||||
"sqlalchemy.sql.selectable" => Some(Self::SqlalchemySqlSelectable),
|
||||
"sqlalchemy.sql._selectable_constructors" => {
|
||||
Some(Self::SqlalchemySqlSelectableConstructors)
|
||||
}
|
||||
"sqlalchemy.orm.attributes" => Some(Self::SqlalchemyOrmAttributes),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -419,6 +441,11 @@ mod tests {
|
||||
let stdlib_search_path = SearchPath::vendored_stdlib();
|
||||
|
||||
for module in KnownModule::iter() {
|
||||
// Third-party modules aren't available in the vendored stdlib
|
||||
if module.is_third_party() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let module_name = module.name();
|
||||
|
||||
assert_eq!(
|
||||
|
||||
@@ -204,48 +204,9 @@ fn definition_expression_type<'db>(
|
||||
/// A [`TypeTransformer`] that is used in `apply_type_mapping` methods.
|
||||
pub(crate) type ApplyTypeMappingVisitor<'db> = TypeTransformer<'db, TypeMapping<'db, 'db>>;
|
||||
|
||||
/// Key type for the `has_relation_to` visitor.
|
||||
///
|
||||
/// For most type comparisons, we use the full `Type` as the key. However, for protocol-to-protocol
|
||||
/// comparisons, we use the underlying `ClassLiteral` (ignoring specialization) to detect infinite
|
||||
/// recursion that occurs with recursive generic protocols.
|
||||
///
|
||||
/// For example, with:
|
||||
/// ```python
|
||||
/// class C[T](Protocol):
|
||||
/// a: 'C[set[T]]'
|
||||
/// ```
|
||||
///
|
||||
/// Checking `C[set[int]] <: C[set[int]]` leads to checking `C[set[set[int]]] <: C[set[set[int]]]`,
|
||||
/// then `C[set[set[set[int]]]] <: C[set[set[set[int]]]]`, etc. Each level has different type
|
||||
/// specializations, so using full types as keys doesn't detect the cycle. By using `ClassLiteral`
|
||||
/// as the key for protocol comparisons, we detect that we're comparing protocol `C` against itself
|
||||
/// regardless of specialization.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub(crate) enum TypeRelationKey<'db> {
|
||||
/// A regular type - used for most comparisons.
|
||||
Type(Type<'db>),
|
||||
/// A protocol class literal (without specialization) - used for protocol-to-protocol comparisons
|
||||
/// to detect recursive generic protocols.
|
||||
ProtocolClass(ClassLiteral<'db>),
|
||||
}
|
||||
|
||||
impl<'db> From<Type<'db>> for TypeRelationKey<'db> {
|
||||
fn from(ty: Type<'db>) -> Self {
|
||||
TypeRelationKey::Type(ty)
|
||||
}
|
||||
}
|
||||
|
||||
/// A [`CycleDetector`] that is used in `has_relation_to` methods.
|
||||
pub(crate) type HasRelationToVisitor<'db> = CycleDetector<
|
||||
TypeRelation<'db>,
|
||||
(
|
||||
TypeRelationKey<'db>,
|
||||
TypeRelationKey<'db>,
|
||||
TypeRelation<'db>,
|
||||
),
|
||||
ConstraintSet<'db>,
|
||||
>;
|
||||
/// A [`PairVisitor`] that is used in `has_relation_to` methods.
|
||||
pub(crate) type HasRelationToVisitor<'db> =
|
||||
CycleDetector<TypeRelation<'db>, (Type<'db>, Type<'db>, TypeRelation<'db>), ConstraintSet<'db>>;
|
||||
|
||||
impl Default for HasRelationToVisitor<'_> {
|
||||
fn default() -> Self {
|
||||
@@ -2012,7 +1973,7 @@ impl<'db> Type<'db> {
|
||||
}
|
||||
|
||||
(Type::TypeAlias(self_alias), _) => {
|
||||
relation_visitor.visit((self.into(), target.into(), relation), || {
|
||||
relation_visitor.visit((self, target, relation), || {
|
||||
self_alias.value_type(db).has_relation_to_impl(
|
||||
db,
|
||||
target,
|
||||
@@ -2025,7 +1986,7 @@ impl<'db> Type<'db> {
|
||||
}
|
||||
|
||||
(_, Type::TypeAlias(target_alias)) => {
|
||||
relation_visitor.visit((self.into(), target.into(), relation), || {
|
||||
relation_visitor.visit((self, target, relation), || {
|
||||
self.has_relation_to_impl(
|
||||
db,
|
||||
target_alias.value_type(db),
|
||||
@@ -2491,7 +2452,7 @@ impl<'db> Type<'db> {
|
||||
) => ConstraintSet::from(false),
|
||||
|
||||
(Type::Callable(self_callable), Type::Callable(other_callable)) => relation_visitor
|
||||
.visit((self.into(), target.into(), relation), || {
|
||||
.visit((self, target, relation), || {
|
||||
self_callable.has_relation_to_impl(
|
||||
db,
|
||||
other_callable,
|
||||
@@ -2503,7 +2464,7 @@ impl<'db> Type<'db> {
|
||||
}),
|
||||
|
||||
(_, Type::Callable(other_callable)) => {
|
||||
relation_visitor.visit((self.into(), target.into(), relation), || {
|
||||
relation_visitor.visit((self, target, relation), || {
|
||||
self.try_upcast_to_callable(db).when_some_and(|callables| {
|
||||
callables.has_relation_to_impl(
|
||||
db,
|
||||
@@ -2538,26 +2499,7 @@ impl<'db> Type<'db> {
|
||||
}
|
||||
|
||||
(_, Type::ProtocolInstance(protocol)) => {
|
||||
// For protocol-to-protocol comparisons, use ClassLiteral keys to detect
|
||||
// infinite recursion with recursive generic protocols (e.g., `class C[T](Protocol): a: C[set[T]]`).
|
||||
// When both types are protocols of the same class, the types may differ due to
|
||||
// different specializations, but comparing them would lead to infinite recursion.
|
||||
let (self_key, target_key) = if let Type::ProtocolInstance(self_protocol) = self {
|
||||
// Both are protocol instances - try to use class literals as keys
|
||||
// for detecting cycles in recursive generic protocols
|
||||
match (self_protocol.class_literal(db), protocol.class_literal(db)) {
|
||||
(Some(self_class), Some(target_class)) => (
|
||||
TypeRelationKey::ProtocolClass(self_class),
|
||||
TypeRelationKey::ProtocolClass(target_class),
|
||||
),
|
||||
// One or both are synthesized protocols - fall back to full types
|
||||
_ => (self.into(), target.into()),
|
||||
}
|
||||
} else {
|
||||
// Source is not a protocol - use full types
|
||||
(self.into(), target.into())
|
||||
};
|
||||
relation_visitor.visit((self_key, target_key, relation), || {
|
||||
relation_visitor.visit((self, target, relation), || {
|
||||
self.satisfies_protocol(
|
||||
db,
|
||||
protocol,
|
||||
@@ -2573,7 +2515,7 @@ impl<'db> Type<'db> {
|
||||
(Type::ProtocolInstance(_), _) => ConstraintSet::from(false),
|
||||
|
||||
(Type::TypedDict(self_typeddict), Type::TypedDict(other_typeddict)) => relation_visitor
|
||||
.visit((self.into(), target.into(), relation), || {
|
||||
.visit((self, target, relation), || {
|
||||
self_typeddict.has_relation_to_impl(
|
||||
db,
|
||||
other_typeddict,
|
||||
@@ -2588,23 +2530,18 @@ impl<'db> Type<'db> {
|
||||
// compatible `Mapping`s. `extra_items` could also allow for some assignments to `dict`, as
|
||||
// long as `total=False`. (But then again, does anyone want a non-total `TypedDict` where all
|
||||
// key types are a supertype of the extra items type?)
|
||||
(Type::TypedDict(_), _) => {
|
||||
relation_visitor.visit((self.into(), target.into(), relation), || {
|
||||
KnownClass::Mapping
|
||||
.to_specialized_instance(
|
||||
db,
|
||||
[KnownClass::Str.to_instance(db), Type::object()],
|
||||
)
|
||||
.has_relation_to_impl(
|
||||
db,
|
||||
target,
|
||||
inferable,
|
||||
relation,
|
||||
relation_visitor,
|
||||
disjointness_visitor,
|
||||
)
|
||||
})
|
||||
}
|
||||
(Type::TypedDict(_), _) => relation_visitor.visit((self, target, relation), || {
|
||||
KnownClass::Mapping
|
||||
.to_specialized_instance(db, [KnownClass::Str.to_instance(db), Type::object()])
|
||||
.has_relation_to_impl(
|
||||
db,
|
||||
target,
|
||||
inferable,
|
||||
relation,
|
||||
relation_visitor,
|
||||
disjointness_visitor,
|
||||
)
|
||||
}),
|
||||
|
||||
// A non-`TypedDict` cannot subtype a `TypedDict`
|
||||
(_, Type::TypedDict(_)) => ConstraintSet::from(false),
|
||||
@@ -2904,7 +2841,7 @@ impl<'db> Type<'db> {
|
||||
// `bool` is a subtype of `int`, because `bool` subclasses `int`,
|
||||
// which means that all instances of `bool` are also instances of `int`
|
||||
(Type::NominalInstance(self_instance), Type::NominalInstance(target_instance)) => {
|
||||
relation_visitor.visit((self.into(), target.into(), relation), || {
|
||||
relation_visitor.visit((self, target, relation), || {
|
||||
self_instance.has_relation_to_impl(
|
||||
db,
|
||||
target_instance,
|
||||
|
||||
@@ -4207,6 +4207,9 @@ pub enum KnownClass {
|
||||
ConstraintSet,
|
||||
GenericContext,
|
||||
Specialization,
|
||||
// sqlalchemy
|
||||
SqlalchemySelect,
|
||||
SqlalchemyInstrumentedAttribute,
|
||||
}
|
||||
|
||||
impl KnownClass {
|
||||
@@ -4315,7 +4318,9 @@ impl KnownClass {
|
||||
| Self::GenericContext
|
||||
| Self::Specialization
|
||||
| Self::ProtocolMeta
|
||||
| Self::TypedDictFallback => Some(Truthiness::Ambiguous),
|
||||
| Self::TypedDictFallback
|
||||
| Self::SqlalchemySelect
|
||||
| Self::SqlalchemyInstrumentedAttribute => Some(Truthiness::Ambiguous),
|
||||
|
||||
Self::Tuple => None,
|
||||
}
|
||||
@@ -4405,7 +4410,9 @@ impl KnownClass {
|
||||
| KnownClass::BuiltinFunctionType
|
||||
| KnownClass::ProtocolMeta
|
||||
| KnownClass::Template
|
||||
| KnownClass::Path => false,
|
||||
| KnownClass::Path
|
||||
| KnownClass::SqlalchemySelect
|
||||
| KnownClass::SqlalchemyInstrumentedAttribute => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4492,7 +4499,9 @@ impl KnownClass {
|
||||
| KnownClass::BuiltinFunctionType
|
||||
| KnownClass::ProtocolMeta
|
||||
| KnownClass::Template
|
||||
| KnownClass::Path => false,
|
||||
| KnownClass::Path
|
||||
| KnownClass::SqlalchemySelect
|
||||
| KnownClass::SqlalchemyInstrumentedAttribute => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4578,7 +4587,9 @@ impl KnownClass {
|
||||
| KnownClass::BuiltinFunctionType
|
||||
| KnownClass::ProtocolMeta
|
||||
| KnownClass::Template
|
||||
| KnownClass::Path => false,
|
||||
| KnownClass::Path
|
||||
| KnownClass::SqlalchemySelect
|
||||
| KnownClass::SqlalchemyInstrumentedAttribute => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4677,7 +4688,9 @@ impl KnownClass {
|
||||
| Self::ProtocolMeta
|
||||
| Self::Template
|
||||
| Self::Path
|
||||
| Self::Mapping => false,
|
||||
| Self::Mapping
|
||||
| Self::SqlalchemySelect
|
||||
| Self::SqlalchemyInstrumentedAttribute => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4766,7 +4779,9 @@ impl KnownClass {
|
||||
| KnownClass::ConstraintSet
|
||||
| KnownClass::GenericContext
|
||||
| KnownClass::Specialization
|
||||
| KnownClass::InitVar => false,
|
||||
| KnownClass::InitVar
|
||||
| KnownClass::SqlalchemySelect
|
||||
| KnownClass::SqlalchemyInstrumentedAttribute => false,
|
||||
KnownClass::NamedTupleFallback | KnownClass::TypedDictFallback => true,
|
||||
}
|
||||
}
|
||||
@@ -4882,6 +4897,8 @@ impl KnownClass {
|
||||
Self::Template => "Template",
|
||||
Self::Path => "Path",
|
||||
Self::ProtocolMeta => "_ProtocolMeta",
|
||||
Self::SqlalchemySelect => "Select",
|
||||
Self::SqlalchemyInstrumentedAttribute => "InstrumentedAttribute",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5203,6 +5220,8 @@ impl KnownClass {
|
||||
| Self::Specialization => KnownModule::TyExtensions,
|
||||
Self::Template => KnownModule::Templatelib,
|
||||
Self::Path => KnownModule::Pathlib,
|
||||
Self::SqlalchemySelect => KnownModule::SqlalchemySqlSelectable,
|
||||
Self::SqlalchemyInstrumentedAttribute => KnownModule::SqlalchemyOrmAttributes,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5291,7 +5310,9 @@ impl KnownClass {
|
||||
| Self::BuiltinFunctionType
|
||||
| Self::ProtocolMeta
|
||||
| Self::Template
|
||||
| Self::Path => Some(false),
|
||||
| Self::Path
|
||||
| Self::SqlalchemySelect
|
||||
| Self::SqlalchemyInstrumentedAttribute => Some(false),
|
||||
|
||||
Self::Tuple => None,
|
||||
}
|
||||
@@ -5383,7 +5404,9 @@ impl KnownClass {
|
||||
| Self::BuiltinFunctionType
|
||||
| Self::ProtocolMeta
|
||||
| Self::Template
|
||||
| Self::Path => false,
|
||||
| Self::Path
|
||||
| Self::SqlalchemySelect
|
||||
| Self::SqlalchemyInstrumentedAttribute => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5489,6 +5512,8 @@ impl KnownClass {
|
||||
"Template" => &[Self::Template],
|
||||
"Path" => &[Self::Path],
|
||||
"_ProtocolMeta" => &[Self::ProtocolMeta],
|
||||
"Select" => &[Self::SqlalchemySelect],
|
||||
"InstrumentedAttribute" => &[Self::SqlalchemyInstrumentedAttribute],
|
||||
_ => return None,
|
||||
};
|
||||
|
||||
@@ -5569,7 +5594,9 @@ impl KnownClass {
|
||||
| Self::Awaitable
|
||||
| Self::Generator
|
||||
| Self::Template
|
||||
| Self::Path => module == self.canonical_module(db),
|
||||
| Self::Path
|
||||
| Self::SqlalchemySelect
|
||||
| Self::SqlalchemyInstrumentedAttribute => module == self.canonical_module(db),
|
||||
Self::NoneType => matches!(module, KnownModule::Typeshed | KnownModule::Types),
|
||||
Self::SpecialForm
|
||||
| Self::TypeAliasType
|
||||
@@ -5924,6 +5951,10 @@ mod tests {
|
||||
source: PythonVersionSource::default(),
|
||||
});
|
||||
for class in KnownClass::iter() {
|
||||
if class.canonical_module(&db).is_third_party() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let class_name = class.name(&db);
|
||||
let class_module =
|
||||
resolve_module_confident(&db, &class.canonical_module(&db).name()).unwrap();
|
||||
@@ -5952,6 +5983,10 @@ mod tests {
|
||||
});
|
||||
|
||||
for class in KnownClass::iter() {
|
||||
if class.canonical_module(&db).is_third_party() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check the class can be looked up successfully
|
||||
class.try_to_class_literal_without_logging(&db).unwrap();
|
||||
|
||||
@@ -5977,6 +6012,7 @@ mod tests {
|
||||
// This makes the test far faster as it minimizes the number of times
|
||||
// we need to change the Python version in the loop.
|
||||
let mut classes: Vec<(KnownClass, PythonVersion)> = KnownClass::iter()
|
||||
.filter(|class| !class.canonical_module(&db).is_third_party())
|
||||
.map(|class| {
|
||||
let version_added = match class {
|
||||
KnownClass::Template => PythonVersion::PY314,
|
||||
|
||||
@@ -1353,6 +1353,10 @@ pub enum KnownFunction {
|
||||
RevealProtocolInterface,
|
||||
/// `ty_extensions.reveal_mro`
|
||||
RevealMro,
|
||||
|
||||
/// `sqlalchemy.select`
|
||||
#[strum(serialize = "select")]
|
||||
SqlalchemySelect,
|
||||
}
|
||||
|
||||
impl KnownFunction {
|
||||
@@ -1425,6 +1429,9 @@ impl KnownFunction {
|
||||
|
||||
Self::TypeCheckOnly => matches!(module, KnownModule::Typing),
|
||||
Self::NamedTuple => matches!(module, KnownModule::Collections),
|
||||
Self::SqlalchemySelect => {
|
||||
matches!(module, KnownModule::SqlalchemySqlSelectableConstructors)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1896,6 +1903,56 @@ impl KnownFunction {
|
||||
|
||||
overload.set_return_type(Type::module_literal(db, file, module));
|
||||
}
|
||||
|
||||
KnownFunction::SqlalchemySelect => {
|
||||
// Try to extract types from InstrumentedAttribute[T] arguments.
|
||||
// If all arguments are InstrumentedAttribute instances, we construct
|
||||
// Select[tuple[T_1, T_2, ...]] where T_i are the inner types.
|
||||
//
|
||||
// We check the class via `class_literal.known(db)` rather than using
|
||||
// `known_specialization` because the class may be re-exported and not
|
||||
// directly importable from its canonical module.
|
||||
let inner_types: Option<Vec<_>> = parameter_types
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(|param_type| {
|
||||
let Type::NominalInstance(instance) = param_type else {
|
||||
return None;
|
||||
};
|
||||
let class = instance.class(db);
|
||||
let (class_literal, specialization) = class.class_literal(db);
|
||||
if class_literal.known(db)
|
||||
!= Some(KnownClass::SqlalchemyInstrumentedAttribute)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
specialization?.types(db).first().copied()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let Some(inner_types) = inner_types else {
|
||||
// Fall back to whatever we infer from the function signature
|
||||
return;
|
||||
};
|
||||
|
||||
if inner_types.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Construct Select[tuple[T1, T2, ...]]
|
||||
// We get the return type's class from the overload rather than looking
|
||||
// it up via try_to_class_literal, since the class may be re-exported.
|
||||
let Type::NominalInstance(return_instance) = overload.return_type() else {
|
||||
return;
|
||||
};
|
||||
let select_class = return_instance.class(db).class_literal(db).0;
|
||||
let tuple_type = Type::heterogeneous_tuple(db, inner_types);
|
||||
let class_type = select_class.apply_specialization(db, |generic_context| {
|
||||
generic_context.specialize(db, vec![tuple_type].into())
|
||||
});
|
||||
overload.set_return_type(Type::instance(db, class_type));
|
||||
}
|
||||
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
@@ -1964,6 +2021,8 @@ pub(crate) mod tests {
|
||||
|
||||
KnownFunction::ImportModule => KnownModule::ImportLib,
|
||||
KnownFunction::NamedTuple => KnownModule::Collections,
|
||||
|
||||
KnownFunction::SqlalchemySelect => continue,
|
||||
};
|
||||
|
||||
let function_definition = known_module_symbol(&db, module, function_name)
|
||||
|
||||
@@ -659,16 +659,6 @@ impl<'db> ProtocolInstanceType<'db> {
|
||||
}
|
||||
}
|
||||
|
||||
/// If this is a class-based protocol, return its class literal (without specialization).
|
||||
///
|
||||
/// Returns `None` for synthesized protocols that don't correspond to a class definition.
|
||||
pub(super) fn class_literal(self, db: &'db dyn Db) -> Option<ClassLiteral<'db>> {
|
||||
match self.inner {
|
||||
Protocol::FromClass(class) => Some(class.class_literal(db).0),
|
||||
Protocol::Synthesized(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the meta-type of this protocol-instance type.
|
||||
pub(super) fn to_meta_type(self, db: &'db dyn Db) -> Type<'db> {
|
||||
match self.inner {
|
||||
|
||||
Reference in New Issue
Block a user