Python

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

Migrate OpenAI SDK Has Fix

Playground Link

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

Playground Link

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

Playground Link

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

Playground Link

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

Playground Link

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

Rule Essentials

Project Setup

Tooling Overview

API Usage

Command Line Interface

Rule Config

Python

Migrate OpenAI SDK Has Fix

Description

YAML

Example

Diff

Contributed by

Prefer Generator Expressions Has Fix

Description

YAML

Example

Diff

Contributed by

Use Walrus Operator in `if` statementHas Fix

Description

YAML

Example

Diff

Contributed by

Remove `async` function Has Fix

Description

YAML

Example

Diff

Contributed by

Refactor pytest fixtures

Description

Example

Renaming Fixtures

Diff

Type Hinting Fixtures

Diff

Python ​

Migrate OpenAI SDK Has Fix ​

Description ​

YAML ​

Example ​

Diff ​

Contributed by ​

Prefer Generator Expressions Has Fix ​

Description ​

YAML ​

Example ​

Diff ​

Contributed by ​

Use Walrus Operator in if statementHas Fix ​

Description ​

YAML ​

Example ​

Diff ​

Contributed by ​

Remove async function Has Fix ​

Description ​

YAML ​

Example ​

Diff ​

Contributed by ​

Refactor pytest fixtures ​

Description ​

Example ​

Renaming Fixtures ​

Diff ​

Type Hinting Fixtures ​

Diff ​

Python

Migrate OpenAI SDK Has Fix

Description

YAML

Example

Diff

Contributed by

Prefer Generator Expressions Has Fix

Description

YAML

Example

Diff

Contributed by

Use Walrus Operator in `if` statementHas Fix

Description

YAML

Example

Diff

Contributed by

Remove `async` function Has Fix

Description

YAML

Example

Diff

Contributed by

Refactor pytest fixtures

Description

Example

Renaming Fixtures

Diff

Type Hinting Fixtures

Diff