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 //[!code --]
from openai import Client //[!code ++]
from flask import Flask, jsonify

app = Flask(__name__)
openai.api_key = os.getenv("OPENAI_API_KEY") //[!code --]
client = Client(os.getenv("OPENAI_API_KEY")) //[!code ++]

@app.route("/chat", methods=("POST"))
def index():
    animal = request.form["animal"]
    response = openai.Completion.create( //[!code --]
    response = client.completions.create( //[!code ++]
      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

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() // [!code --]

if a: // [!code --]
if a := foo(): // [!code ++]
    do_bar()

Contributed by

Inspired by reddit user /u/jackerhack

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