Skip to content

Commit

Permalink
Enable annotation quoting for multi-line expressions (#9142)
Browse files Browse the repository at this point in the history
Given:

```python
x: DataFrame[
    int
] = 1
```

We currently wrap the annotation in single quotes, which leads to a
syntax error:

```python
x: "DataFrame[
    int
]" = 1
```

There are a few options for what to suggest for users here... Use triple
quotes:

```python
x: """DataFrame[
    int
]""" = 1
```

Or, use an implicit string concatenation (which may require
parentheses):

```python
x: ("DataFrame["
    "int"
"]") = 1
```

The solution I settled on here is to use the `Generator`, which
effectively means we write it out on a single line, like:

```python
x: "DataFrame[int]" = 1
```

It's kind of the "least opinionated" solution, but it does mean we'll
expand to a very long line in some cases.

Closes #9136.
  • Loading branch information
charliermarsh authored Dec 15, 2023
1 parent 6c224ce commit d1a7bc3
Show file tree
Hide file tree
Showing 6 changed files with 113 additions and 14 deletions.
2 changes: 1 addition & 1 deletion crates/ruff_diagnostics/src/edit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use ruff_text_size::{Ranged, TextRange, TextSize};

/// A text edit to be applied to a source file. Inserts, deletes, or replaces
/// content at a given location.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Edit {
/// The start location of the edit.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,3 +72,18 @@ def f():

def baz() -> DataFrame | Series:
...


def f():
from pandas import DataFrame, Series

def baz() -> (
DataFrame |
Series
):
...

class C:
x: DataFrame[
int
] = 1
26 changes: 15 additions & 11 deletions crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use ruff_diagnostics::Edit;
use ruff_python_ast::call_path::from_qualified_name;
use ruff_python_ast::helpers::{map_callable, map_subscript};
use ruff_python_ast::{self as ast, Expr};
use ruff_python_codegen::Stylist;
use ruff_python_codegen::{Generator, Stylist};
use ruff_python_semantic::{
Binding, BindingId, BindingKind, NodeId, ResolvedReference, SemanticModel,
};
Expand Down Expand Up @@ -215,6 +215,7 @@ pub(crate) fn quote_annotation(
semantic: &SemanticModel,
locator: &Locator,
stylist: &Stylist,
generator: Generator,
) -> Result<Edit> {
let expr = semantic.expression(node_id).expect("Expression not found");
if let Some(parent_id) = semantic.parent_expression_id(node_id) {
Expand All @@ -224,51 +225,54 @@ pub(crate) fn quote_annotation(
// If we're quoting the value of a subscript, we need to quote the entire
// expression. For example, when quoting `DataFrame` in `DataFrame[int]`, we
// should generate `"DataFrame[int]"`.
return quote_annotation(parent_id, semantic, locator, stylist);
return quote_annotation(parent_id, semantic, locator, stylist, generator);
}
}
Some(Expr::Attribute(parent)) => {
if expr == parent.value.as_ref() {
// If we're quoting the value of an attribute, we need to quote the entire
// expression. For example, when quoting `DataFrame` in `pd.DataFrame`, we
// should generate `"pd.DataFrame"`.
return quote_annotation(parent_id, semantic, locator, stylist);
return quote_annotation(parent_id, semantic, locator, stylist, generator);
}
}
Some(Expr::Call(parent)) => {
if expr == parent.func.as_ref() {
// If we're quoting the function of a call, we need to quote the entire
// expression. For example, when quoting `DataFrame` in `DataFrame()`, we
// should generate `"DataFrame()"`.
return quote_annotation(parent_id, semantic, locator, stylist);
return quote_annotation(parent_id, semantic, locator, stylist, generator);
}
}
Some(Expr::BinOp(parent)) => {
if parent.op.is_bit_or() {
// If we're quoting the left or right side of a binary operation, we need to
// quote the entire expression. For example, when quoting `DataFrame` in
// `DataFrame | Series`, we should generate `"DataFrame | Series"`.
return quote_annotation(parent_id, semantic, locator, stylist);
return quote_annotation(parent_id, semantic, locator, stylist, generator);
}
}
_ => {}
}
}

let annotation = locator.slice(expr);

