From b1c02d8f09dd2c0ecaa2e3d95ce434735cc80854 Mon Sep 17 00:00:00 2001
From: Alex Caza <alex@alexcaza.com>
Date: Sun, 8 Oct 2023 18:47:58 -0400
Subject: [PATCH] Add better types for HeaderColumn

---
 lib/__specs__/helpers.spec.ts | 23 +++++++++++++++++++++++
 lib/helpers.ts                | 30 ++++++++++++++++++++++++------
 lib/types.ts                  | 16 +++++++++++++---
 3 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/lib/__specs__/helpers.spec.ts b/lib/__specs__/helpers.spec.ts
index 410f469..278b3a1 100644
--- a/lib/__specs__/helpers.spec.ts
+++ b/lib/__specs__/helpers.spec.ts
@@ -159,6 +159,19 @@ describe("Helpers", () => {
         expect(dateAndCity).toEqual('"date","city"' + endOfLine);
       });
     });
+
+    describe("pretty header mappings", () => {
+      it("should allow columnHeaders to contain objects with a display label", () => {
+        const config = mkConfig({
+          columnHeaders: ["name", { key: "date", displayLabel: "Date" }],
+        });
+        const nameAndDate = addHeaders(config, [
+          "name",
+          { key: "date", displayLabel: "Date" },
+        ])(mkCsvOutput(""));
+        expect(nameAndDate).toEqual('"name","Date"' + endOfLine);
+      });
+    });
   });
 
   describe("addBody", () => {
@@ -171,6 +184,16 @@ describe("Helpers", () => {
       )(mkCsvOutput(""));
       expect(nameAndDate).toEqual('"rouky","2023-09-02"' + endOfLine);
     });
+
+    it("should build csv body with pretty headers", () => {
+      const config = mkConfig({});
+      const nameAndDate = addBody(
+        config,
+        ["name", { key: "date", displayLabel: "Date" }],
+        [{ name: "rouky", date: "2023-09-02" }],
+      )(mkCsvOutput(""));
+      expect(nameAndDate).toEqual('"rouky","2023-09-02"' + endOfLine);
+    });
   });
 
   describe("formatData", () => {
diff --git a/lib/helpers.ts b/lib/helpers.ts
index 84b6491..21afb45 100644
--- a/lib/helpers.ts
+++ b/lib/helpers.ts
@@ -1,17 +1,34 @@
 import { byteOrderMark, endOfLine } from "./config";
 import { EmptyHeadersError } from "./errors";
 import {
+  ColumnHeader,
   ConfigOptions,
   CsvOutput,
   CsvRow,
+  HeaderDisplayLabel,
+  HeaderKey,
   Newtype,
   WithDefaults,
   mkCsvOutput,
   mkCsvRow,
+  mkHeaderDisplayLabel,
+  mkHeaderKey,
   pack,
   unpack,
 } from "./types";
 
+const getHeaderKey = (columnHeader: ColumnHeader): HeaderKey =>
+  typeof columnHeader === "object"
+    ? mkHeaderKey(columnHeader.key)
+    : mkHeaderKey(columnHeader);
+
+const getHeaderDisplayLabel = (
+  columnHeader: ColumnHeader,
+): HeaderDisplayLabel =>
+  typeof columnHeader === "object"
+    ? mkHeaderDisplayLabel(columnHeader.displayLabel)
+    : mkHeaderDisplayLabel(columnHeader);
+
 export const thread = <T>(initialValue: T, ...fns: Array<Function>): T =>
   fns.reduce((r, fn) => fn(r), initialValue);
 
@@ -41,7 +58,7 @@ export const addFieldSeparator =
     pack<T>(unpack(output) + config.fieldSeparator);
 
 export const addHeaders =
-  (config: WithDefaults<ConfigOptions>, headers: Array<string>) =>
+  (config: WithDefaults<ConfigOptions>, headers: Array<ColumnHeader>) =>
   (output: CsvOutput): CsvOutput => {
     if (!config.showColumnHeaders) {
       return output;
@@ -55,7 +72,8 @@ export const addHeaders =
 
     let row = mkCsvRow("");
     for (let keyPos = 0; keyPos < headers.length; keyPos++) {
-      row = buildRow(config)(row, formatData(config, headers[keyPos]));
+      const header = getHeaderDisplayLabel(headers[keyPos]);
+      row = buildRow(config)(row, formatData(config, header));
     }
 
     row = mkCsvRow(unpack(row).slice(0, -1));
@@ -65,7 +83,7 @@ export const addHeaders =
 export const addBody =
   <T extends Array<{ [k: string]: unknown }>>(
     config: WithDefaults<ConfigOptions>,
-    headers: Array<string>,
+    headers: Array<ColumnHeader>,
     bodyData: T,
   ) =>
   (output: CsvOutput): CsvOutput => {
@@ -73,11 +91,11 @@ export const addBody =
     for (var i = 0; i < bodyData.length; i++) {
       let row = mkCsvRow("");
       for (let keyPos = 0; keyPos < headers.length; keyPos++) {
-        const header = headers[keyPos];
+        const header = getHeaderKey(headers[keyPos]);
         const data =
-          typeof bodyData[i][header] === "undefined"
+          typeof bodyData[i][unpack(header)] === "undefined"
             ? config.replaceUndefinedWith
-            : bodyData[i][header];
+            : bodyData[i][unpack(header)];
         row = buildRow(config)(row, formatData(config, data));
       }
 
diff --git a/lib/types.ts b/lib/types.ts
index bd00e99..07f1fee 100644
--- a/lib/types.ts
+++ b/lib/types.ts
@@ -5,6 +5,8 @@ export type Newtype<URI, A> = {
 
 export type WithDefaults<T> = Required<T>;
 
+export type ColumnHeader = string | { key: string; displayLabel: string };
+
 export type ConfigOptions = {
   filename?: string;
   fieldSeparator?: string;
@@ -16,14 +18,20 @@ export type ConfigOptions = {
   title?: string;
   useTextFile?: boolean;
   useBom?: boolean;
-  columnHeaders?: Array<string>;
+  columnHeaders?: Array<ColumnHeader>;
   useKeysAsHeaders?: boolean;
   boolDisplay?: { true: string; false: string };
   replaceUndefinedWith?: string | boolean | null;
 };
 
-export interface CsvOutput
-  extends Newtype<{ readonly CsvOutput: unique symbol }, string> {}
+export type HeaderKey = Newtype<{ readonly HeaderKey: unique symbol }, string>;
+
+export type HeaderDisplayLabel = Newtype<
+  { readonly HeaderDisplayLabel: unique symbol },
+  string
+>;
+
+export type CsvOutput = Newtype<{ readonly CsvOutput: unique symbol }, string>;
 
 export type CsvRow = Newtype<{ readonly CsvRow: unique symbol }, string>;
 
@@ -37,3 +45,5 @@ export const unpack = <T extends Newtype<any, any>>(newtype: T): T["_A"] =>
 
 export const mkCsvOutput = pack<CsvOutput>;
 export const mkCsvRow = pack<CsvRow>;
+export const mkHeaderKey = pack<HeaderKey>;
+export const mkHeaderDisplayLabel = pack<HeaderDisplayLabel>;