@no_type_check support (#15122)

Co-authored-by: Carl Meyer <[email protected]>

crates/red_knot_python_semantic/resources/mdtest/suppressions/no_type_check.md

@@ -0,0 +1,118 @@
+# `@no_type_check`
+
+> If a type checker supports the `no_type_check` decorator for functions, it should suppress all
+> type errors for the def statement and its body including any nested functions or classes. It
+> should also ignore all parameter and return type annotations and treat the function as if it were
+> unannotated. [source](https://typing.readthedocs.io/en/latest/spec/directives.html#no-type-check)
+
+## Error in the function body
+
+```py
+from typing import no_type_check
+
+@no_type_check
+def test() -> int:
+    return a + 5
+```
+
+## Error in nested function
+
+```py
+from typing import no_type_check
+
+@no_type_check
+def test() -> int:
+    def nested():
+        return a + 5
+```
+
+## Error in nested class
+
+```py
+from typing import no_type_check
+
+@no_type_check
+def test() -> int:
+    class Nested:
+        def inner(self):
+            return a + 5
+```
+
+## Error in preceding decorator
+
+Don't suppress diagnostics for decorators appearing before the `no_type_check` decorator.
+
+```py
+from typing import no_type_check
+
+@unknown_decorator  # error: [unresolved-reference]
+@no_type_check
+def test() -> int:
+    return a + 5
+```
+
+## Error in following decorator
+
+Unlike Pyright and mypy, suppress diagnostics appearing after the `no_type_check` decorator. We do
+this because it more closely matches Python's runtime semantics of decorators. For more details, see
+the discussion on the
+[PR adding `@no_type_check` support](https://github.com/astral-sh/ruff/pull/15122#discussion_r1896869411).
+
+```py
+from typing import no_type_check
+
+@no_type_check
+@unknown_decorator
+def test() -> int:
+    return a + 5
+```
+
+## Error in default value
+
+```py
+from typing import no_type_check
+
+@no_type_check
+def test(a: int = "test"):
+    return x + 5
+```
+
+## Error in return value position
+
+```py
+from typing import no_type_check
+
+@no_type_check
+def test() -> Undefined:
+    return x + 5
+```
+
+## `no_type_check` on classes isn't supported
+
+Red Knot does not support decorating classes with `no_type_check`. The behaviour of `no_type_check`
+when applied to classes is
+[not specified currently](https://typing.readthedocs.io/en/latest/spec/directives.html#no-type-check),
+and is not supported by Pyright or mypy.
+
+A future improvement might be to emit a diagnostic if a `no_type_check` annotation is applied to a
+class.
+
+```py
+from typing import no_type_check
+
+@no_type_check
+class Test:
+    def test(self):
+        return a + 5  # error: [unresolved-reference]
+```
+
+## `type: ignore` comments in `@no_type_check` blocks
+
+```py
+from typing import no_type_check
+
+@no_type_check
+def test():
+    # error: [unused-ignore-comment] "Unused `knot: ignore` directive: 'unresolved-reference'"
+    return x + 5  # knot: ignore[unresolved-reference]
+```

crates/red_knot_python_semantic/src/ast_node_ref.rs

@@ -43,7 +43,7 @@ impl<T> AstNodeRef<T> {
     }
 
     /// Returns a reference to the wrapped node.
-    pub fn node(&self) -> &T {
+    pub const fn node(&self) -> &T {
         // SAFETY: Holding on to `parsed` ensures that the AST to which `node` belongs is still
         // alive and not moved.
         unsafe { self.node.as_ref() }

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -6,10 +6,9 @@ use rustc_hash::{FxHashMap, FxHashSet};
 use ruff_db::files::File;
 use ruff_db::parsed::ParsedModule;
 use ruff_index::IndexVec;
+use ruff_python_ast as ast;
 use ruff_python_ast::name::Name;
 use ruff_python_ast::visitor::{walk_expr, walk_pattern, walk_stmt, Visitor};
-use ruff_python_ast::{self as ast, Pattern};
-use ruff_python_ast::{BoolOp, Expr};
 
 use crate::ast_node_ref::AstNodeRef;
 use crate::module_name::ModuleName;
@@ -289,7 +288,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         constraint
     }
 
-    fn build_constraint(&mut self, constraint_node: &Expr) -> Constraint<'db> {
+    fn build_constraint(&mut self, constraint_node: &ast::Expr) -> Constraint<'db> {
         let expression = self.add_standalone_expression(constraint_node);
         Constraint {
             node: ConstraintNode::Expression(expression),
@@ -408,11 +407,11 @@ impl<'db> SemanticIndexBuilder<'db> {
         let guard = guard.map(|guard| self.add_standalone_expression(guard));
 
         let kind = match pattern {
-            Pattern::MatchValue(pattern) => {
+            ast::Pattern::MatchValue(pattern) => {
                 let value = self.add_standalone_expression(&pattern.value);
                 PatternConstraintKind::Value(value, guard)
             }
-            Pattern::MatchSingleton(singleton) => {
+            ast::Pattern::MatchSingleton(singleton) => {
                 PatternConstraintKind::Singleton(singleton.value, guard)
             }
             _ => PatternConstraintKind::Unsupported,
@@ -1492,8 +1491,8 @@ where
                     if index < values.len() - 1 {
                         let constraint = self.build_constraint(value);
                         let (constraint, constraint_id) = match op {
-                            BoolOp::And => (constraint, self.add_constraint(constraint)),
-                            BoolOp::Or => self.add_negated_constraint(constraint),
+                            ast::BoolOp::And => (constraint, self.add_constraint(constraint)),
+                            ast::BoolOp::Or => self.add_negated_constraint(constraint),
                         };
                         let visibility_constraint = self
                             .add_visibility_constraint(VisibilityConstraint::VisibleIf(constraint));

crates/red_knot_python_semantic/src/semantic_index/symbol.rs

@@ -463,10 +463,7 @@ impl NodeWithScopeKind {
     }
 
     pub fn expect_function(&self) -> &ast::StmtFunctionDef {
-        match self {
-            Self::Function(function) => function.node(),
-            _ => panic!("expected function"),
-        }
+        self.as_function().expect("expected function")
     }
 
     pub fn expect_type_alias(&self) -> &ast::StmtTypeAlias {
@@ -475,6 +472,13 @@ impl NodeWithScopeKind {
             _ => panic!("expected type alias"),
         }
     }
+
+    pub const fn as_function(&self) -> Option<&ast::StmtFunctionDef> {
+        match self {
+            Self::Function(function) => Some(function.node()),
+            _ => None,
+        }
+    }
 }
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]

crates/red_knot_python_semantic/src/types.rs

@@ -3086,13 +3086,16 @@ pub enum KnownFunction {
     Len,
     /// `typing(_extensions).final`
     Final,
+
+    /// [`typing(_extensions).no_type_check`](https://typing.readthedocs.io/en/latest/spec/directives.html#no-type-check)
+    NoTypeCheck,
 }
 
 impl KnownFunction {
     pub fn constraint_function(self) -> Option<KnownConstraintFunction> {
         match self {
             Self::ConstraintFunction(f) => Some(f),
-            Self::RevealType | Self::Len | Self::Final => None,
+            Self::RevealType | Self::Len | Self::Final | Self::NoTypeCheck => None,
         }
     }
 
@@ -3111,6 +3114,9 @@ impl KnownFunction {
             ),
             "len" if definition.is_builtin_definition(db) => Some(KnownFunction::Len),
             "final" if definition.is_typing_definition(db) => Some(KnownFunction::Final),
+            "no_type_check" if definition.is_typing_definition(db) => {
+                Some(KnownFunction::NoTypeCheck)
+            }
             _ => None,
         }
     }

crates/red_knot_python_semantic/src/types/context.rs

@@ -8,14 +8,16 @@ use ruff_db::{
 use ruff_python_ast::AnyNodeRef;
 use ruff_text_size::Ranged;
 
+use super::{binding_ty, KnownFunction, TypeCheckDiagnostic, TypeCheckDiagnostics};
+
+use crate::semantic_index::semantic_index;
+use crate::semantic_index::symbol::ScopeId;
 use crate::{
     lint::{LintId, LintMetadata},
     suppression::suppressions,
     Db,
 };
 
-use super::{TypeCheckDiagnostic, TypeCheckDiagnostics};
-
 /// Context for inferring the types of a single file.
 ///
 /// One context exists for at least for every inferred region but it's
@@ -30,17 +32,21 @@ use super::{TypeCheckDiagnostic, TypeCheckDiagnostics};
 /// on the current [`TypeInference`](super::infer::TypeInference) result.
 pub(crate) struct InferContext<'db> {
     db: &'db dyn Db,
+    scope: ScopeId<'db>,
     file: File,
     diagnostics: std::cell::RefCell<TypeCheckDiagnostics>,
+    no_type_check: InNoTypeCheck,
     bomb: DebugDropBomb,
 }
 
 impl<'db> InferContext<'db> {
-    pub(crate) fn new(db: &'db dyn Db, file: File) -> Self {
+    pub(crate) fn new(db: &'db dyn Db, scope: ScopeId<'db>) -> Self {
         Self {
             db,
-            file,
+            scope,
+            file: scope.file(db),
             diagnostics: std::cell::RefCell::new(TypeCheckDiagnostics::default()),
+            no_type_check: InNoTypeCheck::default(),
             bomb: DebugDropBomb::new("`InferContext` needs to be explicitly consumed by calling `::finish` to prevent accidental loss of diagnostics."),
         }
     }
@@ -68,11 +74,19 @@ impl<'db> InferContext<'db> {
         node: AnyNodeRef,
         message: fmt::Arguments,
     ) {
+        if !self.db.is_file_open(self.file) {
+            return;
+        }
+
         // Skip over diagnostics if the rule is disabled.
         let Some(severity) = self.db.rule_selection().severity(LintId::of(lint)) else {
             return;
         };
 
+        if self.is_in_no_type_check() {
+            return;
+        }
+
         let suppressions = suppressions(self.db, self.file);
 
         if let Some(suppression) = suppressions.find_suppression(node.range(), LintId::of(lint)) {
@@ -112,6 +126,42 @@ impl<'db> InferContext<'db> {
         });
     }
 
+    pub(super) fn set_in_no_type_check(&mut self, no_type_check: InNoTypeCheck) {
+        self.no_type_check = no_type_check;
+    }
+
+    fn is_in_no_type_check(&self) -> bool {
+        match self.no_type_check {
+            InNoTypeCheck::Possibly => {
+                // Accessing the semantic index here is fine because
+                // the index belongs to the same file as for which we emit the diagnostic.
+                let index = semantic_index(self.db, self.file);
+
+                let scope_id = self.scope.file_scope_id(self.db);
+
+                // Inspect all ancestor function scopes by walking bottom up and infer the function's type.
+                let mut function_scope_tys = index
+                    .ancestor_scopes(scope_id)
+                    .filter_map(|(_, scope)| scope.node().as_function())
+                    .filter_map(|function| {
+                        binding_ty(self.db, index.definition(function)).into_function_literal()
+                    });
+
+                // Iterate over all functions and test if any is decorated with `@no_type_check`.
+                function_scope_tys.any(|function_ty| {
+                    function_ty
+                        .decorators(self.db)
+                        .iter()
+                        .filter_map(|decorator| decorator.into_function_literal())
+                        .any(|decorator_ty| {
+                            decorator_ty.is_known(self.db, KnownFunction::NoTypeCheck)
+                        })
+                })
+            }
+            InNoTypeCheck::Yes => true,
+        }
+    }
+
     #[must_use]
     pub(crate) fn finish(mut self) -> TypeCheckDiagnostics {
         self.bomb.defuse();
@@ -131,6 +181,17 @@ impl fmt::Debug for InferContext<'_> {
     }
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
+pub(crate) enum InNoTypeCheck {
+    /// The inference might be in a `no_type_check` block but only if any
+    /// ancestor function is decorated with `@no_type_check`.
+    #[default]
+    Possibly,
+
+    /// The inference is known to be in an `@no_type_check` decorated function.
+    Yes,
+}
+
 pub(crate) trait WithDiagnostics {
     fn diagnostics(&self) -> &TypeCheckDiagnostics;
 }