diff --git a/cuburn/genome/db.py b/cuburn/genome/db.py index 7c1101a..1041486 100644 --- a/cuburn/genome/db.py +++ b/cuburn/genome/db.py @@ -1,6 +1,8 @@ import os import json +import convert + class GenomeDB(object): """ Abstract base class for accessing genomes by ID. This is likely to be @@ -17,6 +19,40 @@ class GenomeDB(object): def stash(self, id, gnm): self.stashed[id] = gnm + def get_anim(self, name, half=False): + """ + Given the identifier of any type of genome that can be converted to an + animation, do so. + + Returns `(gnm, basename)`, where gnm is the animation genome as a + plain Python dict and basename is probably a suitable name for output + files. + """ + basename = os.path.basename(name) + split = basename.rsplit('.', 1) + head, ext = split[0], split[1] if len(split) == 2 else '' + if ext in ('json', 'flam3', 'flame'): + basename = head + + if os.path.isfile(name) and ext in ('flam3', 'flame'): + with open(name) as fp: + gnm_str = fp.read() + flames = convert.XMLGenomeParser.parse(gnm_str) + if len(flames) != 1: + warnings.warn( + '%d flames in file, only using one.' % len(flames)) + gnm = convert.flam3_to_node(flames[0]) + else: + gnm = self.get(name) + + if gnm['type'] == 'node': + gnm = convert.node_to_anim(gnm, half=half) + elif gnm['type'] == 'edge': + gnm = convert.edge_to_anim(self, gnm) + assert gnm['type'] == 'animation', 'Unrecognized genome type.' + + return gnm, basename + class OneFileDB(GenomeDB): def __init__(self, dct): assert dct.get('type') == 'onefiledb', "Doesn't look like a OneFileDB." diff --git a/main.py b/main.py index 7c5c766..6815b71 100755 --- a/main.py +++ b/main.py @@ -42,27 +42,12 @@ def main(args, prof): import pycuda.autoinit gdb = db.connect(args.genomedb) - if os.path.isfile(args.flame) and (args.flame.endswith('.flam3') or - args.flame.endswith('.flame')): - with open(args.flame) as fp: - gnm_str = fp.read() - flames = convert.XMLGenomeParser.parse(gnm_str) - if len(flames) != 1: - warnings.warn('%d flames in file, only using one.' % len(flames)) - gnm = convert.flam3_to_node(flames[0]) - else: - gnm = gdb.get(args.flame) - - if gnm['type'] == 'node': - gnm = convert.node_to_anim(gnm, half=args.half) - elif gnm['type'] == 'edge': - gnm = convert.edge_to_anim(gdb, gnm) - assert gnm['type'] == 'animation', 'Unrecognized genome type.' + gnm, basename = gdb.get_anim(args.flame, args.half) gprof, times = use.wrap_genome(prof, gnm) rmgr = render.RenderManager() - basename = os.path.basename(args.flame).rsplit('.', 1)[0] + '_' + basename += '_' if args.name is not None: basename = args.name prefix = os.path.join(args.dir, basename)