From f5e1b54ffe24169ba78ac357458a222d29606c98 Mon Sep 17 00:00:00 2001
From: Charles Marsh <charlie@springdisc.com>
Date: Thu, 11 Aug 2022 11:29:42 -0400
Subject: [PATCH] Create a dependency graph

---
 resources/test/src/if_tuple.py |   6 ++
 src/bin/generate_graph.rs      | 109 +++++++++++++++++++++++++++++++++
 src/bin/rust_python_linter.rs  |   4 +-
 src/lib.rs                     |   2 +-
 src/linter.rs                  |   4 +-
 5 files changed, 121 insertions(+), 4 deletions(-)
 create mode 100644 src/bin/generate_graph.rs

diff --git a/resources/test/src/if_tuple.py b/resources/test/src/if_tuple.py
index 501a7f8ad2..a68044a651 100644
--- a/resources/test/src/if_tuple.py
+++ b/resources/test/src/if_tuple.py
@@ -1,3 +1,9 @@
+import os
+from multiprocessing import pool as fool
+from spr_platform import module
+from spr_platform.sub_module import other
+import spr_platform.sub_module
+
 if (1, 2):
     pass
 
diff --git a/src/bin/generate_graph.rs b/src/bin/generate_graph.rs
new file mode 100644
index 0000000000..bf982f5620
--- /dev/null
+++ b/src/bin/generate_graph.rs
@@ -0,0 +1,109 @@
+use std::path::PathBuf;
+
+use anyhow::Result;
+use rustpython_parser::ast::{Stmt, StmtKind, Suite};
+use rustpython_parser::parser;
+
+use ::rust_python_linter::fs;
+use ::rust_python_linter::visitor::{walk_stmt, Visitor};
+
+#[allow(dead_code)]
+#[derive(Debug)]
+struct ModuleImport {
+    module_name: Option<String>,
+    remote_name: Option<String>,
+    local_name: Option<String>,
+    lineno: usize,
+    pragma: usize,
+}
+
+#[derive(Default)]
+struct ImportVisitor {
+    imports: Vec<ModuleImport>,
+}
+
+// Inspired by: https://github.com/blais/snakefood/blob/f902c9a099f7c5bb75154a747bf098259211025d/lib/python/snakefood/find.py#L241
+impl Visitor for ImportVisitor {
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        match &stmt.node {
+            StmtKind::Import { names } => {
+                for alias in names {
+                    self.imports.push(ModuleImport {
+                        module_name: Some(alias.name.clone()),
+                        remote_name: None,
+                        local_name: Some(
+                            alias.asname.clone().unwrap_or_else(|| alias.name.clone()),
+                        ),
+                        lineno: stmt.location.row(),
+                        pragma: 0,
+                    })
+                }
+            }
+            StmtKind::ImportFrom {
+                module,
+                names,
+                level,
+            } => {
+                if let Some(module_name) = module {
+                    if module_name == "__future__" {
+                        return;
+                    }
+                }
+
+                for alias in names {
+                    if alias.name == "*" {
+                        self.imports.push(ModuleImport {
+                            module_name: module.clone(),
+                            remote_name: None,
+                            local_name: None,
+                            lineno: stmt.location.row(),
+                            pragma: *level,
+                        })
+                    } else {
+                        self.imports.push(ModuleImport {
+                            module_name: module.clone(),
+                            remote_name: Some(alias.name.clone()),
+                            local_name: Some(
+                                alias.asname.clone().unwrap_or_else(|| alias.name.clone()),
+                            ),
+                            lineno: stmt.location.row(),
+                            pragma: *level,
+                        })
+                    }
+                }
+            }
+            _ => {}
+        }
+        walk_stmt(self, stmt);
+    }
+}
+
+fn collect_imports(python_ast: &Suite) -> Vec<ModuleImport> {
+    python_ast
+        .iter()
+        .flat_map(|stmt| {
+            let mut visitor: ImportVisitor = Default::default();
+            visitor.visit_stmt(stmt);
+            visitor.imports
+        })
+        .collect()
+}
+
+fn main() -> Result<()> {
+    // What else is required here? Map from modules to files.
+    let files = fs::iter_python_files(&PathBuf::from("resources/test/src"));
+    for entry in files {
+        // Read the file from disk.
+        let contents = fs::read_file(entry.path())?;
+
+        // Run the parser.
+        let python_ast = parser::parse_program(&contents)?;
+
+        // Collect imports.
+        let imports = collect_imports(&python_ast);
+        for import in imports {
+            println!("{} imports: {:?}", entry.path().to_string_lossy(), import)
+        }
+    }
+    Ok(())
+}
diff --git a/src/bin/rust_python_linter.rs b/src/bin/rust_python_linter.rs
index 573218a458..500d03a3a6 100644
--- a/src/bin/rust_python_linter.rs
+++ b/src/bin/rust_python_linter.rs
@@ -10,7 +10,7 @@ use notify::{raw_watcher, RecursiveMode, Watcher};
 use rayon::prelude::*;
 use walkdir::DirEntry;
 
-use ::rust_python_linter::fs::iter_python_files;
+use ::rust_python_linter::fs;
 use ::rust_python_linter::linter::check_path;
 use ::rust_python_linter::logging::set_up_logging;
 use ::rust_python_linter::message::Message;
@@ -33,7 +33,7 @@ struct Cli {
 fn run_once(files: &[PathBuf], cache: bool) -> Result<Vec<Message>> {
     // Collect all the files to check.
     let start = Instant::now();
-    let files: Vec<DirEntry> = files.iter().flat_map(iter_python_files).collect();
+    let files: Vec<DirEntry> = files.iter().flat_map(fs::iter_python_files).collect();
     let duration = start.elapsed();
     debug!("Identified files to lint in: {:?}", duration);
 
diff --git a/src/lib.rs b/src/lib.rs
index e15e1a3d7c..536acd9124 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,4 +7,4 @@ pub mod linter;
 pub mod logging;
 pub mod message;
 mod settings;
-mod visitor;
+pub mod visitor;
diff --git a/src/linter.rs b/src/linter.rs
index 89b2707d0f..2a939ae737 100644
--- a/src/linter.rs
+++ b/src/linter.rs
@@ -19,8 +19,10 @@ pub fn check_path(path: &Path, mode: &cache::Mode) -> Result<Vec<Message>> {
     // Read the file from disk.
     let contents = fs::read_file(path)?;
 
-    // Run the linter.
+    // Run the parser.
     let python_ast = parser::parse_program(&contents)?;
+
+    // Run the linter.
     let messages: Vec<Message> = check_ast(&python_ast)
         .into_iter()
         .chain(check_lines(&contents))