ast-grep

Appearance

Python ​

This page curates a list of example ast-grep rules to check and to rewrite Python code.

Migrate OpenAI SDK Has Fix ​

Description ​

OpenAI has introduced some breaking changes in their API, such as using Client to initialize the service and renaming the Completion method to completions . This example shows how to use ast-grep to automatically update your code to the new API.

API migration requires multiple related rules to work together. The example shows how to write multiple rules in a single YAML file. The rules and patterns in the example are simple and self-explanatory, so we will not explain them further.

YAML ​

yaml
id: import-openai
language: python
rule:
  pattern: import openai
fix: from openai import Client
---
id: rewrite-client
language: python
rule:
  pattern: openai.api_key = $KEY
fix: client = Client($KEY)
---
id: rewrite-chat-completion
language: python
rule:
  pattern: openai.Completion.create($$$ARGS)
fix: |-
  client.completions.create(
    $$$ARGS
  )

Example ​

python
import os
import openai
from flask import Flask, jsonify

app = Flask(__name__)
openai.api_key = os.getenv("OPENAI_API_KEY")

@app.route("/chat", methods=("POST"))
def index():
    animal = request.form["animal"]
    response = openai.Completion.create(
        model="text-davinci-003",
        prompt=generate_prompt(animal),
        temperature=0.6,
    )
    return jsonify(response.choices)

Diff ​

python
import os
import openai 
from openai import Client 
from flask import Flask, jsonify

app = Flask(__name__)
openai.api_key = os.getenv("OPENAI_API_KEY") 
client = Client(os.getenv("OPENAI_API_KEY")) 

