[ty] Pull types on synthesized Python files created by mdtest (#18539)

Author: AlexWaygood

Cargo.lock

@@ -50,6 +50,7 @@ strum_macros = { workspace = true }
 [dev-dependencies]
 ruff_db = { workspace = true, features = ["testing", "os"] }
 ruff_python_parser = { workspace = true }
+ty_python_semantic = { workspace = true, features = ["testing"] }
 ty_test = { workspace = true }
 ty_vendored = { workspace = true }
 
@@ -63,6 +64,7 @@ quickcheck_macros = { version = "1.0.0" }
 
 [features]
 serde = ["ruff_db/serde", "dep:serde", "ruff_python_ast/serde"]
+testing = []
 
 [lints]
 workspace = true

crates/ty_python_semantic/Cargo.toml

@@ -139,6 +139,8 @@ x: int = MagicMock()
 
 ## Invalid
 
+<!-- pull-types:skip -->
+
 `Any` cannot be parameterized:
 
 ```py

crates/ty_python_semantic/resources/mdtest/annotations/any.md

@@ -58,6 +58,8 @@ def _(c: Callable[[int, 42, str, False], None]):
 
 ### Missing return type
 
+<!-- pull-types:skip -->
+
 Using a parameter list:
 
 ```py

crates/ty_python_semantic/resources/mdtest/annotations/callable.md

@@ -14,6 +14,8 @@ directly.
 
 ### Negation
 
+<!-- pull-types:skip -->
+
 ```py
 from typing import Literal
 from ty_extensions import Not, static_assert
@@ -371,6 +373,8 @@ static_assert(not is_single_valued(Literal["a"] | Literal["b"]))
 
 ## `TypeOf`
 
+<!-- pull-types:skip -->
+
 We use `TypeOf` to get the inferred type of an expression. This is useful when we want to refer to
 it in a type expression. For example, if we want to make sure that the class literal type `str` is a
 subtype of `type[str]`, we can not use `is_subtype_of(str, type[str])`, as that would test if the
@@ -412,6 +416,8 @@ def f(x: TypeOf) -> None:
 
 ## `CallableTypeOf`
 
+<!-- pull-types:skip -->
+
 The `CallableTypeOf` special form can be used to extract the `Callable` structural type inhabited by
 a given callable object. This can be used to get the externally visibly signature of the object,
 which can then be used to test various type properties.

crates/ty_python_semantic/resources/mdtest/type_api.md

@@ -84,6 +84,8 @@ d.a = 2
 
 ## Too many arguments
 
+<!-- pull-types:skip -->
+
 ```py
 from typing import ClassVar
 

crates/ty_python_semantic/resources/mdtest/type_qualifiers/classvar.md

@@ -45,6 +45,8 @@ reveal_type(FINAL_E)  # revealed: int
 
 ## Too many arguments
 
+<!-- pull-types:skip -->
+
 ```py
 from typing import Final
 

crates/ty_python_semantic/resources/mdtest/type_qualifiers/final.md

@@ -35,6 +35,9 @@ pub mod types;
 mod unpack;
 mod util;
 
+#[cfg(feature = "testing")]
+pub mod pull_types;
+
 type FxOrderSet<V> = ordermap::set::OrderSet<V, BuildHasherDefault<FxHasher>>;
 
 /// Returns the default registry with all known semantic lints.

crates/ty_python_semantic/src/lib.rs

@@ -0,0 +1,134 @@
+//! A utility visitor for testing, which attempts to "pull a type" for ever sub-node in a given AST.
+//!
+//! This is used in the "corpus" and (indirectly) the "mdtest" integration tests for this crate.
+//! (Mdtest uses the `pull_types` function via the `ty_test` crate.)
+
+use crate::{Db, HasType, SemanticModel};
+use ruff_db::{files::File, parsed::parsed_module};
+use ruff_python_ast::{
+    self as ast, visitor::source_order, visitor::source_order::SourceOrderVisitor,
+};
+
+pub fn pull_types(db: &dyn Db, file: File) {
+    let mut visitor = PullTypesVisitor::new(db, file);
+
+    let ast = parsed_module(db.upcast(), file).load(db.upcast());
+
+    visitor.visit_body(ast.suite());
+}
+
+struct PullTypesVisitor<'db> {
+    model: SemanticModel<'db>,
+}
+
+impl<'db> PullTypesVisitor<'db> {
+    fn new(db: &'db dyn Db, file: File) -> Self {
+        Self {
+            model: SemanticModel::new(db, file),
+        }
+    }
+
+    fn visit_target(&mut self, target: &ast::Expr) {
+        match target {
+            ast::Expr::List(ast::ExprList { elts, .. })
+            | ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                for element in elts {
+                    self.visit_target(element);
+                }
+            }
+            _ => self.visit_expr(target),
+        }
+    }
+}
+
+impl SourceOrderVisitor<'_> for PullTypesVisitor<'_> {
+    fn visit_stmt(&mut self, stmt: &ast::Stmt) {
+        match stmt {
+            ast::Stmt::FunctionDef(function) => {
+                let _ty = function.inferred_type(&self.model);
+            }
+            ast::Stmt::ClassDef(class) => {
+                let _ty = class.inferred_type(&self.model);
+            }
+            ast::Stmt::Assign(assign) => {
+                for target in &assign.targets {
+                    self.visit_target(target);
+                }
+                self.visit_expr(&assign.value);
+                return;
+            }
+            ast::Stmt::For(for_stmt) => {
+                self.visit_target(&for_stmt.target);
+                self.visit_expr(&for_stmt.iter);
+                self.visit_body(&for_stmt.body);
+                self.visit_body(&for_stmt.orelse);
+                return;
+            }
+            ast::Stmt::With(with_stmt) => {
+                for item in &with_stmt.items {
+                    if let Some(target) = &item.optional_vars {
+                        self.visit_target(target);
+                    }
+                    self.visit_expr(&item.context_expr);
+                }
+
+                self.visit_body(&with_stmt.body);
+                return;
+            }
+            ast::Stmt::AnnAssign(_)
+            | ast::Stmt::Return(_)
+            | ast::Stmt::Delete(_)
+            | ast::Stmt::AugAssign(_)
+            | ast::Stmt::TypeAlias(_)
+            | ast::Stmt::While(_)
+            | ast::Stmt::If(_)
+            | ast::Stmt::Match(_)
+            | ast::Stmt::Raise(_)
+            | ast::Stmt::Try(_)
+            | ast::Stmt::Assert(_)
+            | ast::Stmt::Import(_)
+            | ast::Stmt::ImportFrom(_)
+            | ast::Stmt::Global(_)
+            | ast::Stmt::Nonlocal(_)
+            | ast::Stmt::Expr(_)
+            | ast::Stmt::Pass(_)
+            | ast::Stmt::Break(_)
+            | ast::Stmt::Continue(_)
+            | ast::Stmt::IpyEscapeCommand(_) => {}
+        }
+
+        source_order::walk_stmt(self, stmt);
+    }
+
+    fn visit_expr(&mut self, expr: &ast::Expr) {
+        let _ty = expr.inferred_type(&self.model);
+
+        source_order::walk_expr(self, expr);
+    }
+
+    fn visit_comprehension(&mut self, comprehension: &ast::Comprehension) {
+        self.visit_expr(&comprehension.iter);
+        self.visit_target(&comprehension.target);
+        for if_expr in &comprehension.ifs {
+            self.visit_expr(if_expr);
+        }
+    }
+
+    fn visit_parameter(&mut self, parameter: &ast::Parameter) {
+        let _ty = parameter.inferred_type(&self.model);
+
+        source_order::walk_parameter(self, parameter);
+    }
+
+    fn visit_parameter_with_default(&mut self, parameter_with_default: &ast::ParameterWithDefault) {
+        let _ty = parameter_with_default.inferred_type(&self.model);
+
+        source_order::walk_parameter_with_default(self, parameter_with_default);
+    }
+
+    fn visit_alias(&mut self, alias: &ast::Alias) {
+        let _ty = alias.inferred_type(&self.model);
+
+        source_order::walk_alias(self, alias);
+    }
+}

crates/ty_python_semantic/src/pull_types.rs

@@ -471,8 +471,10 @@ impl<'db> TypeInference<'db> {
     #[track_caller]
     pub(crate) fn expression_type(&self, expression: ScopedExpressionId) -> Type<'db> {
         self.try_expression_type(expression).expect(
-            "expression should belong to this TypeInference region and \
-            TypeInferenceBuilder should have inferred a type for it",
+            "Failed to retrieve the inferred type for an `ast::Expr` node \
+            passed to `TypeInference::expression_type()`. The `TypeInferenceBuilder` \
+            should infer and store types for all `ast::Expr` nodes in any `TypeInference` \
+            region it analyzes.",
         )
     }