왜 쓸까?

파이썬 코드로 파일 하나를 열어보고 싶다.
어떻게 할까?

file = open("asd.txt", "r")
lines = file.readlines()
...
file.close()

아주 일반적인 방법이다.
허나 파일 작업을 모두 끝낸 후 file.close() 를 명시해야 한다는 불편한 점이 있다.
이럴 때 사용할 수 있는 것이 python의 Context manager다.

위와 똑같은 코드를 context manager를 사용하여 표현한다면,

with open("asd.txt", "r") as file:
    lines = file.readlines()
    ...

끝이다. file.close()를 명시하지 않아도 되고, 코드의 어떤 구간에 파일 작업이 있는지 개행으로 표현되어 알아보기 매우 편하다.

이런 context manager는 많은 방면에서 사용되고 있다.
DB를 읽고 쓰는 작업에 원자성atomicity를 보장하거나,
위처럼 파일 시스템을 이용하는 작업에 사용될 수 있다.

직접 만들기

참고: Python3 contextlib

더 나아가, 우리가 직접 context manager를 만들 수 있다.
크게 두가지 방법이 있다.

.__enter__()와 .__exit__() 메서드를 구현하는 방법과
contextlib.contextmanager() 함수 데코레이터를 사용하는 방법이다.
둘을 같이 사용할 수도 있고, Async context manager를 만드는 방법도 있지만, 여기서는 다루지 않겠다. 후에 때가 된다면...

1. .__enter__() & .__exit__()

예시 코드를 먼저 써보겠다.

class ListManager:

    def __enter__(self):
        self.my_list = []
        return self.my_list

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"my list is {self.my_list}")

with ListManager() as l:
    l.append(1)
    l.append(2)
    l.append(3)

이 파이썬 코드를 실행하면 아래와 같은 출력이 나온다.

my list is [1, 2, 3]

__exit__() 메서드가 저렇게 많은 인자를 받는 이유는 여기를 보면 알 수 있다.
짧게 설명하자면, context manager에서 탈출 시 발생한 에러를 __exit__() 메서드에서 다루기 위함이다.

여기서 contextlib.AbstractContextManager를 ListManager 클래스에서 구현케 하면 __enter__() 메서드와 __exit__() 메서드 구현을 강제시킬 수 있다.

from contextlib import AbstractContextManager

class ListManager(AbstractContextManager):

    def __enter__(self):
        self.my_list = []
        return self.my_list

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"my list is {self.my_list}")

2. contextlib.contextmanager()

1번의 예시 코드와 비슷하게 작성해보겠다.

from contextlib import contextmanager

@contextmanager
def list_manager():
    my_list = []

    try:
        yield my_list
    finally:
        print(my_list)

with list_manager() as l:
    l.append(1)
    l.append(2)
    l.append(3)

1번과 똑같이 my list is [1, 2, 3]가 출력된다.

여담

개인적으로, 공식문서의 이 부분이 조금 헷갈렸다.

필자가 정리한 바로는, 어떤 context manager 객체에close() 메서드가 구현되어있지 않으면 해당 객체 내부에서 사용되는 리소스를 관리-닫거나close 풀어주기release-해줘야 한다.

위 캡쳐에 이어서 나오는 예시 코드는 아래와 같다.

from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwds):
    # Code to acquire resource, e.g.:
    resource = acquire_resource(*args, **kwds)
    try:
        yield resource
    finally:
        # Code to release resource, e.g.:
        release_resource(resource)

>>> with managed_resource(timeout=3600) as resource:
...     # Resource is released at the end of this block,
...     # even if code in the block raises an exception

리소스를 release 해주는 코드가 필수로 있어야 한다는 예시다.

실사용

Machine learning 훈련 상태를 기록하는 DB 테이블을 관리하고자 사용했다.

아래 코드는 django 프레임워크의 model 모듈을 사용한 예시다.

from contextlib import contextmanager 
from django.db import models

class TrainingTask(models.Model):
    ...
    
    @classmethod
    @contextmanager
    def enter_training_context(cls, data):
        task = cls.create_task(data)
        task.status = "STARTED"
            
        try:
            yield task
        except:
            task.status = "FAILED"
            raise
        else:
            task.status = "SUCCESS"
        finally:
            task.save()

모델 객체 사용은 아래와 같다.

def train(training_data):
    parsed_data = parse(training_data)
    ...
    
    with enter_training_context(parsed_data) as task:
        ...
        task.result = result
    
    return

결론

멋진 문법이다.
허나 멋지고 다양하게 쓰일 수 있는 만큼 어디서 어떻게 써야 좋을지 고민해야 한다.
예를 들어, 바로위에 쓴 실사용 예시 코드를 decorator를 사용해서 train() 함수를 감싸주는 방식은 어떨까?
혹은 contextlib.ContextDecorator를 구현하는 방식은 어떨까?

좋은 코드를 짜기위한 좋은 도구를 하나 더 알아간다고 생각하자.