Import management for Python

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

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:

language python

import_from(source="pydantic") => .

from typing import List
from pydantic import BaseModel
from pydantic import More

from typing import List

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:

language python

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

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

Here it changes the parameters:

from litellm import completion

completion(model="gpt-3")

from litellm import completion

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

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

from openai import completion

completion(model="gpt-3")

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:

language python

`completion($params)` where {
  add_import(source="litellm", name="completion")
}

completion(model="gpt-3")

from litellm import completion

completion(model="gpt-3")

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

from openai import other
from litellm import completion

completion(model="gpt-3")

from openai import other
from litellm import completion

completion(model="gpt-3")

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

language python

`completion($params)` => `openai.completion($params)` where {
  add_import(source="openai")
}

completion(model="gpt-3")

import openai

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