Python 고급 문법 3: with 블럭과 @contextmanager

환영합니다, Rolling Ress의 카루입니다.

​

첫 번째 글에서 with 블럭에 대해 소개한 적이 있습니다. 파일 입출력 등을 할 때, with 블럭을 사용하면 try ~ finally 구문을 번거롭게 사용할 필요 없이 자원 해제를 자동으로 해주는 편리한 구문이었죠. 사실 이것도 숨겨진 메서드가 자동으로 호출되는 트릭에 가깝습니다. C++의 소멸자, C#의 IDisposible.Dispose()와 비슷합니다. 파이썬에서는 __enter__()과 __exit__(ty, val, tb) 메서드로 with 문의 동작을 정의할 수 있습니다.

​

'with' Statement (with 문)

f = open("data.txt", "w")
try:
    f.write("Hello")
finally:
    f.close()

정리가 필요한 작업들이 있습니다. 대표적인 예시가 파일 입출력입니다. 파일을 열었으면 반드시 닫아주어야 합니다. 단순히 open()을 실행한 뒤 close()를 실행해도 될 것 같지만, 오류가 나서 프로그램의 제어 흐름이 달라진다면 파일이 닫히지 않을 수도 있습니다. 이를 예방하기 위해 try ~ finally로 파일을 반드시 닫도록 합니다.

with open("data.txt", "w") as f:
    f.write("Hello")

하지만 with 문을 사용한다면 close()를 명시적으로 호출할 필요가 없습니다. 이건 with 문이 자동으로 파일을 닫아주기 때문입니다. 하지만 어떻게 알고 파일을 닫아줄 수 있을까요?

​

컨텍스트 매니저(Context Manager)는 시작할 때 준비 작업을 하고, 끝날 때 뒷정리를 하는 작업을 직접 정의할 수 있게 해줍니다. 지난 시간에 Iterator 클래스와 Generator 함수를 배웠듯, 클래스를 이용하여 구현하는 방법과 제너레이터 함수를 이용하는 방법 두 가지가 존재합니다.

import time

class Timer:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"[{self.name}] 타이머 시작")
        self.start_time = time.time()
        return self.start_time      # with ~ as 뒤에 받을 값 (t)을 반환합니다.

    def __exit__(self, exc_type, exc_val, traceback):
        self.end_time = time.time()
        elapsed = self.end_time - self.start_time
        print(f"[{self.name}] 타이머 종료 (소요 시간: {elapsed:.4f}s)")
        
        # 만약 에러가 발생했다면 exc_type에 에러 정보가 담깁니다.
        # return True를 하면 에러를 삼켜버리고(무시하고) 프로그램을 계속 실행시킵니다.

with Timer("데이터 크롤링") as t:
    print("...열심히 데이터를 수집하는 중...")
    time.sleep(1.5)  # 1.5초가 걸리는 작업이라고 가정
    print("...수집 완료!")

 

t = Timer("데이터 크롤링") 대신 with 문을 사용하여 타이머를 구현하였습니다. 명시적으로 __enter__() 또는 __exit__()을 호출하지 않아도 with문의 시작점과 종료점에서 자동으로 함수를 호출해줍니다.

​

t에는 Timer 객체가 들어갑니다. 이후 __enter__()가 실행되며 start_time을 기록하고, with 문 내부의 로직이 실행됩니다. with문을 벗어날 때 __exit__()이 실행되고 소요 시간을 계산하여 출력합니다.

​

하지만 함수로도 충분히 가능한 일을 클래스까지 도입해서 만들 필요가 있을까요? 그래서, 클래스가 필요 없는 제너레이터 방식을 사용할 수 있습니다.

​

@contextmanager

데코레이터는 이미 배웠습니다. 함수를 변경하지 않고 부가기능을 추가해 준다고 했죠. 파이썬 기본 라이브러리의 데코레이터는 예약어처럼 생각하고 사용하여도 무방합니다.

from contextlib import contextmanager
import time

@contextmanager
def timer(name):
    print(f"[{name}] 타이머 시작") # 1. 시작 (준비 작업) : __enter__ 역할
    start_time = time.time()
    try:
        yield start_time       # 2. 값 전달 및 대기 (as 뒤로 넘어갈 값)
    finally:
        end_time = time.time() # 3. 끝 (뒷정리 작업) : __exit__ 역할
        elapsed = end_time - start_time
        print(f"[{name}] 타이머 종료 (소요 시간: {elapsed:.4f}s)")

with timer("데이터 크롤링") as t:
    print("...열심히 데이터를 수집하는 중...")
    time.sleep(1.5)  # 1.5초가 걸리는 작업이라고 가정
    print("...수집 완료!")

코드가 조금 더 간결해집니다. 약간 헷갈릴 수 있는데, yield를 기준으로 위쪽은 __enter__, 아래쪽은 __exit__으로 볼 수 있습니다. yield를 하는 부분이 t로 들어가 with 문 내부에서 사용할 수 있고, with 문을 나오면 yield 아래 부분이 실행되는 겁니다. 매번 try ~ finally를 길게 쓰는 대신 컨텍스트 매니저를 잘 만들어두면 실수를 차단하고 코드를 깔끔하게 유지할 수 있습니다.