Typescript: How to generate and print an AST based on a data structure

Question

I'm starting a new project, and as part of its interface we have a whole bunch of "tokens", a recursive object with string values, like so:

const colors = {
  accent: '#f90',
  primary: {
    active: '#fff',
    inactive: 'silver'
  }
};

We offer a utility for using those through string-based path (e.g., primary.active for #fff in this case). Extracting all possible paths into an array is easy enough, but what we'd like to offer is better autocomplete for the consumers of this package, so rather than 'string', a union or enum of these possible paths. Does anybody maybe have experience with this? My initial approach would be to write a simple script that takes an array and prints it as a union using a template or some such, but given that we want to do this more often and our use-cases will increase in complexity, I think generating and printing an AST is maybe a better approach. I've written babel and recast codemods before, I'm just looking for some guidance with regards to existing toolsets, examples etc. I've done a quick Google but couldn't find anything. Ideally these will recompile along with my normal "watch" process, but that's a stretch-goal ^_^.

user_125734 · Answer

You can extract the object type and create the union types using the compiler API

import * as ts from 'typescript'
import * as fs from 'fs'

var cmd = ts.parseCommandLine(['test.ts']); // replace with target file
// Create the program
let program = ts.createProgram(cmd.fileNames, cmd.options);


type ObjectDictionary = { [key: string]: string | ObjectDictionary}
function extractAllObjects(program: ts.Program, file: ts.SourceFile): ObjectDictionary {
    let empty = ()=> {};
    // Dummy transformation context
    let context: ts.TransformationContext = {
        startLexicalEnvironment: empty,
        suspendLexicalEnvironment: empty,
        resumeLexicalEnvironment: empty,
        endLexicalEnvironment: ()=> [],
        getCompilerOptions: ()=> program.getCompilerOptions(),
        hoistFunctionDeclaration: empty,
        hoistVariableDeclaration: empty,
        readEmitHelpers: ()=>undefined,
        requestEmitHelper: empty,
        enableEmitNotification: empty,
        enableSubstitution: empty,
        isEmitNotificationEnabled: ()=> false,
        isSubstitutionEnabled: ()=> false,
        onEmitNode: empty,
        onSubstituteNode: (hint, node)=>node,
    };
    let typeChecker =  program.getTypeChecker();

    function extractObject(node: ts.ObjectLiteralExpression): ObjectDictionary {
        var result : ObjectDictionary = {};
        for(let propDeclaration of node.properties){            
            if(!ts.isPropertyAssignment( propDeclaration )) continue;
            const propName = propDeclaration.name.getText()
            if(!propName) continue;
            if(ts.isObjectLiteralExpression(propDeclaration.initializer)) {
                result[propName] = extractObject(propDeclaration.initializer);
            }else{
                result[propName] = propDeclaration.initializer.getFullText()
            }
        }
        return result;
    }
    let foundVariables: ObjectDictionary = {};
    function visit(node: ts.Node, context: ts.TransformationContext): ts.Node {
        if(ts.isVariableDeclarationList(node)) {
            let triviaWidth = node.getLeadingTriviaWidth()
            let sourceText = node.getSourceFile().text;
            let trivia = sourceText.substr(node.getFullStart(), triviaWidth);
            if(trivia.indexOf("Generate_Union") != -1) // Will generate fro variables with a comment Generate_Union above them
            {
                for(let declaration of node.declarations) {
                    if(declaration.initializer && ts.isObjectLiteralExpression(declaration.initializer)){
                        foundVariables[declaration.name.getText()] = extractObject(declaration.initializer)
                    }
                }
            }
        }
        return ts.visitEachChild(node, child => visit(child, context), context);
    }
    ts.visitEachChild(file, child => visit(child, context), context);
    return foundVariables;
}



let result = extractAllObjects(program, program.getSourceFile("test.ts")!); // replace with file name 

function generateUnions(dic: ObjectDictionary) {
    function toPaths(dic: ObjectDictionary) : string[] {
        let result: string[] = []
        function extractPath(parent: string, object: ObjectDictionary) {
            for (const key of  Object.keys(object)) {
                let value = object[key]; 
                if(typeof value === "string") {
                    result.push(parent + key);
                }else{
                    extractPath(key + ".", value);
                }
            }
        }
        extractPath("", dic);
        return result;
    }

    return Object.entries(dic)
        .map(([name, values])=> 
        {
            let paths = toPaths(values as ObjectDictionary)
                .map(ts.createStringLiteral)
                .map(ts.createLiteralTypeNode);

            let unionType = ts.createUnionTypeNode(paths);
            return ts.createTypeAliasDeclaration(undefined, undefined, name + "Paths", undefined, unionType);
        })

}

var source = ts.createSourceFile("d.ts", "", ts.ScriptTarget.ES2015);
source = ts.updateSourceFileNode(source, generateUnions(result));

var printer = ts.createPrinter({ });
let r = printer.printFile(source);
fs.writeFileSync("union.ts", r);

user_7549119 · Answer

I think that you can accomplish what you want with a combination of enums and interfaces/types:

``` 
export enum COLORS {
    accent = '#f90',
    primary_active = '#fff',
    primary_inactive = 'silver',
}

interface ICOLORS {
    [COLORS.accent]: COLORS.accent,
    [COLORS.primary_active]: COLORS.primary_active,
    [COLORS.primary_inactive]: COLORS.primary_inactive
}

export type COLOR_OPTIONS = keyof ICOLORS;

export type PRIMARY_COLOR_OPTIONS = keyof Pick;

export function setColor (color: PRIMARY_COLOR_OPTIONS): void {}

// elsewhere:

import {COLORS, setColor} from 'somewhere';

setColor(COLORS.primary_inactive); // works

setColor(COLORS.accent); //error

```

Typescript: How to generate and print an AST based on a data structure

2 Answers