39.2 이터레이터 만들기
이제 __iter__, __next__ 메서드를 구현해서 직접 이터레이터를 만들어보겠습니다. 간단하게 range(횟수)처럼 동작하는 이터레이터입니다.
class 이터레이터이름: def __iter__(self): 코드 def __next__(self): 코드
iterator.py
class Counter: def __init__(self, stop): self.current = 0 # 현재 숫자 유지, 0부터 지정된 숫자 직전까지 반복 self.stop = stop # 반복을 끝낼 숫자 def __iter__(self): return self # 현재 인스턴스를 반환 def __next__(self): if self.current < self.stop: # 현재 숫자가 반복을 끝낼 숫자보다 작을 때 r = self.current # 반환할 숫자를 변수에 저장 self.current += 1 # 현재 숫자를 1 증가시킴 return r # 숫자를 반환 else: # 현재 숫자가 반복을 끝낼 숫자보다 크거나 같을 때 raise StopIteration # 예외 발생 for i in Counter(3): print(i, end=' ')
실행 결과
0 1 2
실행을 해보면 0 1 2가 나옵니다. 이렇게 0부터 지정된 숫자 직전까지 반복하는 이터레이터 Counter를 정의했습니다.
먼저 클래스로 이터레이터를 작성하려면 __init__ 메서드를 만듭니다. 여기서는 Counter(3)처럼 반복을 끝낼 숫자를 받았으므로 self.stop에 stop을 넣어줍니다. 그리고 반복할 때마다 현재 숫자를 유지해야 하므로 속성 self.current에 0을 넣어줍니다(0부터 지정된 숫자 직전까지 반복하므로 0을 넣습니다).
def __init__(self, stop): self.current = 0 # 현재 숫자 유지, 0부터 지정된 숫자 직전까지 반복 self.stop = stop # 반복을 끝낼 숫자
그리고 __iter__ 메서드를 만드는데 여기서는 self만 반환하면 끝입니다. 이 객체는 리스트, 문자열, 딕셔너리, 세트, range처럼 __iter__를 호출해줄 반복 가능한 객체(iterable)가 없으므로 현재 인스턴스를 반환하면 됩니다. 즉, 이 객체는 반복 가능한 객체이면서 이터레이터입니다.
def __iter__(self): return self # 현재 인스턴스를 반환
그다음에 __next__ 메서드를 만듭니다. __next__에서는 조건에 따라 숫자를 만들어내거나 StopIteration 예외를 발생시킵니다. 현재 숫자 self.current가 반복을 끝낼 숫자 self.stop보다 작을 때는 self.current를 1 증가시키고 현재 숫자를 반환합니다. 이때 1 증가한 숫자를 반환하지 않도록 숫자를 증가시키기 전에 r = self.current처럼 반환할 숫자를 변수에 저장해 놓습니다. 그다음에 self.current가 self.stop보다 크거나 같아질 때는 raise StopIteration으로 예외를 발생시킵니다.
def __next__(self): if self.current < self.stop: # 현재 숫자가 반복을 끝낼 숫자보다 작을 때 r = self.current # 반환할 숫자를 변수에 저장 self.current += 1 # 현재 숫자를 1 증가시킴 return r # 숫자를 반환 else: # 현재 숫자가 반복을 끝낼 숫자보다 크거나 같을 때 raise StopIteration # 예외 발생
for 반복문에 Counter(3)을 지정해서 실행해보면 3번 반복하면서 0, 1, 2가 출력됩니다.
for i in Counter(3): print(i)
지금까지 간단한 이터레이터를 만들어보았습니다. 이터레이터를 만들 때는 __init__ 메서드에서 초깃값, __next__ 메서드에서 조건식과 현재값 부분을 주의해야 합니다. 이 부분이 잘못되면 미묘한 버그가 생길 수 있습니다. 예를 들어서 0, 1, 2와 1, 2, 3처럼 3번 반복하는 것은 같지만 숫자가 1씩 밀려서 나오거나 0, 1, 2, 3처럼 반복 횟수가 달라질 수 있으므로 코드를 작성할 때 꼼꼼히 살펴봐야 합니다.
39.2.1 이터레이터 언패킹
참고로 이터레이터는 언패킹(unpacking)이 가능합니다. 즉, 다음과 같이 Counter()의 결과를 변수 여러 개에 할당할 수 있습니다. 물론 이터레이터가 반복하는 횟수와 변수의 개수는 같아야 합니다.
>>> a, b, c = Counter(3) >>> print(a, b, c) 0 1 2 >>> a, b, c, d, e = Counter(5) >>> print(a, b, c, d, e) 0 1 2 3 4
사실 우리가 자주 사용하는 map도 이터레이터입니다. 그래서 a, b, c = map(int, input().split())처럼 언패킹으로 변수 여러 개에 값을 할당할 수 있습니다.
참고 | 반환값을
_에 저장하는 이유
함수를 호출한 뒤 반환값을 저장할 때 _(밑줄 문자)를 사용하는 경우가 있습니다.
>>> _, b = range(2) >>> b 1
사실 이 코드는 a, b = range(2)와 같습니다. 반환값을 언패킹했을 때 _에 할당하는 것은 특정 순서의 반환값 사용하지 않고 무시하겠다는 관례적 표현입니다. 예를 들어 다음과 같은 코드는 언패킹 했을 때 두 번째 변수는 사용하지 않겠다는 뜻입니다.
>>> a, _, c, d = range(4) >>> a, c, d (0, 2, 3)