[Programming/Python] with as에 사용할 수 있는 클래스 만들기

open으로 파일을 열 때 with as를 사용하여 파일 객체의 close를 자동으로 호출해주었다. 이런 방식으로 with as를 사용하려면 클래스에 __enter__와 __exit__ 메서드를 구현해주면 된다.

class 클래스이름:
    def __enter__(self):
        # 시작할 때 실행할 코드
        
    def __exit__(self, exe_type, exc_val, exc_tb):
        # 종료할 때 실행할 코드

with에 클래스의 인스턴스를 지정하고 as 뒤에 결과를 저장할 변수를 지정한다.

with 클래스() as 변수:
    # 코드

다음은 open('hello.tt', 'w') 처럼 동작하는 OpenHello 클래스이다.

class OpenHello:
    def __enter__(self):
        # 파일 객체를 속성에 저장
        self.file = open('hello.txt', 'w')
        # __enter__에서 값을 반환하면 as에 지정한 변수에 들어감
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        # __exit__에서 파일 객체 닫기
        self.file.close()

with OpenHello() as hello:
    hello.write('Hello, World!')

__enter__ 메서드에서 값을 반환하면 as에 지정한 변수에 들어간다. 여기서는 open으로 파일 객체를 만들어 반환했다. 이때 __exit__ 에서도 파일 객체를 사용할 수 있도록 속성에 저장해준다.

__exit__ 메서드는 with as를 완전히 벗어나면 호출된다. 따라서 여기서는 __exit__에서 파일 객체를 닫았다. 이런 방식으로 __enter__에서 객체를 생성하고 __exit__에서 정리 작업을 하면 된다.