Python: with문 만들어 사용하기

with

contextlib — Utilities for with-statement contexts — Python 3.9.1 documentation

 

1. with문

파이썬에서는 with문을 이용해서 외부 리소스를 간편히 이용할 수 있게 해 준다.

 

with문 없이 파일 열기 (항상 close)

file = open("my.txt", "r+")

firstLine = file.readline()

file.close()

 

with문

with open("my.txt", "r+") as file:
    firstLine = file.readLine()
    
print(firstLine)

 

그럼 나만의 with문을 만들려면 어떻게 해야 할까?

 

2. with문 작동 구조

context manager를 이용하는 법은 보통 아래 방식들을 포함한다.

open - close, lock - release, enter - exit, start - stop...

 

with문은 class를 만들고, __enter__와 __exit__부분을 만들어서 원하는 작업을 하면 된다.

class MyWith:
    def create_save(**args):  # 만들고 저장
        ret = MyWith(**args)
        ret.save()
        return ret
        
    def __enter__(self):  # 클래스 진입
        return self
        
    def __exit__(self, exceptions):  # 클래스 호출 종료
        if not exceptions:
            self.save()

# with 이용
with MyWith(my=-1) as obj:
    # obj 수정

 

또는 더 편하게, contextlib을 이용해서 함수로만 만들 수 있다.

import contextlib

@contextlib.contextmanager
def hello():
    print("Hello")
    yield  # yield 아래에 있는 것들을 호출 종료시 불려옴
    print("World!")
    
with hello():
    print("me")
    
""" 결과
Hello
me
World!
"""

 

3. 커스텀 with문 이용

 

함수 실행 시간

데코레이터와 비슷하게 작동한다고 보면 된다.

import time

@contextlib.contextmanager
def timer():
    # 타이머 시작
    start = time.time()
    yield
    # 타이머 종료
    end = time.time()
    print(f'실행 시간: {end - start} 초')

with timer():
    for i in range(10):
        time.sleep(0.1)
print('완료!')

 

Apache Kafka Producer

flush 수동으로 할 필요 없이 알아서 flush 시킨다.

from contextlib import contextmanager


@contextmanager
def prod(settings):
    p = Producer(settings)
    yield p
    p.flush(100)


settings = {"bootstrap.servers": Config.MY_SERVER}

with prod(settings) as p:
    p.produce(Config.TOPIC_ID, key="key_1", value="Hello")

 

JSON open

json파일을 자주 부르게 된다면 파일 읽어서 바로 json으로 return

from contextlib import contextmanager
from pathlib import Path


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
JSON_PATH = os.path.join(BASE_DIR, "test.json")


@contextmanager
def jsonify(mode):
    f = open(JSON_PATH, mode=mode)
    read = json.load(f)
    yield read
    f.close()
    
with jsonify("r+") as f:
    data = f

 

폴더 안 파일 목록

import os


@contextmanager
def in_dir(path):
    # 현재 작업 경로
    cwd = os.getcwd()
    
    # cd path
    os.chdir(path)
    try:
        yield
    finally:
        # 무조건 다시 홈으로 돌아간다
        os.chdir(cwd)


with in_dir("python"):
    print(os.listdir())

# ['README.md', 'src', 'tests']

 

Lock

import threading

class Lock:
    def __init__(self, lock):
        self.lock = lock

    def __enter__(self):
        self.lock.acquire()

    def __exit__(self, exc_type, exc_value, traceback):
        self.lock.release()

# Example usage
lock = threading.Lock()
with Lock(lock):
    # Critical section of code
    pass

 

Socket

import socket

class NetworkConnection:
    def __init__(self, host, port):
        self.host = host
        self.port = port

    def __enter__(self):
        self.sock = socket.create_connection((self.host, self.port))
        return self.sock

    def __exit__(self, exc_type, exc_value, traceback):
        self.sock.close()

# Example usage
with NetworkConnection('www.example.com', 80) as conn:
    conn.sendall(b'GET / HTTP/1.0\r\n\r\n')
    response = conn.recv(1024)
    # Do something with the response
    pass