Search code examples
pythonmatplotlib

How to type-hint a matplotlib.axes._subplots.AxesSubplots object in python3


I was wondering how is the "best" way to type-hint the axis-object of matplotlib-subplots.

running

from matplotlib import pyplot as plt

f, ax = plt.subplots()
print(type(ax))

returns

<class 'matplotlib.axes._subplots.AxesSubplot'>

and running

from matplotlib import axes
print(type(axes._subplots))
print(type(axes._subplots.AxesSubplot))

yields

<class 'module'>
AttributeError: module 'matplotlib.axes._subplots' has no attribute 'AxesSubplots'

So far a solution for type-hinting that works is as follows:

def multi_rocker(
                 axy: type(plt.subplots()[1]), 
                 y_trues: np.ndarray,
                 y_preds: np.ndarray,
                 ):
  """
  One-Vs-All ROC-curve:
  """
  fpr = dict()
  tpr = dict()
  roc_auc = dict()
  n_classes = y_trues.shape[1]
  wanted = list(range(n_classes))
  for i,x in enumerate(wanted):
    fpr[i], tpr[i], _ = roc_curve(y_trues[:, i], y_preds[:, i])
    roc_auc[i] = round(auc(fpr[i], tpr[i]),2)
  extra = 0
  for i in range(n_classes):
    axy.plot(fpr[i], tpr[i],)
  return

And the problem with it is that it isn't clear enough for code-sharing


Solution

  • As described in Type hints for context manager :

    import matplotlib.pyplot as plt
    
    def plot_func(ax: plt.Axes):
        ...