Centralize animation retrieval and naming

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."

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)