Rewrite type annotations on Python 3.7 when __future__ enabled (#953)

resources/test/fixtures/future_annotations.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import List, Optional
 
 from models import (
     Fruit,
@@ -28,3 +29,12 @@ class Foo:
     @classmethod
     def d(cls) -> Fruit:
         return cls(x=0, y=0)
+
+
+def f(x: int) -> List[int]:
+    y = List[int]()
+    y.append(x)
+    return y
+
+
+x: Optional[int] = None

src/check_ast.rs

@@ -78,6 +78,7 @@ pub struct Checker<'a> {
     in_f_string: Option<Range>,
     in_annotation: bool,
     in_deferred_string_annotation: bool,
+    in_deferred_annotation: bool,
     in_literal: bool,
     in_subscript: bool,
     in_withitem: bool,
@@ -124,6 +125,7 @@ impl<'a> Checker<'a> {
             in_f_string: None,
             in_annotation: false,
             in_deferred_string_annotation: false,
+            in_deferred_annotation: false,
             in_literal: false,
             in_subscript: false,
             in_withitem: false,
@@ -1191,10 +1193,13 @@ where
         // Pre-visit.
         match &expr.node {
             ExprKind::Subscript { value, slice, .. } => {
-                // Ex) typing.List[...]
+                // Ex) Optional[...]
                 if !self.in_deferred_string_annotation
                     && self.settings.enabled.contains(&CheckCode::U007)
-                    && self.settings.target_version >= PythonVersion::Py310
+                    && (self.settings.target_version >= PythonVersion::Py310
+                        || (self.settings.target_version >= PythonVersion::Py37
+                            && self.annotations_future_enabled
+                            && self.in_deferred_annotation))
                 {
                     pyupgrade::plugins::use_pep604_annotation(self, expr, value, slice);
                 }
@@ -1233,7 +1238,10 @@ where
                         // Ex) List[...]
                         if !self.in_deferred_string_annotation
                             && self.settings.enabled.contains(&CheckCode::U006)
-                            && self.settings.target_version >= PythonVersion::Py39
+                            && (self.settings.target_version >= PythonVersion::Py39
+                                || (self.settings.target_version >= PythonVersion::Py37
+                                    && self.annotations_future_enabled
+                                    && self.in_deferred_annotation))
                             && typing::is_pep585_builtin(
                                 expr,
                                 &self.from_imports,
@@ -1268,8 +1276,12 @@ where
             }
             ExprKind::Attribute { attr, .. } => {
                 // Ex) typing.List[...]
-                if self.settings.enabled.contains(&CheckCode::U006)
+                if !self.in_deferred_string_annotation
-                    && self.settings.target_version >= PythonVersion::Py39
+                    && self.settings.enabled.contains(&CheckCode::U006)
+                    && (self.settings.target_version >= PythonVersion::Py39
+                        || (self.settings.target_version >= PythonVersion::Py37
+                            && self.annotations_future_enabled
+                            && self.in_deferred_annotation))
                     && typing::is_pep585_builtin(expr, &self.from_imports, &self.import_aliases)
                 {
                     pyupgrade::plugins::use_pep585_annotation(self, expr, attr);
@@ -2743,7 +2755,9 @@ impl<'a> Checker<'a> {
         while let Some((expr, scopes, parents)) = self.deferred_annotations.pop() {
             self.scope_stack = scopes;
             self.parent_stack = parents;
+            self.in_deferred_annotation = true;
             self.visit_expr(expr);
+            self.in_deferred_annotation = false;
         }
     }
 
@@ -2769,9 +2783,9 @@ impl<'a> Checker<'a> {
             }
         }
         for (expr, (scopes, parents)) in allocator.iter().zip(stacks) {
-            self.in_deferred_string_annotation = true;
             self.scope_stack = scopes;
             self.parent_stack = parents;
+            self.in_deferred_string_annotation = true;
             self.visit_expr(expr);
             self.in_deferred_string_annotation = false;
         }

src/linter.rs

@@ -399,6 +399,7 @@ mod tests {
     use crate::checks::CheckCode;
     use crate::linter::test_path;
     use crate::settings;
+    use crate::settings::types::PythonVersion;
 
     #[test_case(CheckCode::A001, Path::new("A001.py"); "A001")]
     #[test_case(CheckCode::A002, Path::new("A002.py"); "A002")]
@@ -695,4 +696,64 @@ mod tests {
         insta::assert_yaml_snapshot!(checks);
         Ok(())
     }
+
+    #[test]
+    fn future_annotations_pep_585_p37() -> Result<()> {
+        let mut checks = test_path(
+            Path::new("./resources/test/fixtures/future_annotations.py"),
+            &settings::Settings {
+                target_version: PythonVersion::Py37,
+                ..settings::Settings::for_rule(CheckCode::U006)
+            },
+            true,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
+    #[test]
+    fn future_annotations_pep_585_py310() -> Result<()> {
+        let mut checks = test_path(
+            Path::new("./resources/test/fixtures/future_annotations.py"),
+            &settings::Settings {
+                target_version: PythonVersion::Py310,
+                ..settings::Settings::for_rule(CheckCode::U006)
+            },
+            true,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
+    #[test]
+    fn future_annotations_pep_604_p37() -> Result<()> {
+        let mut checks = test_path(
+            Path::new("./resources/test/fixtures/future_annotations.py"),
+            &settings::Settings {
+                target_version: PythonVersion::Py37,
+                ..settings::Settings::for_rule(CheckCode::U007)
+            },
+            true,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
+    #[test]
+    fn future_annotations_pep_604_py310() -> Result<()> {
+        let mut checks = test_path(
+            Path::new("./resources/test/fixtures/future_annotations.py"),
+            &settings::Settings {
+                target_version: PythonVersion::Py310,
+                ..settings::Settings::for_rule(CheckCode::U007)
+            },
+            true,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
 }

src/snapshots/ruff__linter__tests__future_annotations.snap

@@ -7,27 +7,27 @@ expression: checks
       - - models.Nut
       - false
   location:
-    row: 5
+    row: 6
     column: 0
   end_location:
-    row: 8
+    row: 9
     column: 1
   fix:
     patch:
       content: "from models import (\n    Fruit,\n)"
       location:
-        row: 5
+        row: 6
         column: 0
       end_location:
-        row: 8
+        row: 9
         column: 1
 - kind:
     UndefinedName: Bar
   location:
-    row: 25
+    row: 26
     column: 18
   end_location:
-    row: 25
+    row: 26
     column: 21
   fix: ~
 

src/snapshots/ruff__linter__tests__future_annotations_pep_585_p37.snap

@@ -0,0 +1,22 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    UsePEP585Annotation: List
+  location:
+    row: 34
+    column: 17
+  end_location:
+    row: 34
+    column: 21
+  fix:
+    patch:
+      content: list
+      location:
+        row: 34
+        column: 17
+      end_location:
+        row: 34
+        column: 21
+

src/snapshots/ruff__linter__tests__future_annotations_pep_585_py310.snap

@@ -0,0 +1,39 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    UsePEP585Annotation: List
+  location:
+    row: 34
+    column: 17
+  end_location:
+    row: 34
+    column: 21
+  fix:
+    patch:
+      content: list
+      location:
+        row: 34
+        column: 17
+      end_location:
+        row: 34
+        column: 21
+- kind:
+    UsePEP585Annotation: List
+  location:
+    row: 35
+    column: 8
+  end_location:
+    row: 35
+    column: 12
+  fix:
+    patch:
+      content: list
+      location:
+        row: 35
+        column: 8
+      end_location:
+        row: 35
+        column: 12
+

src/snapshots/ruff__linter__tests__future_annotations_pep_604_p37.snap

@@ -0,0 +1,21 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind: UsePEP604Annotation
+  location:
+    row: 40
+    column: 3
+  end_location:
+    row: 40
+    column: 16
+  fix:
+    patch:
+      content: int | None
+      location:
+        row: 40
+        column: 3
+      end_location:
+        row: 40
+        column: 16
+

src/snapshots/ruff__linter__tests__future_annotations_pep_604_py310.snap

@@ -0,0 +1,21 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind: UsePEP604Annotation
+  location:
+    row: 40
+    column: 3
+  end_location:
+    row: 40
+    column: 16
+  fix:
+    patch:
+      content: int | None
+      location:
+        row: 40
+        column: 3
+      end_location:
+        row: 40
+        column: 16
+