Ast walker and modifier

example.js

const walk = require('./walk');
const t = require('@babel/types');
const parser = require('yourastparserhere');

// read your code

const ast = parser.parse(code);

const root = walk(ast);

root.find(t.identifer, {
  name: 'test'
}).forEach(path => {
  path.node.name = path.node.name.reverse();
});

Path.js

class Path {
  constructor(node, parent = null, key = null) {
    this.node = node;
    this.parent = parent;
    this.key = key;
  }

  replace(node) {
    node.leadingComments = [...(this.node.leadingComments || [])];
    node.trailingComments = [...(this.node.trailingComments || [])];
    node.loc = {
      start: this.node.loc.start,
      end: this.node.loc.end,
    };

    if (Array.isArray(this.parent.node[this.key])) {
      const idx = this.parent.node[this.key].indexOf(this.node);

      this.node = node;
      this.parent.node[this.key][idx] = this.node;
    } else {
      this.node = node;
      this.parent.node[this.key] = this.node;
    }
  }

  prune() {
    if (this.parent.node.leadingComments == null) {
      this.parent.node.leadingComments = [];
    }

    if (this.parent.node.trailingComments == null) {
      this.parent.node.trailingComments = [];
    }
    this.parent.node.leadingComments.push(...(this.node.leadingComments || []));
    this.parent.node.trailingComments.push(...(this.node.trailingComments || []));

    delete this.parent.node[this.key];

    this.parent = null;
    this.key = null;
  }
}

module.exports = Path;

walk.js

const Path = require('./Path');

function capitalize(str) {
  return `${str.substr(0, 1).toUpperCase()}${str.substr(1)}`;
}

function matches(a, b) {
  if (typeof b === 'function') {
    return b(a);
  }

  if (typeof a !== typeof b) {
    return false;
  }

  if (Array.isArray(b)) {
    const set = new Set(a);

    return b.every(e => set.has(e));
  }

  if (typeof a === 'object') {
    return Object.keys(b).every(key => matches(a[key], b[key]));
  }

  return a === b;
}

function find(ast, type, selectors) {
  if (typeof type === 'function') {
    type = type.name;
  }

  type = capitalize(type);

  const found = [];
  const queue = [new Path(ast)];

  while (queue.length) {
    const path = queue.pop();
    const { node } = path;

    Object.keys(node).forEach(key => {
      if (node[key] != null && node[key].type != null) {
        queue.push(new Path(node[key], path, key));
      }

      if (Array.isArray(node[key])) {
        node[key].forEach(e => {
          if (e.type != null) {
            queue.push(new Path(e, path, key));
          }
        });
      }
    });

    if (node.type !== type) {
      continue;
    }

    if (matches(node, selectors)) {
      found.push(path);
    }
  }

  return found;
}

function closest(path, type, selectors) {
  if (typeof type === 'function') {
    type = type.name;
  }

  type = capitalize(type);

  const curr = path;

  while (curr.parent != null) {
    const { parent } = curr;

    if (parent.node.type === type && matches(parent.node, selectors)) {
      return path.parent;
    }

    curr = parent;
  }

  return null;
}

function forEach(nodes, fn) {
  nodes.forEach(path => {
    fn(path);
  });

  return source(nodes);
}

function replaceWith(nodes, fn) {
  nodes.forEach(path => {
    path.replace(fn(path));
  });

  return source(nodes);
}

function source(nodes) {
  if (!Array.isArray(nodes)) {
    nodes = [nodes];
  }

  return {
    find(type, selectors = {}) {
      const results = nodes.map(node => find(node, type, selectors));

      return source([].concat.apply([], results));
    },
    forEach(fn) {
      forEach(nodes, fn);

      return source(nodes);
    },
    replaceWith(fn) {
      replaceWith(nodes, fn);

      return source(nodes);
    },
    remove() {
      nodes.forEach(node => node.prune());

      return source(nodes);
    },
    filter(fn) {
      return source(nodes.filter(fn));
    },
    size() {
      return nodes.length;
    },
    closest(type, selectors) {
      return source(nodes.map(path => closest(path, type, selectors)));
    },
    nodes() {
      return nodes;
    },
    map(fn) {
      return source(nodes.map(fn));
    },
  };
}

module.exports = source;