제너레이터, 이터레이터 Cheatsheet

Python 제너레이터는 여러모로 유용하다. <파이썬 코딩의 기술>을 읽다가 꼭 메모해두고 사용해야겠다 싶은 사용법들을 몇개 정리해 본다.

리스트를 반환하기 보다는 제너레이터를 사용하라

문자열과 찾으려는 알파벳이 주어지면 문자열에서 해당 알파벳을 찾아 위치(index)를 리스트로 반환하는 함수를 만들어보자.

def index_alpha(text, alphabet):
    result = []
    for idx, letter in enumerate(text):
        if letter == alphabet:
            result.append(idx)
    return result

알고리즘을 풀 때나 Python으로 개발을 하다 보면 생각보다 이런식으로 코드를 짜는 경우가 많다. 데이터를 처리하고, 리스트에 차곡차곡 순회 결과물을 담아 반환하는 식이다.

이 함수의 단점은 가독성 측면도 있지만, 성능 문제도 있다. 만약 result가 너무 큰 경우 성능 문제도 존재한다.

def index_alpha_iter(text, alphabet):
    for idx, letter in enumerate(text):
        if letter == alphabet:
            yield idx

# 필요할 경우 편리하게 순회할 수 있다.
it = index_alpha_iter('a text is a text', 't')
for idx in it:
    print(idx)

# 이런식으로 언제든 리스트로 변환 가능하다.
result = list(index_alpha_iter('some text', 't'))

이렇게 리팩토링하면 이터레이터를 반환한다. 제너레이터는 경제적이다. 코드의 가독성도 상승했다.

이터레이터를 순회하며 데이터를 처리할 수도 있고, 언제든 리스트로 변환도 가능하다.

한편, 기억해야할 사실이 있다.

이터레이터는 상태를 가진다. 제너레이터가 반환한 이터레이터는 호출 시 재사용이 불가능하다.

인자에 대해 이터레이션할 때는 방어적이 돼라

이터레이터를 다른 함수의 인자로 줄 때는 조심해야 한다.

다음 함수는 리스트의 길이를 반환한다.

이 함수를 이용해 index_alpha_iter를 사용해서 찾은 문자의 개수를 세는 방법을 생각해보자.

def count(letter_indices):
    for i in letter_indices:
        print(i)
    return len(list(letter_indices))

# 문자의 개수 세기
count(index_alpha_iter('some text'), 't')

위 함수에 만약 index_alpha_iter('some text', 't')를 인자로 준다면 어떻게 될까?

iterable은 한번 소비되고 나면 빈값이 된다. 따라서 0이 반환되게 된다.

이터레이터를 여러번 호출해 사용하기 위해서는,

  1. 이터레이터의 값을 list로 복사해 사용하는 방법 (메모리 소비)
  2. 클래스 활용

2가지 방법이 모두 이용 가능하다.

2번을 살펴보자면 다음과 같다.

class LetterIndices:
    def __init__(self, text, alphabet):
        self.text = text
        self.alphabet = alphabet
    
    def __iter__(self):
        for idx, letter in enumerate(self.text):
            if letter = self.alphabet:
                yield idx

# 사용예시
letter_indices = LetterIndices('some string', 's')
count(letter_indices)

마지막으로, count 함수 역시, 이터러블인지를 검증하도록 하면 더욱 안전해진다.

def count(letter_indices):
    if iter(letter_indices) is letter_indices:
        raise TypeError('컨테이너 제공 필요')
    for i in letter_indices:
        print(i)
    return len(list(letter_indices))

혹은 collections.abc에서 Iterator를 import해 isinstance로 비교하는 방법도 있다.

from collections.abc import Iterator

# 이런 방식으로 수정한다.
if isinstance(letter_indices, Iterator):
    raise TypeError('***')
 

긴 리스트 컴프리헨션보다는 제너레이터 식을 사용하라

# 파일이 길 경우 메모리 사용량이 늘고, 프로그램이 중단될 수 있음
****value = [len(line) for line in open('text.txt')]

# 제너레이터 이용: ()로 감싸주면 됨
value = (len(line) for line in open('text.txt'))

참고로 제너레이터는 서로 합성 가능하다.

it = iter([1, 2, 3, 4])
squares = ((n, n**2) for n in it)

이터레이터나 제너레이터를 다룰 때는 itertools를 사용하라

import itertools

# concat
it = itertools.chain(iter([1,2]), iter([3,4]))

# repeat
it = itertools.repeat(True, 10) # True, True...True가 10회

# cycle
it = itertools.cycle([1, 2, 3]) # 1, 2, 3, 1, 2, 3...

# tee
it1, it2, it3 = itertools.tee([0, 1, 2], 3) # it1, it2, it3 총 3개의 이터러블 생성

# islice
val = iter([1, 2, 3, 4, 5, 6])
it = itertools.islice(2, 5, 2) # (시작값, 끝값, 증가값) 끝값, 증가값은 생략 가능

# takewhile: True를 만족하는 원소만 반환
it = itertools.takewhile(lambda x : x < 7, [1, 2, 3, 4, 5, 6, 7, 8, 9]) # 1..6

# dropwhile: False를 반환하는 원소를 찾을 때까지 건너뛰고 False만 반환
it = itertools.dropwhile(lambda x : x < 7, [1, 2, 3, 4, 5, 6, 7, 8, 9]) # 7, 8, 9

# filterfalse: False인 원소만
# filter: True인 원소만

# accumulate: functools.reduce와 유사하나 매 단계마다 결과를 내놓음.
it = itertools.accumulate([1, 2, 3, 4, 5], lambda x: x*2)
# 2, 6, 12, 20, 30

# permutations
it = itertools.permutations([1, 2, 3], 2)
# [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]

# combinations
it = itertools.combinations([1, 2, 3], 2)
# [(1, 2), (1, 3), (2, 3)]