Мне нужно оставить определенные ключи в словаре и удалить ненужные

Question

Мне нужно реализовать генератор, который будет возвращать значения ключей словарного массива через yield.Если один аргумент передается в генератор, выведите только значение переданного ключа.Если более одного аргумента, вернуть словари с ключами, которые были переданы в генератор. Вывести вывод в одну или две строки

Что я пробовал

def field(items, *args):
    assert len(args) > 0
    # Необходимо реализовать генератор
    if (len(args) == 1):
        for es in items:
            if (args[0] in es and args[0] != 'None'):
                yield es[args[0]]
    else:
        for i in range(len(items)):
            en_value = []
            dict = {}
            for j in range(len(args)):
                if args[j] != None:
                    en_value.append(items[i][args[j]])
                    dict[str(args[j])] = en_value[j]
            yield dict

Ожидается:

print(field(goods, 'title'))

выход: «Чехол», «Диван», «Стеллах», «Вечалка»

print(field(goods, 'title', 'price'))

вывод: {'title': 'Cover', 'price': 2000}, {'title': 'Sofa', 'price': 5300}

И это работает, но не работает, если я это сделаю:

for idx in field(goods, None, 'color'):
    print(idx)

Ошибка:

IndexError: список индексов вне диапазона.

Вопрос:

Можете ли вы помочь мне исправить это прекрасное?

Answer 1

Если я правильно понял, что вы хотите, ваша функция field должна выглядеть следующим образом:

def field(items, *args):
    clargs = [x for x in args if x is not None]
    assert len(clargs) > 0
    if len(clargs) == 1:
        for es in items:
            if clargs[0] in es:
                yield es[clargs[0]]
    else:
        for es in items:
            yield {j:es[j] for j in clargs}

Использование списка goods:

In [9]: list(field(goods, 'title'))
Out[9]: ['Cover', 'Sofa', 'Stellach', 'Vechalka']

In [10]: list(field(goods, 'title', 'price'))
Out[10]: 
[{'price': 2000, 'title': 'Cover'},
 {'price': 5300, 'title': 'Sofa'},
 {'price': 7000, 'title': 'Stellach'},
 {'price': 800, 'title': 'Vechalka'}]