bookdata/cli/
extract_graph.rs

1//! Graph extraction program.
2use std::convert::From;
3use std::path::PathBuf;
4
5use crate::graph::{load_graph, save_gml};
6use crate::prelude::*;
7
8/// Extract a subgraph.
9#[derive(Args, Debug)]
10#[command(name = "extract-graph")]
11pub struct ExtractGraph {
12    #[arg(long = "graph-file")]
13    graph_file: Option<PathBuf>,
14
15    #[arg(short = 'c', long = "cluster")]
16    cluster: Option<i32>,
17
18    #[arg(long = "output", short = 'o')]
19    out_file: Option<PathBuf>,
20}
21
22impl Command for ExtractGraph {
23    fn exec(&self) -> Result<()> {
24        let path = match &self.graph_file {
25            Some(p) => p.clone(),
26            None => PathBuf::from("book-links/book-graph.mp.zst"),
27        };
28
29        info!("loading graph from {}", path.to_string_lossy());
30        let mut graph = load_graph(path)?;
31
32        if let Some(c) = &self.cluster {
33            graph.retain_nodes(|g, n| {
34                let book = g.node_weight(n).unwrap();
35                book.cluster == *c
36            });
37        }
38
39        info!("filtered graph to {} nodes", graph.node_count());
40
41        if let Some(outf) = &self.out_file {
42            save_gml(&graph, outf)?;
43        }
44
45        Ok(())
46    }
47}