Don't replace string type annotations when module is only imported with TYPE_CHECKING

[david-zwicker]

pyupgrade aggressively replaces types that are given as strings, which breaks my code. Here's minimal code to demonstrate the problem:

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import sympy

def f(sym: "sympy.Symbol"):
    import sympy
    ...

The idea is that the sympy package is a heavy import and I thus just want to import it lazily when necessary. Yet, I want to give proper types, which in principle is possible using strings. Additionally using the TYPE_CHECKING guard, mypy can still check the types.

The problem is that pyupgrade automatically changes the code to

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import sympy

def f(sym: sympy.Symbol):
    import sympy
    ...

which doesn't run anymore since sympy.Symbol is not available in a normal context (when TYPE_CHECKING is False).
Is there a work-around with which I can prevent this problem?

[asottile]

I assume you closed this because you were wrong?

[david-zwicker]

Yeah, the issue was that my IDE is not smart enough about this and marked an error. I mistakenly thought that python had a problem with the undefined sympy, but apparently it is smarter than I thought :) All good!