Python Other

Context managers (with statement)

Why use context managers?

In certain operations you might want to ensure that when the operation is done there will be an opportunity to clean up after it. Even if decided to end the operation early or if there is an exception in the middle of the operation.

In the following pseudo-code example you can see that cleanup must be called both at the end and before the early-end, but that still leaves the bad-code that raises exception avoiding the cleanup. That forces us to wrap the whole section in a try-block.

def sample():
    start
    do
    do
    do
    do
    cleanup

What is we have some conditions for early termination?

def sample():
    start
    do
    do
    if we are done early:
        cleanup
        return # early-end
    do
    do
    cleanup

What if we might have an exception in the code?

def sample():
    start
    try:
        do
        do
        if we are done early:
            cleanup
            return early-end
        do
        bad-code    (raises exception)
        do
        cleanup
    finally:
        cleanup

It is a lot of unnecessary code duplication and we can easily forget to add it in every location where we early-end our code.

Using Context Manager

with cm_for_sample():
    start
    do
    do
    if we are done early:
        return early-end
    do
    bad-code    (raises exception)
    do

• cleanup happens automatically, it is defined inside the cm_for_sample

Context Manager examples

A few examples where context managers can be useful:

• Opening a file - close it once we are done with it so we don't leak file descriptors.

• Changing directory - change back when we are done.

• Create temporary directory - remove when we are done.

• Open connection to database - close connection.

• Open SSH connection - close connection.

• More information about context managers

cd in a function

• getcwd
• chdir

In this example we have a function in which we change to a directory and then when we are done we change back to the original directory. For this to work first we save the current working directory using the os.getcwd call. Unfortunatelly in the middle of the code there is a conditional call to return. If that condition is True we won't change back to the original directory. We could fix this by calling os.chdir(start_dir) just before calling return. However this would still not solve the problem if there is an exception in the function.

import sys
import os

def do_something(path):
    start_dir = os.getcwd()
    os.chdir(path)

    content = os.listdir()
    number = len(content)
    print(number)
    if number < 15:
        return

    os.chdir(start_dir)

def main():
    if len(sys.argv) != 2:
        exit(f"Usage: {sys.argv[0]} PATH")
    path = sys.argv[1]
    print(os.getcwd())
    do_something(path)
    print(os.getcwd())

main()

$ python no_context_cd.py /tmp/

/home/gabor/work/slides/python-programming/examples/advanced
19
/home/gabor/work/slides/python-programming/examples/advanced

$ python no_context_cd.py /opt/

/home/gabor/work/slides/python-programming/examples/advanced
9
/opt

• In the second example return was called and thus we stayed on the /opt directory.:w

open in function

This is not the recommended way to open a file, but this is how it was done before the introduction of the with context manager. Here we have the same issue. We have a conditional call to return where we forgot to close the file.

import sys
import re

def do_something(filename):
    fh = open(filename)

    while True:
        line = fh.readline()
        if line is None:
            break
        line = line.rstrip("\n")

        if re.search(r'\A\s*\Z', line):
            return
        print(line)

    fh.close()

def main():
    if len(sys.argv) != 2:
        exit(f"Usage: {sys.argv[0]} FILENAME")
    filename = sys.argv[1]
    do_something(filename)

main()

open in for loop

• stat
• os.stat

Calling write does not immediately write to disk. The Operating System provides buffering as an optimization to avoid frequent access to the disk. In this case it means the file has not been saved before we already check its size.

import os

for ix in range(10):
    filename = f'data{ix}.txt'
    fh = open(filename, 'w')
    fh.write('hello')
    if ix == 0:
        break
    fh.close()
stat = os.stat(filename)
print(stat.st_size)    # 0,   the file has not been saved yet

open in function using with

If we open the file in the recommended way using the with statement then we can be sure that the close method of the fh object will be called when we leave the context of the with statement.

import sys
import re

def do_something(filename):
    with open(filename) as fh:

        while True:
            line = fh.readline()
            if line is None:
                break
            line = line.rstrip("\n")

            if re.search(r'\A\s*\Z', line):
                return
            print(line)


