Writing sustainable Python scripts

Vincent Bernat

Python is a great lan­guage to write a stand­alone script. Get­ting to the re­sult can be a mat­ter of a dozen to a few hun­dred lines of code and, mo­ments later, you can for­get about it and fo­cus on your next task.

Six months later, a co-worker asks you why the script fails and you don’t have a clue: no doc­u­men­ta­tion, hard-coded pa­ram­e­ters, noth­ing logged dur­ing the ex­e­cu­tion and no sen­si­ble tests to fig­ure out what may go wrong.

Turn­ing a “quick-and-dirty” Python script into a sus­tain­able ver­sion, which will be easy to use, un­der­stand and sup­port by your co-work­ers and your fu­ture self, only takes some mod­er­ate ef­fort. As an il­lus­tra­tion, let’s start from the fol­low­ing script solv­ing the in­fa­mous Fizz-Buzz test:

import sys
for n in range(int(sys.argv[1]), int(sys.argv[2])):
    if n % 3 == 0 and n % 5 == 0:
        print("fizzbuzz")
    elif n % 3 == 0:
        print("fizz")
    elif n % 5 == 0:
        print("buzz")
    else:
        print(n)

Documentation#

I find it use­ful to write doc­u­men­ta­tion be­fore cod­ing: it makes the de­sign eas­ier and it en­sures I will not post­pone this task in­def­i­nitely. The doc­u­men­ta­tion can be em­bed­ded at the top of the script:

#!/usr/bin/env python3

"""Simple fizzbuzz generator.

This script prints out a sequence of numbers from a provided range
with the following restrictions:

 - if the number is divisible by 3, then print out "fizz,"
 - if the number is divisible by 5, then print out "buzz,"
 - if the number is divisible by 3 and 5, then print out "fizzbuzz."
"""

The first line is a short sum­mary of the script pur­pose. The re­main­ing para­graphs con­tain ad­di­tional de­tails on its action.

Command-line arguments#

The sec­ond task is to turn hard-coded pa­ram­e­ters into doc­u­mented and con­fig­urable val­ues through com­mand-line ar­gu­ments, us­ing the argparse mod­ule. In our ex­am­ple, we ask the user to spec­ify a range and al­low them to mod­ify the mod­ulo val­ues for “fizz” and “buzz.”

import argparse
import sys


class CustomFormatter(argparse.RawDescriptionHelpFormatter,
                      argparse.ArgumentDefaultsHelpFormatter):
    pass


def parse_args(args=sys.argv[1:]):
    """Parse arguments."""
    parser = argparse.ArgumentParser(
        description=sys.modules[__name__].__doc__,
        formatter_class=CustomFormatter)

    g = parser.add_argument_group("fizzbuzz settings")
    g.add_argument("--fizz", metavar="N",
                   default=3,
                   type=int,
                   help="Modulo value for fizz")
    g.add_argument("--buzz", metavar="N",
                   default=5,
                   type=int,
                   help="Modulo value for buzz")

    parser.add_argument("start", type=int, help="Start value")
    parser.add_argument("end", type=int, help="End value")

    return parser.parse_args(args)


options = parse_args()
for n in range(options.start, options.end + 1):
    # ...

The added value of this mod­i­fi­ca­tion is tremen­dous: pa­ram­e­ters are now prop­erly doc­u­mented and are dis­cov­er­able through the --help flag. More­over, the doc­u­men­ta­tion we wrote in the pre­vi­ous sec­tion is also displayed:

$ ./fizzbuzz.py --help
usage: fizzbuzz.py [-h] [--fizz N] [--buzz N] start end

Simple fizzbuzz generator.

This script prints out a sequence of numbers from a provided range
with the following restrictions:

 - if the number is divisible by 3, then print out "fizz",
 - if the number is divisible by 5, then print out "buzz",
 - if the number is divisible by 3 and 5, then print out "fizzbuzz."

positional arguments:
  start         Start value
  end           End value

