重载接受任意数量的位置参数并返回 int 或 Tuple[int] 的函数

Question

我正在尝试输入提示一个函数，该函数接受任意数量的位置参数，如果大小为 1，则返回 int；如果大于 1，则返回相同大小的元组。

这就是我现在拥有的：

@overload
def timestamp(*date: datetime) -> Tuple[int, ...]: ...


@overload
def timestamp(*date: str) -> Tuple[int, ...]: ...


@overload
def timestamp(*date: int) -> Tuple[int, ...]: ...


def timestamp(*date: Union[datetime, str, int]) -> int | Tuple[int, ...]:
    """
    It converts a date to a timestamp

    :param date: The date to be converted to timestamp
    :return: The number of seconds since the epoch.
    """
    if len(date) == 1: return timestamp_(date[0])
    return tuple([timestamp_(d) for d in date])

mypy 并没有抱怨这一点，但我想以正确的方式重载它，因为使用正确的类型使用这个函数会有点困难

Answer 1

使用

flat

kwarg 的示例：

from datetime import datetime
from typing import Literal, overload


def timestamp_(d) -> int:
    return 1


@overload
def timestamp(*date: datetime | str | int, flat: Literal[True]) -> int:
    ...


@overload
def timestamp(
    *date: datetime | str | int, flat: Literal[False] = False
) -> tuple[int, ...]:
    ...


def timestamp(*date: datetime | str | int, flat=False) -> int | tuple[int, ...]:
    """
    It converts a date to a timestamp

    :param date: The date to be converted to timestamp
    :return: The number of seconds since the epoch.
    """
    if flat:
        if len(date) == 1:
            return timestamp_(date[0])
        else:
            raise ValueError
    return tuple([timestamp_(d) for d in date])


def main():
    print(f"{timestamp(1) = }")
    print(f"{timestamp(1, flat=True) = }")
    print(f"{timestamp(1, 2) = }")


if __name__ == "__main__":
    main()

输出：

timestamp(1) = (1,)
timestamp(1, flat=True) = 1
timestamp(1, 2) = (1, 1)