具有通用理解返回类型的函数取决于输入参数？

Question

假设我有一个函数可以转换字符串集合中的字符，例如 -

from typing import Collection

def replace_in_collection(input_collection: Collection[str]) -> Collection[str]:
    return [x.replace('a', 'b') for x in input_collection]

如何以通用方式创建此函数，以便它返回与输入类型相同的集合类型（不仅仅是上面函数示例中的列表）。

所以我想要的是具有以下行为的东西：

my_set = {'ab','ac'}
my_set_result = replace_in_collection(my_set)
type(my_set_result ) # --> set

my_list = ['ab','ac']
my_list_result = replace_in_collection(my_list)
type(my_list_result) # --> list

有没有一种方法可以在不使用单独类型检查的情况下通过理解来创建它？

Answer 1

假设所有类型的集合都接受 iterable 作为 __init__ 参数，那应该可以工作：

def replace_in_collection(input_collection):
    collection_type = type(input_collection)
    return collection_type(x.replace('a', 'b') for x in input_collection)

它适用于 set、list 和 tuple。

正如 Chris Rands 所说，您可能应该检查您的输入是否属于受支持的类型：

if not isinstance(input_collection, (set, list, tuple)): # add other types if needed
    raise TypeError(f"input_collection must be of type set, list or tuple, not {type(input_collection).__name__}")