8.3 객체의 컨텍스트 관리 프로토콜 지원

문제

Context를 관리하는 프로토콜을 만들어야 할 때 적절하게 클래스를 구현하는 방법

해결

with 구문으로 __enter__() __exit__() 메소드를 구현하자

코드

직접 구현

__enter__() : 시작 시 동작 정의

__exit__() : 종료 시 동작 정의

# Connection Class 구현 
from socket import socket, AF_INET, SOCK_STREAM

class LazyConnection:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.sock = None

    def __enter__(self):
        if self.sock is None:
            self.sock = socket(AF_INET, SOCK_STREAM)
            self.sock.connect((self.host, self.port))
        return self.sock

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.sock:
            self.sock.close()
            self.sock = None
            
# Example 1 : Using with statement
host = 'python.org'
port = 80
with LazyConnection(host, port) as sock:
    sock.sendall(b'Hello, world!')
    data = sock.recv(1024)


# Example 2 : When not using with statement
conn = LazyConnection(host, port)
sock = conn.__enter__()
sock.sendall(b'Hello, world!')
data = sock.recv(1024)
conn.__exit__(None, None, None)

 

ContextManager로 구현

contextmanager 함수 역시 내부에는 __enter__()와 __exit__()는 구현되어있다. 

다만 편하게 쓸 수 있는가 여부에 관한 차이이다. 

# Contextmanager로 구현

from contextlib import contextmanager
from socket import socket, AF_INET, SOCK_STREAM

@contextmanager
def lazy_connection(host, port):
    sock = socket(AF_INET, SOCK_STREAM)
    try:
        sock.connect((host, port))
        yield sock
    finally:
        sock.close()

# Example 1 : Using with statement
host = 'python.org'
port = 80

with lazy_connection(host, port) as sock:
    sock.sendall(b'Hello, world!')
    data = sock.recv(1024)

# Example 2 : When not using with statement
conn = lazy_connection(host, port)
sock = conn.__enter__()
sock.sendall(b'Hello, world!')
data = sock.recv(1024)
conn.__exit__(None, None, None)

 

뭐가 좋은가?

한 사이클의 context를 열고 닫기를 간편하게 구현할 수 있다.

시작과 끝을 정의할 수 있는 모든 동작에 적용 가능하다. 

db connection이라든가, 파일의 생성 및 수정, 모델의 학습 등

 

출처: Beazley, D., & Jones, B. K. (2013). *Python cookbook* (3rd ed., pp. 246–248). O’Reilly Media