allow 'transform()' to transform arbitrary nodes.

Author: rbuckton

src/compiler/emitter.ts

@@ -10,7 +10,7 @@ namespace ts {
 
     /*@internal*/
     // targetSourceFile is when users only want one file in entire project to be emitted. This is used in compileOnSave feature
-    export function emitFiles(resolver: EmitResolver, host: EmitHost, targetSourceFile: SourceFile, emitOnlyDtsFiles?: boolean, transformers?: Transformer[]): EmitResult {
+    export function emitFiles(resolver: EmitResolver, host: EmitHost, targetSourceFile: SourceFile, emitOnlyDtsFiles?: boolean, transformers?: TransformerFactory<SourceFile>[]): EmitResult {
         const compilerOptions = host.getCompilerOptions();
         const moduleKind = getEmitModuleKind(compilerOptions);
         const sourceMapDataList: SourceMapData[] = compilerOptions.sourceMap || compilerOptions.inlineSourceMap ? [] : undefined;
@@ -28,7 +28,7 @@ namespace ts {
         const sourceFiles = getSourceFilesToEmit(host, targetSourceFile);
 
         // Transform the source files
-        const transform = transformFiles(resolver, host, sourceFiles, transformers);
+        const transform = transformNodes(resolver, host, compilerOptions, sourceFiles, transformers, /*allowDtsFiles*/ false);
 
         // Create a printer to print the nodes
         const printer = createPrinter(compilerOptions, {

src/compiler/transformer.ts

@@ -13,7 +13,7 @@
 
 /* @internal */
 namespace ts {
-    function getModuleTransformer(moduleKind: ModuleKind): Transformer {
+    function getModuleTransformer(moduleKind: ModuleKind): TransformerFactory<SourceFile> {
         switch (moduleKind) {
             case ModuleKind.ES2015:
                 return transformES2015Module;
@@ -40,7 +40,7 @@ namespace ts {
         const jsx = compilerOptions.jsx;
         const languageVersion = getEmitScriptTarget(compilerOptions);
         const moduleKind = getEmitModuleKind(compilerOptions);
-        const transformers: Transformer[] = [];
+        const transformers: TransformerFactory<SourceFile>[] = [];
 
         addRange(transformers, customTransformers && customTransformers.before);
 
@@ -84,11 +84,13 @@ namespace ts {
      * Transforms an array of SourceFiles by passing them through each transformer.
      *
      * @param resolver The emit resolver provided by the checker.
-     * @param host The emit host.
-     * @param sourceFiles An array of source files
-     * @param transforms An array of Transformers.
+     * @param host The emit host object used to interact with the file system.
+     * @param options Compiler options to surface in the `TransformationContext`.
+     * @param nodes An array of nodes to transform.
+     * @param transforms An array of `TransformerFactory` callbacks.
+     * @param allowDtsFiles A value indicating whether to allow the transformation of .d.ts files.
      */
-    export function transformFiles(resolver: EmitResolver, host: EmitHost, sourceFiles: SourceFile[], transformers: Transformer[]): TransformationResult {
+    export function transformNodes<T extends Node>(resolver: EmitResolver, host: EmitHost, options: CompilerOptions, nodes: T[], transformers: TransformerFactory<T>[], allowDtsFiles: boolean): TransformationResult<T> {
         const enabledSyntaxKindFeatures = new Array<SyntaxKindFeatureFlags>(SyntaxKind.Count);
         let lexicalEnvironmentVariableDeclarations: VariableDeclaration[];
         let lexicalEnvironmentFunctionDeclarations: FunctionDeclaration[];
@@ -104,7 +106,7 @@ namespace ts {
         // The transformation context is provided to each transformer as part of transformer
         // initialization.
         const context: TransformationContext = {
-            getCompilerOptions: () => host.getCompilerOptions(),
+            getCompilerOptions: () => options,
             getEmitResolver: () => resolver,
             getEmitHost: () => host,
             startLexicalEnvironment,
@@ -134,7 +136,9 @@ namespace ts {
         };
 
         // Ensure the parse tree is clean before applying transformations
-        forEach(sourceFiles, disposeEmitNodes);
+        for (const node of nodes) {
+            disposeEmitNodes(getSourceFileOfNode(getParseTreeNode(node)));
+        }
 
         performance.mark("beforeTransform");
 
@@ -144,8 +148,8 @@ namespace ts {
         // prevent modification of transformation hooks.
         state = TransformationState.Initialized;
 
-        // Transform each source file.
-        const transformed = map(sourceFiles, transformSourceFile);
+        // Transform each node.
+        const transformed = map(nodes, allowDtsFiles ? transformation : transformRoot);
 
         // prevent modification of the lexical environment.
         state = TransformationState.Completed;
@@ -160,17 +164,8 @@ namespace ts {
             dispose
         };
 
-        /**
-         * Transforms a source file.
-         *
-         * @param sourceFile The source file to transform.
-         */
-        function transformSourceFile(sourceFile: SourceFile) {
-            if (isDeclarationFile(sourceFile)) {
-                return sourceFile;
-            }
-
-            return transformation(sourceFile);
+        function transformRoot(node: T) {
+            return node && (!isSourceFile(node) || !isDeclarationFile(node)) ? transformation(node) : node;
         }
 
         /**
@@ -366,7 +361,9 @@ namespace ts {
         function dispose() {
             if (state < TransformationState.Disposed) {
                 // Clean up emit nodes on parse tree
-                forEach(sourceFiles, disposeEmitNodes);
+                for (const node of nodes) {
+                    disposeEmitNodes(getSourceFileOfNode(getParseTreeNode(node)));
+                }
 
                 // Release references to external entries for GC purposes.
                 lexicalEnvironmentVariableDeclarations = undefined;

src/compiler/types.ts

@@ -2308,9 +2308,9 @@
 
     export interface CustomTransformers {
         /** Custom transformers to evaluate before built-in transformations. */
-        before?: Transformer[];
+        before?: TransformerFactory<SourceFile>[];
         /** Custom transformers to evaluate after built-in transformations. */
-        after?: Transformer[];
+        after?: TransformerFactory<SourceFile>[];
     }
 
     export interface SourceMapSpan {
@@ -3867,7 +3867,7 @@
          * are emitted by the pretty printer.
          *
          * NOTE: Transformation hooks should only be modified during `Transformer` initialization,
-         * before returning the `FileTransformer` callback.
+         * before returning the `NodeTransformer` callback.
          */
         onSubstituteNode?: (hint: EmitHint, node: Node) => Node;
 
@@ -3888,14 +3888,14 @@
          * the printer emits a node.
          *
          * NOTE: Transformation hooks should only be modified during `Transformer` initialization,
-         * before returning the `FileTransformer` callback.
+         * before returning the `NodeTransformer` callback.
          */
         onEmitNode?: (hint: EmitHint, node: Node, emitCallback: (hint: EmitHint, node: Node) => void) => void;
     }
 
-    export interface TransformationResult {
+    export interface TransformationResult<T extends Node> {
         /** Gets the transformed source files. */
-        transformed: SourceFile[];
+        transformed: T[];
 
         /** Gets diagnostics for the transformation. */
         diagnostics?: Diagnostic[];
@@ -3925,23 +3925,23 @@
     }
 
     /**
-     * A function that is used to initialize and return a `FileTransformer` callback, which in turn
-     * will be used to transform one or more `SourceFile` objects.
+     * A function that is used to initialize and return a `Transformer` callback, which in turn
+     * will be used to transform one or more nodes.
      */
-    export type Transformer = (context: TransformationContext) => FileTransformer;
+    export type TransformerFactory<T extends Node> = (context: TransformationContext) => Transformer<T>;
 
     /**
-     * A function that transforms a `SourceFile` object.
+     * A function that transforms a node.
      */
-    export type FileTransformer = (node: SourceFile) => SourceFile;
-
-    export type VisitResult<T extends Node> = T | T[];
+    export type Transformer<T extends Node> = (node: T) => T;
 
     /**
      * A function that accepts and possible transforms a node.
      */
     export type Visitor = (node: Node) => VisitResult<Node>;
 
+    export type VisitResult<T extends Node> = T | T[];
+
     export interface Printer {
         /**
          * Print a node and its subtree as-is, without any emit transformations.

src/harness/unittests/customTransforms.ts

@@ -47,7 +47,7 @@ namespace ts {
             `
         }];
 
-        const before: Transformer = context => {
+        const before: TransformerFactory<SourceFile> = context => {
             return file => visitEachChild(file, visit, context);
             function visit(node: Node): VisitResult<Node> {
                 switch (node.kind) {
@@ -63,7 +63,7 @@ namespace ts {
             }
         };
 
-        const after: Transformer = context => {
+        const after: TransformerFactory<SourceFile> = context => {
             return file => visitEachChild(file, visit, context);
             function visit(node: Node): VisitResult<Node> {
                 switch (node.kind) {

src/harness/unittests/transform.ts

@@ -3,7 +3,7 @@
 
 namespace ts {
     describe("TransformAPI", () => {
-        function transformsCorrectly(name: string, source: string, transformers: Transformer[]) {
+        function transformsCorrectly(name: string, source: string, transformers: TransformerFactory<SourceFile>[]) {
             it(name, () => {
                 Harness.Baseline.runBaseline(`transformApi/transformsCorrectly.${name}.js`, () => {
                     const transformed = transform(createSourceFile("source.ts", source, ScriptTarget.ES2015), transformers);

src/services/transform.ts

@@ -2,31 +2,16 @@
 /// <reference path="transpile.ts"/>
 namespace ts {
     /**
-     * Transform one or more source files using the supplied transformers.
-     * @param source A `SourceFile` or an array of `SourceFiles`.
-     * @param transformers An array of `Transformer` callbacks used to process the transformation.
+     * Transform one or more nodes using the supplied transformers.
+     * @param source A single `Node` or an array of `Node` objects.
+     * @param transformers An array of `TransformerFactory` callbacks used to process the transformation.
      * @param compilerOptions Optional compiler options.
      */
-    export function transform(source: SourceFile | SourceFile[], transformers: Transformer[], compilerOptions?: CompilerOptions) {
+    export function transform<T extends Node>(source: T | T[], transformers: TransformerFactory<T>[], compilerOptions?: CompilerOptions) {
         const diagnostics: Diagnostic[] = [];
         compilerOptions = fixupCompilerOptions(compilerOptions, diagnostics);
-        const newLine = getNewLineCharacter(compilerOptions);
-        const sourceFiles = isArray(source) ? source : [source];
-        const fileMap = arrayToMap(sourceFiles, sourceFile => sourceFile.fileName);
-        const emitHost: EmitHost = {
-            getCompilerOptions: () => compilerOptions,
-            getCanonicalFileName: fileName => fileName,
-            getCommonSourceDirectory: () => "",
-            getCurrentDirectory: () => "",
-            getNewLine: () => newLine,
-            getSourceFile: fileName => fileMap.get(fileName),
-            getSourceFileByPath: fileName => fileMap.get(fileName),
-            getSourceFiles: () => sourceFiles,
-            isSourceFileFromExternalLibrary: () => false,
-            isEmitBlocked: () => false,
-            writeFile: () => Debug.fail("'writeFile()' is not supported during transformation.")
-        };
-        const result = transformFiles(/*resolver*/ undefined, emitHost, sourceFiles, transformers);
+        const nodes = isArray(source) ? source : [source];
+        const result = transformNodes(/*resolver*/ undefined, /*emitHost*/ undefined, compilerOptions, nodes, transformers, /*allowDtsFiles*/ true);
         result.diagnostics = concatenate(result.diagnostics, diagnostics);
         return result;
     }