Class FunctionCall

frontend/src/utils/python-poet/poet.ts:153–198

}

export class FunctionCall implements PythonCode {
  private args: PythonCode[];
  public name: string;
  public multiLine: boolean;

  constructor(
    name: string,
    args: PythonCode[] | Record<string, PythonCode>,
    multiLine = false,
  ) {
    this.name = name;
    this.multiLine = multiLine;
    this.args = objectToArgs(args);
  }

  toCode(): string {
    if (this.multiLine) {
      if (this.args.length === 0) {
        return `${this.name}()`;
      }
      if (this.args.length === 1) {
        return `${this.name}(${this.args[0].toCode()})`;
      }
      return `${this.name}(${indentList(this.args)})`;
    }
    return `${this.name}(${this.args.map(asString).join(", ")})`;
  }

  addArg(...args: PythonCode[]): FunctionCall {
    return new FunctionCall(this.name, [...this.args, ...args], this.multiLine);
  }

  chain(
    name: string,
    args: PythonCode[] | Record<string, PythonCode>,
  ): FunctionCall {
    args = objectToArgs(args);

    // If the function call is multi-line, we need to add a newline to the name
    const fullName = this.multiLine
      ? `${this.toCode()}\n.${name}`
      : `${this.toCode()}.${name}`;

    return new FunctionCall(fullName, args, this.multiLine);
  }
}

function objectToArgs(
  obj: Record<string, PythonCode> | PythonCode[],