Skip to content

[ty] Smarter semantic token classification for attribute access on union type#23841

Merged
MichaReiser merged 7 commits intoastral-sh:mainfrom
Geo5:semantic-tokens-union
Mar 17, 2026
Merged

[ty] Smarter semantic token classification for attribute access on union type#23841
MichaReiser merged 7 commits intoastral-sh:mainfrom
Geo5:semantic-tokens-union

Conversation

@Geo5
Copy link
Contributor

@Geo5 Geo5 commented Mar 9, 2026

Summary

ref ty/#791

Updates the semantic token classification for attribute access on union types. It now checks if all elements of the union belong to the same type variant and if so, recursively tries to classify that common variant. Imo this produces nicer results especially when calling a method on an object of some union type.

I decided not to specifically check if all union elements are of some fixed set of variants (e.g. only Type::BoundMethod etc), because i felt like duplicating the logic may be brittle. That means, that some hypothetical type of Union[Union[BoundMethod, BoundMethod], Union[something_else, something_else]] may be misclassified as Method in this case, by i am unsure if it is possible to create such a type.

I also did not change behavior for free function classification, because i could not think of a way where this matters, i.e. creating an alias to a function, already changes the classification to `Variable` without involving unions at all:
from random import random

def f(): ...
def g(): ...

f_alias = f
f_alias()

f_or_g = f if random() else g
f_or_g()

https://play.ty.dev/d8acae76-5642-4841-b4cf-8fd6066cea87

I also added a line in the existing semantic_token::test::attribute_classification test, where there was no test for MyClass.prop, which actually classifies as Property.

Test Plan

Added these two unit test to semantic_tokens.rs:

attribute_on_union_1

Playground 1

main.py

from random import random

class Foo:
    CONSTANT = 42

    def method(self):
        return "hello"

    @property
    def prop(self) -> str:
        return "hello"

class Bar:
    CONSTANT = 24

    def method(self, x: int = 1) -> int:
        return 42

    @property
    def prop(self) -> int:
        return self.CONSTANT


foobar = Foo() if random() else Bar()
y = foobar.method                                # method should be method (bound method)
z = foobar.CONSTANT                              # CONSTANT should be variable with readonly modifier
w = foobar.prop                                  # prop should be property
foobar_cls = Foo if random() else Bar
v = foobar_cls.method                            # method should be method (function)
x = foobar_cls.prop                              # prop should be property
attribute_on_union_2

Playground 2

main.py

from random import random

# There is also this way to create union types:
class Baz:
    if random():
        CONSTANT = 42

        def method(self) -> int:
            return 42

        @property
        def prop(self) -> int:
            return 42
    else:
        CONSTANT = "hello"

        def method(self) -> str:
            return "hello"

        @property
        def prop(self) -> str:
            return "hello"

baz = Baz()
s = baz.method      # method should be bound method
t = baz.CONSTANT    # CONSTANT should be variable with readonly
r = baz.prop        # prop should be property
q = Baz.prop        # prop should be property on the class as well
attributes_on_union_3

Playground 3

This one is more of a negative test, where the union does not contain elements of the same type variant and the classification should therefore fallback to just Variable.

main.py

from random import random

class Baz:
    if random():
        CONSTANT = 42

        def method(self) -> int:
            return 42

        @property
        def prop(self) -> int:
            return 42
    else:
        def CONSTANT(self):
            return "hello"

        @property
        def method(self) -> str:
            return "hello"

        prop: str = "hello"

baz = Baz()
s = baz.method 
t = baz.CONSTANT
r = baz.prop
q = Baz.prop

Geo5 added 2 commits March 9, 2026 17:44
access on objects with a union type.
Mirrors the semantic_tokens::tests::attribute_classification
test.

Also fixes that test in regards to `property` classification.
Previously the comment did not match the result.
on union types.
Update snapshot test from previous commit to the desired results.
@astral-sh-bot astral-sh-bot bot requested a review from dhruvmanila March 9, 2026 17:50
@AlexWaygood AlexWaygood removed their request for review March 9, 2026 17:52
@AlexWaygood AlexWaygood added server Related to the LSP server ty Multi-file analysis & type inference labels Mar 9, 2026
@carljm carljm removed their request for review March 9, 2026 23:44
y = obj.method # method should be method (bound method)
z = obj.CONSTANT # CONSTANT should be variable with readonly modifier
w = obj.prop # prop should be property
w = obj.prop # prop should be variable on an instance
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's leave this comment as is. I do think we want to classify this as a property and I find variable on an instance more difficult to understand. We can rephrase it to instance property if you prefer that.

if Self::is_constant_name(attr_name_str) {
modifiers |= SemanticTokenModifier::READONLY;
}
// If we got a union type, check if all elements of the union belong to the same
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The issue with foobar.prop is that ty isn't the type of the value (that is, the type of foobar) but the type of foobar.prop and foobar.pop inferes to Unknown (because the properties lack a return type annotation).

If we change the code snippet like so

from random import random

class Foo:
    @property
    def prop(self) -> int:
        return 4

class Bar:
    @property
    def prop(self) -> int:
        return 3


foobar: Foo | Bar = Foo() if random() else Bar()
w = foobar.prop                                  # prop should be variable on an instance

then ty is int as expected.

