What is Babel AST and how to write a custom Babel plugin to manipulate AST

What is AST?

AST (Abstract Syntax Tree) is a tree representation of source code that structures code into a hierarchy of nodes, where each node represents a construct in the code (such as variable declarations, function calls, etc.).

Babel AST Specification

Babel uses AST based on the ESTree specification, with extensions for JSX, TypeScript, and other syntax support.

AST Node Types

Common Node Types

Node Type	Description	Example
Program	Program root node	Entire file
Identifier	Identifier	Variable names, function names
Literal	Literal	1, "hello", true
VariableDeclaration	Variable declaration	const, let, var
FunctionDeclaration	Function declaration	function foo() {}
CallExpression	Function call	foo()
BinaryExpression	Binary expression	a + b
MemberExpression	Member expression	obj.prop

AST Example

javascript
// Source code
const sum = (a, b) => a + b;

// AST (simplified)
{
  "type": "VariableDeclaration",
  "kind": "const",
  "declarations": [{
    "type": "VariableDeclarator",
    "id": { "type": "Identifier", "name": "sum" },
    "init": {
      "type": "ArrowFunctionExpression",
      "params": [
        { "type": "Identifier", "name": "a" },
        { "type": "Identifier", "name": "b" }
      ],
      "body": {
        "type": "BinaryExpression",
        "operator": "+",
        "left": { "type": "Identifier", "name": "a" },
        "right": { "type": "Identifier", "name": "b" }
      }
    }
  }]
}

Writing Custom Babel Plugins

1. Basic Plugin Structure

javascript
// my-plugin.js
module.exports = function(babel) {
  const { types: t } = babel;
  
  return {
    name: 'my-custom-plugin',
    visitor: {
      // Visitor methods
      Identifier(path) {
        console.log('Found identifier:', path.node.name);
      }
    }
  };
};

2. Practical Plugin Examples

Example 1: Replace console.log

javascript
// remove-console-plugin.js
module.exports = function(babel) {
  const { types: t } = babel;
  
  return {
    name: 'remove-console',
    visitor: {
      CallExpression(path) {
        const { callee } = path.node;
        
        // Check if it's a console.log call
        if (
          t.isMemberExpression(callee) &&
          t.isIdentifier(callee.object, { name: 'console' }) &&
          t.isIdentifier(callee.property, { name: 'log' })
        ) {
          // Remove the node
          path.remove();
        }
      }
    }
  };
};

Example 2: Auto-add Function Names

javascript
// add-function-name-plugin.js
module.exports = function(babel) {
  const { types: t } = babel;
  
  return {
    name: 'add-function-name',
    visitor: {
      FunctionDeclaration(path) {
        const { node } = path;
        
        // If function has no name, add a default one
        if (!node.id) {
          node.id = t.identifier('anonymous');
        }
      }
    }
  };
};

Example 3: Internationalization String Extraction

javascript
// i18n-plugin.js
module.exports = function(babel) {
  const { types: t } = babel;
  const strings = [];
  
  return {
    name: 'i18n-extractor',
    visitor: {
      StringLiteral(path) {
        const { node } = path;
        
        // Collect all strings
        strings.push(node.value);
        
        // Replace with i18n function call
        path.replaceWith(
          t.callExpression(
            t.identifier('t'),
            [t.stringLiteral(node.value)]
          )
        );
      }
    },
    post(state) {
      // Output collected strings
      console.log('Extracted strings:', strings);
    }
  };
};

3. Using Plugins

javascript
// babel.config.js
module.exports = {
  plugins: [
    './remove-console-plugin.js',
    ['./i18n-plugin.js', { /* plugin options */ }]
  ]
};

Path Object Details

Core Path Methods

javascript
// Path object in visitor
visitor: {
  Identifier(path) {
    // Node info
    console.log(path.node);        // AST node
    console.log(path.parent);      // Parent node
    console.log(path.parentPath);  // Parent path
    
    // Node operations
    path.remove();                 // Delete node
    path.replaceWith(newNode);     // Replace node
    path.insertBefore(newNode);    // Insert before
    path.insertAfter(newNode);     // Insert after
    
    // Traversal
    path.traverse({ ... });        // Subtree traversal
    path.skip();                   // Skip subtree
    path.stop();                   // Stop traversal
    
    // Checks
    path.isIdentifier();           // Check node type
    path.findParent((p) => ...);   // Find parent
    path.getFunctionParent();      // Get function parent
    path.getStatementParent();     // Get statement parent
    
    // Scope
    path.scope.hasBinding('name'); // Check binding
    path.scope.rename('old', 'new'); // Rename
    path.scope.generateUid('name'); // Generate unique identifier
  }
}

Advanced Plugin Techniques

1. State Management

javascript
module.exports = function(babel) {
  return {
    name: 'stateful-plugin',
    pre(state) {
      // Initialize state before traversal
      this.counter = 0;
    },
    visitor: {
      Identifier(path) {
        this.counter++;
      }
    },
    post(state) {
      // Output results after traversal
      console.log(`Found ${this.counter} identifiers`);
    }
  };
};

2. Handling JSX

javascript
// Transform <div>Hello</div> to h('div', null, 'Hello')
module.exports = function(babel) {
  const { types: t } = babel;
  
  return {
    name: 'jsx-transform',
    visitor: {
      JSXElement(path) {
        const { openingElement, children } = path.node;
        const tagName = openingElement.name.name;
        
        // Create h() call
        const callExpr = t.callExpression(
          t.identifier('h'),
          [
            t.stringLiteral(tagName),
            t.nullLiteral(),
            ...children.map(child => {
              if (t.isJSXText(child)) {
                return t.stringLiteral(child.value.trim());
              }
              return child;
            })
          ]
        );
        
        path.replaceWith(callExpr);
      }
    }
  };
};

3. Source Map Support

javascript
module.exports = function(babel) {
  return {
    name: 'sourcemap-plugin',
    visitor: {
      Identifier(path) {
        // Preserve original location info
        path.addComment('leading', ` Original: ${path.node.name} `);
      }
    }
  };
};

Debugging Tips

javascript
// View AST
const parser = require('@babel/parser');
const code = 'const a = 1';
const ast = parser.parse(code);
console.log(JSON.stringify(ast, null, 2));

// Use @babel/template to simplify node creation
const template = require('@babel/template').default;
const buildRequire = template(`
  var IMPORT_NAME = require(SOURCE);
`);

const ast2 = buildRequire({
  IMPORT_NAME: t.identifier('myModule'),
  SOURCE: t.stringLiteral('my-module')
});

Best Practices

1. Use path instead of directly manipulating node - Path provides more context
2. Leverage path.scope - Correctly handle variable scopes
3. Use @babel/template - Simplify creation of complex AST nodes
4. Test plugins - Use transformSync from @babel/core for unit testing
5. Reference official plugins - Learn from Babel official plugin implementations