diff --git a/Trajectories/extract_component.py b/Trajectories/extract_component.py index 466ca72dc03643f22e31869b744ba7443a303bd9..ba75f113e71765e0354ffdcab2ab3aae3295a3e2 100755 --- a/Trajectories/extract_component.py +++ b/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])