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

This commit is contained in:
Charlie Marsh 2022-11-28 20:57:38 -05:00 committed by GitHub
parent 82b0b7941a
commit 9944246f98
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 200 additions and 12 deletions

View file

@ -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

View file

@ -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)
&& self.settings.target_version >= PythonVersion::Py39
if !self.in_deferred_string_annotation
&& 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;
}

View file

@ -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(())
}
}

View file

@ -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: ~

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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