Add command to visualize graph

Author: hendrikvanantwerpen

tree-sitter-stack-graphs/src/cli.rs

@@ -70,6 +70,7 @@ pub mod query;
 pub mod status;
 pub mod test;
 mod util;
+pub mod visualize;
 
 pub mod path_loading {
     use std::path::PathBuf;
@@ -87,6 +88,7 @@ pub mod path_loading {
     use crate::cli::r#match::MatchArgs;
     use crate::cli::status::StatusArgs;
     use crate::cli::test::TestArgs;
+    use crate::cli::visualize::VisualizeArgs;
 
     use super::database::DatabaseArgs;
 
@@ -102,6 +104,7 @@ pub mod path_loading {
         Query(Query),
         Status(Status),
         Test(Test),
+        Visualize(Visualize),
     }
 
     impl Subcommands {
@@ -117,6 +120,7 @@ pub mod path_loading {
                 Self::Query(cmd) => cmd.run(default_db_path),
                 Self::Status(cmd) => cmd.run(default_db_path),
                 Self::Test(cmd) => cmd.run(),
+                Self::Visualize(cmd) => cmd.run(default_db_path),
             }
         }
     }
@@ -269,6 +273,22 @@ pub mod path_loading {
             self.test_args.run(loader)
         }
     }
+
+    /// Visualize command
+    #[derive(clap::Parser)]
+    pub struct Visualize {
+        #[clap(flatten)]
+        db_args: DatabaseArgs,
+        #[clap(flatten)]
+        visualize_args: VisualizeArgs,
+    }
+
+    impl Visualize {
+        pub fn run(self, default_db_path: PathBuf) -> anyhow::Result<()> {
+            let db_path = self.db_args.get_or(default_db_path);
+            self.visualize_args.run(&db_path)
+        }
+    }
 }
 
 pub mod provided_languages {
@@ -287,6 +307,7 @@ pub mod provided_languages {
     use crate::cli::r#match::MatchArgs;
     use crate::cli::status::StatusArgs;
     use crate::cli::test::TestArgs;
+    use crate::cli::visualize::VisualizeArgs;
     use crate::loader::LanguageConfiguration;
 
     use super::database::DatabaseArgs;
@@ -303,6 +324,7 @@ pub mod provided_languages {
         Query(Query),
         Status(Status),
         Test(Test),
+        Visualize(Visualize),
     }
 
     impl Subcommands {
@@ -322,6 +344,7 @@ pub mod provided_languages {
                 Self::Query(cmd) => cmd.run(default_db_path),
                 Self::Status(cmd) => cmd.run(default_db_path),
                 Self::Test(cmd) => cmd.run(configurations),
+                Self::Visualize(cmd) => cmd.run(default_db_path),
             }
         }
     }
@@ -482,4 +505,20 @@ pub mod provided_languages {
             self.test_args.run(loader)
         }
     }
+
+    /// Visualize command
+    #[derive(clap::Parser)]
+    pub struct Visualize {
+        #[clap(flatten)]
+        db_args: DatabaseArgs,
+        #[clap(flatten)]
+        visualize_args: VisualizeArgs,
+    }
+
+    impl Visualize {
+        pub fn run(self, default_db_path: PathBuf) -> anyhow::Result<()> {
+            let db_path = self.db_args.get_or(default_db_path);
+            self.visualize_args.run(&db_path)
+        }
+    }
 }

tree-sitter-stack-graphs/src/cli/visualize.rs

@@ -0,0 +1,53 @@
+// -*- coding: utf-8 -*-
+// ------------------------------------------------------------------------------------------------
+// Copyright © 2023, stack-graphs authors.
+// Licensed under either of Apache License, Version 2.0, or MIT license, at your option.
+// Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
+// ------------------------------------------------------------------------------------------------
+
+use clap::Args;
+use clap::ValueHint;
+use stack_graphs::serde::NoFilter;
+use stack_graphs::storage::SQLiteReader;
+use stack_graphs::NoCancellation;
+use std::path::Path;
+use std::path::PathBuf;
+
+/// Visualize database
+#[derive(Args)]
+pub struct VisualizeArgs {
+    /// Source file or directory paths.
+    #[clap(
+        value_name = "SOURCE_PATH",
+        value_hint = ValueHint::AnyPath,
+    )]
+    pub source_paths: Vec<PathBuf>,
+
+    #[clap(
+        long,
+        short = 'o',
+        value_name = "OUTPUT_PATH",
+        value_hint = ValueHint::AnyPath,
+        default_value = "stack-graph.html",
+    )]
+    pub output: PathBuf,
+}
+
+impl VisualizeArgs {
+    pub fn run(self, db_path: &Path) -> anyhow::Result<()> {
+        let cancellation_flag = &NoCancellation;
+        let mut db = SQLiteReader::open(&db_path)?;
+        for source_path in &self.source_paths {
+            let source_path = source_path.canonicalize()?;
+            db.load_graph_for_file_or_directory(&source_path, cancellation_flag)?;
+        }
+        let (graph, partials, db) = db.get();
+        let html = graph.to_html_string("stack-graph", partials, db, &NoFilter)?;
+        if let Some(dir) = self.output.parent() {
+            std::fs::create_dir_all(dir)?;
+        }
+        std::fs::write(&self.output, html)?;
+        println!("Visualization at {}", self.output.display());
+        Ok(())
+    }
+}