파이썬 기초 - Context Manager

Context Manager

Python의 Context Manager는 특정 작업을 수행할 때, 리소스를 할당하고, 사용한 후 자동으로 해제해주는 객체이다.

일반적으로 with문과 함께 사용되며, 리소스 관리, 로깅, 트랜잭션 관리 등에 활용된다.

기본 구문

파이썬의 with 문을 사용하면 파일 입출력, 데이터베이스 연결 등 리소스를 다룰 때 자동으로 정리(clean-up) 작업을 수행할 수 있다.

with 컨텍스트_매니저 as 변수:
    #TODO

파일 다루기

with 블록이 끝나면 Python이 자동으로 파일을 닫는다. 

# Context Manager를 사용하는 방법
with open("example.txt", "r") as file:
    content = file.read()
# 파일이 자동으로 닫힘

 

Context Manager 없이 파일 다루기

# Context manager를 사용하지 않으면, close를 하지 않을 가능성이 생김
file = open("example.txt", "r")
content = file.read()
file.close()

 

Context Manager의 구동 원리

Context Manager는 __enter__와 __exit__이라는 두 가지 메서드를 통해 동작한다.

class MyContextManager:
    def __enter__(self):
        print("Entering the context")
        return self  # 반환된 객체가 with 블록에서 사용됨
    
    def __exit__(self, exc_type, exc_value, traceback):
        print("Exiting the context")
        if exc_type:
            print(f"오류 발생: {exc_value}")
            import traceback as tb
            tb.print_tb(traceback)  # 트레이스백 출력
            # True를 반환하면 예외가 무시됨, False/None이면 예외가 전파됨
            return False

with MyContextManager() as cm:
    print("Inside the context")

위 코드 실행 결과:

Entering the context
Inside the context
Exiting the context

작동 원리

1. with 문이 실행되면 MyContextManager 인스턴스의 __enter__ 메서드가 호출된다.
2. __enter__ 메서드의 반환값이 as 뒤의 변수(여기서는 cm)에 할당된다.
3. with 블록 내부의 코드가 실행된다.
4. 블록이 끝나거나 예외가 발생하면 __exit__ 메서드가 호출된다.
5. __exit__ 메서드는 예외 정보를 매개변수로 받는다.
   • exc_type: 예외 유형 (예외가 없으면 None)
   • exc_value: 예외 값 (예외가 없으면 None)
   • traceback: 예외의 트레이스백 객체 (예외가 없으면 None)

예제: 파일 핸들러 Context Manager

class CustomFileHandler:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.file = None

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_value, traceback):
        if self.file:
            self.file.close()
        if exc_type:
            print(f"오류 발생: {exc_value}")
            import traceback as tb
            tb.print_tb(traceback)
            return False

with CustomFileHandler("sample.txt", "w") as f:
    f.write("Custom context manager example!")

contextlib을 활용한 Context Manager 생성

contextlib 모듈의 contextmanager 데코레이터를 사용하면 클래스를 만들지 않고도 컨텍스트 관리자를 쉽게 정의할 수 있다.

contextlib.contextmanager 데코레이터 사용

from contextlib import contextmanager

@contextmanager
def custom_context(name):
    print(f"{name} Entering the context")
    try:
        yield name
    finally:
        print(f"{name} Exiting the context")

with custom_context("kyeongseo") as res:
    print(f"I am {res}")

실행 결과

kyeongseo Entering the context
I am kyeongseo
kyeongseo Exiting the context

closing() 사용

closing은 __enter__와 __exit__ 메서드가 없는 객체를 컨텍스트 매니저로 변환해주는 래퍼로,  이를 통해 with 문과 함께 사용할 수 있게 되어, 객체가 사용 후 자동으로 닫히도록 한다.

from contextlib import closing
import urllib.request

# closing을 사용하지 않는 경우
url = urllib.request.urlopen('https://example.com')
try:
    content = url.read()
finally:
    url.close()

# closing을 사용하는 경우
with closing(urllib.request.urlopen('https://example.com')) as url:
    content = url.read()