@app.route("/chat", methods=("POST"))
def index():
    animal = request.form["animal"]
    response = openai.Completion.create( 
    response = client.completions.create( 
      model="text-davinci-003",
      prompt=generate_prompt(animal),
      temperature=0.6,
    )
    return jsonify(response.choices)

Contributed by ​

Herrington Darkholme, inspired by Morgante from grit.io

Prefer Generator Expressions Has Fix ​

Description ​

List comprehensions like [x for x in range(10)] are a concise way to create lists in Python. However, we can achieve better memory efficiency by using generator expressions like (x for x in range(10)) instead. List comprehensions create the entire list in memory, while generator expressions generate each element one at a time. We can make the change by replacing the square brackets with parentheses.

YAML ​

yaml
id: prefer-generator-expressions
language: python
rule:
  pattern: $LIST
  kind: list_comprehension
transform:
  INNER:
    substring: {source: $LIST, startChar: 1, endChar: -1 }
fix: ($INNER)

This rule converts every list comprehension to a generator expression. However, not every list comprehension can be replaced with a generator expression. If the list is used multiple times, is modified, is sliced, or is indexed, a generator is not a suitable replacement.

Some common functions like any, all, and sum take an iterable as an argument. A generator function counts as an iterable, so it is safe to change a list comprehension to a generator expression in this context.

yaml
id: prefer-generator-expressions
language: python
rule:
  pattern: $FUNC($LIST)
constraints:
  LIST: { kind: list_comprehension }
  FUNC:
    any:
      - pattern: any
      - pattern: all
      - pattern: sum
      # ...
transform:
  INNER:
    substring: {source: $LIST, startChar: 1, endChar: -1 }
fix: $FUNC($INNER)

Example ​

python
any([x for x in range(10)])

Diff ​

python
any([x for x in range(10)]) 
any(x for x in range(10))

Contributed by ​

Steven Love

Use Walrus Operator in if statementHas Fix ​

Description ​

The walrus operator (:=, Python 3.8+) merges an assignment into an expression. This rule merges an assignment to a simple name with the immediately following if that tests it. Apply both rules in one scan of the original code: one rewrites only the condition, preserving the body and any elif or else branches, and the other removes the assignment.

Apply only when all hold: a single assignment to a simple name (such as title), assignment immediately before the if (adjacent), RHS (right-hand side) is pure (no side effects), test is a bare truth check. Otherwise skip - see non-cases below. The rules do not check for side effects; review the RHS before applying them.

YAML ​

yaml
id: use-walrus-operator
language: python
rule:
  pattern: $VAR
  kind: identifier
  inside:
    kind: if_statement
    field: condition
    follows:
      pattern:
        context: $VAR = $EXPR
        selector: expression_statement
      stopBy: neighbor
constraints:
  VAR:
    kind: identifier
  EXPR:
    not:
      kind: assignment
fix: $VAR := ($EXPR)
---
id: remove-declaration
language: python
rule:
  pattern:
    context: $VAR = $EXPR
    selector: expression_statement
  precedes:
    kind: if_statement
    has:
      field: condition
      pattern: $VAR
      kind: identifier
    stopBy: neighbor
constraints:
  VAR:
    kind: identifier
  EXPR:
    not:
      kind: assignment
fix: ''

The constraints restrict $VAR to a simple name and exclude chained assignments such as a = b = value. Both rules use the same constraints. Parenthesizing $EXPR preserves tuple expressions such as a = 1, 2.

Negated tests (separate rule) ​

if not X needs parentheses. The rules above never match it; hand-applying their shape flips the branch. Example:

python
# before
def get_title(data):
    title = (data.get('title') or '').strip()
    if not title:
        return fallback
    validate(title)
    return title

# wrong — branch inverted
def get_title(data):
    if title := (data.get('title') or '').strip():
        return fallback
    validate(title)
    return title

Correct form is if not (title := ...):. Apply both rules below together so the RHS is evaluated only once.

yaml
id: use-walrus-operator-not
language: python
rule:
  pattern: not $VAR
  inside:
    kind: if_statement
    field: condition
    follows:
      pattern:
        context: $VAR = $EXPR
        selector: expression_statement
      stopBy: neighbor
constraints:
  VAR:
    kind: identifier
  EXPR:
    not:
      kind: assignment
fix: not ($VAR := ($EXPR))
---
id: remove-declaration-not
language: python
rule:
  pattern:
    context: $VAR = $EXPR
    selector: expression_statement
  precedes:
    kind: if_statement
    has:
      field: condition
      pattern: not $VAR
    stopBy: neighbor
constraints:
  VAR:
    kind: identifier
  EXPR:
    not:
      kind: assignment
fix: ''

Example ​

python
a = foo()
if a:
    do_bar()

Diff ​

python
a = foo() 
if a: 
if a := (foo()): 
    do_bar()

Do not apply when ​

  • The target is an attribute, subscript, or unpacking pattern (for example, obj.title, items[0], or a, b). Walrus targets must be simple names.
  • The assignment is chained (for example, a = b = value) or annotated (for example, a: str = value). These rules skip both.
  • Code sits between assignment and if (reordered or skipped).
  • RHS has side effects; this example assumes a pure expression and does not verify purity.
  • The test needs input validation first (for example, external data): truthy does not mean usable. This rewrite preserves the truth test and does not add validation.

Contributed by ​

Inspired by reddit user /u/jackerhack

Remove async function Has Fix ​

Description ​

The async keyword in Python is used to define asynchronous functions that can be awaited.

In this example, we want to remove the async keyword from a function definition and replace it with a synchronous version of the function. We also need to remove the await keyword from the function body.

By default, ast-grep will not apply overlapping replacements. This means await keywords will not be modified because they are inside the async function body.

However, we can use the rewriter to apply changes inside the matched function body.

YAML ​

yaml
id: remove-async-def
language: python
rule:
  # match async function definition
  pattern:
    context: 'async def $FUNC($$$ARGS): $$$BODY'
    selector: function_definition
rewriters:
# define a rewriter to remove the await keyword
  remove-await-call:
    pattern: 'await $$$CALL'
    fix: $$$CALL # remove await keyword
# apply the rewriter to the function body
transform:
  REMOVED_BODY:
    rewrite:
      rewriters: [remove-await-call]
      source: $$$BODY
fix: |-
  def $FUNC($$$ARGS):
    $REMOVED_BODY

Example ​

python
async def main3():
  await somecall(1, 5)

Diff ​

python
async def main3(): 
  await somecall(1, 5) 
def main3(): 
  somecall(1, 5)

Contributed by ​

Inspired by the ast-grep issue #1185

Refactor pytest fixtures ​

Description ​

One of the most commonly used testing framework in Python is pytest. Among other things, it allows the use of fixtures.

Fixtures are defined as functions that can be required in test code, or in other fixtures, as an argument. This means that all functions arguments with a given name in a pytest context (test function or fixture) are essentially the same entity. However, not every editor's LSP is able to keep track of this, making refactoring challenging.

Using ast-grep, we can define some rules to match fixture definition and usage without catching similarly named entities in a non-test context.

First, we define utils to select pytest test/fixture functions.

yaml
utils:
  is-fixture-function:
    kind: function_definition
    follows:
      kind: decorator
      has:
        kind: identifier
        regex: ^fixture$
        stopBy: end
  is-test-function:
    kind: function_definition
    has:
      field: name
      regex: ^test_

Pytest fixtures are declared with a decorator @pytest.fixture. We match the function_definition node that directly follows a decorator node. That decorator node must have a fixture identifier somewhere. This accounts for different location of the fixture node depending on the type of imports and whether the decorator is used as is or called with parameters.

Pytest functions are fairly straightforward to detect, as they always start with test_ by convention.

The next utils builds onto those two to incrementally:

  • Find if a node is inside a pytest context (test/fixture)
  • Find if a node is an argument in such a context
yaml
utils:
  is-pytest-context:
    # Pytest context is a node inside a pytest
    # test/fixture
    inside:
      stopBy: end
      any:
        - matches: is-fixture-function
        - matches: is-test-function
  is-fixture-arg:
    # Fixture arguments are identifiers inside the 
    # parameters of a test/fixture function
    all:
      - kind: identifier
      - inside:
          kind: parameters
      - matches: is-pytest-context

Once those utils are declared, you can perform various refactoring on a specific fixture.

The following rule adds a type-hint to a fixture.

yaml
rule:
  matches: is-fixture-arg
  regex: ^foo$
fix: 'foo: int'

This one renames a fixture and all its references.

yaml
rule:
  kind: identifier
  matches: is-fixture-context
  regex: ^foo$
fix: 'five'

Example ​

Renaming Fixtures ​

python
@pytest.fixture
def foo() -> int:
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo: int) -> str:
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo: int) -> None:
    assert foo == 5

