如何为调用scikit-learn的自定义函数添加类型提示以通过mypy检查？

Let's break down the issues with your original type hints and fix them properly so mypy doesn't throw errors.

What Was Wrong With Your Original Code? #

Your main issues were:

1. Incorrect return type: The f1_score function returns either a single float (for binary/multi-class with average metrics) or a 1D numpy array of floats (for per-class scores). Your return type hint used Union[List[float], Array1D], which doesn't match the actual output (100 * a float is still a float, not a list).
2. Overly restrictive parameter types: Your Array1D only covers numpy arrays and lists of np.float64, but scikit-learn accepts a wider range of 1D array-like inputs (like pandas Series, lists of integers for class labels, etc.).

Fixed Implementation #

Here's the corrected version with proper type hints that align with scikit-learn's API and satisfy mypy:

from typing import Dict, Union
import numpy as np
from numpy.typing import ArrayLike, NDArray
from sklearn.metrics import f1_score

def f1(y_true: ArrayLike, y_pred: ArrayLike) -> Dict[str, Union[float, NDArray[np.float64]]]:
    score = f1_score(y_true, y_pred)
    return {"f1": 100 * score}

Why This Works #

• Parameter Types: ArrayLike is a numpy-provided type alias that encompasses all 1D array-like inputs scikit-learn accepts:
  • Python lists/tuples of integers or floats
  • Numpy arrays
  • Pandas Series/DataFrame columns (since they're array-like)
  • Sparse matrices (though for f1_score, sparse inputs need to be in a compatible format, which ArrayLike still covers)
• Return Type:
  • Union[float, NDArray[np.float64]] accounts for both scenarios:
    • A single float when using average metrics (e.g., average="macro")
    • A 1D numpy array when returning per-class scores (e.g., average=None)
  • Multiplying the score by 100 preserves the original type (float stays float, array stays array), so this hint is accurate.

Alternative (If You Prefer Explicit Types) #

If you want to avoid ArrayLike and be more explicit about supported input types, you can use this version. It's slightly more verbose but still valid:

from typing import Dict, Union, List, Tuple
import numpy as np
from sklearn.metrics import f1_score

def f1(
    y_true: Union[List[Union[int, float]], Tuple[Union[int, float], ...], np.ndarray],
    y_pred: Union[List[Union[int, float]], Tuple[Union[int, float], ...], np.ndarray]
) -> Dict[str, Union[float, np.ndarray]]:
    score = f1_score(y_true, y_pred)
    return {"f1": 100 * score}

This covers lists, tuples, and numpy arrays (with integer or float values, since class labels can be integers).

内容的提问来源于stack exchange，提问作者BioGeek