TCH setting to autofix with conversion to string

[smackesey]

Currently the TCH (flake8-type-checking) rules require annotations to be stored as strings (i.e. not evaluated). Annotations are stored a strings if:

• The annotations are string literals (i.e. are quoted)
• The annotations are in a module with from __future__ import annotations (which causes evaluation of all annotations in the module to be skipped)
• The annotations are of local variables (the Python interpreter does not evaluate local variable annotations)

One approach to getting TCH autofix working is to inject from __future__ import annotations into all your source modules (using I002). This is not ideal because string annotations unavoidably break code that meets these conditions:

1. annotations are being created in a local scope
2. an annotation references a symbol only available in that local scope
3. the annotation must be operated on at runtime

Here is an example of code meeting these conditions:

def make_bar():
    class Bar:
        ...

    def bar(x: "Bar"):
        ...
    return bar

bar = make_bar()

# `bar` has a string annotation-- if we want to operate on it, we need to evaluate it. This requires
# calling `get_type_hints`, but this fails because "Bar" can't be dereferenced since referent `Bar`
# vanished with local scope
get_type_hints(bar)  # => NameError: name 'Bar' is not defined.

We have plenty of code meeting these conditions in Dagster (mostly in tests), so my experiment in injecting from __future__ import annotations to enable TCH autofix has failed. Other projects that use runtime type annotation introspection are likely to encounter the same issues.

Therefore, a significant enhancement to TCH autofix would be to detect where annotations could be stored as strings, and perform a targeted conversion for you using quoting. Example:

from threading import Lock

def foo(lock: Lock):
    ...

After applying TCH autofix:

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from threading import Lock

def foo(lock: "Lock"):
    ...

This would make TCH autofix a lot more useful for us and probably many other projects.

[charliermarsh]

Yeah, I agree this would be useful. I'm not yet certain how hard it is to get right, but I agree that it would be useful... flake8-type-checking has rules for this, which we could look at for reference: https://pypi.org/project/flake8-type-checking/.

annotations are being created in a local scope.

Can you give an example of this, just to illustrate the problem?

[smackesey]

Can you give an example of this, just to illustrate the problem?

I updated OP with an example.

[charliermarsh]

Thanks. I think it will be difficult to reliably determine whether an annotation is required at runtime 😓

[smackesey]

I think it will be difficult to reliably determine whether an annotation is required at runtime

Haven't thought it all the way through, but can't you just see if the only uses of a symbol are inside non-local type annotations, while also applying the existing settings used to special case e.g. Pydantic? Maybe determining scoping is harder than I think.

[charliermarsh]

