时间戳

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

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

问题描述

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

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

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

标签： pythonlist-comprehensiondictionary-comprehensionset-comprehension

解决方案

假设所有类型的集合都接受 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__}")

推荐阅读