Can shape information be preserved by the LSP ?

[Doublonmousse]

Is there a way to preserve the "unmodified" state of type annotations when viewed by the Language Server Protocol (LSP), or is this a fundamental limitation of the package?

I am currently using Visual Studio Code (VSCode), and I observe that the type annotation is translated back to ndarray (which is also what mypy detects, as far as I know).

The desire is for the shape information to remain within the function's type definition, visible in documentation tooltips. This would enhance the module's practicality, as users would have access to shape information while writing function calls.

One potential solution would be to use macros to translate the jaxtyping type annotation into the docstring. Although this might work for a Python module (displaying that information in the LSP), it doesn't seem effective for scripting use. When using a macro to autogenerate docstrings for a function, the LSP in VSCode still shows the old docstring, even if f.__doc__ has changed. It is possible that this issue lies with VSCode's LSP rather than jaxtyping though.

I believe the type annotation remains intact when generating documentation so that advantage is still there.

[patrick-kidger]

So I think this is just a choice that has been made in the implementation of the LSP -- since you're using VSCode I'm assuming it's probably pylance -- to strip away typing.Annotated when displaying these popups. (Annotated is what jaxtyping uses under-the-hood.)

FWIW, it looks like holding down control might be one solution, that shows the full details.

Besides that, I think a toggle for this would be a reasonable feature request to make against pylance. Annotated is typically used to provide additional information for another tool (as it does for us), but that doesn't mean it's not useful to the developer too.

[Doublonmousse]

Trying holding down control on vscode doesn't seem to make it show the full details (maybe that has changed by now), so what is left is jumping of peeking at the definition of the function.

From the various issues on pylance's repository, they seem to agree on not showing Annotated content, as it is not meant as documentation.

The only reasonable option left would be do to add that information either manually or automatically into the docstring for each of the parameters of the function (if the goal is to have that info directly into the tooltips). Having some automation or check would be better though, because the point is to have the shape information be tested at runtime (hence tested and correct) whilst docstrings could linger with the wrong information.

[patrick-kidger]

Hmm, bother.

In terms of docstrings -- I don't think I can see a way to add it automatically. You could probably arrange to have it be checked automatically, although the details of that would depend on your choice of docstring format.

As that depends on the details of the format, though, I don't think that's something we can include in jaxtyping itself -- best to write your own decorator I think.