如何为自定义limit函数编写类型提示以返回同类型集合？

问题背景 #

以下代码可正常运行，但无论如何调整类型提示，PyCharm或mypy总会报错。已尝试将S的类型设为Sized、Iterable和Collection：

import random
from collections.abc import Collection
from typing import TypeVar

T = TypeVar("T")
S = TypeVar("S", bound=Collection[T])


def limit(i: S, n: int) -> S:
    """
    Limits the size of the input iterable to n.  Truncation is randomly chosen.
    """
    if len(i) <= n:
        return i
    return type(i)(random.sample(list(i), n))

需求是实现类似random.sample的功能，满足：

1. 当n > len(i)时，返回原对象而非抛出错误；
2. 输出集合类型与输入完全一致（如输入list[int | str]则输出同类型，输入set[float]也输出同类型）。

类型提示报错原因与修复 #

类型检查器报错的核心原因：Collection包含的子类并非都支持用普通元素列表作为构造参数。比如dict的构造函数要求传入键值对迭代器，而random.sample返回的是单个元素的列表，这会导致类型不兼容，因此类型检查器无法确认type(i)(random.sample(...))的合法性。

修复方案 #

可以通过自定义Protocol来约束支持从元素列表构造的集合类型：

import random
from collections.abc import Collection, Iterable
from typing import TypeVar, Protocol, Type

T = TypeVar("T")

# 定义协议，要求类型可以从元素迭代器构造
class ConstructibleFromIterable(Protocol[T]):
    @classmethod
    def __call__(cls: Type["ConstructibleFromIterable[T]"], iterable: Iterable[T]) -> "ConstructibleFromIterable[T]": ...

# 绑定同时满足Collection和自定义协议的类型
S = TypeVar("S", bound=Collection[T] & ConstructibleFromIterable[T])

def limit(i: S, n: int) -> S:
    if len(i) <= n:
        return i
    sampled_items = random.sample(list(i), n)
    return type(i)(sampled_items)

这样类型检查器就能确认type(i)可以接收采样后的列表生成对应类型的实例。

标准库是否有现成实现 #

Python标准库中没有完全匹配需求的函数：

• random.sample：当n > len(i)时会抛出ValueError，不满足第一个需求；
• itertools.islice：仅能按顺序截断可迭代对象，返回迭代器而非原类型集合，且不支持随机截断；
• 其他相关工具函数要么无法保留原集合类型，要么不处理n > len(i)的场景。

因此需要自行实现该函数，上述修复类型提示后的版本可满足需求。

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