41.3 인덱스로 접근할 수 있는 이터레이터 만들기

지금까지 __iter____next__ 메서드를 구현하는 방식으로 이터레이터를 만들었습니다. 이번에는 __getitem__ 메서드를 구현하여 인덱스로 접근할 수 있는 이터레이터를 만들어보겠습니다.

앞에서 만든 Counter 이터레이터를 인덱스로 접근할 수 있도록 다시 만들어보겠습니다.

class 이터레이터이름:
    def __getitem__(self, 인덱스):
        코드

iterator_getitem.py

class Counter:
    def __init__(self, stop):
        self.stop = stop
 
    def __getitem__(self, index):
        if index < self.stop:
            return index
        else:
            raise IndexError
 
print(Counter(3)[0], Counter(3)[1], Counter(3)[2])
 
for i in Counter(3):
    print(i, end=' ')

실행 결과

0 1 2
0 1 2

Counter(3)[0]을 출력했을 때 0이 나왔습니다. 마찬가지로 Counter(3)[1]은 1, Counter(3)[2]는 2가 나왔습니다. 그리고 for 반복문에 Counter를 사용했을 때도 1, 2, 3이 나왔습니다.

소스 코드를 잘 보면 __init__ 메서드와 __getitem__ 메서드만 있는데도 동작이 잘 됩니다. 클래스에서 __getitem__만 구현해도 이터레이터가 되며 __iter__, __next__는 생략해도 됩니다(초깃값이 없다면 __init__도 생략 가능).

그럼 __init__ 메서드부터 살펴보겠습니다. 여기서는 Counter(3)처럼 반복을 끝낼 숫자를 받았으므로 self.stopstop을 넣어줍니다.

class Counter:
    def __init__(self, stop):
        self.stop = stop             # 반복을 끝낼 숫자

이제 클래스에서 __getitem__ 메서드를 구현하면 인덱스로 접근할 수 있는 이터레이터가 됩니다. 먼저 __getitem__은 매개변수로 인덱스 index를 받습니다. 그리고 지정된 index가 반복을 끝낼 숫자 self.stop보다 작을 때 index를 반환합니다. indexself.stop보다 크거나 같으면 IndexError를 발생시킵니다. 즉, Counter(3)과 같이 반복을 끝낼 숫자가 3이면 인덱스는 2까지 지정할 수 있습니다.

    def __getitem__(self, index):    # 인덱스를 받음
        if index < self.stop:        # 인덱스가 반복을 끝낼 숫자보다 작을 때
            return index             # 인덱스를 반환
        else:                        # 인덱스가 반복을 끝낼 숫자보다 크거나 같을 때
            raise IndexError         # 예외 발생

이렇게 하면 Counter(3)[0]처럼 이터레이터를 인덱스로 접근할 수 있습니다.

참고로 여기서는 반복할 숫자와 인덱스가 같아서 index를 그대로 반환했지만, index와 식을 조합해서 다른 숫자를 만드는 방식으로 활용할 수 있습니다. 예를 들어 index * 10을 반환하면 10, 20, 30처럼 10 단위로 숫자가 나옵니다. 각자 __getitem__ 메서드를 수정해서 다양한 숫자를 만들어보세요.

최근 수정: 2017년 7월 7일, 금요일, 오후 4:48