2

I wanted to have no errors while using VSCode Pylance type checker.

How to type the axs correctly in the following code:

import matplotlib.pyplot as plt
fig, axs = plt.subplots(2, 2)

In the image below, you can see that Pylance on VSCode is detecting an error.

enter image description here

Onyr
  • 769
  • 5
  • 21
  • "strongly typing" is not the correct terminology here, you just mean *typing*, really. either "type hinting" or "type annotating". Strong typing is a loosely-defined term that applies to languages (python is generally considered strongly typed). – juanpa.arrivillaga Jun 16 '22 at 16:50
  • I think you wanted to say that Python is NOT strongly typed. Actually to me, sure, Python by design is not a strongly typed language. However the same way that VSCode Pylance type-checker expresses it, there are several "levels" of typing in python. For instance, stating some `var` is a `list` is type hinting. Saying it is a `list[plt.Axes] | None` is strongly typing. But here we are just debating about terms. In the end, it's only type hinting. Hope you understand why I chose that term here. – Onyr Jun 17 '22 at 08:06
  • 1
    no, I mean that it *is* strongly typed. You seem to be confusing the expression "strong typing" with "static typing". Python is *not* statically typed (although, nowadays it can be with type-hints and third-party static type checkers like `mypy` and `pyright`). But Python has always been considered strongly typed. Not all statically typed languages are strongly typed either, e.g. C is generally considered weakly typed – juanpa.arrivillaga Jun 17 '22 at 17:24
  • Ok you are right. I will remembrer that thanks. – Onyr Jun 20 '22 at 09:42

2 Answers2

1

For reasons I don't aim to understand, the currently accepted answer kept giving warnings, errors, and lacked successful type inference in my setup.

What I did instead was this, which allows slicing, and allowed Pylance to understand that .plot is Axes.plot:

from typing import TypeVar, Generic, cast

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.axes import Axes

# Little helper class, which is only used as a type.
DType = TypeVar("DType")
class Array(np.ndarray, Generic[DType]):
    def __getitem__(self, key) -> DType:
        return super().__getitem__(key) # type: ignore

# Force assign the type, which is correct for most intents and purposes
fig, _axs = plt.subplots(2, 2)
axs = cast(Array[Axes], _axs)

# Use as an ndarray of Axes
axs[0,0].plot(...)
wschella
  • 141
  • 5
  • 1
    Not only `plot`, but this allow Pylance to understand all methods of axes inside the array. This is a more suitable answer. – Diogo Aug 07 '23 at 22:21
0

It turns out that strongly typing the axs variable is not straightforward at all and requires to understant well how to type np.ndarray.

See this question and this question for more details.

The simplest and most powerful solution is to wrap numpy.ndarray with ' characters, in order to avoid the infamous TypeError: 'numpy._DTypeMeta' object is not subscriptable when Python tries to interpret the [] in the expression.

An example:

import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
import seaborn as sns
from typing import cast, Type, Sequence
import typing 

sns.set() 

# Some example data to display
x = np.linspace(0, 2 * np.pi, 400)
y = np.sin(x ** 2)

fig, axs = plt.subplots(
    2, 2, 
    figsize=(12, 10) # set graph size
)

# typechecking operation
NDArrayOfAxes: typing.TypeAlias = 'np.ndarray[Sequence[Sequence[plt.Axes]], np.dtype[np.object_]]'
axs = cast(np.ndarray, axs)

axs[0, 0].plot(x, y)
axs[0, 0].set_title("main")
axs[1, 0].plot(x, y**2)
axs[1, 0].set_title("shares x with main")
axs[1, 0].sharex(axs[0, 0])
axs[0, 1].plot(x + 1, y + 1)
axs[0, 1].set_title("unrelated")
axs[1, 1].plot(x + 2, y + 2)
axs[1, 1].set_title("also unrelated")
fig.tight_layout()

Which is well detected by Pylance and runs correctly:

enter image description here

Onyr
  • 769
  • 5
  • 21
  • what is the use of the line `NDArrayOfAxes: typing.TypeAlias = 'np.ndarray[Sequence[Sequence[plt.Axes]], np.dtype[np.object_]]'`? By removing this line, the result seems to be the same. – Diogo Aug 07 '23 at 19:33
  • This approach removes the warnings. But Pylance does not recognizes `plot`, `set_title`, etc... as `Axes` methods. [wschella's answer](https://stackoverflow.com/a/74197401/9761768) seems to do that. – Diogo Aug 07 '23 at 19:35