Christopher Prohm

Python is a great tool for scientific computing and I use it on a regular basis for a wide array of tasks, ranging from simulation setup, data analysis, to the simulations themselves in some cases. For many of these tasks support for command line arguments greatly increases the productivity: for example, a single script can analyze different simulations. The Python standard library comes with an argument parser (argparse) out of the box. However, it results in a lot if boilerplate code which is especially a problem when checking some idea with only limited lifetime.

By writing the argmap package I tried to streamline the process for my typical work flows. In contrast to the built-in arparse package argmap is purely declarative and therefore often more concise. Further I designed it specifically such that it only uses annotations on the main method of a script. Therefore, the main method remains a pure Python function and can be easily be imported and reused from other modules.

Consider the following little example which overwrites keys in a json file with the given values:

import argparse, json

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("fname")
    parser.add_argument("--set", default=[], action="append", nargs=2)
    args = parser.parse_args()

    set_key_values(args.fname, args.set)

def set_key_values(fname, set):
    try:
        with open(fname, "r") as fobj:
            data = json.load(fobj)
    
    except IOError:
        data = {}

    for key, value in set:
        data[key] = value

    with open(fname, "w") as fobj:
        json.dump(data, fobj, sort_keys=True)

if __name__ == "__main__":
    main()

The heavy lifting is performed in the set_key_values function. The boilerplate main function construct an argument parser, parses the command-line, and calls the set_key_values function. This script can be used as follows:

$ python set_key_values_argparse.py example.json --set foo bar
$ python set_key_values_argparse.py example.json --set hello world \
    --set tick tock
$ cat example.json
{"foo": "bar", "hello": "world", "tick": "tock"}

With argmap the set_key_values function can be exposed directly, without explicitly constructing an argument parser. It only requires that the arguments are correctly annotated:

import argmap, json

@argmap.arg("set", [(str, str)])
def set_key_values(fname, set=[]):
    try:
        with open(fname, "r") as fobj:
            data = json.load(fobj)
    
    except IOError:
        data = {}

    for key, value in set:
        data[key] = value

    with open(fname, "w") as fobj:
        json.dump(data, fobj, sort_keys=True)

if __name__ == "__main__":
    argmap.argmap(set_key_values)

Here the argument type is indicated by the argmap.arg decorator, which takes a type as its second argument - for example an int or an float. Furthermore, argmap supports compound types out of the box. To get a pair of strings, use (str, str). To get a list of string pairs, use [(str, str)].

The fact that argmap only uses annotations greatly simplifies later code reuse. Although a bit nonsensical with this simple example, the function set_key_values can be reused from a second script without any modifications. The function remains a valid Python function and all arguments can be passed in directly. Suppose we wanted to parse the values with json so we are able to set for example numbers. This feat can be accomplished by reusing the previous function:

import argmap, json
from update_json_argmap import set_key_values

@argmap.arg("set", [(str, str)])
def update_json_parsed(fname, set):
    set_key_values(fname, [(key, json.loads(value)) 
                            for (key, value) in set])

if __name__ == "__main__":
    argmap.argmap(update_json_parsed)

The examples so far has been Python 2 / 3 comaptible. If Python 2 compatibility is not required the decorator annotations can be replaced by function annotations:

import argmap, json

def set_key_values(fname, set : [(str,str)] = []):
    try:
        with open(fname, "r") as fobj:
            data = json.load(fobj)
    
    except IOError:
        data = {}
    
    for key, value in set:
        data[key] = value
    
    with open(fname, "w") as fobj:
        json.dump(data, fobj, sort_keys=True)

if __name__ == "__main__":
    argmap.argmap(set_key_values)

You can find the package here. In this repository, all examples from this post and further ones are included.