파이썬의 `contextmanager`

객체를 잠깐 바꿨다가 되돌려놓기

· 2 min read · #python #contextmanager · @python · English

문제

객체를 잠깐 빌려서 설정 하나를 뒤집은 채로 헬퍼 함수를 실행하고, 함수가 어떻게 끝나든(정상 반환, 조기 반환, 예외 발생) 객체를 원래 상태 그대로 돌려놔야 할 때가 있다. try … finally로 손수 하면 되긴 하지만, 코드가 지저분해지고 빼먹기도 쉽다.

context manager

context manager 는 “준비 → 코드 실행 → 정리”를 읽기 좋은 한 블록으로 보장해준다. with 문을 쓰든 contextlib의 데코레이터 형태를 쓰든 마찬가지다.

속성 하나짜리 최소 예제

from contextlib import contextmanager

@contextmanager
def temporarily(obj, attr, new_value):
    """with 블록 안에서만 obj.attr를 new_value로 설정한다."""
    old_value = getattr(obj, attr)
    setattr(obj, attr, new_value)
    try:
        yield obj          # 이제 with 블록 안의 코드가 실행된다
    finally:
        setattr(obj, attr, old_value)   # 항상 복원

사용:

class Printer:
    def __init__(self):
        self.verbose = False
    def show(self, msg):
        if self.verbose:
            print(msg.upper())
        else:
            print(msg)

p = Printer()

with temporarily(p, "verbose", True):
    p.show("inside")   # INSIDE 출력

p.show("outside")      # outside 출력

여러 속성을 한 번에 바꾸기

뒤집을 설정이 하나 이상이라면 매핑을 받아 순회하면 된다:

@contextmanager
def temporarily_many(obj, **changes):
    """obj의 여러 속성을 일시적으로 설정한다."""
    originals = {name: getattr(obj, name) for name in changes}
    try:
        for name, value in changes.items():
            setattr(obj, name, value)
        yield obj
    finally:
        for name, value in originals.items():
            setattr(obj, name, value)

with temporarily_many(p, verbose=True, prefix="[DBG] "):
    ...

(객체에 속성이 아주 많다면, 헬퍼를 하나 더 쓰는 대신 contextlib.ExitStack으로 단일 목적 context manager 여러 개를 조합할 수 있다.)

메서드에 데코레이터를 달아 손쉽게 되돌리기

속성 조작이 객체 자신의 메서드 안쪽에 속한다면, 내부 본문을 감싸면 된다:

class Printer:
    ...

    @contextmanager
    def loud_mode(self):
        yield from temporarily(self, "verbose", True)

    def log_section(self, messages):
        with self.loud_mode():
            for m in messages:
                self.show(m)

이제 log_section을 호출할 때마다, 출력 중 하나가 예외를 던지더라도 verbose가 이전 상태로 돌아오는 것이 보장된다.

클래스 기반 context manager (__enter__ / __exit__)

상태가 깊게 얽힌 경우에는 클래스 전체를 쓰고 싶을 수 있다:

class SwapAttr:
    def __init__(self, obj, attr, new_value):
        self.obj, self.attr, self.new = obj, attr, new_value
    def __enter__(self):
        self.old = getattr(self.obj, self.attr)
        setattr(self.obj, self.attr, self.new)
        return self.obj
    def __exit__(self, exc_type, exc, tb):
        setattr(self.obj, self.attr, self.old)
        # False를 반환하면 예외가 그대로 전파된다

클래스 형태는 부가적인 기록 데이터를 저장하거나, 같은 매니저 인스턴스를 여러 with 블록에서 공유할 수 있게 해준다.