Which makes me question if the fix here is correct. I think what we want is to infer the type of foobar and if that's a union, classify variant.attr for each union variant.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mhm yes, i agree with your analysis, but i don't think the semantic token type changes between Unknown and int, because the same match arm will be taken?

So i think the match arm to check for a property here does not do the correct thing for foobar.prop as you said, but i think that preexists this change and is also happening in the preexisting, other test where i changed the comment.

Regarding your second sentence; i think this is already done somewhat by the type inference anyway? Since at least for foobar.method the inferred type is a union of bound methods, which we can than classify.

Would you be fine with deferring the bugfix for property classification or would you rather have it fixed in this PR? (I feel like that fix could be more involved than the simple loop i added here until now 😅 )

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think the fix requires that many changes. It's probably something like:

let value_ty = value.inferred_type(model);

if let Some(union) = value_ty.as_union() {
	// classify each variant
	// if all classify the same, return early
}

// Old, unchanged logic:
let ty = expr.inferred_type(model);
...

Copy link
Contributor Author

@Geo5 Geo5 Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought a bit more about this and i think this change alone would still not be perfect, consider this case:

from random import random

class Baz:
    if random():
        @property
        def prop(self) -> int:
            return 42
    else:
        @property
        def prop(self) -> str:
            return "hello"

baz = Baz()
x = baz.prop

where the value_ty would not be a union.
(There may actually be a real bug in type inference for the above example: astral-sh/ty#3032)

I think to fix the property classification you would need to look up the attribute name on the class of the object and check if it is defined with @property. Is that something that is already supported by some function you are aware of in the codebase?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll create a follow up issue for this

} else {
None
}
})()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the reason that this is an immediately invoked lambda?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought i could combine it with the outer match somehow, but iiuc r-a transformed it to a "if let guard" which seems to only be stabilized very recently. I'll change it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lmk if you find the functional style even less appealing 😅

@Geo5 Geo5 marked this pull request as draft March 10, 2026 11:07
@Geo5
Copy link
Contributor Author

Geo5 commented Mar 16, 2026

I tried implementing the suggestion from #23841 (comment)
to fix the property classification, but i stopped at something like:

                let value_ty = attr
                    .value
                    .inferred_type(self.model)
                    .unwrap_or(Type::unknown());
                if let Some(class_ty) = value_ty.as_class_literal() {
                    let place = class_ty.class_member(
                        self.model.db(),
                        &attr.attr,
                        MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK,
                    );
                    // Check if the member is a property
                    if let Ok(typ) = place.into_lookup_result(self.model.db()) {
                        if typ.inner_type().is_property_instance() {
                            self.add_token(
                                &attr.attr,
                                SemanticTokenType::Property,
                                SemanticTokenModifier::empty(),
                            );
                        }
                    }
                }

which did not work and required making many more things public.

So all in all i don't feel comfortable making that change as well and opening back up for review.

@Geo5 Geo5 marked this pull request as ready for review March 16, 2026 15:42
@astral-sh-bot astral-sh-bot bot requested a review from MichaReiser March 16, 2026 15:42
@astral-sh-bot
Copy link

astral-sh-bot bot commented Mar 17, 2026

Typing conformance results

No changes detected ✅

Current numbers
The percentage of diagnostics emitted that were expected errors held steady at 85.29%. The percentage of expected errors that received a diagnostic held steady at 78.13%. The number of fully passing files held steady at 64/132.

@MichaReiser
Copy link
Member

Thank you. Your approach makes sense. I simplified it a bit to avoid using discriminant and recursion but this looks good to me. Thank you, and sorry for my over-complicated suggestion earlier.

@MichaReiser MichaReiser enabled auto-merge (squash) March 17, 2026 12:37
@MichaReiser MichaReiser disabled auto-merge March 17, 2026 12:38
@astral-sh-bot
Copy link

astral-sh-bot bot commented Mar 17, 2026

Memory usage report

Memory usage unchanged ✅

@astral-sh-bot
Copy link

astral-sh-bot bot commented Mar 17, 2026

ecosystem-analyzer results

Lint rule Added Removed Changed
invalid-await 40 0 0
invalid-return-type 1 0 0
Total 41 0 0

Full report with detailed diff (timing results)

@MichaReiser MichaReiser merged commit 471fd12 into astral-sh:main Mar 17, 2026
51 checks passed
carljm added a commit that referenced this pull request Mar 17, 2026
* main:
  [ty] Filter out unsatisfiable inference attempts during generic call narrowing (#24025)
  [ty] Avoid inferring intersection types for call arguments (#23933)
  [ty] Pin mypy_primer in `setup_primer_project.py` (#24020)
  [`pycodestyle`] Recognize `pyrefly:` as a pragma comment (`E501`) (#24019)
  Add company AI policy to contributing guide (#24021)
  [ty] Remove the mypy_primer CI workflow (#24016)
  Update prek dependencies (#23980)
  [ty] Smarter semantic token classification for attribute access on union type (#23841)
  [ty] ecosystem-analyzer: Inline diffs and panic messages (#24015)
  [ty] Improve `.toml` support in the ty playground (#23476)
  PEP 639 license information (#19661)
@Geo5
Copy link
Contributor Author

Geo5 commented Mar 18, 2026

no worries at all and your implementation is way cleaner! 😄

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

server Related to the LSP server ty Multi-file analysis & type inference

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants