Check output file suffix at the beginning

Check output file suffix at the beginning of the script. No sense
going through the calculation if the suffix is wrong.

Trajectories/extract_component.py

@@ -10,15 +10,22 @@ if len(sys.argv) != 4:
     sys.exit("Required arguments: input-file node-number output-file")
 
 input_suffix = pathlib.Path(sys.argv[1]).suffix
+output_suffix = pathlib.Path(sys.argv[3]).suffix
+
+if input_suffix not in {".csv", ".gt", ".graphml"}:
+    sys.exit('Bad input file suffix')
+
+if output_suffix not in {".csv", ".gt", ".graphml"}:
+    sys.exit('Bad output file suffix')
+
 
 if input_suffix == ".csv":
     g = graph_tool.load_graph_from_csv(sys.argv[1], directed = True,
                                        csv_options = {'delimiter': ' ',
                                                       'skipinitialspace': True})
-elif input_suffix in {".gt", ".graphml"}:
-    g = graph_tool.load_graph(sys.argv[1])
 else:
-    sys.exit('Bad input file suffix')
+    # assert input_suffix in {".gt", ".graphml"}
+    g = graph_tool.load_graph(sys.argv[1])
 
 g.set_directed(False)
 
@@ -33,8 +40,6 @@ g.purge_vertices()
 g.set_vertex_filter(None)
 g.set_directed(True)
 
-output_suffix = pathlib.Path(sys.argv[3]).suffix
-
 if output_suffix == ".csv":
     with open(sys.argv[3], "w", newline = "") as f:
         writer = csv.writer(f, lineterminator = "\n", delimiter = " ")
@@ -43,7 +48,6 @@ if output_suffix == ".csv":
             v_source = e.source()
             v_target = e.target()
             writer.writerow([g.vp.name[v_source], g.vp.name[v_target]])
-elif output_suffix in {".gt", ".graphml"}:
-    g.save(sys.argv[3])
 else:
-    sys.exit('Bad output file suffix')
+    # assert output_suffix in {".gt", ".graphml"}
+    g.save(sys.argv[3])