Generate XLSX file containing CommCare locations

generate-locations.py

#! /usr/bin/env python
"""Generate XLSX file containing CommCare locations

Requires: pip install openpyxl==2.2.5
https://gist.github.com/millerdev/57a5c1363773a92d9ee6332e79789cd3
"""
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division

import os
import sys
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from collections import OrderedDict

try:
    input = raw_input
except NameError:
    pass

import openpyxl


def main():
    parser = ArgumentParser(
        description=__doc__.strip(),
        formatter_class=RawDescriptionHelpFormatter,
    )
    parser.add_argument("-n", "--ntypes", type=int, default=5,
        help="Number of location types to create (default: 5).")
    parser.add_argument("-r", "--roots", type=int, default=1,
        help="Number of top-level (root) locations (default: 1).")
    parser.add_argument("-t", "--total", type=int, default=1000,
        help="Total number of locations to generate (default: 1000).")
    parser.add_argument("-p", "--path",
        help="Save as file path. A filename will be auto-generated and saved "
            "in the current directory if not specified.")
    args = parser.parse_args()

    if args.roots > args.total:
        parser.error("Number of roots cannot be larger than total.")

    print(
        "{0.ntypes} types, {0.roots} roots -> {0.total} locations..."
        .format(args)
    )
    temp = "locs-{0.ntypes}-{0.roots}-{0.total}.xlsx"
    path = args.path or temp.format(args)
    if os.path.exists(path):
        print("{} exists".format(path))
        if input("overwrite? [y/N] ").lower() not in ("y", "ye", "yes"):
            print("abort.")
            return

    types = [TYPE_HEADINGS]
    data = OrderedDict([("types", types)])
    parent_type = ""
    for n in range(args.ntypes):
        name = "tx{}".format(n)
        types.append([
            name,
            name.upper(),
            parent_type,
            "",
            "no",
            "no",
        ])
        data[name] = [LOCATION_HEADINGS]
        parent_type = name

    counter = iter(range(args.total))
    try:
        type_rows = types[1:]
        parents = []
        while True:
            # depth-first tree traversal: fill all types
            for type_row in type_rows:
                rows = data[type_row[0]]
                if not parents:
                    # create root locations
                    for r in range(args.roots):
                        rows.append(get_row(next(counter)))
                else:
                    # add a new child location for each parent
                    # the number of parents grows exponentially per type level
                    for parent in parents:
                        rows.append(get_row(next(counter), parent[1]))
                parents = rows[1:]
            if type_rows:
                # start over with one less type row
                parents = data[type_rows[0][0]][1:]
                type_rows = type_rows[1:]
            else:
                # we've hit all types, start over at the second type
                parents = data[types[1][0]][1:]
                type_rows = types[2:]
    except StopIteration:
        pass

    levels = [len(data[type_row[0]]) - 1 for type_row in types[1:]]
    print("locations/level:", ", ".join(str(x) for x in levels))

    book = openpyxl.Workbook(optimized_write=True)

    for title, rows in data.items():
        sheet = book.create_sheet()
        sheet.title = title
        for row in rows:
            sheet.append(row)

    book.save(path)
    print("output:", path)


def get_row(index, parent_code=""):
    code = "lx{}".format(index)
    row = [
        "",
        code,           # site_code
        code.upper(),   # name
        parent_code,    # parent_site_code
        code,           # external_id
        "",
        "",
        "",
        "blue" if index % 2 else "green",
        "",
        "",
    ]
    assert len(row) == len(LOCATION_HEADINGS), \
        (len(row), len(LOCATION_HEADINGS), row)
    return row


TYPE_HEADINGS = [
    "code",
    "name",
    "parent_code",
    "Delete(Y/N)",
    "Shares Cases Y/N",
    "View Child Cases (Y/N)",
]

LOCATION_HEADINGS = [
    "location_id",
    "site_code",
    "name",
    "parent_site_code",
    "external_id",
    "latitude",
    "longitude",
    "Delete(Y/N)",
    "data: color",
    "uncategorized_data",
    "Delete Uncategorized Data(Y/N)",
]

if __name__ == "__main__":
    main()