Context Managers and the with Statement

2019-07-29
#python
programming

Context Managers and the with Statement

  • with statement와 context manager에 대해서 알아봅니다. 사실 쓰는 방법만 알고, 안에서는 어떻게 작동되는지 제대로 공부해 보지 않았었는데, 이번에 recommendation system을 구현해보면서 with statement를 썼기 때문에, 공부해 보았습니다. With statement를 제대로 이해하기 위해서는 decorator와 generator에 대한 사전 지식이 있어야 합니다.

With Statement?

with statement는 리소스 관리 패턴 기능들을 추상화 시켜서, 따로 떼어내어 재사용 할 수 있도록 함으로서 리소스 관리를 간단하게 만들어 줍니다. 제가 with statement를 사용한 예시는 아래와 같습니다. for문 안에서 여러 파일을 열고 닫아야 했는데, 간단하게 with statement로 with open(path) as f 사용하여 파일을 열고 닫는 과정을 실행했습니다. 밑의 코드는 복잡하니 간단한 예시를 들어 봅니다. 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# class안에 있던 __iter__ function 입니다 
    def __iter__(self):
        n_files = len(self.paths)
        form = 'iter={}, file={}/{}, sents={}'
        for i_p, path in enumerate(self.paths):
            with open(path) as f:
                for i_doc, doc in enumerate(f):
                    if self.verbose and (i_doc % 10000 == 0):
                        message = form.format(self.n_iter, i_p, n_files, i_doc)
                        print("\r{}".format(message), end='')
                    try:
                        yield self._clean(doc)
                    
                    except Exception as e:
                        print('\rException file={}, sents={}'.format(i_p, i_doc))
                        print(e)
                        break
        if self.verbose:
            message = form.format(self.n_iter, i_p, n_files, i_doc)
        self.n_iter += 1

아래는 with statement로 data_science.txt파일을 만들고, 문장 ‘what have you learned today?’을 적었습니다. 간단하고 깔끔합니다. with을 사용하여 파일을 여는 것이 좋고, 추천되는 이유는 open안에 있는 파일이 자동적으로 with statement가 끝나면 종료되기 때문입니다. 

1
2
3
with open('data_science.txt', 'w') as f:
    f.write('what have you learned today? ')
    f.write('the with statement!')

만약 with statement를 사용하지 않고, 파일을 열고 닫는다면 아래와 같이 해야 하는데, 보다시피 5줄이나 됩니다. 

1
2
3
4
5
6
f = open('data_science.txt', 'w')
try:
    f.write('what have you learned today? ')
    f.write('the with statement!')
finally:
    f.close()

위에서 try와 finally statement는 아주 중요합니다. 그 이유는 f.write 에서 예외/에러가 생긴다면 아래와 같이 쓰는 것만으로는 완전히 파일을 닫혔다는 것을 보장하지 못하기 때문입니다.

1
2
3
4
5
# 아래는 파일이 완전히 닫혔다는 것을 보장할 수 없습니다!
f = open('data_science.txt', 'w')
f.write('what have you learned today? ')
f.write('the with statement!')
f.close()

Make My Own Context Managers

위와 같은 기능을 context managers를 사용하여 직접 만든 클래스나 함수에 입힐 수 있습니다. Context Manager란 with statement를 직접 만든 객체가 지원하기 위해 따라야 하는 약속또는 인터페이스입니다. 간단하게 ‘enter‘ 와 ‘exit‘ 함수를 클래스에 넣어줌으로서 직접 만든 객체가 context manager로 작동 하도록 하게 할 수 있습니다. 아래와 같이 MyResourceManagement 클래스를 만들어 open() context manager가 동작하게 해 봅니다. 주의 해야 할 점은 exit 함수에는 3개의 argument가 필요합니다 : type, value, traceback.

1
2
3
4
5
6
7
8
9
10
11
12
class MyResourceManagement:
    def __init__(self, file_name):
        self.file_name = file_name
        
    def __enter__(self):
        self.file = open(self.file_name, 'w')
        return self.file
    
    # exit함수에는 3개의 argument가 필요합니다 (type, value, traceback)
    def __exit__(self, ex_type, ex_val, ex_tb):
        if self.file:
            self.file.close()
1
2
3
with MyResourceManagement('data_science.txt') as f:
    f.write('what have you learned today? ')
    f.write('the with statement!')

위와 같이 MyResourceManagement 클래스는 context manager 인터페이스를 따랐고, with statement가 지원이 되는 것을 확인할 수 있습니다. Python은 with statement 안에 있는 내용(context) 으로 들어갈 때 (enter), __enter__를 호출하여 리소스를 얻습니다. 그리고 context를 떠날 때, __exit__함수를 호출하여 리소스가 해제 됩니다.

위 처럼 클래스 기반의 context manager를 사용하는 것만이 방법은 아닙니다. Python의 contextlib 모듈의 contextmanager decorator를 사용하면 with statement를 자동적으로 지원하는 generator 기반의 factory function (closure)를 정의할 수 있습니다. 위의 MyResourceManagement context manager를 decorator기반으로 만들어 봅니다. 

1
2
3
4
5
6
7
8
9
from contextlib import contextmanager

@contextmanager
def resource_management(file_name):
    try:
        f = open(file_name, 'w')
        yield f
    finally:
        f.close()
1
2
3
with resource_management('data_science.txt') as f:
    f.write('what have you learned today? ')
    f.write('the with statement!')

위의 경우, resource_management() 함수는 generator이며, 첫번째로 리소스를 얻습니다. 그 후, 실행을 잠시 멈추고, 리소스를 yield 함으로서 호출자가 리소스를 사용할 수 있게 합니다. 호출자가 with안에 있는 내용을 떠나면, generator는 계속 실행이 되어 남아있는 step들이 실행되게 되고, 리소스는 해제가 됩니다. 클래스 기반 방법과 generator 기반 방법의 context manager는 근본적으로는 같기에 선호하는 방법을 사용하면 됩니다.

Implementation Example : Timer!

위에서 배운 내용으로 code block의 실행 시간을 측정하는 context manager를 작성해 봅니다. 두가지 방법, class-based와 generator-based로 만들어 봅니다. 

1
2
3
4
5
6
7
8
9
10
11
12
13
import time

class Timer:    
    def __init__(self, task):
        self.task = task
        
    def __enter__(self):
        self.start = time.time()
        
    def __exit__(self, ex_type, ex_val, ex_tb):
        self.end = time.time()
        self.interval = self.end - self.start
        print(f"{self.task}: {self.end - self.start} seconds")
1
2
with Timer("List Comprehension"):
    s = [x*x for x in range(10000000)]
List Comprehension: 0.6686520576477051 seconds

1
2
3
4
5
6
7
8
@contextmanager
def timing(task):
    start = time.time()
    # yield를 하면 with 안에 있는 내용이 실행됩니다
    yield 
    end = time.time() 
    interval = end - start
    print(f"{task}: {interval} seconds")
1
2
with timing("List Comprehension"):
    s = [x*x for x in range(10000000)]
List Comprehension: 0.6721289157867432 seconds

Summary

  • with statement는 try/finally를 사용한 exception 처리를 캡슐화 하여 context managers란 이름으로 간단하게 만들어 줍니다.
  • 가장 많이 with statement가 쓰이는 경우는 리소스 관리 (획득과 해제)할 때 입니다. 리소스는 with statement와 함께 얻어지고, 자동적으로 with statement를 떠날 때 해제 됩니다.
  • with statement를 효율적으로 사용하면 리소스 누수를 막을 수 있으며, 코드가 쉽게 읽히도록 해줍니다