[flake8-type-checking] Improve flexibility of runtime-evaluated-decorators (#15204)

Co-authored-by: Micha Reiser <micha@reiser.io>

Cargo.lock

@@ -2969,6 +2969,7 @@ dependencies = [
  "rustc-hash 2.1.0",
  "schemars",
  "serde",
+ "smallvec",
 ]
 
 [[package]]

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/module/app.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import fastapi
+from fastapi import FastAPI as Api
+
+if TYPE_CHECKING:
+    import datetime  # TC004
+    from array import array  # TC004
+
+app = fastapi.FastAPI("First application")
+
+class AppContainer:
+    app = Api("Second application")
+
+app_container = AppContainer()
+
+@app.put("/datetime")
+def set_datetime(value: datetime.datetime):
+    pass
+
+@app_container.app.get("/array")
+def get_array() -> array:
+    pass

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/module/routes.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import pathlib  # OK
+from datetime import date  # OK
+
+from module.app import app, app_container
+
+@app.get("/path")
+def get_path() -> pathlib.Path:
+    pass
+
+@app_container.app.put("/date")
+def set_date(d: date):
+    pass

crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs

@@ -98,12 +98,31 @@ fn runtime_required_decorators(
     }
 
     decorator_list.iter().any(|decorator| {
+        let expression = map_callable(&decorator.expression);
         semantic
-            .resolve_qualified_name(map_callable(&decorator.expression))
+            // First try to resolve the qualified name normally for cases like:
+            // ```python
+            // from mymodule import app
+            //
+            // @app.get(...)
+            // def test(): ...
+            //  ```
+            .resolve_qualified_name(expression)
+            // If we can't resolve the name, then try resolving the assignment
+            // in order to support cases like:
+            // ```python
+            // from fastapi import FastAPI
+            //
+            // app = FastAPI()
+            //
+            // @app.get(...)
+            // def test(): ...
+            // ```
+            .or_else(|| analyze::typing::resolve_assignment(expression, semantic))
             .is_some_and(|qualified_name| {
                 decorators
                     .iter()
-                    .any(|base_class| QualifiedName::from_dotted_name(base_class) == qualified_name)
+                    .any(|decorator| QualifiedName::from_dotted_name(decorator) == qualified_name)
             })
     })
 }

crates/ruff_linter/src/rules/flake8_type_checking/mod.rs

@@ -259,6 +259,33 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("module/app.py"))]
+    #[test_case(Rule::TypingOnlyStandardLibraryImport, Path::new("module/routes.py"))]
+    fn decorator_same_file(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!("{}_{}", rule_code.as_ref(), path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("flake8_type_checking").join(path).as_path(),
+            &settings::LinterSettings {
+                flake8_type_checking: super::settings::Settings {
+                    runtime_required_decorators: vec![
+                        "fastapi.FastAPI.get".to_string(),
+                        "fastapi.FastAPI.put".to_string(),
+                        "module.app.AppContainer.app.get".to_string(),
+                        "module.app.AppContainer.app.put".to_string(),
+                        "module.app.app.get".to_string(),
+                        "module.app.app.put".to_string(),
+                        "module.app.app_container.app.get".to_string(),
+                        "module.app.app_container.app.put".to_string(),
+                    ],
+                    ..Default::default()
+                },
+                ..settings::LinterSettings::for_rule(rule_code)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
+
     #[test_case(
         r"
         from __future__ import annotations

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_module__app.py.snap

@@ -0,0 +1,47 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+app.py:9:12: TC004 [*] Move import `datetime` out of type-checking block. Import is used for more than type hinting.
+   |
+ 8 | if TYPE_CHECKING:
+ 9 |     import datetime  # TC004
+   |            ^^^^^^^^ TC004
+10 |     from array import array  # TC004
+   |
+   = help: Move out of type-checking block
+
+ℹ Unsafe fix
+4  4  | 
+5  5  | import fastapi
+6  6  | from fastapi import FastAPI as Api
+   7  |+import datetime
+7  8  | 
+8  9  | if TYPE_CHECKING:
+9     |-    import datetime  # TC004
+10 10 |     from array import array  # TC004
+11 11 | 
+12 12 | app = fastapi.FastAPI("First application")
+
+app.py:10:23: TC004 [*] Move import `array.array` out of type-checking block. Import is used for more than type hinting.
+   |
+ 8 | if TYPE_CHECKING:
+ 9 |     import datetime  # TC004
+10 |     from array import array  # TC004
+   |                       ^^^^^ TC004
+11 | 
+12 | app = fastapi.FastAPI("First application")
+   |
+   = help: Move out of type-checking block
+
+ℹ Unsafe fix
+4  4  | 
+5  5  | import fastapi
+6  6  | from fastapi import FastAPI as Api
+   7  |+from array import array
+7  8  | 
+8  9  | if TYPE_CHECKING:
+9  10 |     import datetime  # TC004
+10    |-    from array import array  # TC004
+11 11 | 
+12 12 | app = fastapi.FastAPI("First application")
+13 13 |

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-standard-library-import_module__routes.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+

crates/ruff_python_ast/src/name.rs

@@ -277,6 +277,14 @@ impl<'a> QualifiedName<'a> {
         inner.push(member);
         Self(inner)
     }
+
+    /// Extends the qualified name using the given members.
+    #[must_use]
+    pub fn extend_members<T: IntoIterator<Item = &'a str>>(self, members: T) -> Self {
+        let mut inner = self.0;
+        inner.extend(members);
+        Self(inner)
+    }
 }
 
 impl Display for QualifiedName<'_> {

crates/ruff_python_semantic/Cargo.toml

@@ -24,6 +24,7 @@ is-macro = { workspace = true }
 rustc-hash = { workspace = true }
 schemars = { workspace = true, optional = true }
 serde = { workspace = true, optional = true }
+smallvec = { workspace = true }
 
 [dev-dependencies]
 ruff_python_parser = { workspace = true }

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -11,6 +11,7 @@ use ruff_python_stdlib::typing::{
     is_typed_dict, is_typed_dict_member,
 };
 use ruff_text_size::Ranged;
+use smallvec::{smallvec, SmallVec};
 
 use crate::analyze::type_inference::{NumberLike, PythonType, ResolvedPythonType};
 use crate::model::SemanticModel;
@@ -983,23 +984,43 @@ fn find_parameter<'a>(
 /// ```
 ///
 /// This function will return `["asyncio", "get_running_loop"]` for the `loop` binding.
+///
+/// This function will also automatically expand attribute accesses, so given:
+/// ```python
+/// from module import AppContainer
+///
+/// container = AppContainer()
+/// container.app.get(...)
+/// ```
+///
+/// This function will return `["module", "AppContainer", "app", "get"]` for the
+/// attribute access `container.app.get`.
 pub fn resolve_assignment<'a>(
     expr: &'a Expr,
     semantic: &'a SemanticModel<'a>,
 ) -> Option<QualifiedName<'a>> {
