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

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

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

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

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.

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

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

Diff

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.

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:

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

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

a = foo()

if a:
    do_bar()

Diff

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

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

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

Diff

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.

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

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.

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

This one renames a fixture and all its references.

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

Example

Renaming Fixtures

@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

@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

@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

@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

• Playground Link

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

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

Example

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

Diff

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

Contributed by

Bede Carroll

Recursive Rewrite Type Has Fix

• Playground Link

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

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

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

Diff

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

• Playground Link

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

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

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

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