bump.py refactor

bump.py

#!/usr/bin/env python

from __future__ import print_function

import argparse
import re
import sys

IS_PY3 = sys.version_info[0] == 3

def parser():
    """Returns an argument parser."""

    parser = argparse.ArgumentParser(
        prog='bump',
        description="Bumps package version numbers")
    parser.add_argument('files', help='Files to update', nargs='+')
    parser.add_argument('-M', dest='major', action='store_true', default=False,
                        help="Bump major version")
    parser.add_argument('-m', dest='minor', action='store_true', default=False,
                        help="Bump minor version")
    parser.add_argument('-b', dest='build', action='store_true', default=True,
                        help="Bump build version")
    return parser

def bump_version(version_string, options):
    """
    Bumps a version number. Returns a new version number string.
    Raises ValueError.
    """
    version = list(map(int, version_string.split('.')))

    while len(version) < 3:
        version += [0]
    if options.major:
        version = version[0] + 1, 0, 0
    elif options.minor:
        version = version[0], version[1] + 1, 0
    elif options.build:
        version = version[0], version[1], version[2] + 1

    return '.'.join(map(str, version))

def get_matches(files, options):
    """Returns a dict with version string and new version string."""

    matches = {}

    for filename in files:
        with open(filename, 'rb') as f:
            match = re.search(
                '\s*[\'"]?version[\'"]?\s*[=:]\s*[\'"]?([^\'",]+)[\'"]?',
                f.read().decode('utf-8'), re.I)

        if not match:
            print("No version definition found in %s" % filename)
            continue

        try:
            version_string = match.group(1)
            new_version_string = bump_version(version_string, options)

            matches[filename] = dict(match=match,
                                     version_string=version_string,
                                     new_version_string=new_version_string)

        except ValueError:
            print("Invalid version string in %s: %s" % (filename, version_string))

    return matches

def main():
    args = parser().parse_args()

    matches = get_matches(args.files, args)

    if len(matches) == 0:
        exit(1)

    for filename, match in matches.items():
        print("%s: %s => %s" %
            (filename, match['version_string'], match['new_version_string']))

    __input = input if IS_PY3 else raw_input

    if __input('Is this ok? y/n ').lower() != 'y':
        print('Cancelled')
        exit(1)

    for filename, match in matches.items():
        new_version_string = match['new_version_string']
        with open(filename, 'wb') as f:
            content = (bytes(match['match'].string, 'utf-8') if IS_PY3 else
                       match['match'].string)
            if IS_PY3:
                new_version_string = bytes(new_version_string, 'utf-8')
            f.write(content[:match['match'].start(1)] +
                    new_version_string + content[match['match'].end(1):])
            print('Updated', filename)

if __name__ == '__main__':
    main()