def main():
    if len(sys.argv) != 2:
        exit(f"Usage: {sys.argv[0]} FILENAME")
    filename = sys.argv[1]
    do_something(filename)

main()

Plain context manager

from contextlib import contextmanager
import sys

param = ''
if len(sys.argv) == 2:
    #exit(f"Usage: {sys.argv[0]} []")
    param = sys.argv[1]

def code_with_context_manager():
    with my_plain_context():
        print("  In plain context")
        if param == "return":
            return
        if param == "die":
            raise Exception("we have a problem")
        print("  More work")


@contextmanager
def my_plain_context():
    print("setup context")
    try:
        yield
    except Exception as err:
        print(f"  We got an exception: {err}")
    print("cleanup context")

print("START")
code_with_context_manager()
print("END")

START
start context
  In plain context
  More work
end context
END

Param context manager

from contextlib import contextmanager

@contextmanager
def my_param_context(name):
   print(f"start {name}")
   yield
   print(f"end {name}")

with my_param_context("foo"):
   print("In param context")

start foo
In param context
end foo

Context manager that returns a value

from contextlib import contextmanager

import time
import random
import os
import shutil


@contextmanager
def my_tempdir():
    print("start return")
    tmpdir = '/tmp/' + str(time.time()) + str(random.random())
    os.mkdir(tmpdir)
    try:
        yield tmpdir
    finally:
        shutil.rmtree(tmpdir)
        print("end return")

import os
from my_tempdir import my_tempdir

with my_tempdir() as tmp_dir:
    print(f"In return context with {tmp_dir}")
    with open(tmp_dir + '/data.txt', 'w') as fh:
        fh.write("hello")
    print(os.listdir(tmp_dir))

print('')
print(tmp_dir)
print(os.path.exists(tmp_dir))

start return
In return context with /tmp/1578211890.49409370.6063140788762365
['data.txt']
end return

/tmp/1578211890.49409370.6063140788762365
False

Use my tempdir - return

import os
from my_tempdir import my_tempdir

def some_code():
    with my_tempdir() as tmp_dir:
        print(f"In return context with {tmp_dir}")
        with open(tmp_dir + '/data.txt', 'w') as fh:
            fh.write("hello")
        print(os.listdir(tmp_dir))
        return

    print('')
    print(tmp_dir)
    print(os.path.exists(tmp_dir))

some_code()

start return
In return context with /tmp/1578211902.3545020.7667694368935928
['data.txt']
end return

Use my tempdir - exception

import os
from my_tempdir import my_tempdir

with my_tempdir() as tmp_dir:
    print(f"In return context with {tmp_dir}")
    with open(tmp_dir + '/data.txt', 'w') as fh:
        fh.write("hello")
    print(os.listdir(tmp_dir))
    raise Exception('trouble')

print('')
print(tmp_dir)
print(os.path.exists(tmp_dir))

start return
In return context with /tmp/1578211921.12552210.9000097350821897
['data.txt']
end return
Traceback (most recent call last):
  File "use_my_tempdir_exception.py", line 9, in <module>
    raise Exception('trouble')
Exception: trouble

cwd context manager

import os
from contextlib import contextmanager

@contextmanager
def cwd(path):
    oldpwd = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(oldpwd)

import sys
import os
from mycwd import cwd

def do_something(path):
    with cwd(path):
        content = os.listdir()
        if len(content) < 10:
            return

def main():
    if len(sys.argv) != 2:
        exit(f"Usage: {sys.argv[0]} PATH")
    path = sys.argv[1]
    print(os.getcwd())
    do_something(path)
    print(os.getcwd())

main()

$ python context_cd.py /tmp
/home/gabor/work/slides/python/examples/context
/home/gabor/work/slides/python/examples/context

$ python context_cd.py /opt
/home/gabor/work/slides/python/examples/context
/home/gabor/work/slides/python/examples/context

tempdir context manager

• contextlib
• contextmanager
• tempfile
• mkdtemp

import os
from contextlib import contextmanager
import tempfile
import shutil

@contextmanager
def tmpdir():
    dd = tempfile.mkdtemp()
    try:
        yield dd
    finally:
        shutil.rmtree(dd)

