[ty] Apply type mappings to functions eagerly

[dcreager]

TypeMapping is no longer cow-shaped.

Before, TypeMapping defined a to_owned method, which would make an owned copy of the type mapping. This let us apply type mappings to function literals lazily. The primary part of a function that you have to apply the type mapping to is its signature. The hypothesis was that doing this lazily would prevent us from constructing the signature of a function just to apply a type mapping; if you never ended up needed the updated function signature, that would be extraneous work.

But looking at the CI for this PR, it looks like that hypothesis is wrong! And this definitely cleans up the code quite a bit. It also means that over time we can consider replacing all of these TypeMapping enum variants with separate TypeTransformer impls.

[github-actions]

Diagnostic diff on typing conformance tests

No changes detected when running ty on typing conformance tests ✅

[github-actions]

mypy_primer results

Changes were detected when running on open source projects

psycopg (https://github.com/psycopg/psycopg)
- psycopg_pool/psycopg_pool/_acompat.py:185:12: error[invalid-return-type] Return type does not match returned value: expected `T@ensure_async`, found `object`
- Found 699 diagnostics
+ Found 698 diagnostics

scikit-build-core (https://github.com/scikit-build/scikit-build-core)
- src/scikit_build_core/build/_wheelfile.py:51:22: error[no-matching-overload] No overload of function `field` matches arguments
- Found 53 diagnostics
+ Found 52 diagnostics

Memory usage changes were detected when running on open source projects

trio (https://github.com/python-trio/trio)
-     struct fields = ~8MB
+     struct fields = ~9MB

sphinx (https://github.com/sphinx-doc/sphinx)
-     struct metadata = ~13MB
+     struct metadata = ~12MB
-     struct fields = ~12MB
+     struct fields = ~13MB

prefect (https://github.com/PrefectHQ/prefect)
-     struct fields = ~30MB
+     struct fields = ~33MB

[codspeed-hq]

CodSpeed WallTime Performance Report

Merging #20596 will improve performances by 6.85%

_{Comparing dcreager/stop-being-lazy (5fb5cc6) with main (3f640da)}

Summary

⚡ 1 improvement
✅ 7 untouched

Benchmarks breakdown

	Benchmark	BASE	HEAD	Change
⚡	large[sympy]	42.9 s	40.1 s	+6.85%

[codspeed-hq]

CodSpeed Instrumentation Performance Report

Merging #20596 will not alter performance

_{Comparing dcreager/stop-being-lazy (5fb5cc6) with main (3f640da)}

Summary

✅ 13 untouched
⏩ 30 skipped¹

Footnotes

1. 30 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[dcreager]

We might also consider consolidating FunctionType and FunctionLiteral into a single type now. (FunctionLiteral would have the extra fields that hold the updated signatures.)

[AlexWaygood]

do we not need to walk the new fields here (updated_signature and updated_last_definition_signature)? Would it be worth adding a comment about why that's not necessary?

[sharkdp]

Took me quite a bit of time to construct an example where this would have any observable effects, since you first need to apply a type mapping to the function literal and then have a type-visitor walk over it:

from typing import cast
from ty_extensions import TypeOf, Unknown


class C[T]:
    def f(self, t: T) -> T | Unknown:
        raise NotImplementedError


cast(TypeOf[C[int]().f], C[int]().f)

I'm first applying a specialization type mapping through C[int]().f, and then a type visitor via any_over_type(..., contains_unknown_or_todo, ..) by attempting a cast. Without walking over the updated_* fields, we get a

Value is already of type `bound method C[int].f(t: int) -> int | Unknown`

diagnostic (which shouldn't happen, if Unknown is part of the type). After applying the following patch, the diagnostic is gone.

@@ -692,6 +692,14 @@ pub(super) fn walk_function_type<'db, V: super::visitor::TypeVisitor<'db> + ?Siz
     visitor: &V,
 ) {
     walk_function_literal(db, function.literal(db), visitor);
+    if let Some(callable_signature) = function.updated_signature(db) {
+        for signature in &callable_signature.overloads {
+            walk_signature(db, signature, visitor);
+        }
+    }
+    if let Some(signature) = function.updated_last_definition_signature(db) {
+        walk_signature(db, signature, visitor);
+    }
 }

[AlexWaygood]

psycopg (https://github.com/psycopg/psycopg)
- psycopg_pool/psycopg_pool/_acompat.py:185:12: error[invalid-return-type] Return type does not match returned value: expected `T@ensure_async`, found `object`
- Found 699 diagnostics
+ Found 698 diagnostics

Is this ecosystem change correct? If so, can we add an mdtest making sure it doesn't regress?

[carljm]

This looks great! (Modulo existing inline comments.) Thank you!

[carljm]

Maybe some doc comments here explaining what these fields are and when/how they become populated?

[sharkdp]

I added some doc comments and made those fields non-public. @dcreager please feel free to change them in a follow-up, if something seems wrong/missing. I'd like to merge this PR to fix astral-sh/ty#1261.

[sharkdp]

Glad I looked at my notifications this morning before continuing to work on astral-sh/ty#1261! The lazy application of type mappings to functions (and the application of type mappings to those lazy type mappings) seemed to be the root cause of that issue, and applying type-mappings eagerly would have been one of the attempts to solve that bug. And indeed, the problem is completely solved with this PR!

Command	Mean [s]	Min [s]	Max [s]	Relative
./ty_main check /home/shark/playground	11.185 ± 0.064	11.117	11.245	385.74 ± 24.39
./ty_20596 check /home/shark/playground	0.029 ± 0.002	0.027	0.031	1.00

[sharkdp]

psycopg (https://github.com/psycopg/psycopg)
- psycopg_pool/psycopg_pool/_acompat.py:185:12: error[invalid-return-type] Return type does not match returned value: expected `T@ensure_async`, found `object`
- Found 699 diagnostics
+ Found 698 diagnostics

Is this ecosystem change correct? If so, can we add an mdtest making sure it doesn't regress?

from typing import Any, ParamSpec, TypeVar
from inspect import isawaitable
from collections.abc import Callable, Coroutine

T = TypeVar("T")
P = ParamSpec("P")

async def ensure_async(
    f: Callable[P, T] ,
    *args: P.args,
    **kwargs: P.kwargs,
) -> T:
    rv = f(*args, **kwargs)
    if isawaitable(rv):
        rv = await rv
    return rv

Yes, this is a false positive that shouldn't be there. I invested quite a bit of time trying to track it down, and while I haven't fully understood what is going on, it feels like it's not worth pursuing further (given that the behavior arguably improves). First of all, it took me a while to even reproduce it, because it only happens when checking on 3.12, not on 3.13 (which I use by default for testing locally). There are a lot of @Todo types involved, both because we create an intersection type in the if isawaitable(rv) narrowing, and because ParamSpec is involved (which we look for via any_over_type, which is affected by this PR). The correct return type from the await rv expression should be T@ensure_async, but for some reason we infer object on main. And we do infer object on this branch, when checking with 3.13. It might also have to do with normalization, because everything works fine when the seemingly irrelevant Callable[P, T] is removed from union in the f parameter annotation.

[AlexWaygood]

Thank you!

I wondered if we should add tests for #20596 (comment) and/or #20596 (comment), but they both seem obscure/strange enough that it would probably be quite low-value to do so :)