Есть ли лучший способ определить этот генератор?

Question

Скажем, я хочу создать генератор, который возвращает фиксированное количество пользователей из списка, и если количество пользователей, которых я хочу, больше, чем длина списка, он просто повторяет последнего пользователя. Я написал генератор, чтобы сделать это ниже:

def gen_users(users, total): 
    num_users = len(users) 
    cur_user = 0 
    while cur_user < total: 
     if cur_user >= num_users: 
      yield users[-1] 
     else: 
      yield users[cur_user] 
     cur_user += 1 

Так что, если у меня был список пользователей users = ['one', 'two', 'three'] и я называю gen_users(users, 5) я вернусь one, two, three, three, three.

Есть ли более питонический способ сделать это? Использует ли генераторы действительный подход к этой проблеме?

Answer 1

Использование itertools.islice:

>>> import itertools 
>>> 
>>> def gen_users(users, total): 
...  # assert users and total > 0 
...  for i, user in enumerate(itertools.islice(users, total)): 
...   yield user 
...  for _ in range(total - i - 1): # i -> last index 
...   yield user 
... 
>>> list(gen_users([1,2,3], 5)) 
[1, 2, 3, 3, 3] 
>>> list(gen_users([1,2,3], 2)) 
[1, 2] 

В качестве альтернативы вы можете использовать users[:total], но это создаст временную последовательность.

UPDATE

Слегка модифицированная версия кода Hans Zauber по:

>>> from itertools import islice, chain, repeat 
>>> 
>>> def gen_users(users, total): 
...  return islice(chain(users, repeat(users[-1])), total) 
... 
>>> list(gen_users([1,2,3], 5)) 
[1, 2, 3, 3, 3] 
>>> list(gen_users([1,2,3], 3)) 
[1, 2, 3] 
>>> list(gen_users([1,2,3], 2)) 
[1, 2] 

Answer 2

Попробуйте Однострочник:

import itertools 
gen_users = lambda users, total: itertools.chain(itertools.islice(users, total), itertools.repeat(users[-1], total-len(users)))