From e075567bacc55faa3e87acabb41977053a13d430 Mon Sep 17 00:00:00 2001 From: Kevin Phoenix Date: Fri, 19 Jan 2024 17:37:47 -0700 Subject: [PATCH] Add contextmanager support to IO class (#68) --- python/binharness/types/io.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/python/binharness/types/io.py b/python/binharness/types/io.py index 80dfa21..6d78dcb 100644 --- a/python/binharness/types/io.py +++ b/python/binharness/types/io.py @@ -1,10 +1,13 @@ """binharness.types.io - Type definition for file-like objects.""" from __future__ import annotations -from typing import AnyStr, Protocol +from typing import TYPE_CHECKING, AnyStr, ContextManager, Protocol +if TYPE_CHECKING: + from types import TracebackType -class IO(Protocol[AnyStr]): + +class IO(ContextManager["IO[AnyStr]"], Protocol[AnyStr]): """A file-like object.""" def close(self: IO[AnyStr]) -> None: @@ -46,3 +49,17 @@ def write(self: IO[AnyStr], s: AnyStr) -> int | None: def writelines(self: IO[AnyStr], lines: list[AnyStr]) -> None: """Write lines to the file.""" + + def __enter__(self: IO[AnyStr]) -> IO[AnyStr]: + """Enter the runtime context related to this object.""" + return self + + def __exit__( + self: IO[AnyStr], + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Exit the runtime context and close the file if it's open.""" + if not self.closed: + self.close()