From bf298249f099463565a23cbfb006e944846cc145 Mon Sep 17 00:00:00 2001
From: Gadi Cohen <dragon@wastelands.net>
Date: Wed, 28 Apr 2021 16:10:22 +0300
Subject: [PATCH] feat(quote): initial returnAs support

---
 schema.json               | 11 +++++++++++
 src/modules/quote.spec.ts | 38 ++++++++++++++++++++++++++++++++++++++
 src/modules/quote.ts      | 26 +++++++++++++++++++++++---
 3 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/schema.json b/schema.json
index 788a2298..1e5cc815 100644
--- a/schema.json
+++ b/schema.json
@@ -4719,6 +4719,9 @@
             "$ref": "#/definitions/QuoteField"
           },
           "type": "array"
+        },
+        "returnAs": {
+          "$ref": "#/definitions/ResultType"
         }
       },
       "type": "object"
@@ -5100,6 +5103,14 @@
       ],
       "type": "string"
     },
+    "ResultType": {
+      "enum": [
+        "array",
+        "object",
+        "map"
+      ],
+      "type": "string"
+    },
     "RevenueEstimate": {
       "additionalProperties": false,
       "properties": {
diff --git a/src/modules/quote.spec.ts b/src/modules/quote.spec.ts
index d7c72031..389ffd27 100644
--- a/src/modules/quote.spec.ts
+++ b/src/modules/quote.spec.ts
@@ -80,4 +80,42 @@ describe("quote", () => {
     expect(result.symbol).toBe("TSLA");
     expect(result.displayName).toBeDefined();
   });
+
+  describe("returnAs", () => {
+    it("array", async () => {
+      const devel = "quote-AAPL-BABA.json";
+      const results = await yf.quote(
+        ["AAPL", "BABA"],
+        { returnAs: "array" },
+        { devel }
+      );
+      expect(results.length).toBe(2);
+      expect(results[0].symbol).toBe("AAPL");
+      expect(results[1].symbol).toBe("BABA");
+    });
+
+    it("object", async () => {
+      const devel = "quote-AAPL-BABA.json";
+      const results = await yf.quote(
+        ["AAPL", "BABA"],
+        { returnAs: "object" },
+        { devel }
+      );
+      expect(Object.keys(results).length).toBe(2);
+      expect(results.AAPL.symbol).toBe("AAPL");
+      expect(results.BABA.symbol).toBe("BABA");
+    });
+
+    it("map", async () => {
+      const devel = "quote-AAPL-BABA.json";
+      const results = await yf.quote(
+        ["AAPL", "BABA"],
+        { returnAs: "map" },
+        { devel }
+      );
+      expect(results.size).toBe(2);
+      expect(results.get("AAPL").symbol).toBe("AAPL");
+      expect(results.get("BABA").symbol).toBe("BABA");
+    });
+  });
 });
diff --git a/src/modules/quote.ts b/src/modules/quote.ts
index 8d323812..c164b682 100644
--- a/src/modules/quote.ts
+++ b/src/modules/quote.ts
@@ -155,8 +155,11 @@ export type QuoteResponse = Quote[];
 
 export type QuoteField = keyof Quote;
 
+export type ResultType = "array" | "object" | "map";
+
 export interface QuoteOptions {
   fields?: QuoteField[];
+  returnAs?: ResultType;
 }
 
 const queryOptionsDefaults = {};
@@ -189,6 +192,7 @@ export default function quote(
   moduleOptions?: ModuleOptions
 ): Promise<any> {
   const symbols = typeof query === "string" ? query : query.join(",");
+  const returnAs = queryOptionsOverrides && queryOptionsOverrides.returnAs;
 
   return this._moduleExec({
     moduleName: "quote",
@@ -202,6 +206,7 @@ export default function quote(
       transformWith(queryOptions: QuoteOptions) {
         // Options validation ensures this is a string[]
         if (queryOptions.fields) queryOptions.fields.join(",");
+        delete queryOptions.returnAs;
         return queryOptions;
       },
     },
@@ -217,8 +222,23 @@ export default function quote(
 
     moduleOptions,
   }).then((results: Quote[]) => {
-    return typeof query === "string"
-      ? (results[0] as Quote)
-      : (results as Quote[]);
+    if (returnAs) {
+      switch (returnAs) {
+        case "array":
+          return results as Quote[];
+        case "object":
+          const object = {} as any;
+          for (let result of results) object[result.symbol] = result;
+          return object; // TODO: type
+        case "map":
+          const map = new Map();
+          for (let result of results) map.set(result.symbol, result);
+          return map; // TODO: type
+      }
+    } else {
+      return typeof query === "string"
+        ? (results[0] as Quote)
+        : (results as Quote[]);
+    }
   });
 }