Can you explain what you mean in terms of this example? (I.e., why x: Bar can't be quoted.)

def make_bar():
    class Bar:
        ...

    def bar(x: Bar):
        ...

    return bar

bar = make_bar()

print(get_type_hints(bar))

[smackesey]

Can you explain what you mean in terms of this example? (I.e., why x: Bar can't be quoted.)

• With x: Bar (unquoted case), bar.__annotations__["x"] stores a direct reference to Bar.
• When x: "Bar" (quoted case), bar.__annotations__["x"] stores the string "Bar".
• When get_type_hints(bar) is applied, Python tries to resolve each annotation
  • If it's a type, just return it.
  • If it's as string, dereference it.
• When Python tries to dereference the string "Bar", it comes up empty and throws a NameError.

I'm not sure why Python fails in exactly this scenario. I guess the defining module for the object holding the annotations is known, so Python is able to look up the referent of string annotations if they exist in the defining module's __dict__.

[charliermarsh]

But, does "whether we can quote it" depend on whether and when get_type_hints is called? Or is there a more robust and general methodology we can follow?

[smackesey]

But, does "whether we can quote it" depend on whether and when get_type_hints is called? Or is there a more robust and general methodology we can follow?

I would say just be conservative and always assume that get_type_hints might be called-- so if it references a local scope symbol, don't quote it.

In practice referencing local scope symbols is going to be pretty rare in 99% of Python code anyway.

[charliermarsh]

Ah I see -- so we would avoid quoting annotations that rely on symbols defined outside of the module scope.

[charliermarsh]

It's also a little less critical for the TCH rules in particular since we're only dealing with imported symbols. But I suppose if the import itself is within a function scope...

[charliermarsh]

I have something that I've been testing against Dagster here: #6001. The outputs look correct to me -- one example:

--- a/examples/assets_smoke_test/assets_smoke_test/pure_python_assets.py
+++ b/examples/assets_smoke_test/assets_smoke_test/pure_python_assets.py
@@ -1,5 +1,9 @@
+from typing import TYPE_CHECKING
+
 from dagster import SourceAsset, TableSchema, asset
-from pandas import DataFrame
+
+if TYPE_CHECKING:
+    from pandas import DataFrame

 raw_country_populations = SourceAsset(
     "raw_country_populations",
@@ -19,7 +23,7 @@


 @asset
-def country_populations(raw_country_populations) -> DataFrame:
+def country_populations(raw_country_populations) -> "DataFrame":
     country_populations = raw_country_populations.copy()
     country_populations["change"] = (
         country_populations["change"]
@@ -32,13 +36,13 @@ def country_populations(raw_country_populations) -> DataFrame:


 @asset
-def continent_stats(country_populations: DataFrame) -> DataFrame:
+def continent_stats(country_populations: "DataFrame") -> "DataFrame":
     result = country_populations.groupby("continent").agg({"pop2019": "sum", "change": "mean"})
     return result


 @asset
-def country_stats(country_populations: DataFrame, continent_stats: DataFrame) -> DataFrame:
+def country_stats(country_populations: "DataFrame", continent_stats: "DataFrame") -> "DataFrame":
     result = country_populations.join(continent_stats, on="continent", lsuffix="_continent")
     result["continent_pop_fraction"] = result["pop2019"] / result["pop2019_continent"]
     return result

But not sure how best to check whether the rule is having any unintended effects on runtime behavior.

[tylerlaprade]

If there's an easy way for me to locally run a preview version of Ruff, I'm happy to run your enhanced rule on my private codebase and verify the runtime behavior. Obviously would be easier if someone else has a large open-source codebase to test on, though.

[charliermarsh]

Thanks! I can kick off a release build dry-run which will generate wheel files that you can download and pip install. I'll post here with the link.

[charliermarsh]

If you'd like, you can click here, download the "Wheels" at the bottom, unzip, and pip install the wheel that matches your platform.

[tylerlaprade]

This looks much more comprehensive than when I made from __future__ import annotations required (which barely modified anything)! I really like the looks of it. However, a runtime error has been introduced that doesn't get detected by my typechecker:

[tylerlaprade]

I did a global find-and-replace to convert "(\w+)" \| None to $1 | None. This fixed all my runtime issues (at least until I re-ran Ruff).

EDIT: The above find-and-replace loses some of the benefits of this rule. Replacing with "$1 | None" is much better for two reasons:

1. It retains the performance benefits of this rule.
2. It does not get undone when re-running Ruff.

[charliermarsh]

Oh yes, that case makes sense, we need handle it. Thanks so much for testing it out!

[smackesey]

The above outputs look good to me, as long as there's a setting to avoid applying the rule to functions with certain decorators (and IIRC this already exists).

But not sure how best to check whether the rule is having any unintended effects on runtime behavior.

Potential unintended runtime effects are unavoidable since people could always be accessing .__annotations__ and performing arbitrary logic with it. But as long as those annotations aren't coming from local scope, they can be easily worked around by changing direct .__annotations__ access to use typing.get_type_hints.

[charliermarsh]

I'm blocked on supporting unions (e.g., when quoting DataFrame in DataFrame | int, you need to quote the entire union, like "DataFrame | int"). There are just a bunch of edge cases and we don't have access to the full AST when constructing the fix (so we need to do manual lexing). This might require some new infrastructure to be fully robust.

[tylerlaprade]

Totally understandable. For now, is there value in shipping it with an exclusion to not add quotes to any types involving unions? It still seemed like a major improvement over from __future__ import annotations when I tested the preview.