From 12ff53d73e5af47321f17f221396f529891cb63e Mon Sep 17 00:00:00 2001
From: Mikhail Bashurov <mishinomb@gmail.com>
Date: Mon, 16 Sep 2019 08:05:31 +0900
Subject: [PATCH] Introduce transformers program support (#879)

---
 README.md    |  2 +-
 src/index.ts | 68 ++++++++++++++++++++++++++++++++--------------------
 2 files changed, 43 insertions(+), 27 deletions(-)

diff --git a/README.md b/README.md
index bed60d1c6..a866f73b3 100644
--- a/README.md
+++ b/README.md
@@ -140,7 +140,7 @@ _Environment variable denoted in parentheses._
 
 ### Programmatic Only Options
 
-* `transformers` An array of transformers to pass to TypeScript
+* `transformers` `_ts.CustomTransformers | ((p: _ts.Program) => _ts.CustomTransformers)` An object with transformers or a function that accepts a program and returns an transformers object to pass to TypeScript. Function isn't available with `transpileOnly` flag
 * `readFile` Custom TypeScript-compatible file reading function
 * `fileExists` Custom TypeScript-compatible file existence function
 
diff --git a/src/index.ts b/src/index.ts
index 82aadb328..610969870 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -71,7 +71,7 @@ export interface Options {
   ignoreDiagnostics?: Array<number | string>
   readFile?: (path: string) => string | undefined
   fileExists?: (path: string) => boolean
-  transformers?: _ts.CustomTransformers
+  transformers?: _ts.CustomTransformers | ((p: _ts.Program) => _ts.CustomTransformers)
 }
 
 /**
@@ -274,32 +274,23 @@ export function register (opts: Options = {}): Register {
   /**
    * Create the basic required function using transpile mode.
    */
-  let getOutput = function (code: string, fileName: string, lineOffset = 0): SourceOutput {
-    const result = ts.transpileModule(code, {
-      fileName,
-      transformers,
-      compilerOptions: config.options,
-      reportDiagnostics: true
-    })
-
-    const diagnosticList = result.diagnostics ?
-      filterDiagnostics(result.diagnostics, ignoreDiagnostics) :
-      []
-
-    if (diagnosticList.length) reportTSError(configDiagnosticList)
-
-    return [result.outputText, result.sourceMapText as string]
-  }
-
-  let getTypeInfo = function (_code: string, _fileName: string, _position: number): TypeInfo {
-    throw new TypeError(`Type information is unavailable without "--type-check"`)
-  }
+  let getOutput: (code: string, fileName: string, lineOffset: number) => SourceOutput
+  let getTypeInfo: (_code: string, _fileName: string, _position: number) => TypeInfo
 
   // Use full language services when the fast option is disabled.
   if (typeCheck) {
     const memoryCache = new MemoryCache(config.fileNames)
     const cachedReadFile = cachedLookup(debugFn('readFile', readFile))
 
+    const getCustomTransformers = () => {
+      if (typeof transformers === 'function') {
+        const program = service.getProgram()
+        return program ? transformers(program) : undefined
+      }
+
+      return transformers
+    }
+
     // Create the compiler host for type checking.
     const serviceHost: _ts.LanguageServiceHost = {
       getScriptFileNames: () => Array.from(memoryCache.fileVersions.keys()),
@@ -331,14 +322,14 @@ export function register (opts: Options = {}): Register {
       getCurrentDirectory: () => cwd,
       getCompilationSettings: () => config.options,
       getDefaultLibFileName: () => ts.getDefaultLibFilePath(config.options),
-      getCustomTransformers: () => transformers
+      getCustomTransformers: getCustomTransformers
     }
 
     const registry = ts.createDocumentRegistry(ts.sys.useCaseSensitiveFileNames, cwd)
     const service = ts.createLanguageService(serviceHost, registry)
 
     // Set the file contents into cache manually.
-    const updateMemoryCache = function (contents: string, fileName: string) {
+    const updateMemoryCache = (contents: string, fileName: string) => {
       const fileVersion = memoryCache.fileVersions.get(fileName) || 0
 
       // Avoid incrementing cache when nothing has changed.
@@ -348,7 +339,7 @@ export function register (opts: Options = {}): Register {
       memoryCache.fileContents.set(fileName, contents)
     }
 
-    getOutput = function (code: string, fileName: string, lineOffset: number = 0) {
+    getOutput = (code: string, fileName: string) => {
       updateMemoryCache(code, fileName)
 
       const output = service.getEmitOutput(fileName)
@@ -379,7 +370,7 @@ export function register (opts: Options = {}): Register {
       return [output.outputFiles[1].text, output.outputFiles[0].text]
     }
 
-    getTypeInfo = function (code: string, fileName: string, position: number) {
+    getTypeInfo = (code: string, fileName: string, position: number) => {
       updateMemoryCache(code, fileName)
 
       const info = service.getQuickInfoAtPosition(fileName, position)
@@ -388,10 +379,35 @@ export function register (opts: Options = {}): Register {
 
       return { name, comment }
     }
+  } else {
+    if (typeof transformers === 'function') {
+      throw new TypeError('Transformers function is unavailable in "--transpile-only"')
+    }
+
+    getOutput = (code: string, fileName: string): SourceOutput => {
+      const result = ts.transpileModule(code, {
+        fileName,
+        transformers,
+        compilerOptions: config.options,
+        reportDiagnostics: true
+      })
+
+      const diagnosticList = result.diagnostics ?
+        filterDiagnostics(result.diagnostics, ignoreDiagnostics) :
+        []
+
+      if (diagnosticList.length) reportTSError(configDiagnosticList)
+
+      return [result.outputText, result.sourceMapText as string]
+    }
+
+    getTypeInfo = () => {
+      throw new TypeError('Type information is unavailable in "--transpile-only"')
+    }
   }
 
   // Create a simple TypeScript compiler proxy.
-  function compile (code: string, fileName: string, lineOffset?: number) {
+  function compile (code: string, fileName: string, lineOffset = 0) {
     const [value, sourceMap] = getOutput(code, fileName, lineOffset)
     const output = updateOutput(value, fileName, sourceMap, getExtension)
     outputCache.set(fileName, output)