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

[Geo5]

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

The first commit adds a unit test, the second commit implements the change and updates the snapshot.

I also added a line in the existing semantic_token::test::attribute_classification test, where previously the comment said that obj.prop should be classified as Property, but was classified as Variable and there was no test for MyClass.prop, which actually classifies as Property. Happy to change this back, if this changes the intention!

Test Plan

Added this unit test to semantic_tokens.rs:

from random import random

class Foo:
    CONSTANT = 42

    def method(self):
        return "hello"

    @property
    def prop(self):
        return self.CONSTANT

class Bar:
    CONSTANT = 24

    def method(self, x: int = 1):
        return "bye"

    @property
    def prop(self):
        return self.CONSTANT


foobar: Foo | Bar = 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 variable on an instance
foobar_cls = Foo if random() else Bar
v = foobar_cls.method                            # method should be method (function)
x = foobar_cls.prop                              # prop should be property

Current behavior: https://play.ty.dev/48bd1df0-42e8-43a3-9164-cd7779300019

The implementation surprisingly only changes the foobar.method and foobar_cls.method classifications, not the foobar_cls.prop classification, which must somehow already be handled by this match arm:

ruff/crates/ty_ide/src/semantic_tokens.rs

Lines 467 to 470 in e4dfddc

	_ if ty.is_property_instance() => {
	// Actual Python property
	(SemanticTokenType::Property, modifiers)
	}

[Geo5]

I could not find any previously existing function that checks for this, so i decided to only use a closure here instead of adding it to UnionType or Type itself

[MichaReiser]

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.

[MichaReiser]

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.

[Geo5]

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 😅 )

[MichaReiser]

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);
...

[Geo5]

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?

[MichaReiser]

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

[Geo5]

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.

[Geo5]

lmk if you find the functional style even less appealing 😅