optional arguments:
  -h, --help    show this help message and exit

fizzbuzz settings:
  --fizz N      Modulo value for fizz (default: 3)
  --buzz N      Modulo value for buzz (default: 5)

The argparse mod­ule is quite pow­er­ful. If you are not fa­mil­iar with it, skim­ming through the doc­u­men­ta­tion is help­ful. I like to use the abil­ity to de­fine sub-com­mands and ar­gu­ment groups.

Logging#

A nice ad­di­tion to a script is to dis­play in­for­ma­tion dur­ing its ex­e­cu­tion. The logging mod­ule is a good fit for this pur­pose. First, we de­fine the logger:

import logging
import logging.handlers
import os
import sys

logger = logging.getLogger(os.path.splitext(os.path.basename(sys.argv[0]))[0])

Then, we make its ver­bosity con­fig­urable: logger.​debug() should out­put some­thing only when a user runs our script with --debug and --silent should mute the logs un­less an ex­cep­tional con­di­tion oc­curs. For this pur­pose, we add the fol­low­ing code in parse_args():

# In parse_args()
g = parser.add_mutually_exclusive_group()
g.add_argument("--debug", "-d", action="store_true",
               default=False,
               help="enable debugging")
g.add_argument("--silent", "-s", action="store_true",
               default=False,
               help="don't log to console")

We add this func­tion to con­fig­ure logging:

def setup_logging(options):
    """Configure logging."""
    root = logging.getLogger("")
    root.setLevel(logging.WARNING)
    logger.setLevel(options.debug and logging.DEBUG or logging.INFO)
    if not options.silent:
        ch = logging.StreamHandler()
        ch.setFormatter(logging.Formatter(
            "%(levelname)s[%(name)s] %(message)s"))
        root.addHandler(ch)

The main body of our script be­comes this:

if __name__ == "__main__":
    options = parse_args()
    setup_logging(options)

    try:
        logger.debug("compute fizzbuzz from {} to {}".format(options.start,
                                                             options.end))
        for n in range(options.start, options.end + 1):
            # ...
    except Exception as e:
        logger.exception("%s", e)
        sys.exit(1)
    sys.exit(0)

If the script may run unattended—e.g. from a crontab, we can make it log to syslog:1

def setup_logging(options):
    """Configure logging."""
    root = logging.getLogger("")
    root.setLevel(logging.WARNING)
    logger.setLevel(options.debug and logging.DEBUG or logging.INFO)
    if not options.silent:
        if not sys.stderr.isatty():
            facility = logging.handlers.SysLogHandler.LOG_DAEMON
            sh = logging.handlers.SysLogHandler(address='/dev/log',
                                                facility=facility)
            sh.setFormatter(logging.Formatter(
                "{0}[{1}]: %(message)s".format(
                    logger.name, os.getpid())))
            root.addHandler(sh)
        else:
            ch = logging.StreamHandler()
            ch.setFormatter(logging.Formatter(
                "%(levelname)s[%(name)s] %(message)s"))
            root.addHandler(ch)

For this ex­am­ple, this is a lot of code just to use logger.​debug() once, but in a real script, this will come in handy to help users un­der­stand how the task is completed.

$ ./fizzbuzz.py --debug 1 3
DEBUG[fizzbuzz] compute fizzbuzz from 1 to 3
1
2
fizz

Tests#

Unit tests are very use­ful to en­sure an ap­pli­ca­tion be­haves as in­tended. It is not com­mon to use them in scripts, but writ­ing a few of them greatly im­proves their re­li­a­bil­ity. Let’s turn the code in the in­ner “for” loop into a func­tion with some in­ter­ac­tive ex­am­ples of us­age in its doc­u­mentation:

