Polish

segments.py

 #!/usr/bin/env python3
 
 """Collapses the input graph into a new graph: the graph of
-segments. An edge of the graph is a segment. The new graph is output
-in the binary format of graph-tool. Input file should be an edgelist
-in CSV format. Output file can have suffix ".gt" or ".graphml".
+segments. A vertex of the new graph is a segment. The new graph is
+output in the binary format of graph-tool. Input file should be an
+edgelist in CSV format. Output file can have suffix ".gt" or
+".graphml".
 
 """
 
 import graph_tool
 import time
 import sys
+import pathlib
 
 ##########################
 # load the edgelist file #
 ##########################
 
 if len(sys.argv) != 3: sys.exit("Required arguments: input-file output-file")
+
+if pathlib.Path(sys.argv[2]).suffix not in {".gt", ".graphml"}:
+    sys.exit('Output file suffix must be ".gt" or ".graphml"')
+
 t0 = time.perf_counter()
 print('Loading edgelist file...')
 g = graph_tool.load_graph_from_csv(sys.argv[1],
@@ -23,6 +29,9 @@ g = graph_tool.load_graph_from_csv(sys.argv[1],
                                    eprop_names = ['nl_cost_function'],
                                    csv_options = {'delimiter': ' ',
                                                   'skipinitialspace': True})
+print("Input graph:")
+print("Number of vertices:", g.num_vertices())
+print("Number of edges:", g.num_edges())
 t1 = time.perf_counter()
 print(f'Loading done in {t1 - t0:.0f} s')
 t0 = t1
@@ -44,7 +53,7 @@ g.ep['nl_cost_function'] = g.new_ep('float')
 # Processing #
 ##############
 
-print('Starting processing')
+print('Collapsing into segments...')
 g.set_fast_edge_removal()
 verts_to_del = []
 
@@ -68,14 +77,15 @@ for v in g.vertices():
         verts_to_del.append(v)
     
 t1 = time.perf_counter()
-print(f'Done processing in {t1 - t0:.0f} s')
+print(f'Done collapsing in {t1 - t0:.0f} s')
 t0 = t1
 print("New graph:")
 print("Number of vertices:", g.num_vertices())
 print("Number of edges:", g.num_edges())
 print("Internal properties:")
 g.list_properties()
-print(f'Empty nodes: {len(verts_to_del)}. Deleting empty nodes...')
+print('Empty nodes:', len(verts_to_del))
+print('Deleting empty nodes...')
 g.remove_vertex(verts_to_del, fast=True)
 t1 = time.perf_counter()
 print(f"Done deleting in {t1 - t0:.0f} s")