Python code templates

From raju

Template1

Self contained python3 code that involves argument parsing, classes, unit tests and regression tests.

Directory Structure

 % tree                                
.
├── readme.txt
├── tests
│   ├── __init__.py
│   └── test_time_calc.py
└── time_calc.py

1 directory, 4 files

The main code is in time_calc.py and it is tested in tests/test_time_calc.py .

The tests/__init__.py is an empty file. It is needed for the tests to be discovered automatically. If the file is present, we can simply do

% python3 -m unittest -v

to run all the tests. If it is not there, the tests have to be specified manually. For example

% python3 -m unittest -v tests/test_time_calc.py

Sample Run

Run the script

 % python3 time_calc.py --start 2:05 --end 5:13
188.0

which shows the number of minutes between two time stamps specified in hours:minutes format.

Run the tests

 % python3 -m unittest -v
test_convert_to_minutes (tests.test_time_calc.ConvertTime) ... ok
test_diff_time (tests.test_time_calc.ConvertTime) ... ok

----------------------------------------------------------------------
Ran 2 tests in 0.002s

OK

File Contents

  • time_calc.py
 % cat time_calc.py 
import sys

"""
Utility script to calculate difference between two time stamps.

Sample usage:
 % python3 time_calc.py --start 2:05 --end 5:13
188.0
"""


def parse_arguments(args):
    import argparse
    parser = argparse.ArgumentParser(
        description='get difference between two timestamps'
        )
    parser.add_argument(
        "--start", action='store',
        default="0:0", dest="start",
        help="Starting timestamp HH:MM")
    parser.add_argument(
        "--end", action='store',
        default="0:0", dest="end",
        help="Ending timestamp HH:MM")
    return parser.parse_args(args)


def diff_time(begin, end):
    hm_begin = HourMinute(begin)
    hm_end = HourMinute(end)
    total_minutes_begin = hm_begin.total_minutes()
    total_minutes_end = hm_end.total_minutes()
    return (total_minutes_end - total_minutes_begin)


class HourMinute:
    hours = 0
    mins = 0

    def __init__(self, hhmm_str):
        tokens = hhmm_str.split(":")
        self.hours = float(tokens[0])
        self.mins = float(tokens[1])

    def total_minutes(self):
        return self.hours*60 + self.mins

    def dump(self):
        print("hours = ", self.hours)
        print("minutes = ", self.mins)

if __name__ == "__main__":
    args = parse_arguments(sys.argv[1:])
    print(diff_time(args.start, args.end))
  • readme.txt
 % cat readme.txt 
To run the unit tests simply do

 % python3 -m unittest -v                 
test_convert_to_minutes (tests.test_time_calc.TestConvertTime) ... ok
test_diff_time (tests.test_time_calc.TestConvertTime) ... ok

----------------------------------------------------------------------
Ran 2 tests in 0.001s

OK
  • tests/test_time_calc.py
 % cat tests/test_time_calc.py 
import unittest

import time_calc


class ConvertTime(unittest.TestCase):
    def test_convert_to_minutes(self):
        hm = time_calc.HourMinute("8:25")
        total_minutes = 505
        result = hm.total_minutes()
        self.assertEqual(result, total_minutes)

        hm = time_calc.HourMinute("14:12")
        total_minutes = 852
        result = hm.total_minutes()
        self.assertEqual(result, total_minutes)

    def test_diff_time(self):
        args = time_calc.parse_arguments(['--start', '2:05', '--end', '5:13'])
        difference = 188
        result = time_calc.diff_time(args.start, args.end)
        self.assertEqual(result, difference)

        args = time_calc.parse_arguments(['--start', '5:30', '--end', '10:05'])
        difference = 275
        result = time_calc.diff_time(args.start, args.end)
        self.assertEqual(result, difference)


if __name__ == '__main__':
    unittest.main()
  • tests/__init__.py

This is just an empty file. It is useful for automatically discovering the tests.

 % ls -al tests/__init__.py
