有什么方法可以定义带有前导可选参数的 Python 函数？

Question

众所周知，可选参数必须在参数列表的末尾，如下所示：

def func(arg1, arg2, ..., argN=default)

我在PyTorch 包中看到了一些例外情况。例如，我们可以在torch.randint 中找到这个问题。如图所示，它的位置参数中有一个前导可选参数！怎么可能？

Docstring:
randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor

我们如何以与上述类似的方式定义函数？

Answer 1

我的另一个答案是关于对火炬库进行逆向工程，但是我想将这个答案专门用于说明如何以一种简单、直接的方式实现类似的机制。

我们有multipledispatch 库：

在 Python 中进行多次调度的一种相对健全的方法。 这种多分派的实现是高效的，大部分是完整的，执行静态分析以避免冲突，并提供可选的命名空间支持。看起来也不错。

所以让我们利用它：

from multipledispatch import dispatch

@dispatch(int, int)
def randint(low, high):
    my_randint(low, high)

@dispatch(int)
def randint(high):
    my_randint(0, high)

def my_randint(low, high):
    print(low, high)

# 0 5
randint(5)

# 2 3
randint(2, 3)

Answer 2

单个函数不允许只有前导可选参数：

8.6. Function definitions

[...] 如果一个参数有一个默认值，那么直到“*”的所有后续参数也必须有一个默认值——这是一个语法限制，没有由语法。

请注意，这不包括仅关键字参数，它们从不按位置接收参数。

---

如果需要，可以通过手动实现参数匹配的参数来模拟这种行为。例如，可以基于 arity 进行调度，或者显式匹配可变参数。

def leading_default(*args):
    # match arguments to "parameters"
    *_, low, high, size = 0, *args
    print(low, high, size)

leading_default(1, 2)     # 0, 1, 2
leading_default(1, 2, 3)  # 1, 2, 3

一种简单的调度形式通过迭代签名并调用第一个匹配的签名来实现函数重载。

import inspect


class MatchOverload:
    """Overload a function via explicitly matching arguments to parameters on call"""
    def __init__(self, base_case=None):
        self.cases = [base_case] if base_case is not None else []

    def overload(self, call):
        self.cases.append(call)
        return self

    def __call__(self, *args, **kwargs):
        failures = []
        for call in self.cases:
            try:
                inspect.signature(call).bind(*args, **kwargs)
            except TypeError as err:
                failures.append(str(err))
            else:
                return call(*args, **kwargs)
        raise TypeError(', '.join(failures))


@MatchOverload
def func(high, size):
    print('two', 0, high, size)


@func.overload
def func(low, high, size):
    print('three', low, high, size)


func(1, 2, size=3)    # three 1 2 3
func(1, 2)            # two 0 1 2
func(1, 2, 3, low=4)  # TypeError: too many positional arguments, multiple values for argument 'low'

Answer 3

你的发现让我着迷，因为在 Python（以及我知道的所有其他语言）中确实是非法的，具有领先的可选参数，这在我们的案例中肯定会出现：

SyntaxError: non-default argument follows default argument

我很怀疑，但我已经搜索了源代码：

我发现，TensorFactories.cpp 的第 566-596 行实际上有几个 (!) randint 的实现：

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ randint ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Tensor randint(int64_t high, IntArrayRef size, const TensorOptions& options) {
  return native::randint(high, size, c10::nullopt, options);
}

Tensor randint(
    int64_t high,
    IntArrayRef size,
    c10::optional<Generator> generator,
    const TensorOptions& options) {
  return native::randint(0, high, size, generator, options);
}

Tensor randint(
    int64_t low,
    int64_t high,
    IntArrayRef size,
    const TensorOptions& options) {
  return native::randint(low, high, size, c10::nullopt, options);
}

Tensor randint(
    int64_t low,
    int64_t high,
    IntArrayRef size,
    c10::optional<Generator> generator,
    const TensorOptions& options) {
  auto result = at::empty(size, options);
  return result.random_(low, high, generator);
}

这种模式在gen_pyi.py 的第 466-471 行再次出现，它为顶级函数生成类型签名：

        'randint': ['def randint(low: _int, high: _int, size: _size, *,'
                    ' generator: Optional[Generator]=None, {}) -> Tensor: ...'
                    .format(FACTORY_PARAMS),
                    'def randint(high: _int, size: _size, *,'
                    ' generator: Optional[Generator]=None, {}) -> Tensor: ...'
                    .format(FACTORY_PARAMS)],

所以，基本上发生的情况是没有“真正的”可选参数，而是几个函数，其中一个存在，另一个不存在。

这意味着，当调用 randint 时不带 low 参数时，它被设置为 0：

Tensor randint(
    int64_t high,
    IntArrayRef size,
    c10::optional<Generator> generator,
    const TensorOptions& options) {
  return native::randint(0, high, size, generator, options);
}

进一步研究，关于如何可能存在多个具有相同名称和不同参数的函数的 OP 请求：

再次返回gen_pyi.py，我们看到这些函数被收集到在第436行定义的unsorted_function_hints，然后在第509-513行用于创建function_hints，最后将function_hints设置为@987654340 @ 在第 670 行。

env 字典用于编写pyi 存根文件。

这些存根文件使用函数/方法重载，如PEP-484 中所述。

函数/方法重载，利用@overload装饰器：

@overload 装饰器允许描述支持多种不同参数类型组合的函数和方法。这种模式在内置模块和类型中经常使用。

这是一个例子：

from typing import overload

class bytes:
    ...
    @overload
    def __getitem__(self, i: int) -> int: ...
    @overload
    def __getitem__(self, s: slice) -> bytes: ...

所以我们基本上定义了同一个函数__getitem__，但参数不同。

还有一个例子：

from typing import Callable, Iterable, Iterator, Tuple, TypeVar, overload

T1 = TypeVar('T1')
T2 = TypeVar('T2')
S = TypeVar('S')

@overload
def map(func: Callable[[T1], S], iter1: Iterable[T1]) -> Iterator[S]: ...
@overload
def map(func: Callable[[T1, T2], S],
        iter1: Iterable[T1], iter2: Iterable[T2]) -> Iterator[S]: ...
# ... and we could add more items to support more than two iterables

这里我们定义了同一个函数map，但参数个数不同。