Implement PEP 604 annotation rewrites (#369)

README.md

@@ -294,6 +294,8 @@ The 🛠 emoji indicates that a rule is automatically fixable by the `--fix` com
 | U003 | TypeOfPrimitive | Use `str` instead of `type(...)` |  | 🛠 |
 | U004 | UselessObjectInheritance | Class `...` inherits from object |  | 🛠 |
 | U005 | NoAssertEquals | `assertEquals` is deprecated, use `assertEqual` instead |  | 🛠 |
+| U006 | UsePEP585Annotation | Use `list` instead of `List` for type annotations |  | 🛠 |
+| U007 | UsePEP604Annotation | Use `X | Y` for type annotations |  | 🛠 |
 | M001 | UnusedNOQA | Unused `noqa` directive |  | 🛠 |
 
 ## Integrations

resources/test/fixtures/U007.py

@@ -0,0 +1,40 @@
+from typing import Optional
+
+
+def f(x: Optional[str]) -> None:
+    ...
+
+
+import typing
+
+
+def f(x: typing.Optional[str]) -> None:
+    ...
+
+
+from typing import Union
+
+
+def f(x: Union[str, int, Union[float, bytes]]) -> None:
+    ...
+
+
+import typing
+
+
+def f(x: typing.Union[str, int]) -> None:
+    ...
+
+
+from typing import Union
+
+
+def f(x: "Union[str, int, Union[float, bytes]]") -> None:
+    ...
+
+
+import typing
+
+
+def f(x: "typing.Union[str, int]") -> None:
+    ...

src/ast.rs

@@ -1,4 +1,5 @@
 pub mod checks;
+pub mod helpers;
 pub mod operations;
 pub mod relocate;
 pub mod types;

src/ast/helpers.rs

@@ -0,0 +1,9 @@
+use rustpython_ast::{Expr, ExprKind};
+
+pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
+    match &expr.node {
+        ExprKind::Attribute { attr, .. } => target == attr,
+        ExprKind::Name { id, .. } => target == id,
+        _ => false,
+    }
+}

src/check_ast.rs

@@ -5,12 +5,14 @@ use std::path::Path;
 use log::error;
 use once_cell::sync::Lazy;
 use regex::Regex;
+use rustpython_ast::Location;
 use rustpython_parser::ast::{
     Arg, Arguments, Constant, Excepthandler, ExcepthandlerKind, Expr, ExprContext, ExprKind,
     KeywordData, Operator, Stmt, StmtKind, Suite,
 };
 use rustpython_parser::parser;
 
+use crate::ast::helpers::match_name_or_attr;
 use crate::ast::operations::{extract_all_names, SourceCodeLocator};
 use crate::ast::relocate::relocate_expr;
 use crate::ast::types::{
@@ -100,14 +102,6 @@ impl<'a> Checker<'a> {
     }
 }
 
-fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
-    match &expr.node {
-        ExprKind::Attribute { attr, .. } => target == attr,
-        ExprKind::Name { id, .. } => target == id,
-        _ => false,
-    }
-}
-
 enum SubscriptKind {
     AnnotatedSubscript,
     PEP593AnnotatedSubscript,
@@ -709,7 +703,14 @@ where
 
         // Pre-visit.
         match &expr.node {
-            ExprKind::Subscript { value, .. } => {
+            ExprKind::Subscript { value, slice, .. } => {
+                // Ex) typing.List[...]
+                if self.settings.enabled.contains(&CheckCode::U007)
+                    && self.settings.target_version >= PythonVersion::Py39
+                {
+                    plugins::use_pep604_annotation(self, expr, value, slice);
+                }
+
                 if match_name_or_attr(value, "Literal") {
                     self.in_literal = true;
                 }
@@ -1605,14 +1606,37 @@ impl<'a> Checker<'a> {
     where
         'b: 'a,
     {
-        while let Some((location, expression)) = self.deferred_string_annotations.pop() {
+        while let Some((range, expression)) = self.deferred_string_annotations.pop() {
+            // HACK(charlie): We need to modify `range` such that it represents the range of the
+            // expression _within_ the string annotation (as opposed to the range of the string
+            // annotation itself). RustPython seems to return an off-by-one start column for every
+            // string value, so we check for double quotes (which are really triple quotes).
+            let contents = self.locator.slice_source_code_at(&range.location);
+            let range = if contents.starts_with("\"\"") || contents.starts_with("\'\'") {
+                Range {
+                    location: Location::new(range.location.row(), range.location.column() + 2),
+                    end_location: Location::new(
+                        range.end_location.row(),
+                        range.end_location.column() - 2,
+                    ),
+                }
+            } else {
+                Range {
+                    location: Location::new(range.location.row(), range.location.column()),
+                    end_location: Location::new(
+                        range.end_location.row(),
+                        range.end_location.column() - 1,
+                    ),
+                }
+            };
+
             if let Ok(mut expr) = parser::parse_expression(expression, "<filename>") {
-                relocate_expr(&mut expr, location);
+                relocate_expr(&mut expr, range);
                 allocator.push(expr);
             } else if self.settings.enabled.contains(&CheckCode::F722) {
                 self.checks.push(Check::new(
                     CheckKind::ForwardAnnotationSyntaxError(expression.to_string()),
-                    self.locate_check(location),
+                    self.locate_check(range),
                 ));
             }
         }

src/checks.rs

@@ -141,6 +141,7 @@ pub enum CheckCode {
     U004,
     U005,
     U006,
+    U007,
     // Meta
     M001,
 }
@@ -230,6 +231,7 @@ pub enum CheckKind {
     NoAssertEquals,
     UselessObjectInheritance(String),
     UsePEP585Annotation(String),
+    UsePEP604Annotation,
     // Meta
     UnusedNOQA(Option<String>),
 }
@@ -322,6 +324,7 @@ impl CheckCode {
             CheckCode::U004 => CheckKind::UselessObjectInheritance("...".to_string()),
             CheckCode::U005 => CheckKind::NoAssertEquals,
             CheckCode::U006 => CheckKind::UsePEP585Annotation("List".to_string()),
+            CheckCode::U007 => CheckKind::UsePEP604Annotation,
             // Meta
             CheckCode::M001 => CheckKind::UnusedNOQA(None),
         }
@@ -401,6 +404,7 @@ impl CheckKind {
             CheckKind::UselessMetaclassType => &CheckCode::U001,
             CheckKind::NoAssertEquals => &CheckCode::U005,
             CheckKind::UsePEP585Annotation(_) => &CheckCode::U006,
+            CheckKind::UsePEP604Annotation => &CheckCode::U007,
             CheckKind::UselessObjectInheritance(_) => &CheckCode::U004,
             // Meta
             CheckKind::UnusedNOQA(_) => &CheckCode::M001,
@@ -603,6 +607,7 @@ impl CheckKind {
                     name,
                 )
             }
+            CheckKind::UsePEP604Annotation => "Use `X | Y` for type annotations".to_string(),
             // Meta
             CheckKind::UnusedNOQA(code) => match code {
                 None => "Unused `noqa` directive".to_string(),
@@ -626,6 +631,7 @@ impl CheckKind {
                 | CheckKind::UselessMetaclassType
                 | CheckKind::UselessObjectInheritance(_)
                 | CheckKind::UsePEP585Annotation(_)
+                | CheckKind::UsePEP604Annotation
         )
     }
 }

src/code_gen.rs

@@ -547,7 +547,7 @@ impl SourceGenerator {
         Ok(())
     }
 
-    fn unparse_expr<U>(&mut self, ast: &Expr<U>, level: u8) -> fmt::Result {
+    pub fn unparse_expr<U>(&mut self, ast: &Expr<U>, level: u8) -> fmt::Result {
         macro_rules! opprec {
             ($opty:ident, $x:expr, $enu:path, $($var:ident($op:literal, $prec:ident)),*$(,)?) => {
                 match $x {

src/linter.rs

@@ -1001,4 +1001,16 @@ mod tests {
         insta::assert_yaml_snapshot!(checks);
         Ok(())
     }
+
+    #[test]
+    fn u007() -> Result<()> {
+        let mut checks = check_path(
+            Path::new("./resources/test/fixtures/U007.py"),
+            &settings::Settings::for_rule(CheckCode::U007),
+            &fixer::Mode::Generate,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
 }

src/plugins.rs

@@ -7,6 +7,7 @@ mod super_call_with_parameters;
 mod type_of_primitive;
 mod unnecessary_abspath;
 mod use_pep585_annotation;
+mod use_pep604_annotation;
 mod useless_metaclass_type;
 mod useless_object_inheritance;
 
@@ -19,5 +20,6 @@ pub use super_call_with_parameters::super_call_with_parameters;
 pub use type_of_primitive::type_of_primitive;
 pub use unnecessary_abspath::unnecessary_abspath;
 pub use use_pep585_annotation::use_pep585_annotation;
+pub use use_pep604_annotation::use_pep604_annotation;
 pub use useless_metaclass_type::useless_metaclass_type;
 pub use useless_object_inheritance::useless_object_inheritance;

src/plugins/use_pep604_annotation.rs

@@ -0,0 +1,77 @@
+use anyhow::{anyhow, Result};
+use rustpython_ast::{Expr, ExprKind};
+
+use crate::ast::helpers::match_name_or_attr;
+use crate::ast::types::Range;
+use crate::autofix::fixer;
+use crate::check_ast::Checker;
+use crate::checks::{Check, CheckKind, Fix};
+use crate::code_gen::SourceGenerator;
+
+pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, slice: &Expr) {
+    if match_name_or_attr(value, "Optional") {
+        let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr));
+        if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply) {
+            let mut generator = SourceGenerator::new();
+            if let Ok(()) = generator.unparse_expr(slice, 0) {
+                if let Ok(content) = generator.generate() {
+                    check.amend(Fix {
+                        content: format!("{} | None", content),
+                        location: expr.location,
+                        end_location: expr.end_location,
+                        applied: false,
+                    })
+                }
+            }
+        }
+        checker.add_check(check);
+    } else if match_name_or_attr(value, "Union") {
+        let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr));
+        if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply) {
+            match &slice.node {
+                ExprKind::Slice { .. } => {
+                    // Invalid type annotation.
+                }
+                ExprKind::Tuple { elts, .. } => {
+                    // Multiple arguments.
+                    let parts: Result<Vec<String>> = elts
+                        .iter()
+                        .map(|expr| {
+                            let mut generator = SourceGenerator::new();
+                            generator
+                                .unparse_expr(expr, 0)
+                                .map_err(|_| anyhow!("Failed to parse."))?;
+                            generator
+                                .generate()
+                                .map_err(|_| anyhow!("Failed to generate."))
+                        })
+                        .collect();
+                    if let Ok(parts) = parts {
+                        let content = parts.join(" | ");
+                        check.amend(Fix {
+                            content,
+                            location: expr.location,
+                            end_location: expr.end_location,
+                            applied: false,
+                        })
+                    }
+                }
+                _ => {
+                    // Single argument.
+                    let mut generator = SourceGenerator::new();
+                    if let Ok(()) = generator.unparse_expr(slice, 0) {
+                        if let Ok(content) = generator.generate() {
+                            check.amend(Fix {
+                                content,
+                                location: expr.location,
+                                end_location: expr.end_location,
+                                applied: false,
+                            });
+                        }
+                    }
+                }
+            }
+        }
+        checker.add_check(check);
+    }
+}

src/snapshots/ruff__linter__tests__f722.snap

@@ -9,6 +9,6 @@ expression: checks
     column: 13
   end_location:
     row: 9
-    column: 17
+    column: 16
   fix: ~
 

src/snapshots/ruff__linter__tests__f821.snap

@@ -45,7 +45,7 @@ expression: checks
     column: 5
   end_location:
     row: 58
-    column: 9
+    column: 8
   fix: ~
 - kind:
     UndefinedName: TOMATO
@@ -81,7 +81,7 @@ expression: checks
     column: 10
   end_location:
     row: 114
-    column: 24
+    column: 23
   fix: ~
 - kind:
     UndefinedName: foo
@@ -90,7 +90,7 @@ expression: checks
     column: 15
   end_location:
     row: 122
-    column: 19
+    column: 18
   fix: ~
 - kind:
     UndefinedName: bar
@@ -99,6 +99,6 @@ expression: checks
     column: 22
   end_location:
     row: 122
-    column: 26
+    column: 25
   fix: ~
 

src/snapshots/ruff__linter__tests__u007.snap

@@ -0,0 +1,133 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind: UsePEP604Annotation
+  location:
+    row: 4
+    column: 10
+  end_location:
+    row: 4
+    column: 23
+  fix:
+    content: str | None
+    location:
+      row: 4
+      column: 10
+    end_location:
+      row: 4
+      column: 23
+    applied: false
+- kind: UsePEP604Annotation
+  location:
+    row: 11
+    column: 10
+  end_location:
+    row: 11
+    column: 30
+  fix:
+    content: str | None
+    location:
+      row: 11
+      column: 10
+    end_location:
+      row: 11
+      column: 30
+    applied: false
+- kind: UsePEP604Annotation
+  location:
+    row: 18
+    column: 10
+  end_location:
+    row: 18
+    column: 46
+  fix:
+    content: "str | int | Union[float, bytes]"
+    location:
+      row: 18
+      column: 10
+    end_location:
+      row: 18
+      column: 46
+    applied: false
+- kind: UsePEP604Annotation
+  location:
+    row: 18
+    column: 26
+  end_location:
+    row: 18
+    column: 45
+  fix:
+    content: float | bytes
+    location:
+      row: 18
+      column: 26
+    end_location:
+      row: 18
+      column: 45
+    applied: false
+- kind: UsePEP604Annotation
+  location:
+    row: 25
+    column: 10
+  end_location:
+    row: 25
+    column: 32
+  fix:
+    content: str | int
+    location:
+      row: 25
+      column: 10
+    end_location:
+      row: 25
+      column: 32
+    applied: false
+- kind: UsePEP604Annotation
+  location:
+    row: 32
+    column: 11
+  end_location:
+    row: 32
+    column: 47
+  fix:
+    content: "str | int | Union[float, bytes]"
+    location:
+      row: 32
+      column: 11
+    end_location:
+      row: 32
+      column: 47
+    applied: false
+- kind: UsePEP604Annotation
+  location:
+    row: 32
+    column: 11
+  end_location:
+    row: 32
+    column: 47
+  fix:
+    content: float | bytes
+    location:
+      row: 32
+      column: 11
+    end_location:
+      row: 32
+      column: 47
+    applied: false
+- kind: UsePEP604Annotation
+  location:
+    row: 39
+    column: 11
+  end_location:
+    row: 39
+    column: 33
+  fix:
+    content: str | int
+    location:
+      row: 39
+      column: 11
+    end_location:
+      row: 39
+      column: 33
+    applied: false
+