파이썬 - 재사용 가능한 TRY, CATCH 동작을 원한다면 contextlib과 with 문을 사용하라

이동욱

2021/09/24

Categories: 파이썬

BATTER WAY 18. 재사용 가능한 try/finally 동작을 원한다면 contextlib과 with 문을 사용하라

파이썬의 with 문은 코드가 특별한 컨텍스트 안에서 실행되는 경우를 표현한다.
예를 들어, 상호 배제 락(뮤텍스)를 with 문 안에서 사용하면 락을 소유 했을 때만, 코드 블록이 실행된다는 것을 의미한다.

from threading import Lock

Lock = Lock()
with lock:
  # 어떤 불변 조건을 유지하면서 작업을 수행한다.
  ...

Lock 클래스가 with 문을 적절히 활성화해주므로 위의 예제는 아래의 코드와 동등하다.

lock.acquire()
try:
  # 어떤 불변 조건을 유지하면서 작업을 수행한다.
  ...
finally:
  lock.release()

위와 같은 경우엔느 with 문을 사용하는 것이 훨씬 더 낫다. 이유는 try/finally 구조를 반복적으로 사용할 필요가 없고, acquire에 대응하는 release 를 실수로 빠뜨리는 경우를 방지할 수 있기 때문이다.
contextlib 내장 모듈을 사용하면 만든 객체나 함수를 with 문에서 쉽게 사용할 수 있다.
contextlib 모듈은 with 문에 쓸 수 있는 함수를 간단히 만들 수 있는 contextmanager 데코레이터를 제공한다.
이 데코레이터를 사용하는 방법이 __enter__, __exit__ 특별 메서드를 사용해서 새롭게 클래스를 정의하는 방법보다 훨씬 쉽다.

사용 예제

예를 들어서, 어떤 코드 영역에서 디버깅 관련 로그를 더 많이 남기고 싶다고 가정을 할 때, 아래 코드는 디버깅 로그를 남기는 함수이다.

def my_function():
   logging.debug('debug data')
   logging.error('error log')
   logging.debug('additional debug data')

프로그램의 기본 로그 수준은 WARNING 이므로, 이 함수를 실행하면 오류 메시지만 화면에 출력된다.
따라서 컨텍스트 매니저를 정의하여 이 함수의 로그 수준을 일시적으로 높혀보겠다.
이 헬퍼 함수는 with 블록을 실행하기 직전에 로그 수준을 높이고, 블록을 실행한 후에 로그 수준을 이전 수준으로 복구 시키는 코드이다.

@contextmanager
def debug_logging(level):
    logger = logging.getLogger()
    old_level = logger.getEffectiveLevel()
    logger.setLevel(level)
    try:
        yield
    finally:
        logger.setLevel(old_level)

yield 식은 with 블록의 내용이 실행되는 부분을 지정한다.
with 블록 안에서 발생한 예외는 어떤 것이든 yield 식에 의해서ㄷ 다시 발생되기 때문에 이 예외를 헬퍼 함수 안에서 잡아 낼 수 있다.

with debug_logging(logging.DEBUG):
    my_function()

위의 사진과 같이 모든 로그가 출력된 것을 확인할 수 있다.

with와 대상 변수 함께 사용하기

with 문에 전달된 컨텍스트 매니저가 객체를 반환할 수 도 있다. 이렇게 반환된 객체는 with 복합문의 일부로 지정된 지역 변수에 대입된다.
이를 통해서 with 블록 안에서 실행되는 코드가 직접 컨텍스트 객체와 상화 작용할 수 있다.
예를 들어서 파일을 작성하고 이 파일이 제대로 닫혔는지 확인하고 싶다고 하자.
with 문에 open을 전달하면 이렇게 할 수 있다. open은 with 문에서 as를 통해 대상으로 지정된 변수에게 파일 핸들을 제공하고, with 블록에서 나갈 때 이 핸들을 닫는다.

with open('my_output.txt', 'w') as handle:
  handle.write('data')

이러한 접근 방법은 파일 핸들을 매번 수동으로 열고 닫는 것보다 파이썬 다운 방식이며, 이 방식을 사용하면 코드 실행이 with 문을 벗어날 때 결국에는 파일이 닫힌다고 확신할 수 있다.
그리고 코드에서 문제가 될 수 있는 부분을 강조함으로써 파일 핸들이 열린 채로 실행되는 코드의 양을 줄이도록 한다. 일반적으로 파일 핸들이 열려 있는 부분을 줄이면 좋다.

정리

with 문을 사용하면 try/finally 블록을 통해서 사용해야하는 로직을 재활용하면서, 가독성을 높힐 수 있다.
contextlib 내장 모듈이 제공하는 contextmanager 데코레이터를 사용하면 여러분이 만든 함수를 with 문에 사용할 수 있다.
컨텍스트 매니저가 yield 하는 값은 with 문의 as 부분에 전달된다. 이를 특별한 컨텍스트 내부에서 실행되는 코드 안에서 직접 그 컨텍스트에 접근할 수 있다.

참고 문헌

파이썬 코딩의 기술

>> Home