numpy.apply_along_axis触发mypy类型不兼容错误的解决方法

2024年10月18日补充：简化复现代码 #

以下是最小化的问题复现代码：

mypy_arg_type.py:

import numpy as np
from numpy.typing import NDArray
import random

def winner(_: NDArray[np.bytes_]) -> bytes | None:
    return b"." if bool(random.randint(0, 1)) else None

board = np.full((2, 2), ".", "|S1")
for w in np.apply_along_axis(winner, 0, board):
    print(w)

运行Python代码输出：

b'.'
None

运行mypy检查报错：

mypy_arg_type.py:9: error: Argument 1 to "apply_along_axis" has incompatible type "Callable[[ndarray[Any, dtype[bytes_]]], bytes | None]"; expected "Callable[[ndarray[Any, dtype[Any]]], _SupportsArray[dtype[Never]] | _NestedSequence[_SupportsArray[dtype[Never]]]]"  [arg-type]
mypy_arg_type.py:9: note: This is likely because "winner" has named arguments: "_". Consider marking them positional-only
Found 1 error in 1 file (checked 1 source file)

原始问题背景 #

我正在实现四子棋的获胜判断功能，棋盘尺寸为6x7，列标从A到G，获胜条件是横、竖、斜向或反斜向连成4枚同色棋子。

示例 #

• 输入：["A_Red", "B_Yellow", "A_Red", "B_Yellow", "A_Red", "B_Yellow", "G_Red", "B_Yellow"]
• 棋盘状态：

R Y . . . . R
R Y . . . . .
R Y . . . . .
. Y . . . . .
. . . . . . .
. . . . . . .

• 获胜者：Yellow

实现代码 #

import itertools
import numpy as np
from numpy.typing import NDArray

def who_is_winner(pieces: list[str]) -> str:
    def parse_board() -> NDArray[np.bytes_]:
        m, n = 6, 7
        indices = [0] * n
        # 单字节编码，字节序为‘|’（不适用）
        board = np.full((m, n), ".", "|S1")
        for p in pieces:
            col = ord(p[0]) - ord("A")
            board[indices[col], col] = p[2]
            indices[col] += 1

        return board

    def winner(arr: NDArray[np.bytes_]) -> np.bytes_ | None:
        i = len(arr)
        xs = next(
            (xs for j in range(i - 3) if (xs := set(arr[j : j + 4])) < {b"R", b"Y"}),
            {None},
        )
        return xs.pop()

    def axis(x: int) -> np.bytes_ | None:
        # Axis 0是按列处理，1是按行处理
        return next(
            (w for w in np.apply_along_axis(winner, x, board) if w is not None), None
        )

    def diag(d: int) -> np.bytes_ | None:
        # 对角线编号相对于主对角线
        b = board if bool(d) else np.fliplr(board)
        return next(
            (w for d in range(-3, 4) if (w := winner(b.diagonal(d))) is not None), None
        )

    board = parse_board()
    match next(
        (
            w
            for f, i in itertools.product((axis, diag), (0, 1))
            if (w := f(i)) is not None
        ),
        None,
    ):
        case b"Y":
            return "Yellow"
        case b"R":
            return "Red"
        case _:
            return "Draw"

mypy报错信息 #

error: Argument 1 to "apply_along_axis" has incompatible type "Callable[[ndarray[Any, dtype[bytes_]]], bytes_ | None]"; expected "Callable[[ndarray[Any, dtype[Any]]], _SupportsArray[dtype[bytes_]] | _NestedSequence[_SupportsArray[dtype[bytes_]]]]"  [arg-type]
note: This is likely because "winner" has named arguments: "arr". Consider marking them positional-only

根据apply_along_axis的文档，该函数允许返回单个值，和我的代码逻辑一致。将winner改为位置参数后，仅移除了提示信息，但错误依然存在。使用环境：Python 3.12.5、mypy 1.11.2。

解决方案 #

方法1：调整winner函数的类型注解 #

mypy对apply_along_axis的类型定义要求输入函数能接受任意dtype的数组，同时返回值需符合数组或嵌套序列的约束。可以放宽输入参数类型，并将返回的单个值包装成数组：

def winner(arr: NDArray[np.generic]) -> NDArray[np.bytes_] | None:
    result = b"." if bool(random.randint(0, 1)) else None
    return np.array(result, dtype="|S1") if result is not None else None

方法2：用typing.cast绕过类型检查 #

若确认逻辑正确，仅需绕过mypy的类型推断，可强制转换函数类型：

from typing import cast, Callable
import numpy as np
from numpy.typing import NDArray, _SupportsArray, dtype

# 在调用处转换类型
winner_cast = cast(Callable[[NDArray[np.generic]], _SupportsArray[dtype[np.bytes_]] | None], winner)
for w in np.apply_along_axis(winner_cast, 0, board):
    print(w)

方法3：升级依赖版本 #

numpy的部分类型定义问题在新版本中已修复，尝试升级到numpy 2.0+和mypy 1.12+，可能自动解决类型不兼容问题。

方法4：手动遍历替代apply_along_axis #

避免使用apply_along_axis，直接遍历数组的行/列，完全控制类型：

# 替代apply_along_axis(winner, 0, board)的循环
for col in board.T:
    w = winner(col)
    if w is not None:
        print(w)

内容的提问来源于stack exchange，提问作者Abhijit Sarkar