Note: phiên bản Tiếng Việt của bài này ở link dưới.

https://duongnt.com/context-manager-vie

Context manager is an effective way to manage resources in Python. It helps abstract away the resource handling logic. But there is no rule that says we can only use it for one purpose. With some creativity, we can use context manager to measure method runtime, format output,…

Primer on context manager

Manually manage resource

Let’s say we have a class to manage Internet connections. It has methods to open and close a connection.

class Connection:
    def __init__(self):
        self._connection = None

    def open(self):
        # code to open connection

    def close(self):
        # code to close connection

    def send(self, data):
        # code to send data

The traditional way to make sure that the connection is always closed after use is to use try-finally.

conn = Connection()
try:
    conn.open()
    conn.send(data)
finally:
    conn.close()

Make the Connection class into a context manager

Instead of writing try-finally ourselves, we can turn the Connection class into a context manager. All we need to do is move the code to open the connection into an __enter__ method; and the code to close the connection into an __exit__ method.

class Connection:
    # ... omitted

    def __enter__(self):
        self.open()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._close()

Then we can use it with the with statement. The close method will always be called, even if an error occurs.

with Connection() as conn:
    conn.send(data)

Use context manager to measure the runtime of methods

A class to measure methods’ runtime

As we can see in the previous section, the statements inside __enter__ will be executed when we enter the with block. And the statements inside __exit__ will be executed when we exit it. We can take advantage of this to measure the runtime of a method. Below is one such context manager.

class Measure:
    def __init__(self):
        self.start = None
        self.end = None

    def __enter__(self):
        self.start = datetime.datetime.now()
        # We don't return anything here because we don't use the class instance anyway

    def __exit__(self, exc_type, exc_value, traceback):
        self.end = datetime.datetime.now()
        diff = (self.end - self.start)
        print(f'Run time: {diff.total_seconds()}s')

Let’s use it to measure the run time of a test method.

def test_target():
    for i in range(10000000):
        j = i + 1

with Measure() as _:
    test_target()

On my machine, the code above prints this to the command line.

Run time: 1.133019s

Context manager without defining a new class

It’s actually possible to create a context manager without defining a new class. We can use the contextmanager decorator from contextlib instead. Below is the version that uses contextmanager.

@contextmanager
def measure_generator():
    try:
        start = datetime.datetime.now()
        yield
    finally:
        end = datetime.datetime.now()
        diff = (end - start)
        print(f'Run time: {diff.total_seconds()}s')

We can see that after recording the start time, we yield control back to the caller. Then when the caller leaves the with block, the remaining code inside the finally block will be executed. Normally, this is where we put the cleanup logic.

Using measure_generator is identical to using the Measure class.

with measure_generator() as _:
    test_target()

Beautify HTML tags with context manager

At first glance, formatting HTML tags seems totally unrelated to context manager. But what if we could write code that looked like this, with the indentation of with statements matching the indentation of HTML tags?

with HtmlTag('html') as _:
    with HtmlTag('head') as _:
        with HtmlTag('title') as title:
            title.print('Duong Blog')
        with HtmlTag('script') as script:
            script.print('https://example.com/script.js')
    with HtmlTag('body') as _:
        with HtmlTag('h1') as header:
            header.print('Awesome header')
        with HtmlTag('p') as section:
            section.print('Lorem ipsum dolor sit amet, consectetur adipiscing elit.')

Create the HtmlTag class

The HtmlTag has a class attribute to record the current tag depth. When a new object is created, it will store that depth in its own instance attribute.

class HtmlTag:
    INDENT = 2 # 2 spaces for each indentation level
    depth = 0 # the current indentation depth
    def __init__(self, tag):
        self.tag = tag
        self.depth = HtmlTag.depth

We can divide our problem into three parts. Writing the opening tag, writing the content, and writing the closing tag. Naturally, the opening tag will be written inside the __enter__ method. We will use self.depth to properly indent it. Also, we need to increase HtmlTag.depth (not self.depth) whenever we enter a new with block.

def __enter__(self):
    print(' ' * HtmlTag.INDENT * self.depth + f'<{self.tag}>')
    HtmlTag.depth += 1
    return self

Similarly, as we exit the with block, the __exit__ method will handle closing the tag and decreasing HtmlTag.depth.

def __exit__(self, exc_type, exc_value, traceback):
    print(' ' * HtmlTag.INDENT * self.depth + f'</{self.tag}>')
    HtmlTag.depth -= 1

Between those tags, we use self.depth and the print method to write the content. Remember that the content needs to be one level deeper than the tags.

def print(self, txt):
    print(' ' * HtmlTag.INDENT * (self.depth + 1) + txt)

You can find the complete code in this link.

Test our formatter

When running the script above, you should see the result below.

An exercise for the reader

Normally, the opening tag, content, and closing tag of title/script/h1/p are written on the same line. How can we modify HtmlTag to do that? The solution is to add a parameter to the constructor of HtmlTag to indicate if we want to write everything on the same line. Please try implementing it yourself before checking my solution in this link.

Your solution should output something like this.

Conclusion

Although the Measure and HtmlTag classes we implemented are not suitable for production, playing around with them is a good way to deepen our understanding of context manager. What about you? What unique use case can you come up with?