schematic for ngrx8

ngrx8-update.ts

import { camelize, dasherize } from '@angular-devkit/core/src/utils/strings';
import { Rule, SchematicContext, Tree, UpdateRecorder } from '@angular-devkit/schematics';
import { tsquery } from '@phenomnomnominal/tsquery';
import { SyntaxKind } from 'typescript';
import * as ts from 'typescript';

// You don't have to export the function as default. You can also have more than one rule factory
// per file.
export function ngrx8Update(_options: any): Rule {
    return (_tree: Tree, _context: SchematicContext) => {
        const workspaceConfigBuffer = _tree.read('angular.json');
        if (workspaceConfigBuffer) {
            // const workspaceConfig = JSON.parse(workspaceConfigBuffer.toString());
        }
        _tree.getDir('./src').visit(path => {
            migrateActions(_tree, path);
            migrateEffects(_tree, path);
            migrateReducers(_tree, path);
            migrateTS(_tree, path);
            migrateRegistryTS(_tree, path);
            migrateModuleTS(_tree, path);
        });
        return _tree;
    };

    function migrateModuleTS(_tree: Tree, path: string) {
        if (path.match(/\.module\.ts/)) {
            const fileContent: Buffer | null = _tree.read(path);
            if (fileContent) {
                const newTS = ts.createSourceFile(getFilename(path), fileContent.toString(), ts.ScriptTarget.Latest, true);
                const declarationRecorder = _tree.beginUpdate(path);
                // const actions = tsquery(newTS, 'ImportClause ImportSpecifier:has(Identifier[name=/\\w+Actions/])');
                const imports = tsquery(newTS, 'ImportClause:has(Identifier[name=/\\w+Actions/])');
                imports.forEach(node => {
                    const newText = node.getText().replace(/\w+Actions,?/, '');
                    replace(declarationRecorder, node, newText);
                });
                const providersNodes = tsquery(newTS, 'PropertyAssignment:has(Identifier[name=providers])');
                providersNodes.forEach(node => {
                    const text = node.getText();
                    if (text.includes('Actions')) {
                        const newText = node.getText().replace(/(?:\.{3})?\w+Actions,?/, '');
                        replace(declarationRecorder, node, newText + ',');
                    }
                });
                _tree.commitUpdate(declarationRecorder);
            }
        }
    }
    function migrateRegistryTS(_tree: Tree, path: string) {
        if (path.match(/\.registry\.ts/)) {
            const fileContent: Buffer | null = _tree.read(path);
            if (fileContent) {
                const newTS = ts.createSourceFile(getFilename(path), fileContent.toString(), ts.ScriptTarget.Latest, true);
                const declarationRecorder = _tree.beginUpdate(path);
                const actions = tsquery(newTS, 'VariableStatement:has(Identifier[name=/\\w+Actions/])');
                if (actions.length) {
                    declarationRecorder.remove(actions[0].getStart(), actions[0].getWidth());
                    _tree.commitUpdate(declarationRecorder);
                }
            }
        }
    }
    function migrateTS(_tree: Tree, path: string) {
        if (path.match(/\.(component|service|guard|resolve)\.ts/)) {
            const fileContent: Buffer | null = _tree.read(path);
            if (fileContent) {
                const newTS = ts.createSourceFile(getFilename(path), fileContent.toString(), ts.ScriptTarget.Latest, true);
                const declarationRecorder = _tree.beginUpdate(path);
                const importedAction = migrateActionImports(newTS, declarationRecorder);
                const injectionTokens = removeActionInjectionFromCtor(newTS, importedAction, declarationRecorder);
                replaceActionMethod(injectionTokens, newTS, declarationRecorder);
                _tree.commitUpdate(declarationRecorder);
            }
        }
    }

    function migrateReducers(_tree: Tree, path: string) {
        if (path.endsWith('.reducer.ts')) {
            const fileContent: Buffer | null = _tree.read(path);
            if (fileContent) {
                const newTS = backupFile(path, fileContent, _tree);
                const declarationRecorder = _tree.beginUpdate(path);
                const importedAction = migrateActionImports(newTS, declarationRecorder);
                const functionNode = tsquery(newTS, 'FunctionDeclaration:has(Identifier[name=/\\w+Reducer/])')[0];
                insertNewReducer(newTS, declarationRecorder, functionNode);
                migrateReducerFunction(functionNode, declarationRecorder);
                declarationRecorder.insertRight(0, "import { createReducer, on, Action } from '@ngrx/store';\r\n");
                _tree.commitUpdate(declarationRecorder);
            }
        }

        function insertNewReducer(newTS: ts.SourceFile, declarationRecorder: UpdateRecorder, functionNode: ts.Node) {
            const caseNodes = tsquery(newTS, 'FunctionDeclaration CaseClause');
            let actions: string[] = [];
            const onStatements = caseNodes
                .map(node => {
                    const action = convertActionType(node.getChildAt(1).getText());
                    actions.push(action);
                    const reStatement = node.getChildAt(node.getChildCount() - 1).getText();
                    if (!reStatement) {
                        // no case statement because it uses next case statement
                        return undefined;
                    }
                    let reText = reStatement;
                    if (!reStatement.startsWith('{')) {
                        reText = `{${reStatement}}`;
                    }
                    let param = 'state';
                    if (reText.includes('action.payload')) {
                        reText = reText.replace(/action\.payload/g, 'payload');
                        param += ', {payload}';
                    }
                    const re = `on( ${actions.join(', ')}, (${param})=> ${reText})`;
                    actions = [];
                    return re;
                })
                .filter(x => !!x);
            const newReducer = `const featureReducer = createReducer( initialState, \r\n ${onStatements.join(', \r\n')});\r\n`;
            declarationRecorder.insertLeft(functionNode.getStart(), newReducer);
        }
    }

    function convertActionType(action: string) {
        const tokens = action.split('.');
        if (tokens.length === 3) {
            return `${tokens[0]}.${camelize(tokens[2].toLowerCase())}`;
        }
        return `${tokens[0]}`;
    }

    function migrateReducerFunction(functionNode: ts.Node, declarationRecorder: UpdateRecorder) {
        // const reducerNameNode = functionNode.getChildAt(2);
        const paramNode = functionNode.getChildAt(4);
        replace(declarationRecorder, paramNode, 'state, action: Action');
        const statement = functionNode.getChildAt(functionNode.getChildCount() - 1);
        replace(declarationRecorder, statement, '{\r\nreturn featureReducer(state, action);\r\n}');
    }

    function backupFile(path: string, fileContent: Buffer, _tree: Tree) {
        const newTS = ts.createSourceFile(getFilename(path), fileContent.toString(), ts.ScriptTarget.Latest, true);
        _tree.rename(path, path + '.old');
        _tree.create(path, newTS.getFullText());
        return newTS;
    }

    function migrateEffects(_tree: Tree, path: string) {
        if (path.endsWith('.effects.ts')) {
            const fileContent: Buffer | null = _tree.read(path);
            if (fileContent) {
                const newTS = backupFile(path, fileContent, _tree);
                const declarationRecorder = _tree.beginUpdate(path);
                const importedAction = migrateActionImports(newTS, declarationRecorder);
                addCreateEffectImport(newTS, declarationRecorder);
                const injectionTokens = removeActionInjectionFromCtor(newTS, importedAction, declarationRecorder);
                useCreateEffectFunction(newTS, declarationRecorder);
                migrateActionType(newTS, declarationRecorder);
                replaceAllActionTypeToAny(newTS, declarationRecorder);
                replaceActionMethod(injectionTokens, newTS, declarationRecorder);
                _tree.commitUpdate(declarationRecorder);
            }
        }

        function addCreateEffectImport(newEffectTS: ts.SourceFile, declarationRecorder: UpdateRecorder) {
            const importEffectNodes = tsquery(newEffectTS, 'ImportDeclaration Identifier[name=Effect]');
            if (importEffectNodes.length === 1) {
                const effectNode = importEffectNodes[0];
                declarationRecorder.remove(effectNode.getStart(), effectNode.getWidth());
                declarationRecorder.insertLeft(effectNode.getStart(), 'createEffect');
            }
        }
        function useCreateEffectFunction(newEffectTS: ts.SourceFile, declarationRecorder: UpdateRecorder) {
            const propertyNodes = tsquery(newEffectTS, 'PropertyDeclaration');
            propertyNodes.forEach(node => {
                const firstToken = node.getChildAt(0) as ts.Node;
                const first = firstToken.getText();
                if (!first.startsWith('@Effect')) {
                    return;
                }
                const isDispatch = !first.match(/dispatch:\s*false/);
                declarationRecorder.remove(firstToken.getStart(), firstToken.getWidth());
                if (node.getChildCount() === 5) {
                    const actionToken = node.getChildAt(3);
                    const actionText = actionToken.getText();
                    if (actionText.startsWith('this._Actions$')) {
                        declarationRecorder.insertLeft(actionToken.getStart(), 'createEffect(() =>\r\n    ');
                        if (isDispatch) {
                            declarationRecorder.insertRight(actionToken.getEnd(), ')');
                        } else {
                            declarationRecorder.insertRight(actionToken.getEnd(), ',\r\n    { dispatch: false }\r\n)');
                        }
                    }
                }
            });
        }
    }

    function replaceActionMethod(
        injectionTokens: { injectToken: string; action: string }[],
        newEffectTS: ts.SourceFile,
        declarationRecorder: UpdateRecorder
    ) {
        injectionTokens.forEach(token => {
            const nodes = tsquery(newEffectTS, `CallExpression Identifier[name=${token.injectToken}]`);
            nodes.forEach(node => {
                const targetNode = node.parent.parent.parent;
                const text = targetNode.getText();
                if (text.startsWith('this.')) {
                    let newText = text.replace('this.', '').replace(token.injectToken, token.action);
                    const re = newText.split(/\(|\)/);
                    if (re.length === 3 && re[1]) {
                        newText = `${re[0]}({payload: ${re[1]}})`;
                    }
                    replace(declarationRecorder, targetNode, newText);
                }
            });
        });
    }

    function replaceAllActionTypeToAny(newEffectTS: ts.SourceFile, declarationRecorder: UpdateRecorder) {
        const allActionsNodes = tsquery(newEffectTS, 'PropertyDeclaration Identifier[name=/All\\w+Actions$/]');
        allActionsNodes.forEach(action => {
            replace(declarationRecorder, action, 'any');
        });
    }

    function migrateActionType(newEffectTS: ts.SourceFile, declarationRecorder: UpdateRecorder) {
        const ofTypeNodes = tsquery(newEffectTS, 'PropertyDeclaration Identifier[name=ofType]');
        ofTypeNodes.forEach(node => {
            const typeNode = node.parent.getChildAt(2);
            const actions = typeNode.getText().split(',');
            const newAction = actions.map(action => convertActionType(action)).join(', ');
            replace(declarationRecorder, typeNode, newAction);
        });
    }

    function removeActionInjectionFromCtor(newEffectTS: ts.SourceFile, modifiedActions: (string | undefined)[], declarationRecorder: UpdateRecorder) {
        const ctorParamNodes = tsquery(newEffectTS, 'Constructor Parameter');
        const injectionTokens: { injectToken: string; action: string }[] = [];
        for (let i = 0; i < ctorParamNodes.length; i++) {
            const node = ctorParamNodes[i];
            const text = node.getText();
            if (modifiedActions.find(action => action && text.includes(action))) {
                const tokens = text.split(' ');
                injectionTokens.push({
                    injectToken: tokens[1].replace(':', ''),
                    action: tokens[2]
                });
                if (i === ctorParamNodes.length - 1) {
                    declarationRecorder.remove(node.getStart(), node.getWidth());
                } else {
                    // +1 is for the , in the end
                    declarationRecorder.remove(node.getStart(), node.getWidth() + 1);
                }
            }
        }
        return injectionTokens;
    }

    function migrateActionImports(newTS: ts.SourceFile, declarationRecorder: UpdateRecorder) {
        const actionImportNodes = tsquery(newTS, 'ImportDeclaration StringLiteral[value=/.+\\.actions$/]');
        const actions = actionImportNodes.map(node => {
            const token = node.parent.getChildAt(1);
            const text = token.getText();
            const actions = text.match(/\w+Actions/g);
            const action = (actions as string[]).find(a => !a.includes('All'));
            replace(declarationRecorder, token, `* as ${action}`);
            return action;
        });
        const actionImportNodes2 = tsquery(newTS, 'ImportSpecifier:has(Identifier[name=/\\w+Actions$/])');
        actionImportNodes2.forEach(node => {
            const text = node.getText();
            if (actions.includes(text) || text.includes('All')) {
                return;
            }
            actions.push(text);
            const importStatement = node.parent.parent.parent;
            const ImportClause = node.parent.parent;
            const newImportClause = node.parent.parent.getText().replace(/\w+Actions,?/g, '');
            replace(declarationRecorder, ImportClause, newImportClause);
            const actionPath =
                importStatement
                    .getChildAt(3)
                    .getText()
                    .slice(0, -1) +
                '/' +
                dasherize(text).replace('-actions', '.actions') +
                "'";
            declarationRecorder.insertRight(importStatement.getEnd(), `\r\nimport * as ${text} from ${actionPath};`);
        });
        return actions;
    }

    function migrateActions(_tree: Tree, path: string) {
        if (path.endsWith('.actions.ts')) {
            let newActionContent = "import { props, createAction } from '@ngrx/store';";
            const fileContent: Buffer | null = _tree.read(path);
            if (fileContent) {
                const filename = getFilename(path);
                const textContent = fileContent.toString();
                const ast = tsquery.ast(textContent);
                newActionContent += getActionImports(ast);
                newActionContent += getActionPrefix(ast, filename);
                const propertyNodes = tsquery(ast, 'PropertyDeclaration');
                const actionInfos = getActionInfo(propertyNodes);
                newActionContent += getActionStatements(actionInfos);
                const newActionTS = ts.createSourceFile(filename, newActionContent, ts.ScriptTarget.Latest, true);
                _tree.rename(path, path + '.old');
                _tree.create(path, newActionTS.getFullText());
            }
        }
    }
    function getFilename(path: string) {
        const pathTokens = path.split('/');
        return pathTokens[pathTokens.length - 1];
    }
    function getActionStatements(actionInfos: IActionInfo[]) {
        if (!actionInfos) {
            return;
        }
        const statements = actionInfos.map(item => {
            if (item.propType) {
                if (item.propType.includes('undefined')) {
                    const propType = item.propType.split('|')[0];
                    return `\r\nexport const ${item.methodName} = createAction( actionTypePrefix + '${item.actionCode}', (input?: { payload: ${propType} }) =>
          input ? { payload: input.payload } : undefined);`;
                }
                return `\r\nexport const ${item.methodName} = createAction( actionTypePrefix + '${item.actionCode}', props<{ payload: ${item.propType} }>());`;
            }
            return `\r\nexport const ${item.methodName} = createAction( actionTypePrefix + '${item.actionCode}');`;
        });
        return statements.join('\r\n');
    }

    function getActionImports(ast: ts.SourceFile) {
        let re = '';
        const importNodes = tsquery(ast, 'ImportDeclaration');
        importNodes.forEach(node => {
            const text = node.getFullText();
            if (text.match(/angular|millerslab/)) {
                return;
            }
            re += text;
        });
        return re;
    }

    function getActionInfo(propertyNodes: ts.Node[]) {
        const actions: IActionInfo[] = [];
        propertyNodes &&
            propertyNodes.forEach(node => {
                // const text = node.getText();
                const nodeChildren = node.getChildren();
                const firstChild = nodeChildren[0].getText();
                if (firstChild === 'static') {
                    return;
                }
                const re: IActionInfo = { methodName: firstChild };
                const cn2 = nodeChildren[2].getChildren();
                if (cn2.length) {
                    cn2.forEach(element => {
                        if (element.kind !== SyntaxKind.SyntaxList) {
                            return;
                        }
                        const text = element.getText();
                        const token = text.split('.');
                        if (token.length === 1) {
                            re.propType = token[0];
                        } else {
                            re.actionCode = token[token.length - 1];
                        }
                    });
                    actions.push(re);
                }
            });
        return actions;
    }

    function getActionPrefix(ast: ts.SourceFile, filename: string) {
        const actionPrefixNodes = tsquery(ast, 'VariableDeclaration:has(Identifier[name=actionTypePrefix]) StringLiteral');
        // const actionPrefixNodes = tsquery(ast, 'VariableStatement');
        let text = '';
        if (actionPrefixNodes.length) {
            text = actionPrefixNodes[0].getText().replace(/'/g, '');
        } else {
            text = filename.split('.')[0];
        }
        return `\r\n\r\nconst actionTypePrefix = '[${text}]'; \r\n`;
    }
    function replace(declarationRecorder: UpdateRecorder, node: ts.Node, content: string) {
        declarationRecorder.remove(node.getStart(), node.getWidth());
        declarationRecorder.insertLeft(node.getStart(), content);
    }
}

export interface IActionInfo {
    methodName?: string;
    actionCode?: string;
    propType?: string;
}