// If the annotation already contains a quote, avoid attempting to re-quote it. For example:
// ```python
// from typing import Literal
//
// Set[Literal["Foo"]]
// ```
if annotation.contains('\'') || annotation.contains('"') {
let text = locator.slice(expr);
if text.contains('\'') || text.contains('"') {
return Err(anyhow::anyhow!("Annotation already contains a quote"));
}

// If we're quoting a name, we need to quote the entire expression.
// Quote the entire expression.
let quote = stylist.quote();
let annotation = format!("{quote}{annotation}{quote}");
Ok(Edit::range_replacement(annotation, expr.range()))
let annotation = generator.expr(expr);

Ok(Edit::range_replacement(
format!("{quote}{annotation}{quote}"),
expr.range(),
))
}
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,7 @@ fn quote_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding])
checker.semantic(),
checker.locator(),
checker.stylist(),
checker.generator(),
))
} else {
None
Expand All @@ -282,7 +283,7 @@ fn quote_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding])
})
.collect::<Result<Vec<_>>>()?;

let mut rest = quote_reference_edits.into_iter().dedup();
let mut rest = quote_reference_edits.into_iter().unique();
let head = rest.next().expect("Expected at least one reference");
Ok(Fix::unsafe_edits(head, rest).isolate(Checker::isolation(
checker.semantic().parent_statement_id(node_id),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -494,6 +494,7 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) ->
checker.semantic(),
checker.locator(),
checker.stylist(),
checker.generator(),
))
} else {
None
Expand All @@ -507,7 +508,7 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) ->
add_import_edit
.into_edits()
.into_iter()
.chain(quote_reference_edits.into_iter().dedup()),
.chain(quote_reference_edits.into_iter().unique()),
)
.isolate(Checker::isolation(
checker.semantic().parent_statement_id(node_id),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,8 @@ quote.py:71:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a typ
73 |- def baz() -> DataFrame | Series:
76 |+ def baz() -> "DataFrame | Series":
74 77 | ...
75 78 |
76 79 |

quote.py:71:35: TCH002 [*] Move third-party import `pandas.Series` into a type-checking block
|
Expand Down Expand Up @@ -251,5 +253,81 @@ quote.py:71:35: TCH002 [*] Move third-party import `pandas.Series` into a type-c
73 |- def baz() -> DataFrame | Series:
76 |+ def baz() -> "DataFrame | Series":
74 77 | ...
75 78 |
76 79 |

quote.py:78:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a type-checking block
|
77 | def f():
78 | from pandas import DataFrame, Series
| ^^^^^^^^^ TCH002
79 |
80 | def baz() -> (
|
= help: Move into type-checking block

Unsafe fix
1 |+from typing import TYPE_CHECKING
2 |+
3 |+if TYPE_CHECKING:
4 |+ from pandas import DataFrame, Series
1 5 | def f():
2 6 | from pandas import DataFrame
3 7 |
--------------------------------------------------------------------------------
75 79 |
76 80 |
77 81 | def f():
78 |- from pandas import DataFrame, Series
79 82 |
80 83 | def baz() -> (
81 |- DataFrame |
82 |- Series
84 |+ "DataFrame | Series"
83 85 | ):
84 86 | ...
85 87 |
86 88 | class C:
87 |- x: DataFrame[
88 |- int
89 |- ] = 1
89 |+ x: "DataFrame[int]" = 1

quote.py:78:35: TCH002 [*] Move third-party import `pandas.Series` into a type-checking block
|
77 | def f():
78 | from pandas import DataFrame, Series
| ^^^^^^ TCH002
79 |
80 | def baz() -> (
|
= help: Move into type-checking block

Unsafe fix
1 |+from typing import TYPE_CHECKING
2 |+
3 |+if TYPE_CHECKING:
4 |+ from pandas import DataFrame, Series
1 5 | def f():
2 6 | from pandas import DataFrame
3 7 |
--------------------------------------------------------------------------------
75 79 |
76 80 |
77 81 | def f():
78 |- from pandas import DataFrame, Series
79 82 |
80 83 | def baz() -> (
81 |- DataFrame |
82 |- Series
84 |+ "DataFrame | Series"
83 85 | ):
84 86 | ...
85 87 |
86 88 | class C:
87 |- x: DataFrame[
88 |- int
89 |- ] = 1
89 |+ x: "DataFrame[int]" = 1


0 comments on commit d1a7bc3

Please sign in to comment.