-rw-r--r-- 1 rajulocal rajulocal 0 Aug 28 22:54 tests/__init__.py

 % cat tests/__init__.py

Finally, check code style by running pep8 on all python programs. A successful run will not print any output.

 % pep8 time_calc.py tests/test_time_calc.py

Template2

In some cases, the test involves computing a dataset first and comparing it with an existing "raw" dataset. The test is presumed to be a success if the computed dataset matches with the raw dataset.

The directory structure in this case can be

% tree
.
|-- mocks
|   `-- raw_data.csv
|-- src
`-- tests

where all the raw datasets are stored in the mocks directory. A sample test function looks like

import unittest
import mymod
import pandas as pd
import pandas.util.testing as tm
import tempfile

class TestMyMod(unittest.TestCase):
    def test_foo(self):
        a = mymod.compute_df(arg1, arg2)

        tmp = tempfile.NamedTemporaryFile()
        a.to_csv(tmp.name, sep=',', index=False)
        df = pd.read_csv(tmp)
        tmp.close()

        raw = pd.read_csv('mocks/raw_data.csv')
        tm.assert_frame_equal(df, raw)
        self.assertTrue(df.equals(raw))
        self.assertTrue(raw.equals(df))

if __name__=='__main__':
    unittest.main()

Template 3

Demonstrates the following

  • Run a python script every night using cron.
  • The python script logs its output into a file using the logger module.
  • It also calls a shell script.
  • The output from shell script is sent to the same log file.
  • Test run the python script in an emulated cron environment.

tags | nightly

cron setup

% crontab -l | stuff.py
MAILTO=""
0 0 * * * $HOME/bin/nightly.py > /dev/null 2>&1

Here stuff.py is a python script that weeds out comments from its input.

The MAILTO line optional. It basically tells cron not to send any emails when it finishes running.

the python script

% cat $HOME/bin/nightly.py
#! /usr/bin/env python3
import os
from datetime import datetime
import logging
import subprocess


def main():
    # See https://docs.python.org/3/howto/logging.html for a tutorial on how to
    # do logging.
    home_dir = os.environ['HOME']
    log_dir = os.path.join(home_dir, 'logs')
    if not os.path.exists(log_dir):
        os.makedirs(log_dir, exist_ok=True)
    log_file_name = "nightly.txt"
    log_file_path = os.path.join(log_dir, log_file_name)
    logging.basicConfig(filename=log_file_path,
                        level=logging.INFO,
                        format='%(asctime)s %(message)s')
    logging.info('Started %s', os.path.abspath(__file__))
    # note_time()
    script_path = os.path.join(home_dir, 'bin', 'nightly.sh')
    run_nightly_sh(script_path, log_file_path)
    logging.info('finished\n')


def note_time():
    a = datetime.now().strftime("%Y%m%d_%H%M%S")
    logging.info('%s', a)


def run_nightly_sh(script_path, log_file_path):
    cmd = os.path.join(script_path) + \
          " >> " + log_file_path + " 2>&1"
    logging.info('%s', cmd)
    subprocess.call(cmd, shell=True)


if __name__ == '__main__':
    main()

the shell script

% cat $HOME/bin/nightly.sh
#!/usr/bin/env dash
set -e -u -x

uptime

Test the script

% run_as_cron.sh $HOME/bin/nightly.py

which will produce something like this in the log file

% tail -n 20 $HOME/logs/nightly.txt
...

2018-03-12 16:01:17,451 Started /home/raju/bin/nightly.py
2018-03-12 16:01:17,451 /home/raju/bin/nightly.sh >> /home/raju/logs/nightly.txt 2>&1
+ uptime
 16:01:17 up 3 days, 22:51,  1 user,  load average: 0.00, 0.00, 0.00
2018-03-12 16:01:17,455 finished

where run_as_cron.sh is a shell script that let's us run any shell script in an emulated cron environment.

Note:- The run_as_cron.sh script requires some initial set up. Check the comments in that script for more details. Basically, we need to dump the cron environment into a file first and then point to it in the script.