I would like to have a context manager for pyplot figures essentially like so:
from contextlib import contextmanager
import matplotlib.pyplot as plt
@contextmanager
def subplots():
(fig, ax) = plt.subplots()
try:
yield (fig, ax)
finally:
plt.close(fig)
Can type hinting be implemented for the returned tuple? The naive
import typing
def subplots() -> typing.Tuple[plt.Figure, plt.Axes]
doesn't work.
Your function is not actually returning a tuple. Rather, it's yielding one -- if you call subplots() without the context manager, you'll get back a generator object. So, you'll need to annotate it as such:
from typing import Tuple, Generator
from contextlib import contextmanager
import matplotlib.pyplot as plt
@contextmanager
def subplots() -> Generator[Tuple[plt.Figure, plt.Axes], None, None]:
(fig, ax) = plt.subplots()
try:
yield (fig, ax)
finally:
plt.close(fig)
You can find more information about Generator in the mypy docs, but in short, it accepts three types: the type of whatever the generator yields, the type of whatever value your generator can be sent, and the type of whatever value your generator finally returns. We don't care about the latter two here, so we leave them as None.
This type does end up being a feeling clunky though. A more concise alternative is to annotate the return type as being an Iterator instead:
from typing import Tuple, Iterator
from contextlib import contextmanager
import matplotlib.pyplot as plt
@contextmanager
def subplots() -> Iterator[Tuple[plt.Figure, plt.Axes]]:
(fig, ax) = plt.subplots()
try:
yield (fig, ax)
finally:
plt.close(fig)
This works because all Generators are actually Iterators -- Generator is a subtype of Iterator. So it's fine to pick the more general return type.