Migration script for vcrpy pull request #73

vcrpy-migration.py

from contextlib import closing
import os
import re
import tempfile

import json


PARTS = [
    'protocol',
    'host',
    'port',
    'path',
]


def build_uri(**parts):
    return "{protocol}://{host}:{port}{path}".format(**parts)


def migrate_json(in_fp, out):
    data = json.load(in_fp)
    for item in data:
        req = item['request']
        uri = {k: req.pop(k) for k in PARTS}
        req['uri'] = build_uri(**uri)
    json.dump(data, out, indent=4)


def migrate_yml(in_f, out):
    uri = dict.fromkeys(PARTS, None)
    for line in in_f:
        for part in uri:
            match = re.match('\s+{}:\s(.*)'.format(part), line)
            if match:
                uri[part] = match.group(1)
                break
        else:
            out.write(line)

        if None not in uri.values():
            out.write("    uri: {}\n".format(build_uri(**uri)))
            uri = dict.fromkeys(PARTS, None)  # reset dict


def migrate(file_path, migration_fn):
    tmp_file = tempfile.NamedTemporaryFile(delete=False)
    with closing(tmp_file) as out:
        with open(file_path) as in_f:
            migration_fn(in_f, out)
    os.rename(tmp_file.name, file_path)


#
#  Usage example
#
migrate('/tmp/in.yml', migrate_yml)
migrate('/tmp/in.json', migrate_json)