[Python] Iterable & Iterator

더 효과적인 반복문 만들기

파이썬 변수 타입 중에 반복문에서 사용할 수 있는 시퀀스 타입들이 있다. 예컨대, 리스트(list), 튜플(tuple) 등이 대표적인 시퀀스 타입들이다. 이렇게 요소를 반복할 수 있는 구조를 가진 객체를 iterable이라고 한다. 아래 예시를 보면 [1,2,3,4]라는 요소를 가진 list인 my_list가 정의되어 있다. list는 iterable 하기 때문에 for문을 통해서 각 원소에 접근할 수 있다. 

그런데 파이썬을 사용하다 보면 이렇게 리스트나 튜플과 같은 자료형 말고, 직접 만든 객체를 반복하고 싶은 경우가 발생한다. 예컨대, TCP 연결에 실패할 때 재시도하는 로직을 추가한다고 가정해 보자. 이때 쉬지 않고 재시도를 하면 부하가 몰릴 수 있기 때문에, 주로 exponential backoff 방식을 사용하여 재시도를 한다. 이를 파이썬으로 구현하면 다음과 같이 구현할 수 있다.

위 RetryManager 클래스에서 마음에 들지 않는 부분은 while문의 조건으로 True가 걸려있는 부분과 실패할 때마다 increase_retry_counter()를 호출하는 부분이다. 왜냐하면 반복 조건이 while 내부에서 need_retry()라는 함수로 별도로 존재하기 때문에, 만약 try 블록 코드가 길어지면 빠져나오는 곳을 정확히 찾기 어렵기 때문이다. 만약 이곳에서 로직이 잘못 고친다면, 영원히 반복문을 나오지 못할 수도 있다.

실제로, 위 로직에서 가장 중요한 부분은 다음 스텝에서 기다릴 시간을 구하는 것과 실제 max_retries만큼만 재시도를 하는 것이다. 이럴 때 간편하게 반복문을 처리할 수 있는 방법이 바로 Iterator를 만드는 것이다.

Iterator는 일종의 프로토콜로, 반복문 처리가 가능한 객체를 의미한다. Iterator를 사용하기 위해서는 두 가지 메서드가 필수적이다. 바로 __iter__()와 __next__()이다. __iter__() 메서드는 iterator 객체 자신을 반환하며, __next__() 메서드는 컬렉션의 다음 요소를 반환한다. 만약 더 이상 반환할 요소가 없을 때는 StopIteration 예외를 발생시켜 반복이 종료되도록 한다. 

Iterator를 만들어서 사용하면, 다음과 같이 간편하게 반복문을 처리할 수 있다. 

RetryManager는 __iter__ 함수를 통해서 내부 클래스인 RetryManagerIterator를 생성하고, 해당 Iterator는 반복문을 돌면서 __next__() 함수를 호출한다. 이렇게 하면, for문이 종료되는 조건을 빠르게 찾을 수도 있고, 반복문이 끝날 때마다 함수를 호출하는 부분을 없앨 수 있다.

Iterator의 또 다른 장점은 모든 데이터를 메모리에 불러오지 않고도 컬렉션의 요소를 하나씩 처리할 수 있다는 점이다. 이는 대용량 데이터를 다룰 때 매우 유용하다. Python에서는 iter() 함수를 사용하여 반복 가능한 객체로부터 iterator를 생성할 수 있다. 예컨대, 리스트에 대해 iter() 함수를 호출하면, 해당 리스트를 순회할 수 있는 iterator 객체가 반환된다. 이후 next() 함수를 이용하여 iterator가 가리키는 다음 요소를 하나씩 접근할 수 있다. 나중에 연산/평가한다고 해서 Lazy Evaluation이라고도 부른다.

메모리 사용량을 확인해 보면 정확하게 알 수 있다. 

천만 개의 숫자가 적힌 리스트를 반복문에 넣는다고 가정해 보자. 가장 쉬운 방법은 리스트로 만들어서 메모리에 올려놓는 것이다. 그렇게 되면 아직 접근하지도 않았는데, 메모리를 천만 개만큼 사용하는 꼴이 된다. 하지만, iterator를 쓰면 접근할 때만 원소를 메모리에 올려서 반환하기 때문에 메모리를 효율적으로 관리할 수 있다. 

그래서 대규모 데이터를 다루는 곳에서는 주로 iterator를 사용해서 메모리를 절약한다. 대표적으로 iterator를 사용하는 곳이 바로 파일 입출력이다. open() 함수는 파일을 모두 읽지 않고 iterator(_io.TextIOWrapper)를 반환한다. 실제 해당 객체에 __iter__()와 __next__()가 있는지 확인해 보면, 둘 다 존재하는 것을 알 수 있다.

__iter__, __next__ 가 있으면 iterator로 사용 가능하다.

파일 iterator를 만들고 나면, for문을 돌면서 한 줄씩 읽어나갈 수 있다. 해당 라인을 읽을 때마다 메모리에 데이터를 올려서 사용하는 방식이므로, 모든 자료를 한 번에 다 읽지 않고도 원하는 작업을 수행할 수 있다.

iterator의 단점은 대체로 반복이 불가하다는 것이다. 이전에도 설명했듯이, iterator는 원본 데이터를 미리 메모리에 넣어두지 않고, 실제 사용할 때 다음 아이템을 메모리에 올리는 구조이다. 따라서 마지막 아이템을 읽은 후에 다시 next() 함수를 호출하게 되면, 당연히 다음 아이템이 없기 때문에 StopIteration Exception이 발생한다.

파일도 마찬가지이다. 이미 한 번 읽어 내려간 파일 스트림을 다시 반복하게 되면, 아무것도 나오지 않는다. 만약 두 번 읽어야 한다면, open()을 두 번 호출하거나 아니면 데이터를 먼저 저장해 놓고 결과물을 반복해서 접근해야 한다. 

따라서 메모리를 아낄 수 있다는 장점이 있다면 무조건 iterator를 사용해서는 안된다. 자주 반복하는 경우라고 하면 미리 연산을 끝내 놓은 결과물을 반복해서 접근하는 게 CPU 소모 관점에서 더 유리할 수도 있다.