diff --git a/generators/python-v2/ast/src/TypeInstantiation.ts b/generators/python-v2/ast/src/TypeInstantiation.ts index 6390a4db673..35762240db6 100644 --- a/generators/python-v2/ast/src/TypeInstantiation.ts +++ b/generators/python-v2/ast/src/TypeInstantiation.ts @@ -6,7 +6,11 @@ import { Reference } from "./Reference"; type InternalTypeInstantiation = Int | Float | Bool | Str | Bytes | List | Set | Tuple | Dict | None | Uuid; -interface StrFormatConfig { +interface IterableConfig { + endWithComma?: boolean; +} + +interface StrConfig { multiline?: boolean; startOnNewLine?: boolean; endWithNewLine?: boolean; @@ -30,7 +34,7 @@ interface Bool { interface Str { type: "str"; value: string; - config?: StrFormatConfig; + config?: StrConfig; } interface Bytes { @@ -41,21 +45,25 @@ interface Bytes { interface List { type: "list"; values: AstNode[]; + config?: IterableConfig; } interface Set { type: "set"; values: AstNode[]; + config?: IterableConfig; } interface Tuple { type: "tuple"; values: AstNode[]; + config?: IterableConfig; } interface Dict { type: "dict"; entries: DictEntry[]; + config?: IterableConfig; } interface DictEntry { @@ -91,7 +99,7 @@ export class TypeInstantiation extends AstNode { public static str( value: string, - config: StrFormatConfig = { + config: StrConfig = { multiline: false, startOnNewLine: false, endWithNewLine: false @@ -104,26 +112,26 @@ export class TypeInstantiation extends AstNode { return new this({ type: "bytes", value }); } - public static list(values: AstNode[]): TypeInstantiation { - const list = new this({ type: "list", values }); + public static list(values: AstNode[], config: IterableConfig = { endWithComma: false }): TypeInstantiation { + const list = new this({ type: "list", values, config }); values.forEach((value) => list.inheritReferences(value)); return list; } - public static set(values: AstNode[]): TypeInstantiation { - const set = new this({ type: "set", values }); + public static set(values: AstNode[], config: IterableConfig = { endWithComma: false }): TypeInstantiation { + const set = new this({ type: "set", values, config }); values.forEach((value) => set.inheritReferences(value)); return set; } - public static tuple(values: AstNode[]): TypeInstantiation { - const tuple = new this({ type: "tuple", values }); + public static tuple(values: AstNode[], config: IterableConfig = { endWithComma: false }): TypeInstantiation { + const tuple = new this({ type: "tuple", values, config }); values.forEach((value) => tuple.inheritReferences(value)); return tuple; } - public static dict(entries: DictEntry[]): TypeInstantiation { - const dict = new this({ type: "dict", entries }); + public static dict(entries: DictEntry[], config: IterableConfig = { endWithComma: false }): TypeInstantiation { + const dict = new this({ type: "dict", entries, config }); entries.forEach((entry) => { dict.inheritReferences(entry.key); dict.inheritReferences(entry.value); @@ -174,48 +182,73 @@ export class TypeInstantiation extends AstNode { case "bytes": writer.write(`b"${this.internalType.value}"`); break; - case "list": + case "list": { + const internalType = this.internalType; writer.write("["); - this.internalType.values.forEach((value, index) => { + internalType.values.forEach((value, index) => { if (index > 0) { writer.write(", "); } value.write(writer); + if (index === internalType.values.length - 1 && internalType.config?.endWithComma) { + writer.write(","); + } }); writer.write("]"); break; - case "set": + } + case "set": { + const internalType = this.internalType; writer.write("{"); - this.internalType.values.forEach((value, index) => { + internalType.values.forEach((value, index) => { if (index > 0) { writer.write(", "); } value.write(writer); + if (index === internalType.values.length - 1 && internalType.config?.endWithComma) { + writer.write(","); + } }); writer.write("}"); break; - case "tuple": + } + case "tuple": { + const internalType = this.internalType; writer.write("("); this.internalType.values.forEach((value, index) => { if (index > 0) { writer.write(", "); } value.write(writer); + if ( + // If the tuple is of length 1, then we must always add a trailing comma + internalType.values.length === 1 || + // Otherwise, check the config that was specified + (index === internalType.values.length - 1 && internalType.config?.endWithComma) + ) { + writer.write(","); + } }); writer.write(")"); break; - case "dict": + } + case "dict": { + const internalType = this.internalType; writer.write("{"); - this.internalType.entries.forEach((entry, index) => { + internalType.entries.forEach((entry, index) => { if (index > 0) { writer.write(", "); } entry.key.write(writer); writer.write(": "); entry.value.write(writer); + if (index === internalType.entries.length - 1 && internalType.config?.endWithComma) { + writer.write(","); + } }); writer.write("}"); break; + } case "none": writer.write("None"); break; @@ -235,7 +268,7 @@ export class TypeInstantiation extends AstNode { }: { writer: Writer; value: string; - } & Pick): void { + } & Pick): void { writer.write('"""'); const lines = value.split("\n"); diff --git a/generators/python-v2/ast/src/__test__/TypeInstantiation.test.ts b/generators/python-v2/ast/src/__test__/TypeInstantiation.test.ts index c0b59b8ab7c..2683f5943c2 100644 --- a/generators/python-v2/ast/src/__test__/TypeInstantiation.test.ts +++ b/generators/python-v2/ast/src/__test__/TypeInstantiation.test.ts @@ -70,31 +70,58 @@ describe("TypeInstantiation", () => { expect(await writer.toStringFormatted()).toMatchSnapshot(); }); - it("list", async () => { - TypeInstantiation.list([ - TypeInstantiation.int(1), - TypeInstantiation.str("two"), - TypeInstantiation.bool(true) - ]).write(writer); - expect(await writer.toStringFormatted()).toMatchSnapshot(); + describe("list", () => { + it("basic", async () => { + TypeInstantiation.list([ + TypeInstantiation.int(1), + TypeInstantiation.str("two"), + TypeInstantiation.bool(true) + ]).write(writer); + expect(await writer.toStringFormatted()).toMatchSnapshot(); + }); }); - it("set", async () => { - TypeInstantiation.set([ - TypeInstantiation.int(1), - TypeInstantiation.str("two"), - TypeInstantiation.bool(true) - ]).write(writer); - expect(await writer.toStringFormatted()).toMatchSnapshot(); + describe("set", () => { + it("basic", async () => { + TypeInstantiation.set([ + TypeInstantiation.int(1), + TypeInstantiation.str("two"), + TypeInstantiation.bool(true) + ]).write(writer); + expect(await writer.toStringFormatted()).toMatchSnapshot(); + }); + + it("should support trailing comma", async () => { + TypeInstantiation.set( + [TypeInstantiation.int(1), TypeInstantiation.str("two"), TypeInstantiation.bool(true)], + { endWithComma: true } + ).write(writer); + expect(await writer.toStringFormatted()).toMatchSnapshot(); + }); }); - it("tuple", async () => { - TypeInstantiation.tuple([ - TypeInstantiation.int(1), - TypeInstantiation.str("two"), - TypeInstantiation.bool(true) - ]).write(writer); - expect(await writer.toStringFormatted()).toMatchSnapshot(); + describe("tuple", () => { + it("basic", async () => { + TypeInstantiation.tuple([ + TypeInstantiation.int(1), + TypeInstantiation.str("two"), + TypeInstantiation.bool(true) + ]).write(writer); + expect(await writer.toStringFormatted()).toMatchSnapshot(); + }); + + it("should support trailing comma", async () => { + TypeInstantiation.tuple( + [TypeInstantiation.int(1), TypeInstantiation.str("two"), TypeInstantiation.bool(true)], + { endWithComma: true } + ).write(writer); + expect(await writer.toStringFormatted()).toMatchSnapshot(); + }); + + it("should handle single-element tuple", async () => { + TypeInstantiation.tuple([TypeInstantiation.int(1)]).write(writer); + expect(await writer.toStringFormatted()).toMatchSnapshot(); + }); }); describe("dict", () => { @@ -115,6 +142,17 @@ describe("TypeInstantiation", () => { expect(await writer.toStringFormatted()).toMatchSnapshot(); expect(dict.getReferences().length).toBe(2); }); + + it("should support trailing comma", async () => { + TypeInstantiation.dict( + [ + { key: TypeInstantiation.str("one"), value: TypeInstantiation.int(1) }, + { key: TypeInstantiation.str("two"), value: TypeInstantiation.bool(true) } + ], + { endWithComma: true } + ).write(writer); + expect(await writer.toStringFormatted()).toMatchSnapshot(); + }); }); it("none", async () => { 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 index 60f2c9e4110..0b22e1efe50 100644 --- a/generators/python-v2/ast/src/__test__/__snapshots__/TypeInstantiation.test.ts.snap +++ b/generators/python-v2/ast/src/__test__/__snapshots__/TypeInstantiation.test.ts.snap @@ -25,6 +25,14 @@ exports[`TypeInstantiation > dict > should correctly write a dict with reference " `; +exports[`TypeInstantiation > dict > should support trailing comma 1`] = ` +"{ + "one": 1, + "two": True, +} +" +`; + exports[`TypeInstantiation > float 1`] = ` "3.14 " @@ -35,7 +43,7 @@ exports[`TypeInstantiation > int 1`] = ` " `; -exports[`TypeInstantiation > list 1`] = ` +exports[`TypeInstantiation > list > basic 1`] = ` "[1, "two", True] " `; @@ -45,11 +53,20 @@ exports[`TypeInstantiation > none 1`] = ` " `; -exports[`TypeInstantiation > set 1`] = ` +exports[`TypeInstantiation > set > basic 1`] = ` "{1, "two", True} " `; +exports[`TypeInstantiation > set > should support trailing comma 1`] = ` +"{ + 1, + "two", + True, +} +" +`; + exports[`TypeInstantiation > str > should render a basic string 1`] = ` ""hello" " @@ -86,11 +103,25 @@ exports[`TypeInstantiation > str > should render a string containing quote 1`] = " `; -exports[`TypeInstantiation > tuple 1`] = ` +exports[`TypeInstantiation > tuple > basic 1`] = ` "(1, "two", True) " `; +exports[`TypeInstantiation > tuple > should handle single-element tuple 1`] = ` +"(1,) +" +`; + +exports[`TypeInstantiation > tuple > should support trailing comma 1`] = ` +"( + 1, + "two", + True, +) +" +`; + exports[`TypeInstantiation > uuid 1`] = ` "UUID("123e4567-e89b-12d3-a456-426614174000") "