Wrap subscripted dicts in parens for f-string conversion (#9238)

Closes #9227.
astral-sh · Dec 21, 2023 · 1e7bc1d · 1e7bc1d
1 parent e241c1c
commit 1e7bc1d
Show file tree

Hide file tree

Showing 3 changed files with 110 additions and 6 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/pyupgrade/UP032_0.py b/crates/ruff_linter/resources/test/fixtures/pyupgrade/UP032_0.py
@@ -243,3 +243,12 @@ async def c():
 ).format(a, b)
 
 ("{}" "{{{}}}").format(a, b)
+
+# The dictionary should be parenthesized.
+"{}".format({0: 1}[0])
+
+# The dictionary should be parenthesized.
+"{}".format({0: 1}.bar)
+
+# The dictionary should be parenthesized.
+"{}".format({0: 1}())
diff --git a/crates/ruff_linter/src/rules/pyupgrade/rules/f_strings.rs b/crates/ruff_linter/src/rules/pyupgrade/rules/f_strings.rs
@@ -137,10 +137,9 @@ enum FormatContext {
     Accessed,
 }
 
-/// Given an [`Expr`], format it for use in a formatted expression within an f-string.
-fn formatted_expr<'a>(expr: &Expr, context: FormatContext, locator: &Locator<'a>) -> Cow<'a, str> {
-    let text = locator.slice(expr);
-    let parenthesize = match (context, expr) {
+/// Returns `true` if the expression should be parenthesized when used in an f-string.
+fn parenthesize(expr: &Expr, text: &str, context: FormatContext) -> bool {
+    match (context, expr) {
         // E.g., `x + y` should be parenthesized in `f"{(x + y)[0]}"`.
         (
             FormatContext::Accessed,
@@ -173,9 +172,44 @@ fn formatted_expr<'a>(expr: &Expr, context: FormatContext, locator: &Locator<'a>
             | Expr::SetComp(_)
             | Expr::DictComp(_),
         ) => true,
+        (_, Expr::Subscript(ast::ExprSubscript { value, .. })) => {
+            matches!(
+                value.as_ref(),
+                Expr::GeneratorExp(_)
+                    | Expr::Dict(_)
+                    | Expr::Set(_)
+                    | Expr::SetComp(_)
+                    | Expr::DictComp(_)
+            )
+        }
+        (_, Expr::Attribute(ast::ExprAttribute { value, .. })) => {
+            matches!(
+                value.as_ref(),
+                Expr::GeneratorExp(_)
+                    | Expr::Dict(_)
+                    | Expr::Set(_)
+                    | Expr::SetComp(_)
+                    | Expr::DictComp(_)
+            )
+        }
+        (_, Expr::Call(ast::ExprCall { func, .. })) => {
+            matches!(
+                func.as_ref(),
+                Expr::GeneratorExp(_)
+                    | Expr::Dict(_)
+                    | Expr::Set(_)
+                    | Expr::SetComp(_)
+                    | Expr::DictComp(_)
+            )
+        }
         _ => false,
-    };
-    if parenthesize && !text.starts_with('(') && !text.ends_with(')') {
+    }
+}
+
+/// Given an [`Expr`], format it for use in a formatted expression within an f-string.
+fn formatted_expr<'a>(expr: &Expr, context: FormatContext, locator: &Locator<'a>) -> Cow<'a, str> {
+    let text = locator.slice(expr);
+    if parenthesize(expr, text, context) && !(text.starts_with('(') && text.ends_with(')')) {
         Cow::Owned(format!("({text})"))
     } else {
         Cow::Borrowed(text)

diff --git a/...inter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP032_0.py.snap b/...inter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP032_0.py.snap
@@ -1141,13 +1141,16 @@ UP032_0.py:240:1: UP032 [*] Use f-string instead of `format` call
     243 |+)
 244 244 | 
 245 245 | ("{}" "{{{}}}").format(a, b)
+246 246 | 
 
 UP032_0.py:245:1: UP032 [*] Use f-string instead of `format` call
     |
 243 | ).format(a, b)
 244 | 
 245 | ("{}" "{{{}}}").format(a, b)
     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP032
+246 | 
+247 | # The dictionary should be parenthesized.
     |
     = help: Convert to f-string
 
@@ -1157,5 +1160,63 @@ UP032_0.py:245:1: UP032 [*] Use f-string instead of `format` call
 244 244 | 
 245     |-("{}" "{{{}}}").format(a, b)
     245 |+(f"{a}" f"{{{b}}}")
+246 246 | 
+247 247 | # The dictionary should be parenthesized.
+248 248 | "{}".format({0: 1}[0])
+
+UP032_0.py:248:1: UP032 [*] Use f-string instead of `format` call
+    |
+247 | # The dictionary should be parenthesized.
+248 | "{}".format({0: 1}[0])
+    | ^^^^^^^^^^^^^^^^^^^^^^ UP032
+249 | 
+250 | # The dictionary should be parenthesized.
+    |
+    = help: Convert to f-string
+
+ℹ Safe fix
+245 245 | ("{}" "{{{}}}").format(a, b)
+246 246 | 
+247 247 | # The dictionary should be parenthesized.
+248     |-"{}".format({0: 1}[0])
+    248 |+f"{({0: 1}[0])}"
+249 249 | 
+250 250 | # The dictionary should be parenthesized.
+251 251 | "{}".format({0: 1}.bar)
+
+UP032_0.py:251:1: UP032 [*] Use f-string instead of `format` call
+    |
+250 | # The dictionary should be parenthesized.
+251 | "{}".format({0: 1}.bar)
+    | ^^^^^^^^^^^^^^^^^^^^^^^ UP032
+252 | 
+253 | # The dictionary should be parenthesized.
+    |
+    = help: Convert to f-string
+
+ℹ Safe fix
+248 248 | "{}".format({0: 1}[0])
+249 249 | 
+250 250 | # The dictionary should be parenthesized.
+251     |-"{}".format({0: 1}.bar)
+    251 |+f"{({0: 1}.bar)}"
+252 252 | 
+253 253 | # The dictionary should be parenthesized.
+254 254 | "{}".format({0: 1}())
+
+UP032_0.py:254:1: UP032 [*] Use f-string instead of `format` call
+    |
+253 | # The dictionary should be parenthesized.
+254 | "{}".format({0: 1}())
+    | ^^^^^^^^^^^^^^^^^^^^^ UP032
+    |
+    = help: Convert to f-string
+
+ℹ Safe fix
+251 251 | "{}".format({0: 1}.bar)
+252 252 | 
+253 253 | # The dictionary should be parenthesized.
+254     |-"{}".format({0: 1}())
+    254 |+f"{({0: 1}())}"