Функция с универсальным типом возвращаемого значения в зависимости от входного аргумента? - PullRequest
Новая
       39

Функция с универсальным типом возвращаемого значения в зависимости от входного аргумента?

1 голос
/ 24 сентября 2019

Предполагается, что у меня есть функция, которая преобразует символ в коллекцию строк, например - 

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

Есть ли способ создать это с пониманием, не используя индивидуальную проверку типов?

1 Ответ

3 голосов
/ 24 сентября 2019

Предполагая, что все типы коллекций принимают итеративный параметр __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.

И, как заметил Крис Рэндс, вам, вероятно, следует проверить, поддерживается ли ваш ввод: 

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__}")
...