функция перегрузки mypy с numpy ndarray и pandas dataframe (типы параметров подписи такие же или более широкие)

Question

У меня есть функция, которая выполняет некоторые арифметические c вещи (квантильную нормализацию) по массиву numpy или pandas фрейму данных. Когда вы вставляете ndarray, вы должны вернуть ndarray, а когда вы вставляете pandas dataframe, вы должны возвращать dataframe:

from typing import Union, overload


@overload
def quantile_normalize(data: pd.DataFrame) -> pd.DataFrame: ...
@overload
def quantile_normalize(data: np.ndarray) -> np.ndarray: ...


def quantile_normalize(data: Union[pd.DataFrame, np.ndarray]) -> Union[pd.DataFrame, np.ndarray]:
    pass

Однако, когда я пытаюсь проверить это с помощью mypy it жалуется:

qnorm/quantile_normalize.py:72: error: Overloaded function signature 2 will never be matched: signature 1's parameter type(s) are the same or broader

Проблемы и ответы, которые я нашел до сих пор, связанные с этой проблемой, похоже, связаны с дополнительным вводом / выводом и типами None. Фрейм данных pandas и массив numpy связаны, однако они должны различаться mypy.

Answer 1

Если в библиотеке отсутствуют подсказки типов, каждый импорт будет разрешен до Any. И numpy, и pandas не соответствуют PEP 526 (не предлагают никаких подсказок типа) и не имеют заглушек в typeshed, поэтому оба pd.DataFrame и np.ndarray будут преобразованы в Any, поэтому обе перегрузки разрешаются на def quantile_normalize(data: Any) -> Any: .... Чтобы решить эту проблему, добавьте заглушки для numpy и pandas.

Либо используйте заглушки существующих типов - я использую data-science-types ( PyPI , GitHub ) которые предлагают заглушки для numpy, pandas и matplotlib:

$ pip install data-science-types

Теперь pd.DataFrame и np.ndarray будут правильно разрешены при запуске mypy. Это также обеспечит вам лучшее завершение кода в каждой среде IDE, которая поддерживает PEP 526 (например, Visual Studio Code или WingIDE) бесплатно.

Или, если вы не можете / не хотите добавлять пакет-заглушку, напишите свои собственные минимальные заглушки, например

# _typeshed/pandas/__init__.pyi

from typing import Any


def __getattr__(name: str) -> Any: ...  # incomplete

class DataFrame:
    def __getattr__(self, name: str) -> Any: ...  # incomplete

и

# _typeshed/numpy/__init__.pyi

from typing import Any


def __getattr__(name: str) -> Any: ...  # incomplete

class ndarray:
    def __getattr__(self, name: str) -> Any: ...  # incomplete

и запустите MYPYPATH=_typeshed mypy ....