[Python] Generator

yield 라는 키워드가 사용되는 것을 generator라고 부른다. 아래 예시 코드를 통해 이해할 수 있다.

def frange(start, stop, increment):
    x = start
    while x < stop:
        yield x
        x += increment

for n in frange(0, 4, 0.5):
    print(n , end=' ')

0 0.5 1.0 1.5 2.0 2.5 3.0 3.5

위 코드는 특정 시작 값부터 종료 값까지 일정 증가값 간격으로 출력하는 코드다. 기본적으로 python에선 range() 함수를 통해서 사용 가능하지만, 소수점에 대해선 지원하지 않는다. frange는 소수점을 지원하는 함수다.

yield x 에 도달하면 x의 값을 return하지만 함수는 종료되지 않고 계속해서 iteration을 수행하게 된다. 결국 while 조건에 부합하지 않을 때 함수가 종료된다. yield는 함수 수행이 종료됐을 때 불필요하게 수행되는 것을 막기 위해 사용된다고 한다. 아래 예시 코드를 보면 더 이해가 쉽다.

def countdown(n):
    print('Starting to count from',n)
    while n > 0:
        yield n
        n -= 1
    print('Done!')
    

c = countdown(3)
c

<generator object countdown at 0x00000223E79CF4C8>

위 코드를 수행했을 때 Starting to count from 가 출력되지 않는다. 왜냐하면 이는 함수가 아니라 yield 가 들어가는 순간 generator가 된다, 즉 iterator처럼 동작하게 된다. 실제 main 문과 함수 구문이 서로 번갈아가면서 수행되게 되기 때문이다.

next(c) # Starting to count from 3
next(c) # 2
next(c) # 1
next(c) # Done!

3 2 1 Done!

위 결과를 보면 이해가 가기 더 쉽다. 처음 countdown(3)이 호출되면 함수가 수행되는게 아니라 단순하게 generator 객체를 return하는 것을 의미한다.
실제 수행은 next()가 호출 돼야지 수행된다.

Advanced Generator

Iterator와 generator를 활용한 코드를 소개한다.

class Node:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children)

    def depth_first(self):
        yield self
        for c in self:
            yield from c.depth_first()

if __name__ == '__main__':
    root = Node(0)
    child1 = Node(1)
    child2 = Node(2)
    root.add_child(child1)
    root.add_child(child2)
    child1.add_child(Node(3))
    child1.add_child(Node(4))
    child2.add_child(Node(5))

    for ch in root.depth_first():
        print(ch, end=' ')

Node(0) Node(1) Node(3) Node(4) Node(2) Node(5)

위 코드는 tree 구조처럼 생성되는데, main 함수에서 for 문을 다시 기술해보면 다음과 같이 된다.

# Previous
    for ch in root.depth_first():
        print(ch, end=' ')
# Alternative
    g = root.depth_first()
    while True:
        ch = next(g)

처음은 root의 generator 함수인 depth_first()를 호출하게 되는데, 실제로 next()가 호출 될 때 수행이 된다. depth_first() 의 yield까지 수행되고 print(ch, end=' ')를 수행하면서 Node(0)을 출력하게 된다.
그리고 다음 next()가 call이 되면서 depth_first() 내부에 for문을 수행하게 된다. 해당 for문의 self는 본인의 iter()를 의미하게 되고 __iter__ special function을 수행한다. 내부엔 본인의 children의 iter를 호출하게 된다. 참고로 for – in 문의 뒤에 선언되는 것은 내부 __iter__ 를 호출한다. 코드를 풀어서 써보자면 아래처럼 될 것이다.

# Previous
    for ch in root.depth_first():
        print(ch, end=' ')
# Alternative
    g = root.depth_first()
    while True:
        ch = next(g)
        print(ch, end=' ')
        it = iter(self)
        while True:
            c = next(it)
            print(c, end=' ')

Reverse Iteration

class Countdown:
    def __init__(self, start):
        self.start = start

    def __iter__(self):
        n = self.start
        while n > 0:
            yield n
            n -= 1

    def __reversed__(self):
        n = 1
        while n <= self.start:
            yield n
            n += 1

c = Countdown(5)
print("Forward:")
for x in c:
    print(x, end=' ')
print('\n')
print("Reverse:")
for x in reversed(c):
    print(x, end=' ')

Forward: 5 4 3 2 1 Reverse: 1 2 3 4 5

Python에선 reverse iteration을 __reversed__ special function을 구현해서 사용 가능하다.

Iterator 일부 얻기

def count(n):
    while True:
        yield n
        n += 1
        
c = count(0)
c[5:10]

TypeError Traceback (most recent call last) <ipython-input-80-0e5ab4786dd3> in 5 6 c = count(0) —-> 7 c[5:10] TypeError: ‘generator’ object is not subscriptable

위 코드는 Python에서 generator를 통해서 생성한 list를 접근을 하려고 하는 코드다. 그러나 실행하면 error를 발생하는데 이는 generator를 통해서 생성하기 위해선 next() 를 수행해야지 생성이 되기 때문이다. 이를 수행시킬 수 있도록 돕는 tool이 있다.

Definition

islice([Input], [Start Index], [End Index])

Example

import itertools

for x in itertools.islice(c, 10, 20):
    print(x)

10 11 12 13 14 15 16 17 18 19

itertools에는 islice (iterator slice)라는 method가 있는데 이를 통해서 특정 부분을 쉽게 참조 할 수 있다.