subset_cft_input.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import yaml
import argparse
import os

parser = argparse.ArgumentParser(description='Process some info.yaml.')
parser.add_argument('infile', type=str,
                    help='input info.yaml')
parser.add_argument('output', type=str,
                    help='output info.yaml')
parser.add_argument('--samples', type=lambda x: x.split(':'),
                    help='samples to process')
parser.add_argument('--seeds', type=lambda x: x.split(':'),
                    help='seeds to process')
parser.add_argument('--no-other-partitions', action='store_true',
                    help='remove other-partitions')


if __name__ == '__main__':
    args = parser.parse_args()
    with open(args.infile) as fp:
        d = yaml.load(fp)
    if args.samples:
        d['samples'] = {k: v for k,v in d['samples'].items() if k in args.samples}
    if args.seeds:
        for sample in d['samples']:
            d['samples'][sample]['seeds'] = {seed_id: seed_obj for seed_id,seed_obj in d['samples'][sample]['seeds'].items() if seed_id in args.seeds and len(seed_obj.items()) > 0}
    for sample in d['samples'].values():
        if sample.get('other-partitions') and args.no_other_partitions:
            del sample['other-partitions']
    if not os.path.exists(os.path.dirname(args.output)):
        os.mkdir(os.path.dirname(args.output))
    with open(args.output, 'w') as out:
        yaml.dump(d, out)