Diff ​

python
@pytest.fixture
def foo() -> int: 
def five() -> int: 
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo: int) -> str: 
def some_fixture(five: int) -> str: 
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo: int) -> None: 
def test_code(five: int) -> None: 
    assert foo == 5
    assert five == 5

Type Hinting Fixtures ​

python
@pytest.fixture
def foo() -> int:
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo) -> str:
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo) -> None:
    assert foo == 5

Diff ​

python
@pytest.fixture
def foo() -> int:
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo) -> str: 
def some_fixture(foo: int) -> str: 
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo) -> None: 
def test_code(foo: int) -> None: 
    assert foo == 5

Rewrite Optional[Type] to Type | None Has Fix ​

Description ​

PEP 604 recommends that Type | None is preferred over Optional[Type] for Python 3.10+.

This rule performs such rewriting. Note Optional[$T] alone is interpreted as subscripting expression instead of generic type, we need to use pattern object to disambiguate it with more context code.

YAML ​

yaml
id: optional-to-none-union
language: python
rule:
  pattern:
    context: 'a: Optional[$T]'
    selector: generic_type
fix: $T | None

Example ​

py
def a(arg: Optional[int]): pass

Diff ​

py
def a(arg: Optional[int]): pass
def a(arg: int | None): pass

Contributed by ​

Bede Carroll

Recursive Rewrite Type Has Fix ​

Description ​

Suppose we want to transform Python's Union[T1, T2] to T1 | T2 and Optional[T] to T | None.

By default, ast-grep will only fix the outermost node that matches a pattern and will not rewrite the inner AST nodes inside a match. This avoids unexpected rewriting or infinite rewriting loop.

So if you are using non-recursive rewriter like this, Optional[Union[int, str]] will only be converted to Union[int, str] | None. Note the inner Union[int, str] is not enabled. This is because the rewriter optional matches Optional[$TYPE] and rewrite it to $TYPE | None. The inner $TYPE is not processed.

However, we can apply rewriters to inner types recursively. Take the optional rewriter as an example, we need to apply rewriters, optional and unions, recursively to $TYPE and get a new variable $NT.

YAML ​

