Skip to content

Commit

Permalink
[red-knot] move standalone expression_ty to TypeInferenceBuilder::fil…
Browse files Browse the repository at this point in the history
…e_expression_ty (#14879)

## Summary

Per suggestion in
#14802 (comment)

This is a bit less error-prone and allows us to handle both expressions
in the current scope or a different scope. Also, there's currently no
need for this method outside of `TypeInferenceBuilder`, so no reason to
expose it in `types.rs`.

## Test Plan

Pure refactor, no functional change; existing tests pass.

---------

Co-authored-by: Alex Waygood <[email protected]>
  • Loading branch information
carljm and AlexWaygood authored Dec 9, 2024
1 parent c62ba48 commit 533e8a6
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 18 deletions.
12 changes: 0 additions & 12 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -225,18 +225,6 @@ fn definition_expression_ty<'db>(
}
}

/// Get the type of an expression from an arbitrary scope.
///
/// Can cause query cycles if used carelessly; caller must be sure that type inference isn't
/// currently in progress for the expression's scope.
fn expression_ty<'db>(db: &'db dyn Db, file: File, expression: &ast::Expr) -> Type<'db> {
let index = semantic_index(db, file);
let file_scope = index.expression_scope_id(expression);
let scope = file_scope.to_scope_id(db, file);
let expr_id = expression.scoped_expression_id(db, scope);
infer_scope_types(db, scope).expression_ty(expr_id)
}

/// Infer the combined type of an iterator of bindings.
///
/// Will return a union if there is more than one binding.
Expand Down
35 changes: 29 additions & 6 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,6 @@ use crate::unpack::Unpack;
use crate::util::subscript::{PyIndex, PySlice};
use crate::Db;

use super::expression_ty;
use super::string_annotation::parse_string_annotation;

/// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope.
Expand Down Expand Up @@ -407,13 +406,37 @@ impl<'db> TypeInferenceBuilder<'db> {

/// Get the already-inferred type of an expression node.
///
/// PANIC if no type has been inferred for this node.
/// ## Panics
/// If the expression is not within this region, or if no type has yet been inferred for
/// this node.
#[track_caller]
fn expression_ty(&self, expr: &ast::Expr) -> Type<'db> {
self.types
.expression_ty(expr.scoped_expression_id(self.db, self.scope()))
}

/// Get the type of an expression from any scope in the same file.
///
/// If the expression is in the current scope, and we are inferring the entire scope, just look
/// up the expression in our own results, otherwise call [`infer_scope_types()`] for the scope
/// of the expression.
///
/// ## Panics
///
/// If the expression is in the current scope but we haven't yet inferred a type for it.
///
/// Can cause query cycles if the expression is from a different scope and type inference is
/// already in progress for that scope (further up the stack).
fn file_expression_ty(&self, expression: &ast::Expr) -> Type<'db> {
let file_scope = self.index.expression_scope_id(expression);
let expr_scope = file_scope.to_scope_id(self.db, self.file);
let expr_id = expression.scoped_expression_id(self.db, expr_scope);
match self.region {
InferenceRegion::Scope(scope) if scope == expr_scope => self.expression_ty(expression),
_ => infer_scope_types(self.db, expr_scope).expression_ty(expr_id),
}
}

/// Infers types in the given [`InferenceRegion`].
fn infer_region(&mut self) {
match self.region {
Expand Down Expand Up @@ -1081,9 +1104,9 @@ impl<'db> TypeInferenceBuilder<'db> {
} = parameter_with_default;
let default_ty = default
.as_ref()
.map(|default| expression_ty(self.db, self.file, default));
.map(|default| self.file_expression_ty(default));
if let Some(annotation) = parameter.annotation.as_ref() {
let declared_ty = expression_ty(self.db, self.file, annotation);
let declared_ty = self.file_expression_ty(annotation);
let inferred_ty = if let Some(default_ty) = default_ty {
if default_ty.is_assignable_to(self.db, declared_ty) {
UnionType::from_elements(self.db, [declared_ty, default_ty])
Expand Down Expand Up @@ -1127,7 +1150,7 @@ impl<'db> TypeInferenceBuilder<'db> {
definition: Definition<'db>,
) {
if let Some(annotation) = parameter.annotation.as_ref() {
let _annotated_ty = expression_ty(self.db, self.file, annotation);
let _annotated_ty = self.file_expression_ty(annotation);
// TODO `tuple[annotated_ty, ...]`
let ty = KnownClass::Tuple.to_instance(self.db);
self.add_declaration_with_binding(parameter.into(), definition, ty, ty);
Expand All @@ -1152,7 +1175,7 @@ impl<'db> TypeInferenceBuilder<'db> {
definition: Definition<'db>,
) {
if let Some(annotation) = parameter.annotation.as_ref() {
let _annotated_ty = expression_ty(self.db, self.file, annotation);
let _annotated_ty = self.file_expression_ty(annotation);
// TODO `dict[str, annotated_ty]`
let ty = KnownClass::Dict.to_instance(self.db);
self.add_declaration_with_binding(parameter.into(), definition, ty, ty);
Expand Down

0 comments on commit 533e8a6

Please sign in to comment.