1 files changed, 18 insertions, 10 deletions
diff --git a/view.py b/view.py
index e4ebaee..c827442 100755
--- a/view.py
+++ b/view.py
@@ -175,6 +175,7 @@ if __name__ == '__main__':
 
     from stl import *
     from materials import *
+    import detectors
 
     parser = optparse.OptionParser('%prog filename.stl')
     parser.add_option('-b', '--bits', type='int', dest='bits',
@@ -191,13 +192,20 @@ if __name__ == '__main__':
         geometry = Geometry()
         geometry.add_solid(Solid(read_stl(args[0]), vacuum, vacuum))
         geometry.build(options.bits)
-
-    elif ext.lower() == '.pkl':
-        import pickle
-        from detectors import *
-
-        f = open(args[0], 'rb')
-        geometry = pickle.load(f)
-        f.close()
-
-    view(geometry, tail)
+        view(geometry, tail)
+    else:
+        import inspect
+        import chroma.geometry
+
+        members = dict(inspect.getmembers(detectors))
+
+        if args[0] in members:
+            if issubclass(members[args[0]], Geometry) or \
+                    issubclass(members[args[0]], chroma.geometry.Geometry):
+                geometry = members[args[0]]()
+                geometry.build(options.bits)
+                view(geometry, args[0])
+            else:
+                sys.exit("%s doesn't inherit from a Geometry object" % args[0])
+        else:
+            sys.exit("couldn't find a geometry with name %s" % args[0])