fix: Support generators in f-strings

snok · Apr 22, 2023 · cbbc4f7 · cbbc4f7
1 parent 94938e5
commit cbbc4f7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/main.rs b/src/main.rs
@@ -272,6 +272,8 @@ mod tests {
             TestCase { input: "logger.error(f'{ {\"foo\": str(e) for errors in all_errors for e in errors} }')".to_string(), expected_output: "logger.error('%s', {'foo': str(e) for errors in all_errors for e in errors})".to_string() },
             // Call containing list comprehension
             TestCase { input: "logger.error(f'{\", \".join([str(e) for e in errors for errors in all_errors])}')".to_string(), expected_output: "logger.error('%s', ', '.join([str(e) for e in errors for errors in all_errors],))".to_string() },
+            // Generator
+            TestCase { input: "logger.exception(f'{\", \".join(b for b in bs)}')".to_string(), expected_output: "logger.exception('%s', ', '.join([b for b in bs],))".to_string() },
         ]
     }
 

diff --git a/src/parse_fstring.rs b/src/parse_fstring.rs
@@ -126,7 +126,7 @@ pub fn parse_formatted_value(value: &Expr, postfix: String, in_call: bool) -> Re
                 quotes.char()
             )
         }
-        ExprKind::ListComp { elt, generators } => {
+        ExprKind::ListComp { elt, generators } | ExprKind::GeneratorExp { elt, generators } => {
             let mut s = format!("[{}", parse_formatted_value(elt, postfix.clone(), true)?,);
             for generator in generators {
                 s.push_str(&format!(