script to add level numbering to alias of EA requirements

requirements.py

""" script to add level numbering to alias of EA requirements

    get csv from EA with the query below:
    select Object_ID, ParentID, Name, Alias, Version, Effort, Phase, TPos
    from t_object where Object_Type="Requirement"
"""

import csv
import sys

def read_reqs(csvfile, reqs):
    """ read requirements from file """
    req_reader = csv.DictReader(csvfile, delimiter=';')
    for row in req_reader:
        row['subs'] = []
        reqs[row['Object_ID']] = row
    return req_reader.fieldnames

def register_parent(req, reqs, tree):
    """ register a requirement with its parent """
    parent_id = req['ParentID']
    if parent_id == '0':
        tree.append(req)
    else:
        reqs[parent_id]['subs'].append(req)

def assign_numbers(reqs, parent_number):
    """ assign a level number to a tree of requirements """
    reqs.sort(key=lambda x: int(x['TPos']))

    num = 0
    for req in reqs:
        num += 1
        if req['ParentID'] == '0':
            req['Alias'] = num
        else:
            req['Alias'] = "%s.%s" % (parent_number, num)
        assign_numbers(req['subs'], req['Alias'])

def write_tree(reqs, writer):
    """  write a requirements tree """
    for req in reqs:
        subs = req['subs']
        del req['subs']
        writer.writerow(req)
        write_tree(subs, writer)

def main(file_name='requirements.csv'):
    """ this is where it all happens """
    reqs = dict()
    fieldnames = None
    with open(file_name, 'rb') as csvfile:
        fieldnames = read_reqs(csvfile, reqs)

    tree = []
    for req in reqs.itervalues():
        register_parent(req, reqs, tree)
    assign_numbers(tree, '')

    req_writer = csv.DictWriter(sys.stdout, fieldnames=fieldnames, lineterminator='\n')
    req_writer.writeheader()
    write_tree(tree, req_writer)

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print "Usage: %s <csv source file>" % sys.argv[0]
        sys.exit(0)
    main(sys.argv[1])