def fizzbuzz(n, fizz, buzz):
    """Compute fizzbuzz nth item given modulo values for fizz and buzz.

    >>> fizzbuzz(5, fizz=3, buzz=5)
    'buzz'
    >>> fizzbuzz(3, fizz=3, buzz=5)
    'fizz'
    >>> fizzbuzz(15, fizz=3, buzz=5)
    'fizzbuzz'
    >>> fizzbuzz(4, fizz=3, buzz=5)
    4
    >>> fizzbuzz(4, fizz=4, buzz=6)
    'fizz'

    """
    if n % fizz == 0 and n % buzz == 0:
        return "fizzbuzz"
    if n % fizz == 0:
        return "fizz"
    if n % buzz == 0:
        return "buzz"
    return n

pytest can en­sure the re­sults are correct:2

$ python3 -m pytest -v --log-level=debug --doctest-modules ./fizzbuzz.py
============================ test session starts =============================
platform linux -- Python 3.7.4, pytest-3.10.1, py-1.8.0, pluggy-0.8.0 -- /usr/bin/python3
cachedir: .pytest_cache
rootdir: /home/bernat/code/perso/python-script, inifile:
plugins: xdist-1.26.1, timeout-1.3.3, forked-1.0.2, cov-2.6.0
collected 1 item

fizzbuzz.py::fizzbuzz.fizzbuzz PASSED                                  [100%]

========================== 1 passed in 0.05 seconds ==========================

In case of an er­ror, pytest dis­plays a mes­sage de­scrib­ing the lo­ca­tion and the na­ture of the failure:

$ python3 -m pytest -v --log-level=debug --doctest-modules ./fizzbuzz.py -k fizzbuzz.fizzbuzz
============================ test session starts =============================
platform linux -- Python 3.7.4, pytest-3.10.1, py-1.8.0, pluggy-0.8.0 -- /usr/bin/python3
cachedir: .pytest_cache
rootdir: /home/bernat/code/perso/python-script, inifile:
plugins: xdist-1.26.1, timeout-1.3.3, forked-1.0.2, cov-2.6.0
collected 1 item

fizzbuzz.py::fizzbuzz.fizzbuzz FAILED                                  [100%]

================================== FAILURES ==================================
________________________ [doctest] fizzbuzz.fizzbuzz _________________________
100
101     >>> fizzbuzz(5, fizz=3, buzz=5)
102     'buzz'
103     >>> fizzbuzz(3, fizz=3, buzz=5)
104     'fizz'
105     >>> fizzbuzz(15, fizz=3, buzz=5)
106     'fizzbuzz'
107     >>> fizzbuzz(4, fizz=3, buzz=5)
108     4
109     >>> fizzbuzz(4, fizz=4, buzz=6)
Expected:
    fizz
Got:
    4

/home/bernat/code/perso/python-script/fizzbuzz.py:109: DocTestFailure
========================== 1 failed in 0.02 seconds ==========================

We can also write unit tests as code. Let’s sup­pose we want to test the fol­low­ing function:

def main(options):
    """Compute a fizzbuzz set of strings and return them as an array."""
    logger.debug("compute fizzbuzz from {} to {}".format(options.start,
                                                         options.end))
    return [str(fizzbuzz(i, options.fizz, options.buzz))
            for i in range(options.start, options.end+1)]

At the end of the script,3 we add the fol­low­ing unit tests, lever­ag­ing pytest’s parametrized test functions:

# Unit tests
import pytest                   # noqa: E402
import shlex                    # noqa: E402


@pytest.mark.parametrize("args, expected", [
    ("0 0", ["fizzbuzz"]),
    ("3 5", ["fizz", "4", "buzz"]),
    ("9 12", ["fizz", "buzz", "11", "fizz"]),
    ("14 17", ["14", "fizzbuzz", "16", "17"]),
    ("14 17 --fizz=2", ["fizz", "buzz", "fizz", "17"]),
    ("17 20 --buzz=10", ["17", "fizz", "19", "buzz"]),
])
def test_main(args, expected):
    options = parse_args(shlex.split(args))
    assert main(options) == expected

