파이썬 제너레이터(generator), yield Python

서론

파이썬에서는 yield 키워드를 통해 iterate하는 걸 볼 수 있다. 한 번에 하나의 요소를 통해 루프할 수 있는 방법이다. 특히 모든 것을 한 번에 메모리에 로드하고 싶지 않은 경우 한 번에 하나씩 데이터를 쉽게 처리할 수 있도록 하는 것이 목적이다. 쉽게 말해, 모든 데이터를 훑지 않고, 사용자가 원할 때 하나씩 값을 처리하는 것이다. 이렇게 값을 처리하는 방법을 Lazy Evaluation이라고 한다. 한 번에 하나씩 하는게 뭔말인지 아래 예제를 통해서 알아보자.

---

yield 란?

내가 아는 yield는 일본드라마인데(허허) 싶겠지만, 예제에 앞서 yield 키워드를 알아야 된다. return과 같은 역할을 수행하는데, return 자리에 yield가 있다면 그건 generator 함수라고 보면 된다. yield가 호출될 때마다 함수 작동을 멈추고 그 때의 상태를 저장한다. 예를들어 1,2,3하다가 멈추면 3은 리턴하되 3까지 했다는 건 저장해 놓는 것. 1에서 1000까지 iterate하려는데 100에서 멈춘 상태라고 가정해보면, 나머지 900은 훑지 않았다. 100에서 멈춘 뒤 다시 사용하는 게 아니라 다시 필요하면 101부터 훑을것이기에 메모리 절약에 우수하다고 하고, 메모리 사용이 큰 데이터 분석에서 주로 사용되는 것이다. 이제 진짜 예제를 통해 알아보자.

def number_generator():
    for i in range(1, 6):
        yield i

맨 마지막에 yield i가 쓰였다. 그렇다면 이건 generator가 되는 것이다. 어떠한 이유로 이 함수(=number_generator()) 가 멈추게 되면 이 함수는 작동을 멈추고 현재 상태를 저장과 함께 그 값을 반환한다. 일단 아래 처럼 for문을 사용해 number_generator를 돌려보자.

gen = number_generator()

for number in gen:
    print(number)
    
###결과###
1
2
3
4
5

 대충보면 return i를 했을 때랑 다를 것이 없어보인다. 하지만 다르다! 우선 gen = number_generator()이다. 만약 number_generator()를 return으로 사용했다면 저 함수는 동작하지 않는다. 하지만 yield를 사용했기 때문에 1한 것을 gen에 담았고 -> for문이 하나씩 gen을 호출 -> gen에서 i가 하나하나 print된다. 호출 할 때마다 yield i도 바뀐 것이다. 이걸 표로 정리해보면 아래와 같다.

	yield	return
방식	값을 내보내고 정지	값 반환 시 종료
상태	내부 상태 유지 됨-> 재호출시 다음 거 계산	내부 상태 유지 안 됨 ->  재 호출시 처음부터
계산 방식	lazy	eager

---

next()를 사용해보자

for문으로도 호출 할 수 있지만 next함수를 이용하면 더 직관적으로 generator를 사용할 수 있다.

gen = number_generator()
print(next(gen))  
print(next(gen)) 
####결과####
1
2

 

---

Generator의 활용

서론에서 말했지만 이러한 하나씩 훑으면서 값을 처리하는 구조상 빅데이터 분석에서 많이 쓰인다고 했다. 이 밖에 API 나 센서와 같이 실시간으로 데이터가 들어올 때 쓰이기도 한다. API를 사용하는 AI Agent에도 쓰일 것이다. 혹은 무한 루프를 만들고 싶을때도 쓰인다. 1일때-> 계산 -> 2일때 -> 계산-> 3일 때-> 계산->... -> 해서 무한 수열을 돌릴지라도 계속 output을 만들 수 있는 것이다. 범위가 정해져 있으면 그만큼만 가능하다는 것이다. 아무튼 무한 generator의 예제는 아래와 같다. print(next(gen))하는 만큼 i가 나오는 함수이다. 예제에서는 2개만 썻지만 123981293개 쓰면 123981293개의 i가 나올 것이다. 

def infinite_generator():
    i = 1
    while True:
        yield i
        i += 1

gen = infinite_generator()
print(next(gen))
print(next(gen))

 

---

for문 돌리고 next무한으로 하는거랑 머가 다름?

본문의 경우 1부터 5까지 for문을 도는 내용이었는데, 만약 next를 사용해 5 이후의 숫자를 원해서 계속 next(gen)을 하면 어떻게 될까? StopIteration exception이 생긴다. 하지만 for문을 써서 무한루프를 돌려 출력하고 싶다면... 된다! 아래 코드 처럼 for _ in iter(int,1):을 써서 무한루프 돌린다음 호출하면 우리가 예상했던 1\n2\n3\n 이 나온다.

def number_generator():
    cnt = 0
    for _ in iter(int, 1):  # for문 무한 루프
        cnt += 1
        yield cnt

gen = number_generator()
print(next(gen))
print(next(gen))
print(next(gen))