Plantacerium
Blog Software Engineer
Blog Software Engineer

Guide to Using AST Grep

Introduction

ast-grep is a powerful, command-line tool for searching and rewriting code based on its abstract syntax tree (AST). This allows for more precise and powerful code modifications than traditional text-based search and replace tools.

This guide provides an overview of ast-grep and how to use it in a production environment, covering common, real-world, and edge use cases.

Key Concepts

  • Abstract Syntax Tree (AST): An AST is a tree representation of the syntactic structure of source code. ast-grep uses ASTs to understand the code's structure, rather than just treating it as plain text.

  • Patterns: You can write patterns that look like the code you want to find. These patterns can include wildcards to match different parts of the code.

  • Rewriting: ast-grep can rewrite code by replacing the matched patterns with a new code structure.

Installation

You can install ast-grep using various package managers:

Using npm

npm install -g @ast-grep/cli

Using pip

pip install ast-grep-cli

Using Homebrew

brew install ast-grep

Real-World Use Cases

Here are some real-world examples of how ast-grep can be used in a production environment:

1. Large-Scale Refactoring

When a library or API changes, you often need to update your entire codebase. ast-grep can automate this process.

Scenario: You are upgrading a testing library and need to replace the old assertion style with a new one.

Old code:

assert.equal(actual, expected);

New code:

expect(actual).toEqual(expected);

ast-grep command:

sg -p 'assert.equal($A, $B)' -r 'expect($A).toEqual($B)' -l js

2. Enforcing Code Standards

You can use ast-grep to enforce coding standards and best practices across your organization.

Scenario: You want to ensure that all fetch calls have a timeout and appropriate error handling.

Pattern to find fetch calls without a timeout:

# rules/no-fetch-without-timeout.yml
id: no-fetch-without-timeout
language: JavaScript
rule:
  pattern: fetch($URL)
  not:
    has:
      pattern: fetch($URL, { timeout: $T })

ast-grep command to lint:

sg scan -r rules/no-fetch-without-timeout.yml

3. Security Vulnerability Patching

ast-grep can be used to find and patch security vulnerabilities at scale.

Scenario: A security vulnerability is discovered in a commonly used function, and you need to replace it with a safer alternative.

Vulnerable code:

const html = unsafeFunction(userInput);

Patched code:

const html = safeFunction(userInput);

ast-grep command:

sg -p 'unsafeFunction($ARG)' -r 'safeFunction($ARG)' -l js

Common Use Cases

Here are some common use cases for ast-grep:

  • API Migration: When a framework or library that you use updates its API, you can use ast-grep to automatically update your code.

  • Code Cleanup: Remove deprecated functions, unused variables, and other code smells.

  • Boilerplate Generation: Automate the creation of boilerplate code for new components or modules.

Edge Use Cases

Here are some advanced and less-common use cases for ast-grep:

1. Analyzing Code for Complex Patterns

You can use ast-grep to find complex code patterns that might be difficult to express with other tools.

Scenario: Find all functions that are called with a specific argument and then immediately awaited.

Pattern:

sg -p 'await $FN($ARG)' -l js

2. Multi-language Support

ast-grep is a polyglot tool and it can be used on a variety of different programming languages.

Scenario: Rewrite Python f-strings to use the str.format() method.

ast-grep command:

sg -p 'f"Hello, {$NAME}"' -r '"Hello, {}".format($NAME)' -l py

3. Custom

Linting Rules

You can create your own custom linting rules to enforce project-specific conventions.

Scenario: Ensure that all database queries are made through a specific helper function.

Pattern to find direct database queries:

# rules/no-direct-db-queries.yml
id: no-direct-db-queries
language: Python
rule:
  pattern: db.execute($QUERY)

ast-grep command:

sg scan -r rules/no-direct-db-queries.yml

Conclusion

ast-grep is a versatile and powerful tool that can significantly improve your workflow for code maintenance, refactoring, and analysis. By understanding its core concepts and exploring its various use cases, you can leverage ast-grep to write cleaner, safer, and more consistent code.