Skip to content

Commit

Permalink
feat(python): Config for Trailing Commas in Iterable Types (#5486)
Browse files Browse the repository at this point in the history 
* Add iterable config

* Update writer logic

* Add test cases

* Handle tuple of length 1
  • Loading branch information
@noanflaherty
noanflaherty authored Dec 26, 2024
1 parent 192e49e commit 82c9d11
Show file tree
Hide file tree
Showing 3 changed files with 145 additions and 43 deletions.
71 changes: 52 additions & 19 deletions generators/python-v2/ast/src/TypeInstantiation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -30,7 +34,7 @@ interface Bool {
interface Str {
type: "str";
value: string;
config?: StrFormatConfig;
config?: StrConfig;
}

interface Bytes {
Expand All @@ -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 {
Expand Down Expand Up @@ -91,7 +99,7 @@ export class TypeInstantiation extends AstNode {

public static str(
value: string,
config: StrFormatConfig = {
config: StrConfig = {
multiline: false,
startOnNewLine: false,
endWithNewLine: false
Expand All @@ -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);
Expand Down Expand Up @@ -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;
Expand All @@ -235,7 +268,7 @@ export class TypeInstantiation extends AstNode {
}: {
writer: Writer;
value: string;
} & Pick<StrFormatConfig, "startOnNewLine" | "endWithNewLine">): void {
} & Pick<StrConfig, "startOnNewLine" | "endWithNewLine">): void {
writer.write('"""');
const lines = value.split("\n");

Expand Down
80 changes: 59 additions & 21 deletions generators/python-v2/ast/src/__test__/TypeInstantiation.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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", () => {
Expand All @@ -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 () => {
Expand Down
37 changes: 34 additions & 3 deletions generators/python-v2/ast/src/__test__/__snapshots__/TypeInstantiation.test.ts.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
"
Expand All @@ -35,7 +43,7 @@ exports[`TypeInstantiation > int 1`] = `
"
`;

exports[`TypeInstantiation > list 1`] = `
exports[`TypeInstantiation > list > basic 1`] = `
"[1, "two", True]
"
`;
Expand All @@ -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"
"
Expand Down Expand Up @@ -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")
"
Expand Down

0 comments on commit 82c9d11

Please sign in to comment.