yml
id: recursive-rewrite-types
language: python
rewriters:
# rewrite Optional[T] to T | None
- id: optional
  rule:
    any:
    - pattern:
        context: 'arg: Optional[$TYPE]'
        selector: generic_type
    - pattern: Optional[$TYPE]
  # recursively apply rewriters to $TYPE
  transform:
    NT:
      rewrite:
        rewriters: [optional, unions]
        source: $TYPE
  # use the new variable $NT
  fix: $NT | None

# similar to Optional, rewrite Union[T1, T2] to T1 | T2
- id: unions
  language: Python
  rule:
    pattern:
      context: 'a: Union[$$$TYPES]'
      selector: generic_type
  transform:
    UNIONS:
      # rewrite all types inside $$$TYPES
      rewrite:
        rewriters: [ rewrite-unions ]
        source: $$$TYPES
        joinBy: " | "
  fix: $UNIONS
- id: rewrite-unions
  rule:
    pattern: $TYPE
    kind: type
  # recursive part
  transform:
    NT:
      rewrite:
        rewriters: [optional, unions]
        source: $TYPE
  fix: $NT

# find all types
rule:
  kind: type
  pattern: $TPE
# apply the recursive rewriters
transform:
  NEW_TYPE:
    rewrite:
      rewriters: [optional, unions]
      source: $TPE
# output
fix: $NEW_TYPE

Example ​

python
results:  Optional[Union[List[Union[str, dict]], str]]

Diff ​

python
results:  Optional[Union[List[Union[str, dict]], str]] 
results:  List[str | dict] | str | None

Contributed by ​

Inspired by steinuil

Rewrite SQLAlchemy with Type Annotations Has Fix ​

Description ​

SQLAlchemy 2.0 recommends using type annotations with Mapped type for modern declarative mapping. The mapped_column() construct can derive its configuration from PEP 484 type annotations.

This rule helps migrate legacy SQLAlchemy code that explicitly uses String type and nullable=True to the modern type annotation approach using Mapped[str | None].

The key technique demonstrated here is using rewriters to selectively filter arguments. The rewriter:

  1. Matches each argument inside the argument_list
  2. Excludes the String type argument
  3. Excludes the nullable=True keyword argument
  4. Keeps all other arguments

YAML ​

yaml
id: remove-nullable-arg
language: python
rule:
  pattern: $X = mapped_column($$$ARGS)
  any:
    - pattern: $X = mapped_column($$$BEFORE, String, $$$MID, nullable=True, $$$AFTER)
    - pattern: $X = mapped_column($$$BEFORE, String, $$$MID, nullable=True)
rewriters:
- id: filter-string-nullable
  rule:
    pattern: $ARG
    inside:
      kind: argument_list
    all:
    - not:
        pattern: String
    - not:
        pattern:
          context: a(nullable=True)
          selector: keyword_argument
  fix: $ARG

transform:
  NEWARGS:
    rewrite:
      rewriters: [filter-string-nullable]
      source: $$$ARGS
      joinBy: ', '
fix: |-
  $X: Mapped[str | None] = mapped_column($NEWARGS)

Example ​

python
message = mapped_column(String, default="hello", nullable=True)

message = mapped_column(String, nullable=True)

_message = mapped_column("message", String, nullable=True)

message = mapped_column(String, nullable=True, unique=True)

message = mapped_column(
  String, index=True, nullable=True, unique=True)

# Should not be transformed
message = mapped_column(String, default="hello")

message = mapped_column(String, default="hello", nullable=False)

message = mapped_column(Integer, default="hello")

Diff ​

python
message = mapped_column(String, default="hello", nullable=True) 
message: Mapped[str | None] = mapped_column(default="hello") 

message = mapped_column(String, nullable=True) 
message: Mapped[str | None] = mapped_column() 

_message = mapped_column("message", String, nullable=True) 
_message: Mapped[str | None] = mapped_column("message") 

message = mapped_column(String, nullable=True, unique=True) 
message: Mapped[str | None] = mapped_column(unique=True) 

message = mapped_column( 
  String, index=True, nullable=True, unique=True) 
message: Mapped[str | None] = mapped_column( 
  index=True, unique=True)

Contributed by ​

Inspired by discussion #2319

Made with ❤️ with Rust

Copyright © 2022-present Herrington Darkholme