Add return type hints to improve code completion suggestions

[binste]

Current behaviour

Some autocomplete engines such as the one used by VS Code heavily use type hints. Without this PR, after using a method on a chart once, VS Code no longer knows that it is still a Chart (or FacetChart, etc.) object:

Autocompletion still works for the first method:

But is lost afterwards:

VS Code states that the return type of mark_bar is Any, meaning that it does not know what it returns.

New behaviour introduced in this PR

This PR adds return type hints to all top-level chart methods I could find. These are usually "self" type hints which tell a static type checker that an object is returned which has the same type as the object on which the method was called. The approach is explained in the "Motivation" section in PEP 673 - Self Type which introduces a new way of doing this but that was only introduced in Python 3.11 so I still use the backwards compatible approach here with bound TypeVar.

The return type now correctly displays as "Chart":

This also works if a method such as repeat changes the chart to a RepeatChart. Calling configure afterwards gives the correct return type hint:

[mattijn]

Impressive work @binste! I'm not very familiar with type hinting.

I was wondering what your opinion is on using https://github.com/python/typing_extensions as dependency for Python < 3.11. It contains a back-port for the Self-type as described in PEP673. I agree with the PEP that it is a bit more readable in comparison to the TypeVar definitions.

I know it only needs to be defined a few times in the generate_schema_wrapper.py so it's not a huge deal, but I like to check with you what you think is best.

[binste]

Thanks @mattijn! :)

Regarding typing_extensions, I see the argument for code readability and typing_extensions is probably a very reliable dependency to add so I'd be onboard with it as well. However, I'm not familiar enough with it to guarantee that there are not some rare cases where users might have a need to pin typing_extensions to an older version. We'd need at least 4.0.1. Probably not but to make sure to not break anything, I opted for this approach, also because it felt like a bit much to add a dependency just to get rid of maybe 5 lines of code. How about proceeding without typing_extensions but reconsider adding it in case we'd add more type hints to Altair and could need some extra features from it?

[mattijn]

I understand that adding another dependency just to reduce 5 lines of code is probably not worth it. I also agree that there may be other recent type-related changes in the core of Python that makes it something to consider if we look a bit broader than this PR alone.

We previously had two types of discussions about types in Altair.

• One is about improving the existing Python code in the /tools folder in the Altair repository, such as this issue (Type hints #2493) and related PR (Add type hints to tools.schemapi.schemapi #2614, btw I now notice that this PR also uses the TypeVar bound approach).

• The other is about including types in the generated Python code of Altair, originally (add typehint Any to Undefined object #2670) adapted to (use 'UndefinedLike = Any' as the type hint #2681) and then reverted back to the original in (Reverting UndefinedLike? #2714).

Especially the result of the second bullet has set the type Any to anything that is Undefined finishing this discussion on, 'what is the type of Undefined?'.

However, with your impressive PRs regarding jsonschema and types, I was wondering, since types are included in the vega-lite jsonschema, if these types can also be inferred for each argument? Maybe recent type-related changes in the core of Python can make this easier? It is merely an open question from which I don't know enough if it is worthwhile to consider.

[binste]

Thanks for writing this down! This is very helpful. It should be possible to create Python type hints based on the vl schema, especially given that this parsing is already done for the docstrings, see e.g. https://github.com/altair-viz/altair/blob/master/altair/vegalite/v5/schema/channels.py#L12509 where the stack argument to alt.X is documented as:

stack : anyOf(:class:`StackOffset`, None, boolean)

There are probably also synergies with replacing the vega-lite types with Python types in all error messages (suggested by you in #2842 (comment)) and docstrings.

As more type hints shouldn't break anyones code, I don't think we would need to necessarily include it in the release candidate but I'm definitely intrigued in giving this a try. I really like the development experience you get for a library with well-defined type hints such as better autocompletion as well as early warnings if you pass a wrong value to an argument before even executing the code, for example if you'd pass an integer to stack. The same linters could then of course be used in CI/CD pipelines.

I'll try to get a draft PR going in the coming weeks so we can further explore how well it works.

[mattijn]

I've opened #2854 to track this feature request.
This branch now has some conflicts since there was also a bit of reordering in the generate_schema_wrapper in #2829. If you can resolve these, than we can merge this as well.

[joelostblom]

This is looking very useful, thank you @binste ! Do you know if jupyterlab autocompletion would benefit from this change as well, or are they not using the same approach as VScode here?

[binste]

Just tested it and JupyterLab can already provide good autocompletion without this PR:

I haven't kept up with the JupyterLab developments but I think it uses some autocomplete engine (maybe jedi?) which relies on inspection of objects instead of type hints. For example the above auto completion only works once you executed the first cell. VS Code can provide it without execution based on type hints alone. Maybe some language servers, see https://github.com/jupyter-lsp/jupyterlab-lsp, could benefit as well or other IDEs such as PyCharm but I haven't tested that.

[joelostblom]

Ah ok so this PR brings VS Code up to the same level as JupyterLab is already at, got it. And you're correct, IPython uses Jedi. I have tested the lsp extension and I think there was talk about integrating it in to core in a future version of jupyterlab so maybe it will be there soon.

[mattijn]

I've the feeling that with the last commit something got mixed up. Currently the following definitions are defined two times.

MARK_METHOD = ...
CONFIG_METHOD = ...
CONFIG_PROP_METHOD = ...

[binste]

Thanks for spotting it! Indeed, there was a merge conflict as I modified those variables in this PR and they were also moved around in the PR which removed vega 5. I made a mistake when resolving that conflict. It's fixed now and I ran generate_schema_wrapper.py to make sure that it produces the expected scripts.

Need to wait for #2869 before black test passes.

[mattijn]

With #2814 we check for top level variables in the __init__.py file. The changes introduced in this PR will currently make the following test fail:

FAILED tests/test_toplevel.py::test_completeness_of__all__ - AssertionError: assert ['Aggregate',...tString', ...] == ['Aggregate',...tString', ...]
[699](https://github.com/altair-viz/altair/actions/runs/4152594080/jobs/7183718320#step:6:700)
  At index 410 diff: 'TOPLEVEL_ONLY_KEYS' != 'TChart'
[700](https://github.com/altair-viz/altair/actions/runs/4152594080/jobs/7183718320#step:6:701)
  Right contains 10 more items,

The new elements added to the __all__ within __init__.py are:

-  'TChart',
-  'TConcatChart',
-  'TEncodingMixin',
-  'TFacetChart',
-  'THConcatChart',
-  'TLayerChart',
-  'TRepeatChart',
-  'TTopLevelMixin',
-  'TVConcatChart',
-  'TypeVar',

If this is expected than I think the generate_schema_wrapper.py should be re-run?

[mattijn]

Can these type hints also be defined in hidden modus? E.g. as such:

_TChart = TypeVar("_TChart", bound="Chart")

[binste]

I like the approach of marking them as private with the "_" prefix so they don't even get imported all the way to the top-level. Could give us also more flexibility when doing #2854 as no-one should rely on the current naming and that these typevars are defined like this. Pushed this change.

[mattijn]

Tests pass, looks better now👍