CODEWITHSUNDEEP
pythonpandasmatplotlibpython-3.8plot-annotations
Jon Franco
Jon Franco

Reputation: 153

Type hint for matplotlib axes

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.

Upvotes: 8

Views: 3164

Answers (1)

mmdanziger
mmdanziger

Reputation: 4658

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.

Upvotes: 2

Related Questions