extract_component.py

#!/usr/bin/env python3

"""This script reads a graph from a graph-tool file or an edge list,
finds the weakly connected component containing a given node, and
writes this component to a new file, either in graph-tool format or as
an edge list.

"""

import csv
import pathlib
import sys

import graph_tool
from graph_tool import topology

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})
else:
    # assert input_suffix in {".gt", ".graphml"}
    g = graph_tool.load_graph(sys.argv[1])

g.set_directed(False)

for v in g.vertices():
    if g.vp.name[v] == int(sys.argv[2]): break
else:
    sys.exit("Node not found")

label = topology.label_out_component(g, v)
g.set_vertex_filter(label)
g.purge_vertices()
g.set_vertex_filter(None)
g.set_directed(True)

if output_suffix == ".csv":
    with open(sys.argv[3], "w", newline = "") as f:
        writer = csv.writer(f, lineterminator = "\n", delimiter = " ")

        for e in g.edges():
            v_source = e.source()
            v_target = e.target()
            writer.writerow([g.vp.name[v_source], g.vp.name[v_target]])
else:
    # assert output_suffix in {".gt", ".graphml"}
    g.save(sys.argv[3])