Skip to content

Commit

Permalink
[red-knot] support deferred evaluation of type expressions (#13131)
Browse files Browse the repository at this point in the history
Prototype deferred evaluation of type expressions by deferring
evaluation of class bases in a stub file. This allows self-referential
class definitions, as occur with the definition of `str` in typeshed
(which inherits `Sequence[str]`).

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
  • Loading branch information
carljm and AlexWaygood authored Aug 28, 2024
1 parent c6023c0 commit 770ef2a
Show file tree
Hide file tree
Showing 3 changed files with 207 additions and 40 deletions.
56 changes: 45 additions & 11 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
use ruff_db::files::File;
use ruff_python_ast::name::Name;
use ruff_python_ast as ast;

use crate::builtins::builtins_scope;
use crate::semantic_index::definition::Definition;
use crate::semantic_index::ast_ids::HasScopedAstId;
use crate::semantic_index::definition::{Definition, DefinitionKind};
use crate::semantic_index::symbol::{ScopeId, ScopedSymbolId};
use crate::semantic_index::{
global_scope, semantic_index, symbol_table, use_def_map, DefinitionWithConstraints,
Expand All @@ -14,7 +15,8 @@ use crate::{Db, FxOrderSet};
pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
pub(crate) use self::diagnostic::TypeCheckDiagnostics;
pub(crate) use self::infer::{
infer_definition_types, infer_expression_types, infer_scope_types, TypeInference,
infer_deferred_types, infer_definition_types, infer_expression_types, infer_scope_types,
TypeInference,
};

mod builder;
Expand Down Expand Up @@ -88,6 +90,24 @@ pub(crate) fn definition_ty<'db>(db: &'db dyn Db, definition: Definition<'db>) -
inference.definition_ty(definition)
}

/// Infer the type of a (possibly deferred) sub-expression of a [`Definition`].
///
/// ## Panics
/// If the given expression is not a sub-expression of the given [`Definition`].
pub(crate) fn definition_expression_ty<'db>(
db: &'db dyn Db,
definition: Definition<'db>,
expression: &ast::Expr,
) -> Type<'db> {
let expr_id = expression.scoped_ast_id(db, definition.scope(db));
let inference = infer_definition_types(db, definition);
if let Some(ty) = inference.try_expression_ty(expr_id) {
ty
} else {
infer_deferred_types(db, definition).expression_ty(expr_id)
}
}

/// Infer the combined type of an array of [`Definition`]s, plus one optional "unbound type".
///
/// Will return a union if there is more than one definition, or at least one plus an unbound
Expand Down Expand Up @@ -243,7 +263,7 @@ impl<'db> Type<'db> {
/// us to explicitly consider whether to handle an error or propagate
/// it up the call stack.
#[must_use]
pub fn member(&self, db: &'db dyn Db, name: &Name) -> Type<'db> {
pub fn member(&self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
match self {
Type::Any => Type::Any,
Type::Never => {
Expand Down Expand Up @@ -314,7 +334,7 @@ impl<'db> Type<'db> {
#[salsa::interned]
pub struct FunctionType<'db> {
/// name of the function at definition
pub name: Name,
pub name: ast::name::Name,

/// types of all decorators on this function
decorators: Vec<Type<'db>>,
Expand All @@ -329,19 +349,33 @@ impl<'db> FunctionType<'db> {
#[salsa::interned]
pub struct ClassType<'db> {
/// Name of the class at definition
pub name: Name,
pub name: ast::name::Name,

/// Types of all class bases
bases: Vec<Type<'db>>,
definition: Definition<'db>,

body_scope: ScopeId<'db>,
}

impl<'db> ClassType<'db> {
/// Return an iterator over the types of this class's bases.
///
/// # Panics:
/// If `definition` is not a `DefinitionKind::Class`.
pub fn bases(&self, db: &'db dyn Db) -> impl Iterator<Item = Type<'db>> {
let definition = self.definition(db);
let DefinitionKind::Class(class_stmt_node) = definition.node(db) else {
panic!("Class type definition must have DefinitionKind::Class");
};
class_stmt_node
.bases()
.iter()
.map(move |base_expr| definition_expression_ty(db, definition, base_expr))
}

/// Returns the class member of this class named `name`.
///
/// The member resolves to a member of the class itself or any of its bases.
pub fn class_member(self, db: &'db dyn Db, name: &Name) -> Type<'db> {
pub fn class_member(self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
let member = self.own_class_member(db, name);
if !member.is_unbound() {
return member;
Expand All @@ -351,12 +385,12 @@ impl<'db> ClassType<'db> {
}

/// Returns the inferred type of the class member named `name`.
pub fn own_class_member(self, db: &'db dyn Db, name: &Name) -> Type<'db> {
pub fn own_class_member(self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
let scope = self.body_scope(db);
symbol_ty_by_name(db, scope, name)
}

pub fn inherited_class_member(self, db: &'db dyn Db, name: &Name) -> Type<'db> {
pub fn inherited_class_member(self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
for base in self.bases(db) {
let member = base.member(db, name);
if !member.is_unbound() {
Expand Down
Loading

0 comments on commit 770ef2a

Please sign in to comment.