Добавьте пользовательскую перегрузку в метод Python с существующими подсказками типов. ⇐ Python

[Anonymous]

Можно ли расширить существующие подсказки типов для пакета Python в коде моего приложения (т. е. не трогая подсказки восходящих типов), добавив пользовательскую перегрузку для метода? (См. также репозиторий с минимальным воспроизводимым примером.)
Например, возьмите подсказки типов для 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, так и Freezeset, и они были исправлены.
Мой вопрос есть ли способ расширить существующие подсказки типов для DataFrame.apply в моем собственном коде, добавив @overload, который позволяет возвращать set из примененной функции?
Я знаю, что могу переопределить весь Frame.pyi, создав typings/pandas-stubs/core/frame.pyi в моем собственном проекте на основе копии исходной версии, куда я бы добавил set в объединение возвращаемых типов в соответствующей перегрузке. Но для этого потребуется синхронизировать остальную часть файла с восходящим потоком по мере его изменений.
Вместо этого я бы хотел продолжать использовать исходный файлframe.pyi, просто расширив введите подсказки для DataFrame.apply в моем специальном коде, чтобы разрешить возврат set из прикладной функции.
Claude LLM предлагает следующий подход:< /p>

Код: Выделить всё

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
)
)

Как упоминалось изначально, я также создал минимальный воспроизводимый пример в этом репозитории, который включает приведенный выше код вместе с зависимостями для удобства экспериментирования.

Подробнее здесь: https://stackoverflow.com/questions/792 ... type-hints