Add import insertion support to autofix capabilities

[charliermarsh]

Summary

We often have rules (and violations) that could in theory be autofixed, but that rely on the presence of an imported symbol that may or may not be available in scope.

The sys-exit-alias rule is a good example. It recommends replacing calls to the builtin exit and quit methods with calls to sys.exit. That's "easy" to autofix... as long as sys.exit is in-scope.

This PR enables rules to import symbols into the current scope in a fairly general way, enabling autofix support for rules like sys-exit-alias. Below, I've also extended autofix support for lru-cache-with-maxsize-none, which relies on making functools.cache available.

How it works

From a rule's perspective, the end-user API leans on a single get_or_import_symbol method, which returns (1) an Edit (to "insert" the symbol into the scope) and (2) the bound name of the resolved symbol (e.g., exit, if we added from sys import exit, or sys.exit, if we added import sys).

Under-the-hood, that method performs a series of lookups:

1. First, it checks to see if sys.exit is already available in any form (e.g., from sys import exit as foo, import sys as foo); if so, it returns a reference to it.
2. Otherwise, it checks to see if there's an import ... from for the relevant module (e.g., from sys import path). If so, it modifies that import to include the necessary symbol (from sys import path, exit). If exit is already bound to something else, it aborts and logs an error.
3. Otherwise, it generates an edit to insert an import sys, and returns sys.exit as the bound reference.

The logic for inserting imports is crude and follows LibCST's own logic: if there are no imports, we insert at top-of-file; otherwise, we insert after the last-seen top-level import statement.

Limitations

There are a few limitations to the current approach:

1. We don't attempt to retain import order at all. We leave that up to the isort rules (or isort itself, if you're using isort).

2. We only respect top-level imports, so here, we'd still insert an import sys:

def f():
  from sys import path

  quit()

3. Every individual fix (that relies on import insertion) has to go through this process. We may want to implement some sort of cache here.

Closes #835.

[github-actions]

PR Check Results

Ecosystem

✅ ecosystem check detected no changes.

Benchmark

Linux

group                                      main                                   pr
-----                                      ----                                   --
linter/all-rules/large/dataset.py          1.00     17.1±0.34ms     2.4 MB/sec    1.04     17.8±0.36ms     2.3 MB/sec
linter/all-rules/numpy/ctypeslib.py        1.00      4.3±0.12ms     3.9 MB/sec    1.03      4.5±0.16ms     3.7 MB/sec
linter/all-rules/numpy/globals.py          1.00   572.2±15.38µs     5.2 MB/sec    1.00   573.5±14.79µs     5.1 MB/sec
linter/all-rules/pydantic/types.py         1.00      7.3±0.25ms     3.5 MB/sec    1.03      7.6±0.16ms     3.4 MB/sec
linter/default-rules/large/dataset.py      1.00      8.6±0.13ms     4.7 MB/sec    1.09      9.4±0.32ms     4.3 MB/sec
linter/default-rules/numpy/ctypeslib.py    1.00  1977.1±55.79µs     8.4 MB/sec    1.03      2.0±0.08ms     8.2 MB/sec
linter/default-rules/numpy/globals.py      1.00    235.8±7.58µs    12.5 MB/sec    1.03   242.3±11.21µs    12.2 MB/sec
linter/default-rules/pydantic/types.py     1.00      4.0±0.09ms     6.4 MB/sec    1.06      4.3±0.11ms     6.0 MB/sec

Windows

group                                      main                                   pr
-----                                      ----                                   --
linter/all-rules/large/dataset.py          1.00     13.7±0.17ms     3.0 MB/sec    1.00     13.7±0.18ms     3.0 MB/sec
linter/all-rules/numpy/ctypeslib.py        1.00      3.4±0.03ms     4.9 MB/sec    1.00      3.4±0.04ms     4.9 MB/sec
linter/all-rules/numpy/globals.py          1.00    358.0±6.67µs     8.2 MB/sec    1.00    357.8±3.90µs     8.2 MB/sec
linter/all-rules/pydantic/types.py         1.01      5.7±0.06ms     4.4 MB/sec    1.00      5.7±0.06ms     4.5 MB/sec
linter/default-rules/large/dataset.py      1.01      7.3±0.07ms     5.6 MB/sec    1.00      7.2±0.07ms     5.7 MB/sec
linter/default-rules/numpy/ctypeslib.py    1.01  1508.0±13.70µs    11.0 MB/sec    1.00  1494.8±18.25µs    11.1 MB/sec
linter/default-rules/numpy/globals.py      1.00    156.1±1.85µs    18.9 MB/sec    1.00    156.3±3.88µs    18.9 MB/sec
linter/default-rules/pydantic/types.py     1.01      3.2±0.03ms     7.9 MB/sec    1.00      3.2±0.04ms     8.0 MB/sec

[charliermarsh]

GitHub isn't picking it up, but everything from here down is moved from rules/isort/rules/add_required_import.rs (which now uses this Importer struct).

[JonathanPlasse]

What happens if sysis shadowed?
Do you import sys with an alias like sys_ with as many _ necessary to avoid the shadowing?

[charliermarsh]

No, we avoid fixing if the name is already bound to something else. In practice, I think that's the least surprising behavior.

[MichaReiser]

Nice!

Can we add a test case that demonstrates what happens if we have

if test:
	import path


exit(1)

or

import path

if test:
	pass

import numpy

// ...

exit(1)

[konstin]

why not from functools import lru_cache? that's what i'd expect and afaik that's what pycharm generally does

[charliermarsh]

Yeah I wasn't sure what the default should be here. I was thinking that it should be rule-dependent (e.g., it makes sense to do from typing import List rather than import typing; typing.List if we need to inject from typing).

Where does PyCharm do this? When you have @cache and cache isn't imported? (But in that case, the user has already expressed that they want the member, not the fully-qualified name.)

[konstin]

yes, this:

@cache
def foo():
    ...

quck fix -> enter -> enter:

from functools import cache


@cache
def foo():
    ...

[charliermarsh]

Yeah but in that case, the user has already written @cache.

Given:

@functools.cache
def f(x):
    return x

Then the same action yields:

import functools

@functools.cache
def f(x):
    return x

Here, we don't know which the user prefers ahead-of-time.