[Python] iterator 와 generator에 대해 알아보기

이번 포스팅에서는 Python에서 중요한 이터레이터와 제네레이터에 대해 알아보고 예시 코드도 살펴보기로 하자.

 

---

1. for loop는 어떻게 만들어질까?

파이썬에서 for를 사용한 loop는 매우 빈번하게 사용된다. 그런데 for loop는 어떤 매직 메소드를 활용해서 구현되는 것일까? 우선 아래와 같은 코드가 있다고 해보자.

 

string = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
# for loop는 iterable한 객체를 자동으로 iterator로 만들어줌
for s in string:
    print('for loop->', s)

 

위 처럼 문자열과 같은 iterable한 객체에 for loop를 사용하게 되면 for loop는 자동으로 iterable한 객체(위 코드에선 문자열)를 iterator로 내부적으로 자동으로 만들어준다. iterable한 객체들은 애초에 내부적으로 __iter__ 매직 메소드를 가지고 있다. 그래서 for loop가 이를 활성화(?) 시켜준다고 생각하면 될 듯 하다. 다시 말해, for를 사용하면 iter(string) 하여 이터레이터로 자동으로 만들어주고 내부적으로 next()를 호출하면서 하나씩 원소를 반환(출력)하게 된다. 그래서 for 문을 사용하지 않고 for 문과 동일한 결과를 만들어주려면 아래와 같이 만들어줄 수 있다.

 

# for loop를 직접 구현해보기
string_iter = iter(string)  # 1.__iter__로 iterator 만들기
while True:
    try:
        print('직접 구현->', next(string_iter))  # 2. __next__로 하나씩 원소 반환
    except StopIteration:
        print('There is no more elements')
        break

 

문자열 string을 iter() 라는 매직 메소드를 call 해서 이터레이터로 만들어준다. 그리고 이 이터레이터를 무한 반복해서 StopIteration 에러가 발생할 때까지 next() 매직 메소드를 호출해 원소를 하나씩 출력한다. 이렇게 하게 되면 for loop 문보다 코드가 길어지고 심지어 예외처리까지 추가해줘야 하는 등 귀찮음(?)이 많이 발생하기 때문에 iter() 와 next()를 내부적으로 자동으로 구현해주는 일반적인 for loop를 보통 사용하는 것이다.

2.  __iter__ 매직 메소드를 직접 바꾸어주어 제네레이터 스타일로 구현하자

자, 문자열이 입력으로 주어졌을 때, 공백으로 문자열을 구분하는 splitter를 만든다고 가정해보자. 우선 __next__ 매직 메소드를 우리가 구현해서 next() 함수를 하나씩 호출할 때마다 공백으로 구분된 문자를 하나씩 반환하도록 해보자.

 

# 문장을 공백으로 구분하는 클래스 만들기 -> 1. next를 활용
class WordSplitter:
    def __init__(self, text):
        self.idx = 0
        self.text = text.split(' ')

    def __next__(self):
        try:
            word = self.text[self.idx]
        except IndexError:
            raise StopIteration('There is no more elements')
        self.idx += 1
        return word

    def __repr__(self):
        return f'Word Splitter {self.text}'

ws = WordSplitter('가지 많은 나무에 바람 잘 날 없다')
print(next(ws))
print(next(ws))
print(next(ws))
print(next(ws))
print(next(ws))
print(next(ws))

 

하지만 위처럼 하게 되면 코드가 매우 길어지게 된다. 그래서 우리는 __iter__ 즉, 이터레이터로 만드는 함수를 제네레이터로 만드는 함수로 변환해서 문자열에 __iter__ 함수를 call 하게 되면 제네레이터 함수가 반환되도록 만들어 줄 수 있다. 그리고 이 만들어진 제네레이터 함수에 for loop를 사용할 수 있다. 이렇게 되면 제네레이터를 사용함으로써 데이터 양이 매우 증가할 때 메모리를 절약할 수 있는 장점도 존재하게 된다.

 

# 문장을 공백으로 구분하는 클래스 만들기 -> 2. Generator 스타일(__iter__ 매직메소드를 변경)
class WordSplitterGen:
    def __init__(self, text):
        self.text = text.split(' ')

    def __iter__(self):
        for word in self.text: # for loop가 자동으로 iterator로 만들어줌
            yield word

    def __repr__(self):
        return f'Word Splitter {self.text}'

ws2 = WordSplitterGen('하늘을 우러러 보자')
ws2_iter = iter(ws2)

for i in ws2_iter:
    print(i)

 

물론 for loop를 사용하지 않고 next() 함수를 계속적으로 호출해서 원소를 하나씩 반환할 수도 있다.  위 코드에서 제네레이터로 만들 때 yield 또는 Python 3.7부터는 await로 반환을 해주는데, 위 단어 구분자를 예시로 들면 yield 또는 await의 역할이 매번 다음에 리턴할 단어의 위치값 즉, 다음에 나올 단어의 index값을 내부적으로 기억한다고 보면 된다.

 

또 한가지 참고로 특정 객체가 어떤 속성을 갖고 있는지 체크하는 3가지 방법이 있다. 첫 번째는 dir() 함수를 사용하는 것, 두 번째는 hasattr() 함수를 사용하는 것, 세 번째는 collections의 abc 라이브러리를 활용하는 방법이 있다. 

 

string = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
string_iter = iter(string)

# 특정 객체가 특정 속성을 갖고 있는지 체크하는 방법
print(dir(string_iter))

print(hasattr(string_iter, '__next__'))

from collections import abc
print(isinstance(string_iter, abc.Callable))