diff --git a/srcxray.py b/srcxray.py
index f881bd4..f37edde 100755
--- a/srcxray.py
+++ b/srcxray.py
@@ -407,7 +407,7 @@ def syscalls():
 # srcxray.py "write_dot(remove_loops(read_dot2('reduced.dot')), 'no-loops.dot')"
 
 def cleanup(a):
-    dg = read_dot2(a)
+    g = to_dg(a)
     print(dg.number_of_edges())
     dg.remove_nodes_from(black_list)
     print(dg.number_of_edges())
@@ -649,6 +649,13 @@ def read_dot2(dot):
     return dg
 
 
+def to_dg(a):
+    if isinstance(a, nx.DiGraph):
+        return a
+    if os.path.isfile(a):
+        return read_dot2(a)
+
+
 def cflow_linux():
     dirs = ('init kernel kernel/time '
             'fs fs/ext4 block '