suppress() 사용

특정 예외를 무시하고 싶을 때 사용한다.

from contextlib import suppress
import os

# FileNotFoundError 예외를 무시하고 계속 실행
with suppress(FileNotFoundError):
    os.remove("존재하지_않는_파일.txt")
    
print("파일이 없어도 프로그램은 계속 실행된다.")

redirect_stdout() 사용

이 함수는 표준 출력(stdout)을 일시적으로 다른 파일 객체로 리디렉션할 수 있게 해준다.

from contextlib import redirect_stdout
import io

# 문자열 버퍼로 stdout 리디렉션하기
f = io.StringIO()
with redirect_stdout(f):
    print("Hello, World!")
    print("redirected")
    
# 캡처된 출력을 가져오기
output = f.getvalue()
print(f"Captured output: {output}")

# 파일로 stdout 리디렉션하기
with open('output.txt', 'w', encoding='utf-8') as f:
    with redirect_stdout(f):
        print("콘솔에 출력되는 대신 파일에 작성됩니다.")

Context Manager 활용 예제

데이터베이스 연결 관리

데이터베이스 연결도 Context Manager를 사용하면 안전하게 관리할 수 있다..

import sqlite3
from contextlib import contextmanager

@contextmanager
def db_connection(db_name):
    conn = sqlite3.connect(db_name)
    try:
        yield conn
    finally:
        conn.close()

with db_connection("test.db") as conn:
    cursor = conn.cursor()
    cursor.execute("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
    conn.commit()

시간 측정하기

특정 코드 블록의 실행 시간을 측정하는 Context Manager

import time
from contextlib import contextmanager

@contextmanager
def timer(작업명="작업"):
    시작 = time.time()
    try:
        yield
    finally:
        종료 = time.time()
        print(f"{작업명} 실행 시간: {종료 - 시작:.4f}초")

# 사용 예
with timer("반복문"):
    # 시간을 측정할 코드
    total = 0
    for i in range(1000000):
        total += i
    print(f"계산 결과: {total}")

임시 설정 변경하기

작업 디렉토리를 임시로 변경하는 Context Manager

import os
from contextlib import contextmanager

@contextmanager
def temporary_working_directory(path):
    """임시로 작업 디렉토리를 변경하는 컨텍스트 매니저"""
    origin_path = os.getcwd()  # 현재 디렉토리 저장
    os.chdir(path)  # 새 디렉토리로 변경
    try:
        yield
    finally:
        os.chdir(origin_path)  # 원래 디렉토리로 복구

# 사용 예
print(f"현재 디렉토리: {os.getcwd()}")

with temporary_working_directory('C:/Users'):
    print(f"임시 디렉토리: {os.getcwd()}")

print(f"다시 원래 디렉토리: {os.getcwd()}")

로그 레벨 임시 변경

로깅 레벨을 임시로 변경하는 Context Manager

import logging
from contextlib import contextmanager

@contextmanager
def temporary_log_level(logger, level):
    """로거의 레벨을 임시로 변경하는 컨텍스트 매니저"""
    origin_level = logger.level
    logger.setLevel(level)
    try:
        yield
    finally:
        logger.setLevel(origin_level)

# 로거 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()

# 기본 레벨(INFO)에서는 DEBUG 메시지가 출력되지 않음
logger.debug("이 메시지는 보이지 않습니다.")
logger.info("이 메시지는 보입니다.")

# 임시로 로그 레벨 변경
with temporary_log_level(logger, logging.DEBUG):
    logger.debug("이제 DEBUG 메시지도 보입니다")
    logger.info("INFO 메시지도 계속 보입니다.")

# 다시 원래 레벨로 돌아감
logger.debug("이 메시지는 다시 보이지 않습니다.")

여러 Context Manager 동시에 사용하기 

Context Manager는 여러 개를 한 번에 사용할 수 있다.

with open("input.txt", "r") as f1, open("output.txt", "w") as f2:
    data = f1.read()
    f2.write(data.upper())  # 대문자로 변환하여 저장