Structural Code Search With ast-grep

ast-grep is an amazing tool for structural code search. With it, you can search for code by it’s AST shape, rather than just a simple string or regex search. I’ve been using it a fair bit at work, and I developed a few scripts that replaced the core logic of the Flashlight project I built last year.

Find JSX tags

This script searches for JSX tags in your code. You can specify a tag name, and optionally an attribute/value where only the tags containing the attribute or attribute/value pair will be listed.

/usr/local/bin/tags

#!/bin/bash

cwd="$(pwd)"
format="rich"

_print_help() {
	if [ -n "$1" ]; then
		printf "%s\n\n" "$1"
	fi

	echo "Usage: tags [options] <name> [attribute] [value]

Options:
  -h, --help             Display this help message and exit
  -c, --cwd <dir>        Set the current working directory to <dir>
  -f, --format <format>  Set the output format (short|long)

Arguments:
  <name>                 The tag name to search for
  [attribute]            Filter by attribute
  [value]                Filter by attribute value"
}

_parse_args() {
	while [[ $# -gt 0 ]]; do
		case "$1" in
		-h | --help)
			_print_help "Find JSX tags"
			exit 0
			;;
		-c | --cwd)
			cwd="$2"
			shift 2
			;;
		-f | --format)
			if [[ "$2" == 'short' ]]; then
				format="short"
			elif [[ "$2" == 'long' ]]; then
				format="rich"
			else
				_print_help "Invalid format: $2"
				exit 1
			fi

			shift 2
			;;
		*)
			if [ -z "$name" ]; then
				name="$1"
			elif [ -z "$attribute" ]; then
				attribute="$1"
			elif [ -z "$value" ]; then
				value="$1"
			else
				_print_help "Unexpected argument: $1"
				exit 1
			fi

			shift
			;;
		esac
	done
}

_split() {
	echo "$1" | while IFS='.' read -ra parts; do
		for part in "${parts[@]}"; do
			echo "$part"
		done
	done
}

_get_rule() {
	namespaced_tag=($(_split "$name"))
	head="
id: tags
language: TSX
utils:
  is-tag:
    any:
      - kind: jsx_opening_element
      - kind: jsx_self_closing_element
    has:
      any:
        - kind: identifier
          regex: '^$name$'
        - kind: member_expression
          all:
            - has:
                kind: identifier
                field: object
                regex: '^${namespaced_tag[0]}$'
            - has:
                kind: property_identifier
                field: property
                regex: '^${namespaced_tag[1]}$'
"

	if [ -n "$attribute" ]; then
		echo "
$head
rule:
  kind: jsx_attribute
  all:
    - has:
        kind: property_identifier
        regex: '^$attribute$'
    $(if [ -n "$value" ]; then
      # TODO: Support boolean, boolean shorthand, null, undefined, and template literals
      echo "
    - has:
        any:
          - kind: string
            has:
              kind: string_fragment
              regex: '^$value$'
          - has:
              kind: number
              regex: '^$value$'
"
		fi)
  inside:
    matches: is-tag
"
	else
		echo "
$head
rule:
  matches: is-tag
"
	fi
}

_find_tags() {
	if [ -z "$name" ]; then
		_print_help "Missing name argument"
		exit 1
	fi

	sg scan \
    --config "$HOME/.sgconfig.yml" \
    --inline-rules "$(_get_rule)" \
		--report-style="$format" \
		"$cwd"
}

_parse_args "$@"
_find_tags

For this script to work, we need to specify a global config file to instruct ast-grep to search for all js, jsx, ts, and tsx file extensions since by default it will only search for the specified language.

~/.sgconfig.yml

ruleDirs: []
languageGlobs:
  tsx: ["*.js", "*.jsx", "*.ts", "*.tsx"]

To use this script, we can run the following command to find code like this: <MyButton onClick={...}>

tags MyButton onClick

Find imports

This script searches for imports in your code. You can specify an import source (e.g. react), and optionally an import specifier (e.g. useState).

/usr/local/bin/imports

#!/bin/bash

cwd="$(pwd)"
format="rich"

_print_help() {
	if [ -n "$1" ]; then
		printf "%s\n\n" "$1"
	fi

	echo "Usage: imports [options] <source> [specifier]

Options:
  -h, --help             Display this help message and exit
  -c, --cwd <dir>        Set the current working directory to <dir>
  -f, --format <format>  Set the output format (short|long)

Arguments:
  <source>               The import source to search for
  [specifier]            Only return imports containing this specifier"
}

_parse_args() {
	while [[ $# -gt 0 ]]; do
		case "$1" in
		-h | --help)
			_print_help "Find JavaScript imports"
			exit 0
			;;
		-c | --cwd)
			cwd="$2"
			shift 2
			;;
		-f | --format)
			if [[ "$2" == 'short' ]]; then
				format="short"
			elif [[ "$2" == 'long' ]]; then
				format="rich"
			else
				_print_help "Invalid format: $2"
				exit 1
			fi

			shift 2
			;;
		*)
			if [ -z "$source" ]; then
				source="$1"
			elif [ -z "$specifier" ]; then
				specifier="$1"
			else
				_print_help "Unexpected argument: $1"
				exit 1
			fi

			shift
			;;
		esac
	done
}

_to_regex() {
	# If already in regex format, return as is removing the leading and traililng slash
	if [[ "$1" == /* && "$1" == */ ]]; then
		echo "${1:1:${#1}-2}"
	else
		echo "^$1$"
	fi
}

_get_rule() {
	head="
id: imports
language: TSX
utils:
  is-import:
    kind: import_statement
    has:
      kind: string
      field: source
      has:
        kind: string_fragment
        regex: '$(_to_regex "$1")'
  "

	# If a specifier is provided, we search for the specifier within the
	# import given the specified source.
	if [ -n "$2" ]; then
		echo "
$head
rule:
  kind: import_specifier
  regex: '$(_to_regex "$2")'
  inside:
    kind: named_imports
    inside:
      kind: import_clause
      inside:
        matches: is-import
"
	else
		echo "
$head
rule:
  matches: is-import
"
	fi
}

_find_imports() {
	if [ -z "$source" ]; then
		_print_help "Missing source argument"
		exit 1
	fi

	sg scan \
		--config "$HOME/.sgconfig.yml" \
		--inline-rules "$(_get_rule "$source" "$specifier")" \
		--report-style="$format" \
		"$cwd"
}

_parse_args "$@"
_find_imports

To use this script, we can run the following command to find code like this: import { useState } from 'react':

imports react useState

It even supports searching by regex! If we want to search for all React hook imports from any module, we can use this command:

imports '/.*/' '/^use.*/'