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 (:=) introduced in Python 3.8 allows you to assign values to variables as part of an expression. This rule aims to simplify code by using the walrus operator in if statements.

This first part of the rule identifies cases where a variable is assigned a value and then immediately used in an if statement to control flow.

yaml
id: use-walrus-operator
language: python
rule:
  pattern: "if $VAR: $$$B"
  follows:
    pattern:
      context: $VAR = $$$EXPR
      selector: expression_statement
fix: |-
  if $VAR := $$$EXPR:
    $$$B

The pattern clause finds an if statement that checks the truthiness of $VAR. If this pattern follows an expression statement where $VAR is assigned $$$EXPR, the fix clause changes the if statements to use the walrus operator.

The second part of the rule:

yaml
id: remove-declaration
rule:
  pattern:
    context: $VAR = $$$EXPR
    selector: expression_statement
  precedes:
    pattern: "if $VAR: $$$B"
fix: ''

This rule removes the standalone variable assignment when it directly precedes an if statement that uses the walrus operator. Since the assignment is now part of the if statement, the separate declaration is no longer needed.

By applying these rules, you can refactor your Python code to be more concise and readable, taking advantage of the walrus operator's ability to combine an assignment with an expression.

YAML ​

yaml
id: use-walrus-operator
language: python
rule:
  follows:
    pattern:
      context: $VAR = $$$EXPR
      selector: expression_statement
  pattern: "if $VAR: $$$B"
fix: |-
  if $VAR := $$$EXPR:
    $$$B
---
id: remove-declaration
language: python
rule:
  pattern:
    context: $VAR = $$$EXPR
    selector: expression_statement
  precedes:
    pattern: "if $VAR: $$$B"
fix: ''

Example ​

python
a = foo()

if a:
    do_bar()

Diff ​

python
a = foo() 

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

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 straghtforward 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

Made with ❤️ with Rust

Copyright © 2022-present Herrington Darkholme