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 __enter__() and __exit__() methods.

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[1], "r") as file:
    for line in file:
        print(line.strip())

When __enter__ is executed on with line the result gets written into file variable.

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 __exit__ is True, an exception, thrown by file generator, would be eaten up. If 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

If 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 LineIterator 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.