Making a Python Context Manager
I’ve decided to get myself up-to-date with some Python features I rarely extend but often use.
One of them is the mighty context manager.
Context managers allow cleanup of resources that need some kind of closure after being used like files and threads(for those familiar with Java something similar is AutoCloseable interface).
More formally PEP 343 defines them as objects that implement a “context management protocol” which
is composed of
As a developer the best way to wrap my head about something is simply to go and implement it.
There are several options for implementing a context manager. As an OOP fan I went the class way. A very minimal context manager with no error handling looks as follows:
import sys from io import IOBase from typing import Tuple, List class SimpleContextManager(object): def __init__(self, file_name: str, method: str): self.file_obj = open(file_name, method) def __enter__(self) -> IOBase: return self.file_obj def __exit__(self, type, value, traceback): self.file_obj.close() return True with SimpleContextManager(sys.argv, "r") as file: for line in file: print(line.strip())
__enter__ is executed on
with line the result gets written into
Any exception that gets thrown before
__enter__ returns is up to the caller to solve.
If an exception gets thrown by
file line generator the exception is eaten inside the context manager.
This is why
__exit__ has type, value, traceback arguments (think exc_type, exc_value, exc_traceback).
If return value of
True, an exception, thrown by file generator, would be eaten up.
False exception propagates out of
with body and caller needs to handler the error.
Now that we know a bit about making a simple context manager how about we create something a bit more useful.
I decided to make a basic CSV parser that can skip header without additional if statements.
Instead of returning file object from context manager I can return a custom iterator
that skips the header. Since
__enter__ can return any object I can also return the header as a tuple.
I’ve first implemented an iterator over lines of file that skips header if configured:
class LineIterator(object): def __init__(self, file_object: IOBase, has_header: bool = True, separator: str = ";"): self.file_object = file_object self.has_header = has_header self.separator = separator self.first_line = True def __next__(self) -> List[str]: line = next(self.file_object) if self.first_line: next(self.file_object) self.first_line = False if line: return line.strip().split(self.separator) def __iter__(self) -> LineIterator: return self
If the iterator receives
has_header as True it will skip the header line.
The final context manager then looks like the following:
class CsvFile(object): def __init__(self, file_name: str, method: str, has_header: bool = True, separator: str = ";"): self.has_header = has_header self.separator = separator self.file_name = file_name self.method = method self.file_obj = None def __enter__(self) -> Tuple[List[str], LineIterator]: if not self.file_obj: try: self.file_obj = open(self.file_name, self.method) except Exception as e: # if this is rethrown caller needs to handle the exception logging.error("Failed opening file") header =  if self.has_header: first_line = next(self.file_obj, None) if first_line: header = first_line.split(self.separator) lines = LineIterator(self.file_obj, self.has_header, self.separator) return header, lines def __exit__(self, type, value, traceback): if value: logging.error("Failed processing CSV", exc_info=(type, value, traceback)) if self.file_obj: self.file_obj.close() return True
has_header is passed in as True this context manager will on
__enter__ return a tuple of headers and
a line iterator which will only iterate through non-header lines.
CsvFile can be used in a
with statement as follows:
with CsvFile("test.csv", "r", has_header=True, separator=';') as (csv_header, csv_file): for line in csv_file: print(line)
Note that if the file doesn’t exist
CsvFile won’t throw. An exception will be generated inside the
and logged at
logging.error("Failed processing CSV", exc_info=(type, value, traceback)).
And that wraps up context managers. Full code samples are available on my GitHub.
Note: There are plenty of good CSV libraries out there so please avoid making your own.