Improve autocompletion for arguments to objects wrapped with use_signature

[binste]

Resolves #2884, see that issue for more context. All methods and functions which are wrapped with use_signature profit from this PR. The PR can be split into three main changes which were all necessary to make this work.

Full function signature:

Keyword completion (which still does not work for me in Jupyter, just VS Code):

Long method-chaining still works:

Function signature and docstrings still show up in Jupyter as before.

I wasn't able to also get docstrings working in VS Code. They still require inspection at runtime and therefore only work in Jupyter. I don't think this is too bad, I find the function signature and autocompletion on keywords much more important as I usually tend to look up this information in the documentation anyway which is better formatted.

Type hints on utils.use_signature

Using the example:

@utils.use_signature(core.FacetedEncoding)
def encode(self, *args, **kwargs) -> Self:

the new type hints on use_signature tell a type checker that encode has the same parameters as core.FacetedEncoding but keeps its original return type Self which is whatever the type of self is. This is the main part. The other changes are necessary to make sure this also works with method chaining.

As typing.ParamSpec was introduced in Python 3.10 I added typing_extensions as a new requirement to Altair. Already briefly discussed this with @mattijn here. It's a very stable dependency and I think now it's worth adding it.

Switching from self-bound TypeVars to Self

Now that we have typing_extensions we can switch to Self. See linked comment above for the discussion. This is also necessary as the new type hints on use_signature did not work with the self-bound TypeVars.

Move utils.use_signature from class to __init__ definitions

In the instances where use_signature was wrapping a class I had to move it to wrapping __init__ instead as else method chaining such as alt.Chart().repeat()... did not work. Reason is that e.g. repeat should show FacetChart as return type but without this change it shows Any as the type hints on use_signature no longer work when a class is wrapped.

This should not have any averse effects on the functionality of Altair. Only change I found is that in the documentation, the docstrings are now included in __init__ instead of the top-level class which I don't think is bad as Python allows to document a class either on the top-level class or in __init__ and IDEs usually treat this as the same. Old:

New:

For consistency, I could move the manually defined docstrings on classes such as alt.Chart also to __init__. Let me know if you'd prefer this in the same PR or if we want to postpone/don't do this change and merge the PR as is. I don't think this alignment is very important but a nice-to-have.

[mattijn]

Nice that you've introduced the typing_extension, it mades the required type hints easier to read. I wish I have the time to do one big review, but for now just add some comments when I find some time to test.

• I noticed the self is now also included in the signature (in Jupyter at least):

---

• Also this is probably a good moment to test if type hinting will improve in Jupyter with the Language Server Protocol (this), as it is lost now:

---

• As it won't give type hints two . deep:

---

• Not sure exactly what happened in my VSCode, since it only gives me thunder methods when I do the following:

[mattijn]

I tried the python language server in Jupyter and it does change things, not sure if they will help us completely.
Type hints for encode are presented as encode(*args, **kwargs) (screenshot top), but doing a shift+tab it gives me a nice signature

By now, I feel the IDEs and language servers are not yet in-sync with each other, so it is hard to implement a solution that suits all. Its already a big step forward as such!

Only question remains from my side is: why is the long chaining method for you still working and for me not anymore? (I'm on VScode 1.75.1)

[binste]

Thanks! I can replicate the issue in Python 3.9, didn't notice it during development as I was on 3.11. I thought the try except ImportError statements are the correct way to use typing_extensions but turns out that this does not work well. The mypy docs (and this GH issue) recommend sys.version_info instead. Same works for VS Code, it now even highlights which part of the code is never reached (this is on Python 3.9):

With this adjustment, method chaining and keyword completion works for me in both Python 3.9 and 3.11. Could you try again? Maybe it also helps the language server in Jupyter.

[binste]

Tests work locally (Python 3.11) but fail in workflow due to the removal of VL 5.2 in the newest version of altair_viewer. I agree that it was reasonable to remove it there so it doesn't get distributed as no official Altair version will use it. Not sure what the current progress with the upgrade is so this PR can also wait. Interim solution could be to pin altair_viewer installation to the penultimate commit and then remove the restriction once #2871 is merged. Let me know if I can help.

[mattijn]

Will try to get #2871 in later today

[mattijn]

All is fine now for long chained methods using VSCode:

---

Using the language server with Jupyter, I now get the following for long chained methods:
Using Altair 4.2:

---

And using the branch of this PR (python -m pip install git+https://github.com/binste/altair.git@type_hint_use_signature):

It seems to be lost after .encode() using this branch, which seems to be a regression to what the language server can decipher using Altair 4.2.

[binste]

I can replicate this and it feels like a bug in python-lsp-server or whichever package is responsible for doing the static code analysis. The below is a minimal reproducible example which I'm running on Python 3.9 with newest Jupyterlab, Jupyterlsp, and python-lsp-server[all]:

from typing import Callable
from typing_extensions import Self

def use_signature(Obj):
    def decorate(f: Callable):
        return f
    return decorate

class SignatureClass:
    def __init__(self, arg1, arg2):
        self.arg1 = arg1

class MainClass:
    @use_signature(SignatureClass)
    def returns_self(self, *args, **kwargs) -> Self:
        return self
    
    def another_method(self):
        pass

This replicates the use_signature usage that we have in Altair. I removed all type hints except for Self and Callable. return_self correctly shows Self as return type:

but it cannot complete `another_method`

When removing Callable from the decorate function it still shows Self as return type and can now complete another_method:

I currently can't find a reason why annotating a function with Callable would throw off a type checker like this, especially as it still is aware of the return type hint Self. I will look at it again tomorrow morning with fresh eyes. If we can't find an explanation for this behaviour I'd suggest to merge for now and I would open an issue in the python-lsp-server repo with this example to get more inputs.

[mattijn]

I've the same feeling that this a limitation of jupyter-lsp and not intended. I agree that we should raise an issue at the relevant repo and merge this in if you don't change your mind after some sleep.
Btw, main repo is now on vl5.6.1, so probably you need to rerun the generate_schema_wrapper and resolve a few conflicts..

[binste]

Opened this issue in the jedi repo.

Great to see 5.6.1 in master! Merged it into this branch and reran the schema generation. Changes still work in VS Code for me. If tests pass, it's good to go from my side.

[mattijn]

All tests pass and since:

davidhalter commented 6 hours ago
Self is currently not supported.

Can merge this in and the jupyter ecosystem can benefit of this PR once jedi supports the typehint Self