sync_epic_roster.py

from __future__ import print_function
import os
import re
from collections import Counter
import itertools
import argparse

import json
import requests


def pt_api_endpoint(path_template, **kwargs):
    return "https://www.pivotaltracker.com/services/v5" + path_template.format(**kwargs)


def project_epic_endpoint(project_id, epic):
    return pt_api_endpoint('/projects/{project_id}/epics/{epic_id}',
                           project_id=project_id,
                           epic_id=epic['id'])


def project_epics_endpoint(project_id):
    return pt_api_endpoint('/projects/{project_id}/epics/', project_id=project_id)


def project_memberships_endpoint(project_id):
    return pt_api_endpoint('/projects/{project_id}/memberships', project_id=project_id)


def project_stories_endpoint(project_id):
    return pt_api_endpoint('/projects/{project_id}/stories/', project_id=project_id)


def get_epics(session, project_id):
    return session.get(project_epics_endpoint(project_id)).json()


def get_epic_features(session, project_id, epic):
    params = {'with_label': epic['label']['name'], 'with_story_type': 'feature', 'fields': 'owner_ids,current_state'}
    return session.get(project_stories_endpoint(project_id), params=params).json()


def get_members(session, project_id):
    return {m['person']['id']: m['person']['initials']
            for m in session.get(project_memberships_endpoint(project_id)).json()}


def update_epic(session, project_id, epic, **attrs):
    # support older version of python-requests which doesn't have json support
    data = json.dumps(attrs)
    headers = {'Content-type': 'application/json'}
    return session.put(project_epic_endpoint(project_id, epic),
                       data=data,
                       headers=headers)


def update_owners_in_epic_name(epic_name, owners):
    if not owners:
        return epic_name
    return u'{0} ({1})'.format(bare_epic_name(epic_name), owners)


def bare_epic_name(epic_name):
    return re.sub(r' ?\([^()]*\)$', '', epic_name)


def update_owners_in_epic_description(description, owners, separator='\n\n----\n\n'):
    if not owners:
        return description
    preamble = 'Story owners: (automatically updated)\n' + owners + separator
    maybe_owners_and_rest = description.split(separator, 1)
    if len(maybe_owners_and_rest) == 2:
        return preamble + maybe_owners_and_rest[1]
    elif len(maybe_owners_and_rest) == 1:
        return preamble + maybe_owners_and_rest[0]
    else:
        return preamble


def format_owners(owner_ids, members):
    if not len(owner_ids):
        return ''
    comparator = lambda (name1, count1), (name2, count2): cmp(count2, count1) if count1 != count2 else cmp(name1, name2)
    member_impacts = sorted(((members[owner_id], count) for owner_id, count in Counter(owner_ids).items()),
                            cmp=comparator)
    return u', '.join(u'{0}:{1}'.format(name, count) for name, count in member_impacts)


def main(args):
    session = requests.Session()
    session.headers.update({'X-TrackerToken': os.environ['PT_TOKEN']})
    project_id = os.environ['PT_PROJECT_ID']
    members = get_members(session, project_id)
    for epic in get_epics(session, project_id):
        features = get_epic_features(session, project_id, epic)
        completed = all(feature['current_state'] in ('accepted', 'unscheduled')
                        for feature in features)
        if completed and args.skip_completed:
            continue
        owner_ids = list(itertools.chain.from_iterable(
            itertools.imap(lambda feature: feature["owner_ids"], features)))
        owners_string = format_owners(owner_ids, members)
        old_description = epic.get('description', '')
        new_description = update_owners_in_epic_description(old_description, owners_string)
        print(epic['name'])
        print('-> status:', 'completed' if completed else 'in progress')
        print(u'-> owners:', owners_string.encode('utf8'))
        if new_description != old_description:
            print('-> updating')
            if not args.dry_run:
                print('->', update_epic(session, project_id, epic, description=new_description).status_code)


def test():
    owner_ids = [1, 1, 1, 2, 2, 3, 1, 4, 5, 4, 3]
    members = {1: 'tk', 2: 'mo', 3: 'mg', 4: 'av', 5: 'sg'}
    epic_names = [
        ('Teacher Results', 'Teacher Results'),
        ('Teacher Results ()', 'Teacher Results'),
        ('Teacher Results (tk:1)', 'Teacher Results'),
        ('Teacher Results (PD) (av:2)', 'Teacher Results (PD)'),
    ]
    assert 'tk:4, av:2, mg:2, mo:2, sg:1' == format_owners(owner_ids, members)
    for epic_name, expected_prefix in epic_names:
        updated = update_owners_in_epic_name(epic_name, 'hello')
        assert updated == '{0} (hello)'.format(expected_prefix)
    epic_descriptions = [
        ('', 'hello\n--\n'),
        ('abc', 'hello\n--\nabc'),
        ('abc--\n--\ndef', 'hello\n--\ndef'),
        ('abc--\n--\ndef\nghi', 'hello\n--\ndef\nghi'),
    ]
    intro = 'Story owners: (automatically updated)\n'
    for epic_description, expected in epic_descriptions:
        updated = update_owners_in_epic_description(epic_description, 'hello', separator='\n--\n')
        assert updated == (intro + expected)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--skip-completed', action='store_true', default=False)
    parser.add_argument('--dry-run', action='store_true', default=False)
    args = parser.parse_args()

    test()
    main(args)