[ty] Implement non-stdlib stub mapping for classes and functions #19471

Gankra · 2025-07-21T18:28:05Z

This implements mapping of definitions in stubs to definitions in the "real" implementation using the approach described in astral-sh/ty#788 (comment)

I've tested this with goto-definition in vscode with code that uses colorama and types-colorama (will put up a screen recording and hopefully tests after lunch).

Notably this implementation does not add support for stub-mapping stdlib modules, which can be done as an essentially orthogonal followup in the implementation of resolve_real_module.

Partially fixes astral-sh/ty#788

crates/ty_python_semantic/src/module_resolver/resolver.rs

Gankra · 2025-07-21T18:30:40Z

crates/ty_python_semantic/src/types/ide_support.rs

+            NodeWithScopeKind::TypeAlias(_)
+            | NodeWithScopeKind::ClassTypeParameters(_)
+            | NodeWithScopeKind::FunctionTypeParameters(_)
+            | NodeWithScopeKind::TypeAliasTypeParameters(_)
+            | NodeWithScopeKind::Lambda(_)
+            | NodeWithScopeKind::ListComprehension(_)
+            | NodeWithScopeKind::SetComprehension(_)
+            | NodeWithScopeKind::DictComprehension(_)
+            | NodeWithScopeKind::GeneratorExpression(_) => {
+                return Err(());


I simply haven't even tried to think about these.

Gankra · 2025-07-21T18:30:55Z

crates/ty_python_semantic/src/types/ide_support.rs

+            DefinitionKind::TypeAlias(_)
+            | DefinitionKind::Import(_)
+            | DefinitionKind::ImportFrom(_)
+            | DefinitionKind::StarImport(_)
+            | DefinitionKind::NamedExpression(_)
+            | DefinitionKind::Assignment(_)
+            | DefinitionKind::AnnotatedAssignment(_)
+            | DefinitionKind::AugmentedAssignment(_)
+            | DefinitionKind::For(_)
+            | DefinitionKind::Comprehension(_)
+            | DefinitionKind::VariadicPositionalParameter(_)
+            | DefinitionKind::VariadicKeywordParameter(_)
+            | DefinitionKind::Parameter(_)
+            | DefinitionKind::WithItem(_)
+            | DefinitionKind::MatchPattern(_)
+            | DefinitionKind::ExceptHandler(_)
+            | DefinitionKind::TypeVar(_)
+            | DefinitionKind::ParamSpec(_)
+            | DefinitionKind::TypeVarTuple(_) => {
+                // Not yet implemented
+                return Err(());


Similarly, didn't even try to think about these.

crates/ty_python_semantic/src/types/ide_support.rs

github-actions · 2025-07-21T18:32:21Z

`mypy_primer` results

No ecosystem changes detected ✅
No memory usage changes detected ✅

Gankra · 2025-07-21T18:33:06Z

crates/ty_python_semantic/src/types/ide_support.rs

+                    // Not yet implemented -- in this case we want to recover something like a Definition
+                    // and build a Definition Path, but this input is a bit too abstract for now.
+                    trace!("Found arbitrary FileWithRange by stub mapping, giving up");
+                    None


I believe this path is currently only used for going to a function keyword argument or something like that?

Looks like it -- but those also have a Definition, so it should be possible to stop using FileWithRange for that case.

It seems like maybe FileWithRange should be more finely targeted to only the "entire module" case -- in every other case I'd think we should have a Definition.

It actually used to just be Module(File) but then i pulled main and it had turned into FileWithRange with that extra case. Not sure what's up with that.

I changed it to FileWithRange when I added support for keyword arguments. @carljm, are you sure those have definitions? I found definitions for keyword parameters but not for keyword arguments.

If I'm reading the code correctly in definitions_for_keyword_argument, the source is an argument, but the ResolvedDefinition resolves to the corresponding parameter in the callable's signature. That's what the doc comment says, and what the code appears to do.

It is correct that there's no Definition for the argument, but there is one for the parameter. Which seems like it should mean that we can return a ResolvedDefinition::Definition there.

crates/ty_python_semantic/src/module_resolver/resolver.rs

github-actions · 2025-07-21T18:38:56Z

`ruff-ecosystem` results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

Formatter (stable)

ℹ️ ecosystem check encountered format errors. (no format changes; 1 project error)

openai/openai-cookbook (error)

warning: Detected debug build without --no-cache.
error: Failed to parse examples/mcp/databricks_mcp_cookbook.ipynb:12:1:8: Simple statements must be separated by newlines or semicolons

Formatter (preview)

ℹ️ ecosystem check encountered format errors. (no format changes; 1 project error)

openai/openai-cookbook (error)

ruff format --preview

warning: Detected debug build without --no-cache.
error: Failed to parse examples/mcp/databricks_mcp_cookbook.ipynb:12:1:8: Simple statements must be separated by newlines or semicolons

AlexWaygood · 2025-07-21T18:46:13Z

It looks like for something like this, where there isn't a separate *-stubs package but instead the stub lives alongside the runtime files in the same package, we still jump to the stub file in go-to definition on this branch:

main.py:

from module import Foo

module.py:

class Foo: ...

module.pyi:

class Foo: ...

E.g. I ran the playground locally with ^those files on this branch, and used go-to definition on the Foo symbol in main.py, and it jumped to module.pyi rather than module.py.

That said, it's something of a distinct task from the "separate *-stubs package" task, so I don't think it should block this PR being merged!

AlexWaygood · 2025-07-21T18:48:04Z

crates/ty_python_semantic/src/types/ide_support.rs

+        let stub_module = file_to_module(db, stub_file)?;
+        trace!("Found stub module: {}", stub_module.name());
+        let real_module = resolve_real_module(db, stub_module.name())?;
+        trace!("Found real module: {}", real_module.name());


I'm choosing to take the insinuation that stubs are somehow "fake" or "inauthentic" as a personal affront, since I'm a typeshed maintainer 😆

Only in the artifice can we find perfection

crates/ty_ide/src/goto_definition.rs

Gankra · 2025-07-21T21:01:21Z

Tests added

AlexWaygood · 2025-07-21T21:32:06Z

Here's a slightly pathological situation where this branch still has... questionable? behaviour (but it's an edge case, so maybe not hugely important):

main.py:

from module import Foo

module/__init__.py:

# empty file

module/Foo.py:

# empty file

module.pyi:

class Foo: ...

If you go-to-definition on Foo in main.py, we currently jump to the class Foo: ... definition in module.pyi on this branch. But arguably the "runtime definition" of Foo is the submodule module/Foo.py (that's what will actually be imported by the statement from module import Foo at runtime!).

Gankra · 2025-07-21T21:57:37Z

Not sure how you were testing that -- in the cursor tests it actually yields "no goto (definition) target found". If you created this is an IDE and e.g. cmd+clicked, you may have just unknowingly triggered goto-declaration (if there's no results for goto-definition, goto-declaration is used by the IDE).

I've pushed up a commit adding a test for this.

Gankra · 2025-07-21T22:02:27Z

It looks like for something like this, where there isn't a separate *-stubs package but instead the stub lives alongside the runtime files in the same package, we still jump to the stub file in go-to definition on this branch:

I also just added a test for this case -- it's fixed now. In an earlier revision I just wasn't propagating ModuleResolveMode deep enough so simple cases still were grabbing .pyi results.

carljm

Fantastic! More to do here, but this is a great start and covers the cases people will care most about.

crates/ty_ide/src/stub_mapping.rs

carljm · 2025-07-21T22:00:43Z

crates/ty_python_semantic/src/types/ide_support.rs

+                    // Not yet implemented -- in this case we want to recover something like a Definition
+                    // and build a Definition Path, but this input is a bit too abstract for now.
+                    trace!("Found arbitrary FileWithRange by stub mapping, giving up");
+                    None


Looks like it -- but those also have a Definition, so it should be possible to stop using FileWithRange for that case.

It seems like maybe FileWithRange should be more finely targeted to only the "entire module" case -- in every other case I'd think we should have a Definition.

crates/ty_python_semantic/src/types/ide_support.rs

carljm · 2025-07-21T22:10:35Z

crates/ty_ide/src/goto_definition.rs

+            .source(
+                "mymodule/__init__.py",
+                r#"
+# empty file


IMO we should not support the case as shown here, because the stub is kind of just lying. There is no MyClass in the namespace of mymodule (by default -- there would be if someone else happened to have done an import mymodule.MyClass somewhere, but this is an unreliable side effect we intentionally don't model.) In other words, I don't think we should eagerly go fetch submodules and treat them as attributes of the parent module for this purpose.

If, on the other hand, there were a from . import MyClass line right here, then I think it would make sense to find that as the definition of MyClass.

If we are going to include this test, I would suggest adding that import to it.

IMO we should not support the case as shown here, because the stub is kind of just lying.

I can't see the full test that you were commenting on here, but note that there are many places where the stub will be "kind of just lying", where it will be important for us to be able to figure out where the actual, truthful runtime definition of something is. For example, if you do "goto declaration" for typing.Iterable in Pylance, it takes you to https://github.com/python/typeshed/blob/ddd8434a42b83923b770450386ee12672ab3c604/stdlib/typing.pyi#L501-L504, but if you do "goto definition", it takes you to https://github.com/python/cpython/blob/c933a6bb329bb97bc7e448388dad1b74f7ca4baa/Lib/typing.py#L2714. The definition is a class in the stub but a variable assignment at runtime -- but being able to jump to the runtime definition is very useful if you're trying to figure out why the runtime behaviour of these aliases is different for some reason!

Similarly, typeshed will usually pretend that types.SimpleNamespace instances are instances of custom classes rather than instances of SimpleNamespace, since attributes on SimpleNamespace instances are untypable

I'm just gonna remove this test for now.

crates/ty_ide/src/goto_definition.rs

crates/ty_python_semantic/src/module_resolver/resolver.rs

MichaReiser · 2025-07-22T06:25:45Z

crates/ty_python_semantic/src/types/ide_support.rs

+        }
+
+        // It's definitely a stub, so now rerun module resolution but with stubs disabled.
+        let stub_module = file_to_module(db, stub_file)?;


This is something we discussed on Friday and I'm fine deferring this to another PR but file_to_module will fail for every resolved real module. That means that import resolution won't work for those files. We may want to change file_to_module to always return the module name if file matches the real_module (and not only the resolved_module)

AlexWaygood · 2025-07-22T11:17:03Z

Not sure how you were testing that -- in the cursor tests it actually yields "no goto (definition) target found". If you created this is an IDE and e.g. cmd+clicked, you may have just unknowingly triggered goto-declaration (if there's no results for goto-definition, goto-declaration is used by the IDE).

I've pushed up a commit adding a test for this.

Oh, yeah, I was interactively testing by running the playground locally using your branch. (You can do this by running cd playground; npm start --workspace ty-playground from the root of the Ruff repository, then opening http://localhost:5173/ in a browser window. It's a great way to experiment with IDE features!) So it might well have just fallen back to goto-declaration!

* main: (76 commits) Move fix suggestion to subdiagnostic (#19464) [ty] Implement non-stdlib stub mapping for classes and functions (#19471) [ty] Disallow illegal uses of `ClassVar` (#19483) [ty] Disallow `Final` in function parameter/return-type annotations (#19480) [ty] Extend `Final` test suite (#19476) [ty] Minor change to diagnostic message for invalid Literal uses (#19482) [ty] Detect illegal non-enum attribute accesses in Literal annotation (#19477) [ty] Reduce size of `TypeInference` (#19435) Run MD tests for Markdown-only changes (#19479) Revert "[ty] Detect illegal non-enum attribute accesses in Literal annotation" [ty] Detect illegal non-enum attribute accesses in Literal annotation [ty] Added semantic token support for more identifiers (#19473) [ty] Make tuple subclass constructors sound (#19469) [ty] Pass down specialization to generic dataclass bases (#19472) [ty] Garbage-collect reachability constraints (#19414) [ty] Implicit instance attributes declared `Final` (#19462) [ty] Expansion of enums into unions of literals (#19382) [ty] Avoid rechecking the entire project when changing the opened files (#19463) [ty] Add warning for unknown `TY_MEMORY_REPORT` value (#19465) [ty] Sync vendored typeshed stubs (#19461) ...

* main: [ty] Use `ThinVec` for sub segments in `PlaceExpr` (#19470) [ty] Splat variadic arguments into parameter list (#18996) [`flake8-pyi`] Skip fix if all `Union` members are `None` (`PYI016`) (#19416) Skip notebook with errors in ecosystem check (#19491) [ty] Consistent use of American english (in rules) (#19488) [ty] Support iterating over enums (#19486) Fix panic for illegal `Literal[…]` annotations with inner subscript expressions (#19489) Move fix suggestion to subdiagnostic (#19464) [ty] Implement non-stdlib stub mapping for classes and functions (#19471) [ty] Disallow illegal uses of `ClassVar` (#19483) [ty] Disallow `Final` in function parameter/return-type annotations (#19480) [ty] Extend `Final` test suite (#19476) [ty] Minor change to diagnostic message for invalid Literal uses (#19482) [ty] Detect illegal non-enum attribute accesses in Literal annotation (#19477) [ty] Reduce size of `TypeInference` (#19435) Run MD tests for Markdown-only changes (#19479) Revert "[ty] Detect illegal non-enum attribute accesses in Literal annotation" [ty] Detect illegal non-enum attribute accesses in Literal annotation [ty] Added semantic token support for more identifiers (#19473) [ty] Make tuple subclass constructors sound (#19469)

* main: (28 commits) [ty] highlight the argument in `static_assert` error messages (astral-sh#19426) [ty] Infer single-valuedness for enums based on `int`/`str` (astral-sh#19510) [ty] Restructure submodule query around `File` dependency [ty] Make `Module` a Salsa ingredient [ty] Reachability analysis for `isinstance(…)` branches (astral-sh#19503) [ty] Normalize single-member enums to their instance type (astral-sh#19502) [ty] Invert `ty_ide` and `ty_project` dependency (astral-sh#19501) [ty] Implement mock language server for testing (astral-sh#19391) [ty] Detect enums if metaclass is a subtype of EnumType/EnumMeta (astral-sh#19481) [ty] perform type narrowing for places marked `global` too (astral-sh#19381) [ty] Use `ThinVec` for sub segments in `PlaceExpr` (astral-sh#19470) [ty] Splat variadic arguments into parameter list (astral-sh#18996) [`flake8-pyi`] Skip fix if all `Union` members are `None` (`PYI016`) (astral-sh#19416) Skip notebook with errors in ecosystem check (astral-sh#19491) [ty] Consistent use of American english (in rules) (astral-sh#19488) [ty] Support iterating over enums (astral-sh#19486) Fix panic for illegal `Literal[…]` annotations with inner subscript expressions (astral-sh#19489) Move fix suggestion to subdiagnostic (astral-sh#19464) [ty] Implement non-stdlib stub mapping for classes and functions (astral-sh#19471) [ty] Disallow illegal uses of `ClassVar` (astral-sh#19483) ... # Conflicts: # crates/ty_ide/src/goto.rs

Implement stub mapping for classes and functions

e3bce74

Gankra requested review from AlexWaygood, MichaReiser, carljm, dcreager and sharkdp as code owners July 21, 2025 18:28

Gankra added server Related to the LSP server ty Multi-file analysis & type inference labels Jul 21, 2025

Gankra commented Jul 21, 2025

View reviewed changes

crates/ty_python_semantic/src/module_resolver/resolver.rs Outdated Show resolved Hide resolved

Gankra commented Jul 21, 2025

View reviewed changes

crates/ty_python_semantic/src/types/ide_support.rs Outdated Show resolved Hide resolved

Gankra commented Jul 21, 2025

View reviewed changes

AlexWaygood reviewed Jul 21, 2025

View reviewed changes

crates/ty_python_semantic/src/module_resolver/resolver.rs Outdated Show resolved Hide resolved

AlexWaygood reviewed Jul 21, 2025

View reviewed changes

Gankra added 2 commits July 21, 2025 16:05

fixups

c0bc99a

add goto_definition tests

449de62

Gankra commented Jul 21, 2025

View reviewed changes

crates/ty_ide/src/goto_definition.rs Show resolved Hide resolved

more properly propagate ModuleResolveMode

69157c5

add test for Alex

d691f50

add another alex test

603d4c2

carljm approved these changes Jul 21, 2025

View reviewed changes

cleanup and lots more tests

37c10ff

Gankra commented Jul 21, 2025

View reviewed changes

crates/ty_ide/src/goto_definition.rs Show resolved Hide resolved

Gankra commented Jul 21, 2025

View reviewed changes

crates/ty_ide/src/goto_definition.rs Show resolved Hide resolved

carljm reviewed Jul 21, 2025

View reviewed changes

crates/ty_ide/src/goto_definition.rs Outdated Show resolved Hide resolved

crates/ty_ide/src/goto_definition.rs Outdated Show resolved Hide resolved

Gankra added 2 commits July 21, 2025 18:56

cleanup comments

1a8cf85

dedup stub-mapped definitions

a10bb30

MichaReiser reviewed Jul 22, 2025

View reviewed changes

fixups

b48e773

Gankra enabled auto-merge (squash) July 22, 2025 12:39

Gankra merged commit c82fa94 into main Jul 22, 2025
36 checks passed

Gankra deleted the gankra/remod branch July 22, 2025 12:42

Gankra mentioned this pull request Jul 22, 2025

file_to_module will misbehave on "real modules" astral-sh/ty#869

Closed

AxillV mentioned this pull request Aug 10, 2025

Surface the "source definition" over the stub definition astral-sh/ty#788

Closed

[ty] Implement non-stdlib stub mapping for classes and functions #19471

[ty] Implement non-stdlib stub mapping for classes and functions #19471

Uh oh!

Conversation

Gankra commented Jul 21, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

mypy_primer results

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

ruff-ecosystem results

Linter (stable)

Linter (preview)

Formatter (stable)

Formatter (preview)

Uh oh!

AlexWaygood commented Jul 21, 2025

Uh oh!

AlexWaygood Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Gankra commented Jul 21, 2025

Uh oh!

AlexWaygood commented Jul 21, 2025

Uh oh!

Gankra commented Jul 21, 2025

Uh oh!

Gankra commented Jul 21, 2025

Uh oh!

carljm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AlexWaygood commented Jul 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

github-actions bot commented Jul 21, 2025 •

edited

Loading

`mypy_primer` results

github-actions bot commented Jul 21, 2025 •

edited

Loading

`ruff-ecosystem` results

AlexWaygood Jul 21, 2025 •

edited

Loading

AlexWaygood commented Jul 22, 2025 •

edited

Loading