From 16535cbd26c81e56962a85ff1fd6fb76a7ce7489 Mon Sep 17 00:00:00 2001
From: Patrick McLaughlin <patrickmclaughlin@bitgo.com>
Date: Wed, 6 Sep 2023 14:20:21 -0400
Subject: [PATCH] feat: allow overriding internal custom codecs

---
 packages/openapi-generator/src/codec.ts       |  4 +-
 packages/openapi-generator/src/project.ts     |  7 +-
 .../openapi-generator/test/project.test.ts    | 95 +++++++++++++++++++
 .../openapi-generator/test/testProject.ts     |  5 +-
 4 files changed, 106 insertions(+), 5 deletions(-)
 create mode 100644 packages/openapi-generator/test/project.test.ts

diff --git a/packages/openapi-generator/src/codec.ts b/packages/openapi-generator/src/codec.ts
index 9327208f..7c687471 100644
--- a/packages/openapi-generator/src/codec.ts
+++ b/packages/openapi-generator/src/codec.ts
@@ -32,7 +32,7 @@ function codecIdentifier(
     } else if (imp.type === 'star') {
       return E.left(`Tried to use star import as codec ${id.value}`);
     }
-    const knownImport = project.knownImports[imp.from]?.[imp.importedName];
+    const knownImport = project.resolveKnownImport(imp.from, imp.importedName);
     if (knownImport !== undefined) {
       return E.right({ type: 'codec', schema: knownImport });
     }
@@ -67,7 +67,7 @@ function codecIdentifier(
     }
 
     const name = id.property.value;
-    const knownImport = project.knownImports[objectSym.from]?.[name];
+    const knownImport = project.resolveKnownImport(objectSym.from, name);
     if (knownImport !== undefined) {
       return E.right({ type: 'codec', schema: knownImport });
     }
diff --git a/packages/openapi-generator/src/project.ts b/packages/openapi-generator/src/project.ts
index c40f2d61..227b8eb5 100644
--- a/packages/openapi-generator/src/project.ts
+++ b/packages/openapi-generator/src/project.ts
@@ -10,7 +10,7 @@ import { parseSource, type SourceFile } from './sourceFile';
 const readFile = promisify(fs.readFile);
 
 export class Project {
-  readonly knownImports: Record<string, Record<string, KnownCodec>>;
+  private readonly knownImports: Record<string, Record<string, KnownCodec>>;
 
   private files: Record<string, SourceFile>;
 
@@ -90,4 +90,9 @@ export class Project {
       }
     }
   }
+
+  resolveKnownImport(path: string, name: string): KnownCodec | undefined {
+    const baseKey = path.startsWith('.') ? '.' : path;
+    return this.knownImports[baseKey]?.[name];
+  }
 }
diff --git a/packages/openapi-generator/test/project.test.ts b/packages/openapi-generator/test/project.test.ts
new file mode 100644
index 00000000..906c00cb
--- /dev/null
+++ b/packages/openapi-generator/test/project.test.ts
@@ -0,0 +1,95 @@
+import * as E from 'fp-ts/lib/Either';
+import assert from 'node:assert';
+import test from 'node:test';
+
+import { TestProject } from './testProject';
+import { parsePlainInitializer, type Schema } from '../src';
+import { KNOWN_IMPORTS, type KnownImports } from '../src/knownImports';
+
+async function testCase(
+  description: string,
+  src: string,
+  knownImports: KnownImports,
+  expected: Record<string, Schema>,
+  expectedErrors: string[] = [],
+) {
+  test(description, async () => {
+    const project = new TestProject(
+      { '/index.ts': src },
+      { ...KNOWN_IMPORTS, ...knownImports },
+    );
+    await project.parseEntryPoint('/index.ts');
+    const sourceFile = project.get('/index.ts');
+    if (sourceFile === undefined) {
+      throw new Error('Source file not found');
+    }
+
+    const actual: Record<string, Schema> = {};
+    const errors: string[] = [];
+    for (const symbol of sourceFile.symbols.declarations) {
+      if (symbol.init !== undefined) {
+        const result = parsePlainInitializer(project, sourceFile, symbol.init);
+        if (E.isLeft(result)) {
+          errors.push(result.left);
+        } else {
+          if (symbol.comment !== undefined) {
+            result.right.comment = symbol.comment;
+          }
+          actual[symbol.name] = result.right;
+        }
+      }
+    }
+
+    assert.deepStrictEqual(errors, expectedErrors);
+    assert.deepStrictEqual(actual, expected);
+  });
+}
+
+const EXTERNAL_CUSTOM_CODEC: KnownImports = {
+  foo: {
+    bar: () => E.right({ type: 'primitive', value: 'string' }),
+  },
+};
+
+const EXTERNAL_CUSTOM_CODEC_SRC = `
+import * as f from 'foo';
+
+export const FOO = f.bar;
+`;
+
+testCase(
+  'External custom codecs are parsed',
+  EXTERNAL_CUSTOM_CODEC_SRC,
+  EXTERNAL_CUSTOM_CODEC,
+  {
+    FOO: { type: 'primitive', value: 'string' },
+  },
+);
+
+const INTERNAL_CODEC_OVERRIDE: KnownImports = {
+  '.': {
+    bar: () => E.right({ type: 'primitive', value: 'string' }),
+  },
+};
+
+const INTERNAL_CODEC_OVERRIDE_SRC = `
+import * as t from 'io-ts';
+import { bar } from './bar';
+
+export const FOO = t.type({ bar: bar });
+`;
+
+testCase(
+  'Internal codec overrides are parsed',
+  INTERNAL_CODEC_OVERRIDE_SRC,
+  INTERNAL_CODEC_OVERRIDE,
+  {
+    FOO: {
+      type: 'object',
+      properties: {
+        bar: { type: 'primitive', value: 'string' },
+      },
+      required: ['bar'],
+    },
+  },
+);
diff --git a/packages/openapi-generator/test/testProject.ts b/packages/openapi-generator/test/testProject.ts
index cd48a2c0..235489f3 100644
--- a/packages/openapi-generator/test/testProject.ts
+++ b/packages/openapi-generator/test/testProject.ts
@@ -4,12 +4,13 @@ import resolve from 'resolve';
 import { promisify } from 'util';
 
 import { Project } from '../src';
+import type { KnownImports } from '../src/knownImports';
 
 export class TestProject extends Project {
   private volume: ReturnType<(typeof Volume)['fromJSON']>;
 
-  constructor(files: NestedDirectoryJSON) {
-    super();
+  constructor(files: NestedDirectoryJSON, knownImports?: KnownImports) {
+    super({}, knownImports);
     this.volume = Volume.fromNestedJSON(files, '/');
   }