fix: improve function defs and f-strings

jvllmr · jvllmr · commit 64d9bc810192 · 2025-01-03T21:31:50.000+01:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,3 +16,4 @@ rustpython-literal = ">=0.4.0"
 rustpython-parser = "0.4.0"
 rustpython-ast = { version = "0.4.0", features = ["fold"] }
 rand = "0.8.5"
+pretty_assertions = "1.4.1"
diff --git a/src/unparser.rs b/src/unparser.rs
@@ -116,6 +116,13 @@ impl Unparser {
         self.source += str_
     }
 
+    fn write_type_comment(&mut self, type_comment: &Option<String>) {
+        if let Some(str_) = type_comment {
+            self.write_str("  # type: ignore");
+            self.write_str(str_);
+        }
+    }
+
     fn block<F>(&mut self, f: F)
     where
         F: FnOnce(&mut Self),
@@ -233,6 +240,7 @@ impl Unparser {
             self.unparse_expr(&returns);
         }
         self.write_str(":");
+        self.write_type_comment(&node.type_comment);
         self.block(|block_self| {
             for value in &node.body {
                 block_self.unparse_stmt(&value);
@@ -268,6 +276,7 @@ impl Unparser {
             self.unparse_expr(&returns);
         }
         self.write_str(":");
+        self.write_type_comment(&node.type_comment);
         self.block(|block_self| {
             for value in &node.body {
                 block_self.unparse_stmt(&value);
@@ -359,9 +368,7 @@ impl Unparser {
         }
         self.write_str(" = ");
         self.unparse_expr(&node.value);
-        if node.type_comment.is_some() {
-            self.write_str(&node.type_comment.as_ref().unwrap());
-        }
+        self.write_type_comment(&node.type_comment);
     }
 
     fn unparse_stmt_type_alias(&mut self, node: &StmtTypeAlias<TextRange>) {
@@ -408,10 +415,7 @@ impl Unparser {
         self.write_str(" in ");
         self.unparse_expr(&node.iter);
         self.write_str(":");
-        if let Some(type_comment) = &node.type_comment {
-            self.write_str(" #type: ignore");
-            self.write_str(type_comment);
-        }
+        self.write_type_comment(&node.type_comment);
         self.block(|block_self| {
             for value in &node.body {
                 block_self.unparse_stmt(value);
@@ -432,10 +436,7 @@ impl Unparser {
         self.write_str(" in ");
         self.unparse_expr(&node.iter);
         self.write_str(":");
-        if let Some(type_comment) = &node.type_comment {
-            self.write_str(" #type: ignore");
-            self.write_str(type_comment);
-        }
+        self.write_type_comment(&node.type_comment);
         self.block(|block_self| {
             for value in &node.body {
                 block_self.unparse_stmt(value);
@@ -992,7 +993,8 @@ impl Unparser {
             match expr {
                 Expr::Constant(ExprConstant { value, .. }) => {
                     if let Constant::Str(str_) = value {
-                        inner_unparser.write_str(str_);
+                        let escaped = str_.replace('{', "{{").replace('}', "}}");
+                        inner_unparser.write_str(&escaped);
                     } else {
                         unreachable!()
                     }
@@ -1202,16 +1204,13 @@ impl Unparser {
 
             if posonly_iter.peek().is_some() {
                 self.write_str(", ");
-            } else if args_iter.peek().is_some()
-                || node.vararg.is_some()
-                || node.vararg.is_some()
-                || kw_iter.peek().is_some()
-                || node.kwarg.is_some()
-            {
-                self.write_str(", /, ");
             }
         }
 
+        if node.posonlyargs.len() > 0 {
+            self.write_str(", /,");
+        }
+
         while let Some(arg) = args_iter.next() {
             self.unparse_arg(arg.as_arg());
             if let Some(default) = &arg.default {
@@ -1230,13 +1229,16 @@ impl Unparser {
         if let Some(vararg) = &node.vararg {
             self.write_str("*");
             self.write_str(&vararg.arg);
+
             if let Some(annotation) = &vararg.annotation {
                 self.write_str(": ");
                 self.unparse_expr(annotation);
             }
             if kw_iter.peek().is_some() || node.kwarg.is_some() {
                 self.write_str(", ");
             }
+        } else if node.kwonlyargs.len() > 0 {
+            self.write_str("*, ");
         }
 
         while let Some(kw) = kw_iter.next() {
diff --git a/test_files/function_def.py b/test_files/function_def.py
@@ -20,3 +20,32 @@ async def name_5[*name_2](  # type: ignore
     name_5: str = "",
 ) -> str:
     return ""
+
+
+def type_comment() -> str:  # type: ignore[misc]
+    return ""
+
+
+async def type_comment_async():  # type: ignore[misc]
+    return ""
+
+
+def name_8(
+    name_5: name_4,  # type: ignore
+    name_3,
+    name_4,
+    /,
+    *,
+    name_1=name_2 <= name_1,  # type: ignore # noqa
+    name_0=~name_5,  # type: ignore
+):
+    pass
+
+
+async def name_5[**name_5, *name_3, **name_2](  # type: ignore
+    name_1: name_1,  # type: ignore # noqa
+    name_2,
+    name_0: name_2,  # type: ignore
+    /,  # type: ignore
+) -> name_3:  # type: ignore
+    name_1
diff --git a/test_files/simple_f_string.py b/test_files/simple_f_string.py
@@ -3,5 +3,6 @@
 f"{world!a:}None{empty_f_string!s:}None"
 answer = 42.000001
 f"{answer:.03f}"
+f"'\"'''\"\"\"{{}}\\"  # noqa
 if __name__ == "__main__":
     print(f"Hello {world}!")