-    let name = expr.as_name_expr()?;
+    // Resolve any attribute chain.
+    let mut head_expr = expr;
+    let mut reversed_tail: SmallVec<[_; 4]> = smallvec![];
+    while let Expr::Attribute(ast::ExprAttribute { value, attr, .. }) = head_expr {
+        head_expr = value;
+        reversed_tail.push(attr.as_str());
+    }
+
+    // Resolve the left-most name, e.g. `foo` in `foo.bar.baz` to a qualified name,
+    // then append the attributes.
+    let name = head_expr.as_name_expr()?;
     let binding_id = semantic.resolve_name(name)?;
     let statement = semantic.binding(binding_id).statement(semantic)?;
     match statement {
-        Stmt::Assign(ast::StmtAssign { value, .. }) => {
+        Stmt::Assign(ast::StmtAssign { value, .. })
-            let ast::ExprCall { func, .. } = value.as_call_expr()?;
+        | Stmt::AnnAssign(ast::StmtAnnAssign {
-            semantic.resolve_qualified_name(func)
-        }
-        Stmt::AnnAssign(ast::StmtAnnAssign {
             value: Some(value), ..
         }) => {
             let ast::ExprCall { func, .. } = value.as_call_expr()?;
-            semantic.resolve_qualified_name(func)
+
+            let qualified_name = semantic.resolve_qualified_name(func)?;
+            Some(qualified_name.extend_members(reversed_tail.into_iter().rev()))
         }
         _ => None,
     }

crates/ruff_workspace/src/options.rs

@@ -1819,6 +1819,21 @@ pub struct Flake8TypeCheckingOptions {
     ///
     /// Common examples include Pydantic's `@pydantic.validate_call` decorator
     /// (for functions) and attrs' `@attrs.define` decorator (for classes).
+    ///
+    /// This also supports framework decorators like FastAPI's `fastapi.FastAPI.get`
+    /// which will work across assignments in the same module.
+    ///
+    /// For example:
+    /// ```python
+    /// import fastapi
+    ///
+    /// app = FastAPI("app")
+    ///
+    /// @app.get("/home")
+    /// def home() -> str: ...
+    /// ```
+    ///
+    /// Here `app.get` will correctly be identified as `fastapi.FastAPI.get`.
     #[option(
         default = "[]",
         value_type = "list[str]",

ruff.schema.json

@@ -1387,7 +1387,7 @@
           }
         },
         "runtime-evaluated-decorators": {
-          "description": "Exempt classes and functions decorated with any of the enumerated decorators from being moved into type-checking blocks.\n\nCommon examples include Pydantic's `@pydantic.validate_call` decorator (for functions) and attrs' `@attrs.define` decorator (for classes).",
+          "description": "Exempt classes and functions decorated with any of the enumerated decorators from being moved into type-checking blocks.\n\nCommon examples include Pydantic's `@pydantic.validate_call` decorator (for functions) and attrs' `@attrs.define` decorator (for classes).\n\nThis also supports framework decorators like FastAPI's `fastapi.FastAPI.get` which will work across assignments in the same module.\n\nFor example: ```python import fastapi\n\napp = FastAPI(\"app\")\n\n@app.get(\"/home\") def home() -> str: ... ```\n\nHere `app.get` will correctly be identified as `fastapi.FastAPI.get`.",
           "type": [
             "array",
             "null"