Infer the type of an implementation function based on usage in `rule`

[JohnnyMorganz]

First of all, I just wanted to say that this project is awesome. It works really well for me, thank you for taking the time to do this!

I don't know how your type system is set up, but I was curious whether the following was possible.

Given a starlark file of

def _example_rule_impl(ctx):
    pass

example_rule = rule(
    implementation = _example_rule_impl,
    attrs = {
        "srcs": attr.label_list(allow_files = True)
    }
)

Right now, the type of ctx in _example_rule_impl is Unknown. This is resolved by adding the comment # type: (ctx) -> Unknown.

Since we can see that this function is used in rule later, do you reckon it's possible to automatically infer that ctx should be a ctx type?

A step further (and this is the main reason why I ask this). If the above is possible, it would be awesome if the ctx type could understand the attrs defined in the rule. This would mean e.g. ctx.files.srcs in the above function would autocomplete and typecheck appropriately.

[withered-magic]

Thanks so much for using starpls, and glad to hear it's working well for you!

I've actually considered this before, and unfortunately doing it "correctly" is probably pretty far out of the capabilities of the current type system. It seems like some Hindley-Milner sorta system with unification would probably be the "real" way to solve this, but the type system as currently implemented is much, much dumber than that and basically just infers types of variables based on the last seen declaration, which actually works well enough for Bazel since most files aren't actually that complicated. (Although I do have an in-progress PR for introducing a form of data-flow analysis, which will eventually enable stuff like type guards, type narrowing, etc.)

However, I do think such functionality would fit nicely in starpls, and given that we already use heuristics to drive some features (e.g. Go to Definition for targets, targets in Document Symbols), I think coming up with another heuristic to drive this feature should work fine as well (probably gated behind a startup option)

Off the top of my head, we could probably start off with something simple like keeping track of all rule invocations that have implementation arguments; once we're conducting type inference on the impl function, we can check for a rule/repository_rule/module_extension invocation that has the impl function's name as the implementation argument, and, if we find one, assign the corresponding attributes to ctx.attr. Would something like this cover your use case? (this functionality would be opt-in anyways so it wouldn't cause any weird behaviors to happen by default)

[JohnnyMorganz]

This sounds reasonable to me, and I think it would cover the use cases I have in mind. Thank you!

Also, a thought that might be useful temporarily. For ctx.files, ctx.file, ctx.outputs and ctx.executable, is it possible to give them a type equivalent to something like typescript's index signatures for now? So that ctx.files.foo would get a type of List[File]. Would help with ergonomics when using those params, but of-course wouldn't catch e.g. typos when there is no label list attr of name foo. Don't think it would work for the generic ctx.attr though. And maybe it is not so simple.

[withered-magic]

Yep, that should be possible! structs are supported in the type system already, so it would probably just be a modification to the struct type to store an enum that can be either the known fields or the type of unknown fields, e.g.

enum StructData {
    KnownFields(Vec<Field>),
    MemberTy(Ty)
}

[withered-magic]

Also, a thought that might be useful temporarily. For ctx.files, ctx.file, ctx.outputs and ctx.executable, is it possible to give them a type equivalent to something like typescript's index signatures for now? So that ctx.files.foo would get a type of List[File].

@JohnnyMorganz This is on main now!

[JohnnyMorganz]

Awesome, thank you so much! I shall try this out

[withered-magic]

Thanks! You'll still need the # type: (ctx) -> Unknown comment since I'm still working on that bit

[withered-magic]

Support for this feature is on main now! Make sure to update your editor config if you want to try it out, the functionality is gated behind the --experimental_infer_ctx_attributes flag under the starpls server subcommand!

[JohnnyMorganz]

Just had a chance to try it out with the new release, and it works pretty well! Thank you for taking the time to implement it so quickly :)

The experimental option is able to infer ctx.attr well, but is it meant to also infer ctx.files etc?

[withered-magic]

No problem! And ah I didn't implement it for those just yet since it proved to be a bit trickier, I'll follow up in a separate issue!