numpy 函数返回 pandas Series 或 DataFrames 时的类型提示答案

【问题标题】：Type hinting when numpy functions return pandas Series or DataFramesnumpy 函数返回 pandas Series 或 DataFrames 时的类型提示
【发布时间】：2021-01-14 20:35:48
【问题描述】：

当使用 Python 类型提示时，我如何处理在应用于此类时返回 pandas 对象的 numpy 函数？

比如我写的时候

import numpy as np
import pandas as pd
def clip_probability(probability: pd.Series) -> pd.Series:
    return np.clip(
        probability,
        0,
        1,
    )

我得到Expected type 'Series', got 'ndarray' instead 即使np.clip 在第一个参数是一个时确实返回了一个系列。

如果有特定于 IDE 的解决方案，我使用的是 PyCharm 2020。

【问题讨论】：

在这种特定情况下，np.clip 可以替换为pd.Series.clip，IIUC。

标签： python pandas numpy type-hinting

【解决方案1】：

您是否尝试过强制转换函数的返回值？ Numpy 和 Pandas 通常擅长在它们的类型之间进行转换。像这样：

return pd.Series(np.clip(
        probability,
        0,
        1,
    )
)

【讨论】：

我希望有一个无需更改代码的解决方案（虽然这是一个功能性无操作，但它仍然意味着额外的函数调用）。请注意，np.clip 的输出已经是pd.Series

【解决方案2】：

我会使用typing.cast:

from typing import cast

def clip_probability(probability: pd.Series) -> pd.Series:
    out = np.clip(
        probability,
        0,
        1,
    )
    return cast(pd.Series, out)

这在运行时没有任何作用，但会向类型检查器发出信号，out 确实是 Pandas 系列。

【讨论】：

只是出于好奇，在解释器级别的支持下，这真的在运行时什么都不做，还是仍然转换为少量字节码指令？
（我暂时不接受答案，以防万一有人出现并显示猴子补丁类型提示的干净方法。这样我每次打电话@987654326时都不必做这个演员表@/np.exp 等，但每个只有一次。）
这是一个立即返回的附加函数调用：github.com/python/typing/blob/master/src/typing.py#L1436-L1444。我不认为你可以让你的熊猫系列小到足以让这件事变得重要。
我同意这一点，我只是好奇。无论如何，每次调用此类函数时额外的源代码仍然让我有些困扰。如果没有人有更简洁的，我会重新接受

【解决方案3】：

np.clip 返回一个ndarray（参见https://numpy.org/doc/stable/reference/generated/numpy.clip.html）。

您首先需要将这个 ndarray 转换为 Pandas 系列：

return pd.Series(np.clip(
        probability,
        0,
        1,
    ))

【讨论】：

import numpy as np; import pandas as pd; type(np.clip(pd.Series([0,2]), 0, 1))