The test func­tion runs once for each of the pro­vided pa­ram­e­ters. The args part is used as in­put for the parse_​args() func­tion to get the ap­pro­pri­ate op­tions we need to pass to the main() func­tion. The expected part is com­pared to the re­sult of the main() func­tion. When ev­ery­thing works as ex­pected, pytest says:

python3 -m pytest -v --log-level=debug --doctest-modules ./fizzbuzz.py
============================ test session starts =============================
platform linux -- Python 3.7.4, pytest-3.10.1, py-1.8.0, pluggy-0.8.0 -- /usr/bin/python3
cachedir: .pytest_cache
rootdir: /home/bernat/code/perso/python-script, inifile:
plugins: xdist-1.26.1, timeout-1.3.3, forked-1.0.2, cov-2.6.0
collected 7 items

fizzbuzz.py::fizzbuzz.fizzbuzz PASSED                                  [ 14%]
fizzbuzz.py::test_main[0 0-expected0] PASSED                           [ 28%]
fizzbuzz.py::test_main[3 5-expected1] PASSED                           [ 42%]
fizzbuzz.py::test_main[9 12-expected2] PASSED                          [ 57%]
fizzbuzz.py::test_main[14 17-expected3] PASSED                         [ 71%]
fizzbuzz.py::test_main[14 17 --fizz=2-expected4] PASSED                [ 85%]
fizzbuzz.py::test_main[17 20 --buzz=10-expected5] PASSED               [100%]

========================== 7 passed in 0.03 seconds ==========================

When an er­ror oc­curs, pytest pro­vides a use­ful as­sess­ment of the situation:

$ python3 -m pytest -v --log-level=debug --doctest-modules ./fizzbuzz.py
[…]
================================== FAILURES ==================================
__________________________ test_main[0 0-expected0] __________________________

args = '0 0', expected = ['0']

    @pytest.mark.parametrize("args, expected", [
        ("0 0", ["0"]),
        ("3 5", ["fizz", "4", "buzz"]),
        ("9 12", ["fizz", "buzz", "11", "fizz"]),
        ("14 17", ["14", "fizzbuzz", "16", "17"]),
        ("14 17 --fizz=2", ["fizz", "buzz", "fizz", "17"]),
        ("17 20 --buzz=10", ["17", "fizz", "19", "buzz"]),
    ])
    def test_main(args, expected):
        options = parse_args(shlex.split(args))
>       assert main(options) == expected
E       AssertionError: assert ['fizzbuzz'] == ['0']
E         At index 0 diff: 'fizzbuzz' != '0'
E         Full diff:
E         - ['fizzbuzz']
E         + ['0']

fizzbuzz.py:160: AssertionError
----------------------------- Captured log call ------------------------------
fizzbuzz.py                125 DEBUG    compute fizzbuzz from 0 to 0
===================== 1 failed, 6 passed in 0.05 seconds =====================

The call to logger.​debug() is in­cluded in the out­put. This is an­other good rea­son to use the log­ging fea­ture! If you want to know more about the won­der­ful fea­tures of pytest, have a look at “Test­ing net­work soft­ware with pytest and Linux namespaces.”

To sum up, en­hanc­ing a Python script to make it more sus­tain­able can be done in four steps:

  1. add doc­u­men­ta­tion at the top,
  2. use the argparse mod­ule to doc­u­ment the dif­fer­ent pa­rameters,
  3. use the logging mod­ule to log de­tails about progress, and
  4. add some unit tests.

You can find the com­plete ex­am­ple on GitHub and use it as a template!

Up­date (2019-06)

There are some in­ter­est­ing threads about this ar­ti­cle on Lob­sters and Red­dit. While the ad­di­tion of doc­u­men­ta­tion and com­mand-line ar­gu­ments seems to be well-re­ceived, logs and tests are some­times re­ported as too ver­bose. Dan Con­nolly wrote “Prac­ti­cal pro­duc­tion python scripts” as an an­swer to this post.