123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219
import { concat, indexOfNeedle } from "./deps.ts";import {  dummyAsyncIterable,  getUint8Array,  hasPrefixFrom,  isAsyncIterable,  makeAsyncIterable,  SyncAsyncIterable,} from "./utils.ts";
/** Options for CSV writer */export interface CSVWriterOptions {  columnSeparator: string | Uint8Array;  lineSeparator: string | Uint8Array;  quote: string | Uint8Array;}
/** Options for `CSVWriter.writeCell` */export interface CSVWriteCellOptions {  forceQuotes: boolean;}
const defaultCSVWriterOptions = {  columnSeparator: ",",  lineSeparator: "\n",  quote: '"',};
/** Class for manual CSV writing: * *       const writer = new CSVWriter(f, { *         columnSeparator: "\t", *         lineSeparator: "\r\n", *       }); *       await writer.writeCell("a\nb"); *       await writer.nextLine(); *       await writer.writeCell('1"2'); */export class CSVWriter {  private writer: Deno.Writer;  private columnSeparator: Uint8Array;  private lineSeparator: Uint8Array;  private quote: Uint8Array;  private firstColumn: boolean;
  constructor(writer: Deno.Writer, options?: Partial<CSVWriterOptions>) {    this.writer = writer;    this.columnSeparator = getUint8Array(      (options && options.columnSeparator) ||        defaultCSVWriterOptions.columnSeparator,    );    this.lineSeparator = getUint8Array(      (options && options.lineSeparator) ||        defaultCSVWriterOptions.lineSeparator,    );    this.quote = getUint8Array(      (options && options.quote) || defaultCSVWriterOptions.quote,    );    this.firstColumn = true;  }
  public writeCell(    str: string | Uint8Array | AsyncIterable<Uint8Array>,    options?: Partial<CSVWriteCellOptions>,  ): Promise<void> {    if (isAsyncIterable(str)) {      return this._writeCellAsyncIterable(str, { wrap: true });    }
    const arr = getUint8Array(str);    const wrap = options?.forceQuotes ||      indexOfNeedle(arr, this.quote) >= 0 ||      indexOfNeedle(arr, this.columnSeparator) >= 0 ||      indexOfNeedle(arr, this.lineSeparator) >= 0;
    return this._writeCellAsyncIterable(dummyAsyncIterable(arr), {      wrap,    });  }
  private async _writeCellAsyncIterable(    iterable: AsyncIterable<Uint8Array>,    options: { wrap: boolean },  ): Promise<void> {    const { quote } = this;    const { wrap } = options;
    const iterator = iterable[Symbol.asyncIterator]();
    let inputBuffer = new Uint8Array();    let inputBufferEmpty = false;    let inputBufferIndex = 0;
    if (this.firstColumn) {      this.firstColumn = false;    } else {      await this.writer.write(this.columnSeparator);    }
    if (wrap) {      await this.writer.write(this.quote);    }
    while (true) {      const inputBufferUnprocessed = inputBuffer.length - inputBufferIndex;
      if (inputBufferEmpty && inputBufferUnprocessed === 0) {        break;      }
      if (!inputBufferEmpty && inputBufferUnprocessed < quote.length) {        const { done, value } = await iterator.next();        if (done) {          inputBufferEmpty = true;        } else {          inputBuffer = concat(inputBuffer, value);        }        continue;      }
      if (wrap && hasPrefixFrom(inputBuffer, quote, inputBufferIndex)) {        await this.writer.write(quote);        await this.writer.write(quote);        inputBufferIndex += quote.length;        continue;      }
      if (inputBufferUnprocessed > 0) {        await this.writer.write(          inputBuffer.subarray(inputBufferIndex, inputBufferIndex + 1),        );        inputBufferIndex++;        continue;      }
      throw new Error("unexpected");    }
    if (wrap) {      await this.writer.write(this.quote);    }  }
  public async nextLine() {    this.firstColumn = true;    await this.writer.write(this.lineSeparator);  }}
/** Write CSV with sync or async row iterators: * *       await writeCSV(f, [["a", "b"], ["1", "2"]]); * *       const asyncRowGenerator = async function*() { *         yield ["a", "b"]; *         yield ["1", "2"]; *       } *       await writeCSV(f, asyncRowGenerator()); */export async function writeCSV(  writer: Deno.Writer,  iter: SyncAsyncIterable<    SyncAsyncIterable<string | Uint8Array | AsyncIterable<Uint8Array>>  >,  options?: Partial<CSVWriterOptions & CSVWriteCellOptions>,) {  const csv = new CSVWriter(writer, options);
  let firstLine = true;
  for await (const row of makeAsyncIterable(iter)) {    if (firstLine) {      firstLine = false;    } else {      await csv.nextLine();    }
    for await (const cell of makeAsyncIterable(row)) {      await csv.writeCell(cell, options);    }  }}
/** Write CSV with sync or async object iterators: * *       await writeCSVObjects(f, [{a: "1"}, {a: "2"}], { header: ["a"] }); * *       const asyncObjectsGenerator = async function*() { *         yield { a: "1", b: "2", c: "3" }; *         yield { a: "4", b: "5", c: "6" }; *       } *       await writeCSVObjects(f, asyncObjectsGenerator(), { header: ["a", "b", "c"] }); */export async function writeCSVObjects(  writer: Deno.Writer,  iter: SyncAsyncIterable<{ [key: string]: string }>,  options: Partial<CSVWriterOptions & CSVWriteCellOptions> & {    header: string[];  },) {  const { header } = options;
  const row = function* (obj: { [key: string]: string }) {    for (const key of header) {      yield obj[key];    }  };
  const rows = async function* () {    yield header;
    for await (const obj of makeAsyncIterable(iter)) {      yield row(obj);    }  };
  await writeCSV(writer, rows(), options);}