Python3.x 如何向 scikitlearn 函数添加类型提示

发布于08月25日

我有以下简单的函数:

def f1(y_true, y_pred):
    return {"f1": 100 * sklearn.metrics.f1_score(y_true, y_pred)}

根据SCRICKIT-LEARN文档，f1_score的参数可以有以下类型:

y_true:1D数组，或标签指示器数组/稀疏矩阵
y_pred:1D数组，或标签指示器数组/稀疏矩阵

并且输出的类型为:

浮点或浮点数组，Shape=[n_Unique_Labels]

如何向该函数添加类型提示，以便mypy不会出现错误？

我try 了以下不同的版本:

Array1D = NewType('Array1D', Union[np.ndarray, List[np.float64]])

def f1(y_true: Union[List[float], Array1D], y_pred: Union[List[float], Array1D]) -> Dict[str, Union[List[float], Array1D]]:
    return {"f1": 100 * sklearn.metrics.f1_score(y_true, y_pred)}

但这造成了错误.

推荐答案

这是我用来避免类似的输入问题的方法.它利用了1.20中引入的numpy typing.ArrayLike类型涵盖List[float]，因此不必担心显式涵盖它.

在上面运行带有NumPy v1.23.1的mypy v0.971没有显示任何问题.

from typing import List, Dict
import numpy as np
import numpy.typing as npt
import sklearn.metrics


def f1(y_true: npt.ArrayLike, y_pred: npt.ArrayLike) -> Dict[str, npt.ArrayLike]:
    return {"f1": 100 * sklearn.metrics.f1_score(y_true, y_pred)}

y_true_list: List[float] = [1, 0, 1, 0]
y_pred_list: List[float] = [1, 0, 1, 1]
y_true_np: npt.ArrayLike = np.array(y_true_list)
y_pred_np: npt.ArrayLike = np.array(y_pred_list)

assert f1(y_true_list, y_pred_list) == f1(y_true_np, y_pred_np)