使用现有类型提示向 Python 方法添加自定义重载

Question

是否可以通过为方法添加自定义重载来扩展 Python 包在我的应用程序代码中的现有类型提示（即不触及上游类型提示）？ （另请参阅 repo 和最小的可重复示例。）

例如，采用

pandas.DataFrame.apply

 中定义的 

pandas-stubs/core/frame.pyi 的类型提示。根据定义，它们不允许应用的函数返回一个集合，如下所示：

import pandas as pd

df = pd.DataFrame(
    {
        "a": list("abc"),
        "b": list("def"),
        "c": list("aaa"),
    }
)
print(df.apply(lambda row: set(x for x in row if not pd.isna(x)), axis=1))

安装

pandas-stubs

后，Pyright 报告：

/Users/david/repos/eva/pandas-apply-set/reprex.py
  /Users/david/repos/eva/pandas-apply-set/reprex.py:10:7 - error: No overloads for "apply" match the provided arguments (reportCallIssue)
  /Users/david/repos/eva/pandas-apply-set/reprex.py:10:16 - error: Argument of type "(row: Unknown) -> set[Unknown]" cannot be assigned to parameter "f" of type "(...) -> Series[Any]" in function "apply"
    Type "(row: Unknown) -> set[Unknown]" is not assignable to type "(...) -> Series[Any]"
      Function return type "set[Unknown]" is incompatible with type "Series[Any]"
        "set[Unknown]" is not assignable to "Series[Any]" (reportArgumentType)
2 errors, 0 warnings, 0 informations

但是在运行时一切实际上都很好——脚本打印出：

0       {a, d}
1    {e, a, b}
2    {c, f, a}
dtype: object

所以这可能是一个错误，并且类型提示过于严格。已针对

pandas.Series.apply

的类型提示报告了类似的错误，涉及

set

和

frozenset

，并已修复。

我的问题是：有没有办法可以在我自己的代码中扩展

DataFrame.apply

的现有类型提示，添加一个

@overload

允许从应用的函数返回

set

？

我知道我可以通过在我自己的项目中基于上游版本的副本创建

frame.pyi

来覆盖整个

typings/pandas-stubs/core/frame.pyi

，我将在适当的重载中将

set

添加到返回类型的并集。但这需要在文件发生变化时保持文件的其余部分与上游同步。

相反，我想继续使用上游

frame.pyi

，只需在自定义代码中扩展

DataFrame.apply

的类型提示，以允许从应用的函数返回

set

。

克劳德法学硕士建议采用以下方法：

import typing as t
import pandas as pd

class DataFrameExt(pd.DataFrame):
    @t.overload
    def apply(
        self,
        f: t.Callable[..., set[t.Any]],
        raw: bool = ...,
        result_type: None = ...,
        args: t.Any = ...,
        *,
        axis: t.Literal[1],
        **kwargs: t.Any,
    ) -> pd.Series[t.Any]: ...

pd.DataFrame.apply = DataFrameExt.apply

但是这个特定版本会产生自己的错误。

有可能让它发挥作用吗？或者我应该将

DataFrame

调用站点上的

.apply

转换为自定义的不相关类，该类具有

.apply

的正确类型提示？比如：

import typing as t
import pandas as pd

class DataFrameApplyOverride:
    def apply(
        self,
        f: t.Callable[..., set[t.Any]],
        raw: bool = ...,
        result_type: None = ...,
        args: t.Any = ...,
        *,
        axis: t.Literal[1],
        **kwargs: t.Any,
    ) -> pd.Series[t.Any]: ...

# And then, at the call site of .apply:
print(
    t.cast(DataFrameApplyOverride, df).apply(
        lambda row: set(x for x in row if not pd.isna(x)), axis=1
    )
)

正如最初提到的，我还在 this repo 中设置了一个最小的可重现示例，其中包括上面的代码以及易于实验的依赖项。

Answer 1

作为框架挑战，不要尝试修改

DataFrame.apply

的类型存根，而是修改函数以使其返回预期的数据类型之一。

from __future__ import annotations

import pandas as pd

from typing import TypeVar

T = TypeVar("T")

df = pd.DataFrame(
    {
        "a": list("abc"),
        "b": list("def"),
        "c": list("aaa"),
    }
)

def to_set(row: pd.Series[T]) -> pd.Series[T]:
    return pd.Series(
        data=tuple(set(x for x in row if not pd.isna(x))),
        dtype=row.dtype,
        name=row.name,
    )

print(df.apply(to_set, axis=1))