Import management for Python

Python pattern

Grit includes standard patterns for declaratively finding, adding, and updating imports in Python.

import_from($source) pattern

The import_from pattern is used to find an import statement. The $source metavariable can be used to specify the module that is being imported. This pattern will match any import statement that imports from the specified module.

For example, you can use the following pattern to remove all imports from the pydantic module:

GRIT 
language python

import_from(source="pydantic") => .
BEFORE 
from typing import List
from pydantic import BaseModel
from pydantic import More
AFTER 
from typing import List

imported_from($source) pattern

The imported_from($source) pattern is used to filter an identifier to cases that are imported from a specific module $source. This is useful for narrowing commonly used names.

For example, you can use the following pattern to replace the model parameter for completion calls, but only when it is imported from the litellm module:

GRIT 
language python

`$completion($params)` where {
  $completion <: imported_from(source="litellm"),

  $completion <: `completion`,
  $params <: contains `model=$_` => `model="gpt-4-turbo"`,
}

Here it changes the parameters:

BEFORE 
from litellm import completion

completion(model="gpt-3")
AFTER 
from litellm import completion

completion(model="gpt-4-turbo")

But if completion is imported from another module, it will not be changed:

PYTHON 
from openai import completion

completion(model="gpt-3")

add_import($source, $name) predicate

The add_import($source, $name) predicate can be used inside a where clause to add an import statement to the top of the file. If $name isn't already imported from $source, the import statement will be added.

Note this is idempotent, so it will not add the import if it is already present and you can safely call it multiple times.

For example, this pattern can be used to add a completion import from the litellm package:

GRIT 
language python

`completion($params)` where {
  add_import(source="litellm", name="completion")
}
BEFORE 
completion(model="gpt-3")
AFTER 
from litellm import completion

completion(model="gpt-3")

If the import is already present, the pattern will not change the file.

BEFORE 
from openai import other
from litellm import completion

completion(model="gpt-3")
AFTER 
from openai import other
from litellm import completion

completion(model="gpt-3")

Bare imports

If you want to add a bare import (e.g. import openai), use add_import($source) without specifying a name:

GRIT 
language python

`completion($params)` => `openai.completion($params)` where {
  add_import(source="openai")
}
BEFORE 
completion(model="gpt-3")
AFTER 
import openai

openai.completion(model="gpt-3")

remove_import($source) predicate

The remove_import($source) predicate can be used inside a where clause to remove an import statement, if it is present.

For example, you can use the following pattern to remove all imports from the pydantic module:

GRIT 
language python

import_from(source="pydantic") => . where {
  remove_import(source="pydantic")
}
BEFORE 
from typing import List
from pydantic import BaseModel
from pydantic import More
AFTER 
from typing import List