diff --git a/crates/ty_python_semantic/resources/mdtest/function/return_type.md b/crates/ty_python_semantic/resources/mdtest/function/return_type.md index 1ca2b10dca..1dfe9eec62 100644 --- a/crates/ty_python_semantic/resources/mdtest/function/return_type.md +++ b/crates/ty_python_semantic/resources/mdtest/function/return_type.md @@ -127,6 +127,80 @@ def f(x: int | str): return x ``` +### In `if TYPE_CHECKING` block + +Inside an `if TYPE_CHECKING` block, we allow "stub" style function definitions with empty bodies, +since these functions will never actually be called. + +```py +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + def f() -> int: ... + +else: + def f() -> str: + return "hello" + +reveal_type(f) # revealed: def f() -> int + +if not TYPE_CHECKING: + ... +elif True: + def g() -> str: ... + +else: + def h() -> str: ... + +if not TYPE_CHECKING: + def i() -> int: + return 1 + +else: + def i() -> str: ... + +reveal_type(i) # revealed: def i() -> str + +if False: + ... +elif TYPE_CHECKING: + def j() -> str: ... + +else: + def j_() -> str: ... # error: [invalid-return-type] + +if False: + ... +elif not TYPE_CHECKING: + def k_() -> str: ... # error: [invalid-return-type] + +else: + def k() -> str: ... + +class Foo: + if TYPE_CHECKING: + def f(self) -> int: ... + +if TYPE_CHECKING: + class Bar: + def f(self) -> int: ... + +def get_bool() -> bool: + return True + +if TYPE_CHECKING: + if get_bool(): + def l() -> str: ... + +if get_bool(): + if TYPE_CHECKING: + def m() -> str: ... + +if TYPE_CHECKING: + if not TYPE_CHECKING: + def n() -> str: ... +``` + ## Conditional return type ```py diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index 32ba361e1d..ef273f2b14 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -90,6 +90,8 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> { /// Flags about the file's global scope has_future_annotations: bool, + /// Whether we are currently visiting an `if TYPE_CHECKING` block. + in_type_checking_block: bool, // Used for checking semantic syntax errors python_version: PythonVersion, @@ -130,6 +132,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { try_node_context_stack_manager: TryNodeContextStackManager::default(), has_future_annotations: false, + in_type_checking_block: false, scopes: IndexVec::new(), place_tables: IndexVec::new(), @@ -248,6 +251,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { node_with_kind, children_start..children_start, reachability, + self.in_type_checking_block, ); let is_class_scope = scope.kind().is_class(); self.try_node_context_stack_manager.enter_nested_scope(); @@ -719,7 +723,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { // since its the pattern that introduces any constraints, not the body.) Ideally, that // standalone expression would wrap the match arm's pattern as a whole. But a standalone // expression can currently only wrap an ast::Expr, which patterns are not. So, we need to - // choose an Expr that can “stand in” for the pattern, which we can wrap in a standalone + // choose an Expr that can "stand in" for the pattern, which we can wrap in a standalone // expression. // // See the comment in TypeInferenceBuilder::infer_match_pattern for more details. @@ -1498,6 +1502,17 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { let mut last_predicate = self.record_expression_narrowing_constraint(&node.test); let mut last_reachability_constraint = self.record_reachability_constraint(last_predicate); + + let is_outer_block_in_type_checking = self.in_type_checking_block; + + let if_block_in_type_checking = is_if_type_checking(&node.test); + + // Track if we're in a chain that started with "not TYPE_CHECKING" + let mut is_in_not_type_checking_chain = is_if_not_type_checking(&node.test); + + self.in_type_checking_block = + if_block_in_type_checking || is_outer_block_in_type_checking; + self.visit_body(&node.body); let mut post_clauses: Vec = vec![]; @@ -1516,6 +1531,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { // if there's no `else` branch, we should add a no-op `else` branch Some((None, Default::default())) }); + for (clause_test, clause_body) in elif_else_clauses { // snapshot after every block except the last; the last one will just become // the state that we merge the other snapshots into @@ -1538,12 +1554,34 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { self.record_reachability_constraint(last_predicate); } + // Determine if this clause is in type checking context + let clause_in_type_checking = if let Some(elif_test) = clause_test { + if is_if_type_checking(elif_test) { + // This block has "TYPE_CHECKING" condition + true + } else if is_if_not_type_checking(elif_test) { + // This block has "not TYPE_CHECKING" condition so we update the chain state for future blocks + is_in_not_type_checking_chain = true; + false + } else { + // This block has some other condition + // It's in type checking only if we're in a "not TYPE_CHECKING" chain + is_in_not_type_checking_chain + } + } else { + is_in_not_type_checking_chain + }; + + self.in_type_checking_block = clause_in_type_checking; + self.visit_body(clause_body); } for post_clause_state in post_clauses { self.flow_merge(post_clause_state); } + + self.in_type_checking_block = is_outer_block_in_type_checking; } ast::Stmt::While(ast::StmtWhile { test, @@ -2711,3 +2749,18 @@ impl ExpressionsScopeMapBuilder { ExpressionsScopeMap(interval_map.into_boxed_slice()) } } + +/// Returns if the expression is a `TYPE_CHECKING` expression. +fn is_if_type_checking(expr: &ast::Expr) -> bool { + matches!(expr, ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING") +} + +/// Returns if the expression is a `not TYPE_CHECKING` expression. +fn is_if_not_type_checking(expr: &ast::Expr) -> bool { + matches!(expr, ast::Expr::UnaryOp(ast::ExprUnaryOp { op, operand, .. }) if *op == ruff_python_ast::UnaryOp::Not + && matches!( + &**operand, + ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING" + ) + ) +} diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs index f8962a6183..81b650f19e 100644 --- a/crates/ty_python_semantic/src/semantic_index/place.rs +++ b/crates/ty_python_semantic/src/semantic_index/place.rs @@ -525,10 +525,20 @@ impl FileScopeId { #[derive(Debug, salsa::Update, get_size2::GetSize)] pub struct Scope { + /// The parent scope, if any. parent: Option, + + /// The node that introduces this scope. node: NodeWithScopeKind, + + /// The range of [`FileScopeId`]s that are descendants of this scope. descendants: Range, + + /// The constraint that determines the reachability of this scope. reachability: ScopedReachabilityConstraintId, + + /// Whether this scope is defined inside an `if TYPE_CHECKING:` block. + in_type_checking_block: bool, } impl Scope { @@ -537,12 +547,14 @@ impl Scope { node: NodeWithScopeKind, descendants: Range, reachability: ScopedReachabilityConstraintId, + in_type_checking_block: bool, ) -> Self { Scope { parent, node, descendants, reachability, + in_type_checking_block, } } @@ -573,6 +585,10 @@ impl Scope { pub(crate) fn reachability(&self) -> ScopedReachabilityConstraintId { self.reachability } + + pub(crate) fn in_type_checking_block(&self) -> bool { + self.in_type_checking_block + } } #[derive(Copy, Clone, Debug, PartialEq, Eq)] diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index af0f27eb08..15c280f92f 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -2135,6 +2135,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { if self.in_function_overload_or_abstractmethod() { return; } + if self.scope().scope(self.db()).in_type_checking_block() { + return; + } if let Some(class) = self.class_context_of_current_method() { enclosing_class_context = Some(class); if class.is_protocol(self.db()) {