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