[Python] Tip - 리스트 대신 제너레이터를 반환하도록 하자

예를 들어 문자열에 있는 모든 단어의 인덱스를 출력하는 예제를 만든다고 합니다. 다음 코드에서는 append 메서드로 리스트에 결과들을 누적하고 함수가 끝날 떄 해당 리스트를 반환합니다.

def index_words(text): result = [] if text: result.append(0) for index, letter in enumerate(text): if letter == ' ': result.append(index + 1) return result address = 'Four score and seven years ago...' result = index_words(address) print(result[:3]) # 결과 # [0, 5, 11]

여기서 index_words 함수에는 두 가지 문제가 있습니다.

1. 코드가 약간 복잡하고 깔끔하지 않음

   1. 새로운 결과가 나올 때마다 append 메서드를 호출해야 합니다. 메서드 호출(result.append)이 많아서 리스트에 추가하는 값(index + 1)이 덜 중요해 보입니다. 결과 리스트를 생성하는데 한 줄이 필요하고 그 값을 반환하는 데도 한 줄이 필요합니다. 함수 몸체에 문자가 130개 가량 (공백 제외)있지만 그 중에서 중요한 문자는 약 75개 입니다.
      
      이 함수를 작성하는 더 좋은 방법은 제너레이터(generator)를 사용하는 것입니다. 제너레이터는 yield 표현식을 사용하는 함수입니다. 제너레이터 함수는 호출되면 실제로 실행하지 않고 바로 이터레이터(iterator)를 반환합니다.
      내장 함수 next를 호출할 때마다 이터레이터는 제너레이터가 다음 yield 표현식으로 진행하게 합니다. 제너레이터에서 yield에 전달한 값을 이터레이터가 호출하는 쪽에 반환합니다.
      

      def index_words_iter(text): if text: yield 0 for index, letter in enumerate(text): if letter == ' ': yield index +1 result = list(index_words_iter(address))

      결과 리스트와 연동하는 부분이 모두 사라져서 훨씬 이해하기 쉽습니다. 결과는 리스트가 아닌 yield 표현식으로 전달됩니다. 제너레이터 호출로 반환되는 이터레이터를 내장함수 list에 전달하면 손쉽게 리스트로 변환할 수 있습니다.
      
      

2. 반환되기 전에 모든 결과를 리스트에 저장해야 함

   1. 입력이 매우 많으면 프로그램 실행 중에 메모리가 고갈되어 동작이 멈추는 원인이 됩니다. 반면 제너레이터로 작성한 버전은 다양한 길이의 입력에도 쉽게 이용할 수 있습니다.
      다음은 파일에서 입력을 한 번에 한 줄씩 읽어서 한 번에 한 단어씩 출력을 내어주는 제너레이터입니다. 이 함수가 동작할 때 사용하는 메모리는 입력 한 줄의 최대길이입니다.
      

      from itertools import islice def index_file(handle): offset = 0 for line in handle: if line: yield offset for letter in line: offset += 1 if letter == ' ': yield offset with open('address.txt', 'r') as f: it = index_file(f) result = islice(it, 0, 3) print(list(result)) # 결과 [0, 5, 11]

       이와 같은 제너레이터를 정의할 때 알아둬야 할 단 하나는 반환되는 이터레이터에 상태가 있고 재사용할 수 없다는 사실을 호출하는 쪽에서 알아야 한다는 점입니다.

요약

제너레이터를 사용하는 방법이 누적된 결과의 리스트를 반환하는 방법보다 이해하기 명확함 제너레이터에서 반환한 이터레이터는 제너레이터 함수의 본문에 있는 yield 표현식에 전달된 값들의 집합 제너레이터는 모든 입력과 출력을 메모리에 저장하지 않으므로 입력값의 양을 알기 어려울 때도 연속된 출력을 만들 수 있음