[`flake8-annotations`] Fix return type annotations to handle shadowed builtin symbols (`ANN201`, `ANN202`, `ANN204`, `ANN205`, `ANN206`) (#20612)

## Summary

Fixes #20610

---------

Co-authored-by: Brent Westbrook <brentrwestbrook@gmail.com>

crates/ruff_linter/resources/test/fixtures/flake8_annotations/shadowed_builtins.py

@@ -0,0 +1,21 @@
+from collections import UserString as str
+from typing import override
+
+def foo():
+    return "!"
+
+def _foo():
+    return "!"
+
+class C:
+    @override
+    def __str__(self):
+        return "!"
+    
+    @staticmethod
+    def foo():
+        return "!"
+    
+    @classmethod
+    def bar(cls):
+        return "!"

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

@@ -142,23 +142,25 @@ impl AutoPythonType {
                 });
                 Some((expr, vec![no_return_edit]))
             }
-            AutoPythonType::Atom(python_type) => {
+            AutoPythonType::Atom(python_type) => type_expr(python_type, checker, at),
-                let expr = type_expr(python_type)?;
-                Some((expr, vec![]))
-            }
             AutoPythonType::Union(python_types) => {
                 if target_version >= PythonVersion::PY310 {
                     // Aggregate all the individual types (e.g., `int`, `float`).
+                    let mut all_edits = Vec::new();
                     let names = python_types
                         .iter()
                         .sorted_unstable()
-                        .map(|python_type| type_expr(*python_type))
+                        .map(|python_type| {
+                            let (expr, mut edits) = type_expr(*python_type, checker, at)?;
+                            all_edits.append(&mut edits);
+                            Some(expr)
+                        })
                         .collect::<Option<Vec<_>>>()?;
 
                     // Wrap in a bitwise union (e.g., `int | float`).
                     let expr = pep_604_union(&names);
 
-                    Some((expr, vec![]))
+                    Some((expr, all_edits))
                 } else {
                     let python_types = python_types
                         .into_iter()
@@ -167,7 +169,7 @@ impl AutoPythonType {
 
                     match python_types.as_slice() {
                         [python_type, PythonType::None] | [PythonType::None, python_type] => {
-                            let element = type_expr(*python_type)?;
+                            let (element, mut edits) = type_expr(*python_type, checker, at)?;
 
                             // Ex) `Optional[int]`
                             let (optional_edit, binding) = checker
@@ -175,12 +177,18 @@ impl AutoPythonType {
                                 .import(at)
                                 .ok()?;
                             let expr = typing_optional(element, Name::from(binding));
-                            Some((expr, vec![optional_edit]))
+                            edits.push(optional_edit);
+                            Some((expr, edits))
                         }
                         _ => {
+                            let mut all_edits = Vec::new();
                             let elements = python_types
                                 .into_iter()
-                                .map(type_expr)
+                                .map(|python_type| {
+                                    let (expr, mut edits) = type_expr(python_type, checker, at)?;
+                                    all_edits.append(&mut edits);
+                                    Some(expr)
+                                })
                                 .collect::<Option<Vec<_>>>()?;
 
                             // Ex) `Union[int, str]`
@@ -189,7 +197,8 @@ impl AutoPythonType {
                                 .import(at)
                                 .ok()?;
                             let expr = typing_union(&elements, Name::from(binding));
-                            Some((expr, vec![union_edit]))
+                            all_edits.push(union_edit);
+                            Some((expr, all_edits))
                         }
                     }
                 }
@@ -199,26 +208,44 @@ impl AutoPythonType {
 }
 
 /// Given a [`PythonType`], return an [`Expr`] that resolves to that type.
-fn type_expr(python_type: PythonType) -> Option<Expr> {
+///
-    fn name(name: &str) -> Expr {
+/// If the [`Expr`] relies on importing any external symbols, those imports will be returned as
-        Expr::Name(ast::ExprName {
+/// additional edits.
-            id: name.into(),
+pub(crate) fn type_expr(
+    python_type: PythonType,
+    checker: &Checker,
+    at: TextSize,
+) -> Option<(Expr, Vec<Edit>)> {
+    fn name(name: &str, checker: &Checker, at: TextSize) -> Option<(Expr, Vec<Edit>)> {
+        let (edit, binding) = checker
+            .importer()
+            .get_or_import_builtin_symbol(name, at, checker.semantic())
+            .ok()?;
+        let expr = Expr::Name(ast::ExprName {
+            id: binding.into(),
             range: TextRange::default(),
             node_index: ruff_python_ast::AtomicNodeIndex::NONE,
             ctx: ExprContext::Load,
-        })
+        });
+        Some((expr, edit.map_or_else(Vec::new, |edit| vec![edit])))
     }
 
     match python_type {
-        PythonType::String => Some(name("str")),
+        PythonType::String => name("str", checker, at),
-        PythonType::Bytes => Some(name("bytes")),
+        PythonType::Bytes => name("bytes", checker, at),
-        PythonType::Number(number) => match number {
+        PythonType::Number(number) => {
-            NumberLike::Integer => Some(name("int")),
+            let symbol = match number {
-            NumberLike::Float => Some(name("float")),
+                NumberLike::Integer => "int",
-            NumberLike::Complex => Some(name("complex")),
+                NumberLike::Float => "float",
-            NumberLike::Bool => Some(name("bool")),
+                NumberLike::Complex => "complex",
-        },
+                NumberLike::Bool => "bool",
-        PythonType::None => Some(name("None")),
+            };
+            name(symbol, checker, at)
+        }
+        PythonType::None => {
+            let expr = Expr::NoneLiteral(ast::ExprNoneLiteral::default());
+            Some((expr, vec![]))
+        }
         PythonType::Ellipsis => None,
         PythonType::Dict => None,
         PythonType::List => None,

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

@@ -229,4 +229,20 @@ mod tests {
         assert_diagnostics!(diagnostics);
         Ok(())
     }
+
+    #[test]
+    fn shadowed_builtins() -> Result<()> {
+        let diagnostics = test_path(
+            Path::new("flake8_annotations/shadowed_builtins.py"),
+            &LinterSettings::for_rules(vec![
+                Rule::MissingReturnTypeUndocumentedPublicFunction,
+                Rule::MissingReturnTypePrivateFunction,
+                Rule::MissingReturnTypeSpecialMethod,
+                Rule::MissingReturnTypeStaticMethod,
+                Rule::MissingReturnTypeClassMethod,
+            ]),
+        )?;
+        assert_diagnostics!(diagnostics);
+        Ok(())
+    }
 }

crates/ruff_linter/src/rules/flake8_annotations/rules/definition.rs

@@ -4,13 +4,14 @@ use ruff_python_ast::identifier::Identifier;
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::{self as ast, Expr, Stmt};
 use ruff_python_semantic::Definition;
+use ruff_python_semantic::analyze::type_inference::{NumberLike, PythonType};
 use ruff_python_semantic::analyze::visibility;
 use ruff_python_stdlib::typing::simple_magic_return_type;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::{Checker, DiagnosticGuard};
 use crate::registry::Rule;
-use crate::rules::flake8_annotations::helpers::auto_return_type;
+use crate::rules::flake8_annotations::helpers::{auto_return_type, type_expr};
 use crate::rules::ruff::typing::type_hint_resolves_to_any;
 use crate::{Edit, Fix, FixAvailability, Violation};
 
@@ -825,11 +826,31 @@ pub(crate) fn definition(
                     },
                     function.identifier(),
                 );
-                if let Some(return_type) = return_type {
+                if let Some(return_type_str) = return_type {
-                    diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
+                    // Convert the simple return type to a proper expression that handles shadowed builtins
-                        format!(" -> {return_type}"),
+                    let python_type = match return_type_str {
-                        function.parameters.end(),
+                        "str" => PythonType::String,
-                    )));
+                        "bytes" => PythonType::Bytes,
+                        "int" => PythonType::Number(NumberLike::Integer),
+                        "float" => PythonType::Number(NumberLike::Float),
+                        "complex" => PythonType::Number(NumberLike::Complex),
+                        "bool" => PythonType::Number(NumberLike::Bool),
+                        "None" => PythonType::None,
+                        _ => return, // Unknown type, skip
+                    };
+
+                    if let Some((expr, edits)) =
+                        type_expr(python_type, checker, function.parameters.start())
+                    {
+                        let return_type_expr = checker.generator().expr(&expr);
+                        diagnostic.set_fix(Fix::unsafe_edits(
+                            Edit::insertion(
+                                format!(" -> {return_type_expr}"),
+                                function.parameters.end(),
+                            ),
+                            edits,
+                        ));
+                    }
                 }
                 diagnostics.push(diagnostic);
             }

crates/ruff_linter/src/rules/flake8_annotations/snapshots/ruff_linter__rules__flake8_annotations__tests__shadowed_builtins.snap

@@ -0,0 +1,124 @@
+---
+source: crates/ruff_linter/src/rules/flake8_annotations/mod.rs
+---
+ANN201 [*] Missing return type annotation for public function `foo`
+ --> shadowed_builtins.py:4:5
+  |
+2 | from typing import override
+3 |
+4 | def foo():
+  |     ^^^
+5 |     return "!"
+  |
+help: Add return type annotation: `builtins.str`
+1 | from collections import UserString as str
+2 | from typing import override
+3 + import builtins
+4 | 
+  - def foo():
+5 + def foo() -> builtins.str:
+6 |     return "!"
+7 | 
+8 | def _foo():
+note: This is an unsafe fix and may change runtime behavior
+
+ANN202 [*] Missing return type annotation for private function `_foo`
+ --> shadowed_builtins.py:7:5
+  |
+5 |     return "!"
+6 |
+7 | def _foo():
+  |     ^^^^
+8 |     return "!"
+  |
+help: Add return type annotation: `builtins.str`
+1  | from collections import UserString as str
+2  | from typing import override
+3  + import builtins
+4  | 
+5  | def foo():
+6  |     return "!"
+7  | 
+   - def _foo():
+8  + def _foo() -> builtins.str:
+9  |     return "!"
+10 | 
+11 | class C:
+note: This is an unsafe fix and may change runtime behavior
+
+ANN204 [*] Missing return type annotation for special method `__str__`
+  --> shadowed_builtins.py:12:9
+   |
+10 | class C:
+11 |     @override
+12 |     def __str__(self):
+   |         ^^^^^^^
+13 |         return "!"
+   |
+help: Add return type annotation: `str`
+1  | from collections import UserString as str
+2  | from typing import override
+3  + import builtins
+4  | 
+5  | def foo():
+6  |     return "!"
+--------------------------------------------------------------------------------
+10 | 
+11 | class C:
+12 |     @override
+   -     def __str__(self):
+13 +     def __str__(self) -> builtins.str:
+14 |         return "!"
+15 |     
+16 |     @staticmethod
+note: This is an unsafe fix and may change runtime behavior
+
+ANN205 [*] Missing return type annotation for staticmethod `foo`
+  --> shadowed_builtins.py:16:9
+   |
+15 |     @staticmethod
+16 |     def foo():
+   |         ^^^
+17 |         return "!"
+   |
+help: Add return type annotation: `builtins.str`
+1  | from collections import UserString as str
+2  | from typing import override
+3  + import builtins
+4  | 
+5  | def foo():
+6  |     return "!"
+--------------------------------------------------------------------------------
+14 |         return "!"
+15 |     
+16 |     @staticmethod
+   -     def foo():
+17 +     def foo() -> builtins.str:
+18 |         return "!"
+19 |     
+20 |     @classmethod
+note: This is an unsafe fix and may change runtime behavior
+
+ANN206 [*] Missing return type annotation for classmethod `bar`
+  --> shadowed_builtins.py:20:9
+   |
+19 |     @classmethod
+20 |     def bar(cls):
+   |         ^^^
+21 |         return "!"
+   |
+help: Add return type annotation: `builtins.str`
+1  | from collections import UserString as str
+2  | from typing import override
+3  + import builtins
+4  | 
+5  | def foo():
+6  |     return "!"
+--------------------------------------------------------------------------------
+18 |         return "!"
+19 |     
+20 |     @classmethod
+   -     def bar(cls):
+21 +     def bar(cls) -> builtins.str:
+22 |         return "!"
+note: This is an unsafe fix and may change runtime behavior