[MNT]: Improve return type of ioff and ion to improve Pyright analysis of bound variables

[FeldrinH]

Summary

The following code causes Pyright (the type checker used by Pylance, default type checker for Python in VSCode) to emit an error:

import matplotlib.pyplot as plt

with plt.ioff():
    fig, ax = plt.subplots()

ax.plot([...], [...]) # <-- Pyright error: "ax" is possibly unbound

The root cause of the issue seems to be that ioff and ion return contextlib.ExitStack, which can in theory supress excpetions, as indicated by the return type of ErrorStack.__exit__ being bool rather than Literal[False] or None. (See microsoft/pyright#7009)

In practice though, the way matplotlib uses ExitStack, it can't swallow exceptions. It would be convenient for typed codebases if the return type of ion and ioff reflected that.

Proposed fix

A simple context manager like this could be used instead to make it clear to type checkers that exceptions will never be supressed (note that the return type of __exit__ is None):

class _InteractiveContextManager(contextlib.AbstractContextManager):
    def __init__(self, callback: Callable):
        self._callback = callback

    def __enter__(self) -> None:
        return None
    
    def __exit__(self, __exc_type, __exc_value, __traceback) -> None:
        self._callback()

This also has the added benefit that by returning None from __enter__ we can better enforce that the context manager is not stored or accessed by the user.

[tacaswell]

We still need to be able to use plt.ion() and plt.ioff() as stand-alone function (that is not as context managers). So long as that still works, I do not see the harm in making this change.

[FeldrinH]

We still need to be able to use plt.ion() and plt.ioff() as stand-alone function (that is not as context managers).

Yes, that should continue to work exactly as it does right now.

As far as I know the only effect this change will have on runtime behavior is that if someone is using the return value of ion or ioff under the assumption that it is specifically a contextlib.ExitStack then this will break their code. ion/ioff documentation explicitly recommends against doing that, so I hope that the number of people affected would be fairly small.

[anntzer]

It would seem wrong to replace a simple, short implementation of ion/ioff by a lengthier version just to appease type checkers. In fact this topic has been explicitly discussed e.g. at https://discuss.python.org/t/add-an-else-clause-to-with-statements-to-detect-early-termination/38031/1, and I'd guess https://discuss.python.org/t/add-an-else-clause-to-with-statements-to-detect-early-termination/38031/54 would be a more appealing solution.

[FeldrinH]

It would seem wrong to replace a simple, short implementation of ion/ioff by a lengthier version just to appease type checkers. In fact this topic has been explicitly discussed e.g. at https://discuss.python.org/t/add-an-else-clause-to-with-statements-to-detect-early-termination/38031/1, and I'd guess https://discuss.python.org/t/add-an-else-clause-to-with-statements-to-detect-early-termination/38031/54 would be a more appealing solution.

I agree that a general solution, like the one proposed in the discussion you linked, would be more appealing. However, as far as I can tell that solution is currently stalled by other type system limitations and will not land in the near future. Also, if it is made dependent on PEP 696, as has been proposed here, then it will probably require Python >= 3.13.

Given all that, I think solving this now with a small change in matplotlib is a more practical approach.

I think that adding 7 lines of readable, simple code is well worth it to have things "just work" in projects that use matplotlib with type checking. I exclusively work with type checking enabled and if the type checker does not consider some code correct then that, to me, is a fairly major concern.

[tacaswell]

On one hand I agree it is less than ideal to complicate the codebase to cope with limitation of the type checkers. However I do not think this is significantly different than version gating we do to work around bugs in our dependencies and support multiple versions.

This is not the right time or place to litigate the usefulness of typing in Python. We ship type stubs and some of our users do find it useful so we have to accept some complexity to make them work nicely.

If we get a PR we should merge it (but it should include a comment that it is extra complex due to limitations of the type system).

[FeldrinH]

I realized that this specific issue with Pyright can actually be fixed with a much more minimal change: by simply changing the return types of ion and ioff to AbstractContextManager Pyright no longer complains about unbound variables.

[anntzer]

This actually seems right from an API PoV: we only guarantee that what ion/ioff returns can be used as a CM; them being ExitStacks is really an implementation detail.

[FeldrinH]

I have created a pull request. Also, given that the new proposed fix has no effect on runtime behavior, is there a possibility that this could be included in the next patch release?