from mytmpdir import tmpdir
import os

with tmpdir() as temp_dir:
    print(temp_dir)
    with open( os.path.join(temp_dir, 'some.txt'), 'w') as fh:
        fh.write("hello")
    print(os.path.exists(temp_dir))
    print(os.listdir(temp_dir))

print(os.path.exists(temp_dir))

/tmp/tmprpuywa3_
True
['some.txt']
False

Context manager with class

• enter
• exit

class MyCM:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f'__enter__ {self.name}')
        return self

    def __exit__(self, exception_type, exception, traceback):
        print(f'__exit__  {self.name}')

    def something(self):
        print(f'something {self.name}')

def main():
    with MyCM('Foo') as cm:
        print(cm.name)
        cm.something()
        #raise Exception('nono')
    print('in main - after')

main()
print('after main')

Context managers with class

• enter
• exit

Even if there was en exception in the middle of the process, the exit methods of each object will be called.

class MyCM:
    def __init__(self, n):
        self.name = n

    def __enter__(self):
        print('__enter__', self.name)

    def __exit__(self, exception_type, exception, traceback):
        print('__exit__ ', self.name)

    def something(self):
        print('something', self.name)

def main():
    a = MyCM('a')
    b = MyCM('b')
    with a, b:
        a.partner = b
        b.partner = a
        a.something()
        raise Exception('nono')
        b.something()
    print('in main - after')

main()
print('after main')

__enter__ a
__enter__ b
something a
__exit__  b
__exit__  a
Traceback (most recent call last):
  File "context-managers.py", line 27, in <module>
    main()
  File "context-managers.py", line 23, in main
    raise Exception('nono')
Exception: nono

Context manager: with for file

• with

import sys

if len(sys.argv) != 2:
    sys.stderr.write('Usage: {} FILENAME\n'.format(sys.argv[0]))
    exit()

file = sys.argv[1]
print(file)
with open(file) as f:
    for line in f:
        val = 30/int(line)

print('done')

With - context managers

• with

class WithClass:
    def __init__(self, name='default'):
        self.name = name

    def __enter__(self):
        print('entering the system')
        return self.name

    def __exit__(self, exc_type, exc_value, traceback):
        print('exiting the system')

    def __str__(self):
        return 'WithObject:'+self.name

x = WithClass()
with x as y:
    print(x,y)

Exercise: Context manager

Create a few CSV file likes these:

a11,a12
a21,a22

b13,b14
b23,b24

c15,c16
c25,c26

Merge them horizontally to get this:

a11,a12,b13,b14,c15,c16
a21,a22,b23,b24,c25,c26

• Do it without your own context manager
• Create a context manager called myopen that accepts N filenames. It opens the first one to write and the other N-1 to read

with myopen(outfile, infile1, infile2, infile3) as out, ins:
    ...

Exercise: Tempdir on Windows

Make the tempdir context manager example work on windows as well. Probably need to cd out of the directory.

Solution: Context manager

import sys
from contextlib import contextmanager

if len(sys.argv) < 3:
    exit(f"Usage: {sys.argv[0]} OUTFILE INFILEs")

outfile = sys.argv[1]
infiles = sys.argv[2:]
#print(outfile)
#print(infiles)

@contextmanager
def myopen(outfile, *infiles):
    #print(len(infiles))
    out = open(outfile, 'w')
    ins = []
    for filename in infiles:
        ins.append(open(filename, 'r'))
    try:
        yield out, ins
    except Exception as ex:
        print(ex)
        pass
    finally:
        out.close()
        for fh in ins:
            fh.close()


with myopen(outfile, *infiles) as (out_fh, input_fhs):
    #print(out_fh.__class__.__name__)
    #print(len(input_fhs))
    while True:
        row = ''
        done = False
        for infh in (input_fhs):
            line = infh.readline()
            #print(f"'{line}'")
            if not line:
                done = True
                break
            if row:
                row += ','
            row += line.rstrip("\n")
        if done:
            break
        out_fh.write(row)
        out_fh.write("\n")