Type hint for matplotlib axes

Question

I'm writing type annotations for some functions and I couldn't find the best way to annotate an argument which is expected to be of type matplotlib.axes._axes.Axes.

I'm not even sure if I'm doing the right thing for an argument expected to be of type pandas.core.frame.DataFrame.

def myfunc(
        df: pd.DataFrame, # is this correct?
        ax: matplotlib.axes._axes.Axes # do I have to use this?
        ) -> None:
   
    # body of my function

Some insights?

Thanks.

Answer 1

Don't reference the .pyi file in the typehint, just use the actual class name. Your IDE or mypy will figure it out.

If you want mypy to work with matplotlib version <3.8 you will need data-science-types. data-science-types only needs to be installed, not imported.

See this matplotlib issue. Related to this answer.

The following has been tested with mypy after running pip install data-science-types.

import pandas as pd
import matplotlib.axes
from matplotlib import pyplot as plt
import numpy as np


def myfunc(
    df: pd.DataFrame,
    ax: matplotlib.axes.Axes,
) -> None:
    ax.plot(df.values.mean(), df.values.std())


if __name__ == "__main__":
    f, axis = plt.subplots()
    df = pd.DataFrame(data=np.random.rand(50, 10))
    myfunc(df, axis)

Running mypy returns Success: no issues found in 1 source file.