feat(python): Multi-line Strings (#5105)

* Add multi line string support * Update tests * Update snapshots * Fix multi-line string writes * Support config
fern-api · Nov 17, 2024 · b3922cc · b3922cc
1 parent b02a296
commit b3922cc
Show file tree

Hide file tree

Showing 7 changed files with 129 additions and 9 deletions.
diff --git a/generators/commons/src/ast/AbstractWriter.ts b/generators/commons/src/ast/AbstractWriter.ts
@@ -38,7 +38,10 @@ export class AbstractWriter {
      * @param text
      */
     public writeNoIndent(text: string): void {
-        this.writeInternal(text);
+        const currIndentLevel = this.indentLevel;
+        this.indentLevel = 0;
+        this.write(text);
+        this.indentLevel = currIndentLevel;
     }
 
     /**

diff --git a/generators/python-v2/ast/src/Field.ts b/generators/python-v2/ast/src/Field.ts
@@ -48,7 +48,8 @@ export class Field extends AstNode {
         }
 
         if (this.initializer !== undefined) {
-            writer.write(` = ${this.initializer}`);
+            writer.write(" = ");
+            this.initializer.write(writer);
         }
 
         if (this.docs != null) {

diff --git a/generators/python-v2/ast/src/TypeInstantiation.ts b/generators/python-v2/ast/src/TypeInstantiation.ts
@@ -24,6 +24,9 @@ interface Bool {
 interface Str {
     type: "str";
     value: string;
+    config?: {
+        multiline?: boolean;
+    };
 }
 
 interface Bytes {
@@ -82,8 +85,8 @@ export class TypeInstantiation extends AstNode {
         return new this({ type: "bool", value });
     }
 
-    public static str(value: string): TypeInstantiation {
-        return new this({ type: "str", value });
+    public static str(value: string, config = { multiline: false }): TypeInstantiation {
+        return new this({ type: "str", value, config });
     }
 
     public static bytes(value: string): TypeInstantiation {
@@ -143,7 +146,13 @@ export class TypeInstantiation extends AstNode {
                 }
                 break;
             case "str":
-                writer.write(`"${this.internalType.value}"`);
+                if (this.internalType.config?.multiline) {
+                    this.writeStringWithTripleQuotes({ writer, value: this.internalType.value });
+                } else {
+                    writer.write(
+                        `"${this.escapeString(this.internalType.value, { doubleQuote: true, newline: true })}"`
+                    );
+                }
                 break;
             case "bytes":
                 writer.write(`b"${this.internalType.value}"`);
@@ -200,6 +209,30 @@ export class TypeInstantiation extends AstNode {
                 assertNever(this.internalType);
         }
     }
+
+    private writeStringWithTripleQuotes({ writer, value }: { writer: Writer; value: string }): void {
+        writer.write('"""');
+        const parts = value.split("\n");
+        const head = parts[0] + "\n";
+        const tail = parts.slice(1).join("\n");
+        writer.write(this.escapeString(head));
+        writer.writeNoIndent(this.escapeString(tail));
+        writer.write('"""');
+    }
+
+    private escapeString(
+        value: string,
+        config: { doubleQuote?: boolean; newline?: boolean } = { doubleQuote: true, newline: false }
+    ): string {
+        let escapedValue = value;
+        if (config.doubleQuote) {
+            escapedValue = escapedValue.replaceAll('"', '\\"');
+        }
+        if (config.newline) {
+            escapedValue = escapedValue.replaceAll("\n", "\\n");
+        }
+        return escapedValue;
+    }
 }
 
 export { Type };
diff --git a/generators/python-v2/ast/src/__test__/Class.test.ts b/generators/python-v2/ast/src/__test__/Class.test.ts
@@ -125,4 +125,20 @@ describe("class", () => {
 
         expect(clazz.getReferences().length).toBe(2);
     });
+
+    it("class with de-indented multi-line string", async () => {
+        const clazz = python.class_({
+            name: "MyClass"
+        });
+        clazz.add(
+            python.field({
+                name: "multiline_string",
+                type: python.Type.str(),
+                initializer: python.TypeInstantiation.str("Hello\nWorld", { multiline: true })
+            })
+        );
+
+        clazz.write(writer);
+        expect(await writer.toStringFormatted()).toMatchSnapshot();
+    });
 });
diff --git a/generators/python-v2/ast/src/__test__/TypeInstantiation.test.ts b/generators/python-v2/ast/src/__test__/TypeInstantiation.test.ts
@@ -31,9 +31,38 @@ describe("TypeInstantiation", () => {
         });
     });
 
-    it("str", async () => {
-        TypeInstantiation.str("hello").write(writer);
-        expect(await writer.toStringFormatted()).toMatchSnapshot();
+    describe("str", () => {
+        it("should render a basic string", async () => {
+            TypeInstantiation.str("hello").write(writer);
+            expect(await writer.toStringFormatted()).toMatchSnapshot();
+        });
+
+        it("should render a string containing quote", async () => {
+            TypeInstantiation.str('She said "hello!"').write(writer);
+            expect(await writer.toStringFormatted()).toMatchSnapshot();
+        });
+
+        it("should render a string containing escaped newline characters", async () => {
+            TypeInstantiation.str("\n\n####\n\n").write(writer);
+            expect(await writer.toStringFormatted()).toMatchSnapshot();
+        });
+
+        it("should render a multiline string", async () => {
+            TypeInstantiation.str("\n\n####\n\n", { multiline: true }).write(writer);
+            expect(await writer.toStringFormatted()).toMatchSnapshot();
+        });
+
+        it("should render a multiline string containing escaped quotes", async () => {
+            TypeInstantiation.str('She said "Hi"\nHe said "bye"\nShe said "okay then"', { multiline: true }).write(
+                writer
+            );
+            expect(await writer.toStringFormatted()).toMatchSnapshot();
+        });
+
+        it("should render a string containing escaped newline characters and quotes", async () => {
+            TypeInstantiation.str('She said "Hi"\nHe said "bye"\nShe said "okay then"').write(writer);
+            expect(await writer.toStringFormatted()).toMatchSnapshot();
+        });
     });
 
     it("bytes", async () => {

diff --git a/generators/python-v2/ast/src/__test__/__snapshots__/Class.test.ts.snap b/generators/python-v2/ast/src/__test__/__snapshots__/Class.test.ts.snap
@@ -13,6 +13,13 @@ class MyDataClass:
 "
 `;
 
+exports[`class > class with de-indented multi-line string 1`] = `
+"class MyClass:
+    multiline_string: str = """Hello
+World"""
+"
+`;
+
 exports[`class > class with generic parent reference 1`] = `
 "class MyClass(MyParentClass[MyParentType]):
     pass

diff --git a/generators/python-v2/ast/src/__test__/__snapshots__/TypeInstantiation.test.ts.snap b/generators/python-v2/ast/src/__test__/__snapshots__/TypeInstantiation.test.ts.snap
@@ -50,11 +50,42 @@ exports[`TypeInstantiation > set 1`] = `
 "
 `;
 
-exports[`TypeInstantiation > str 1`] = `
+exports[`TypeInstantiation > str > should render a basic string 1`] = `
 ""hello"
 "
 `;
 
+exports[`TypeInstantiation > str > should render a multiline string 1`] = `
+""""
+
+####
+
+"""
+"
+`;
+
+exports[`TypeInstantiation > str > should render a multiline string containing escaped quotes 1`] = `
+""""She said \\"Hi\\"
+He said \\"bye\\"
+She said \\"okay then\\" """
+"
+`;
+
+exports[`TypeInstantiation > str > should render a string containing escaped newline characters 1`] = `
+""\\n\\n####\\n\\n"
+"
+`;
+
+exports[`TypeInstantiation > str > should render a string containing escaped newline characters and quotes 1`] = `
+"'She said "Hi"\\nHe said "bye"\\nShe said "okay then"'
+"
+`;
+
+exports[`TypeInstantiation > str > should render a string containing quote 1`] = `
+"'She said "hello!"'
+"
+`;
+
 exports[`TypeInstantiation > tuple 1`] = `
 "(1, "two", True)
 "