有沒有更好的方法來定義這個發生器？

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。

有沒有更pythonic的方式來做到這一點？正在使用生成器來解決這個問題嗎？

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

的漢斯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)))