looks good

ImportGraph/Cli.lean

@@ -8,6 +8,7 @@ import Batteries.Lean.IO.Process
 import ImportGraph.CurrentModule
 import ImportGraph.Imports
 import ImportGraph.Lean.Name
+import ImportGraph.Unused
 
 open Cli
 
@@ -17,15 +18,15 @@ open ImportGraph
 /-- Write an import graph, represented as a `NameMap (Array Name)` to the ".dot" graph format.
     Nodes in the `unused` set will be shaded light gray. -/
 def asDotGraph
-    (graph : NameMap (Array Name)) (unused : NameSet := {}) (header := "import_graph") :
+    (graph : NameMap (Array Name)) (used : NameSet := {}) (header := "import_graph") :
     String := Id.run do
   let mut lines := #[s!"digraph \"{header}\" " ++ "{"]
-  -- Add node styling for unused nodes
-  lines := lines.push "  node [style=filled];"
-  -- Add nodes and edges
   for (n, is) in graph do
-    if unused.contains n then
-      lines := lines.push s!"  \"{n}\" [fillcolor=lightgray];"
+    if used.contains n then
+      lines := lines.push s!"  \"{n}\";"
+    else
+      lines := lines.push s!"  \"{n}\" [style=filled, fillcolor=lightgray];"
+    -- Then add edges
     for i in is do
       lines := lines.push s!"  \"{i}\" -> \"{n}\";"
   lines := lines.push "}"
@@ -45,10 +46,10 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
   searchPathRef.set compile_time_search_path%
   let dotFile ← try unsafe withImportModules #[{module := to}] {} (trustLevel := 1024) fun env => do
     let mut graph := env.importGraph
-    let ctx := { options := {}, fileName := "<input>", fileMap := default }
-    let state := { env }
-    let used ← Prod.fst <$> (CoreM.toIO (env.transitivelyRequiredModules to) ctx state)
-    let unused := graph.fold (fun acc n _ => if used.contains n then acc else acc.insert n) NameSet.empty
+    let used ← do
+      let ctx := { options := {}, fileName := "<input>", fileMap := default }
+      let state := { env }
+      Prod.fst <$> (CoreM.toIO (env.transitivelyRequiredModules to) ctx state)
     if let Option.some f := from? then
       graph := graph.downstreamOf (NameSet.empty.insert f)
     if ¬(args.hasFlag "include-deps") then
@@ -65,7 +66,7 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
       graph := graph.filterGraph filterMathlibMeta (replacement := `«Mathlib.Tactics»)
     if args.hasFlag "reduce" then
       graph := graph.transitiveReduction
-    return asDotGraph graph (unused := unused)
+    return asDotGraph graph (used := used)
   catch err =>
     -- TODO: try to build `to` first, so this doesn't happen
     throw <| IO.userError <| s!"{err}\nIf the error above says `unknown package`, " ++

ImportGraph/RequiredModules.lean

@@ -72,6 +72,7 @@ and identifies all modules contains constants which are transitively required by
 def Environment.transitivelyRequiredModules (env : Environment) (module : Name) : CoreM NameSet := do
   let mut requiredModules : NameSet := {}
   for (_, ci) in env.constants.map₁.toList do
+    if !ci.name.isInternal then
     if let some m := env.getModuleFor? ci.name then
       if m = module then
         for r in (← ci.name.transitivelyRequiredModules env) do

ImportGraph/Unused.lean

@@ -0,0 +1,3 @@
+/-!
+This is a dummy file to test the `unused` flag in the import graph.
+-/

ImportGraphTest.lean

@@ -0,0 +1,2 @@
+import ImportGraphTest.Imports
+import ImportGraphTest.Dot

ImportGraphTest/.gitignore

@@ -0,0 +1 @@
+produced.dot

ImportGraphTest/Dot.lean

@@ -0,0 +1,29 @@
+import ImportGraph.Cli
+
+def readFile (path : System.FilePath) : IO String :=
+  IO.FS.readFile path
+
+def runGraphCommand : IO Unit := do
+  let _ ← IO.Process.output {
+    cmd := "lake"
+    args := #["exe", "graph", "--to", "ImportGraph.Cli", "ImportGraphTest/produced.dot"]
+  }
+
+def compareOutputs (expected : String) (actual : String) : IO Bool := do
+  let expectedLines := expected.splitOn "\n" |>.filter (·.trim.length > 0) |>.map String.trim
+  let actualLines := actual.splitOn "\n" |>.filter (·.trim.length > 0) |>.map String.trim
+  pure (expectedLines == actualLines)
+
+/-- info: Test passed: The graph command output matches the expected.dot file. -/
+#guard_msgs in
+#eval show IO Unit from do
+  runGraphCommand
+  let expectedOutput ← readFile "ImportGraphTest/expected.dot"
+  let actualOutput ← readFile "ImportGraphTest/produced.dot"
+  let isEqual ← compareOutputs expectedOutput actualOutput
+  if isEqual then
+    IO.println "Test passed: The graph command output matches the expected.dot file."
+  else
+    IO.println "Test failed: The graph command output does not match the expected.dot file."
+    IO.println s!"Expected:\n{expectedOutput}"
+    IO.println s!"Actual:\n{actualOutput}"

test/Imports.lean → ImportGraphTest/Imports.lean

@@ -22,7 +22,17 @@ ImportGraph.RequiredModules
 #guard_msgs in
 #find_home importTest
 
-open Elab Command in
+open Elab Command
+
+elab "#transitivelyRequiredModules_test" : command => do
+  let env ← getEnv
+  let unused ← liftCoreM <| env.transitivelyRequiredModules `ImportGraph.RequiredModules
+  logInfo s!"{unused.contains `Init.Data.Option.Lemmas}"
+
+/-- info: true -/
+#guard_msgs in
+#transitivelyRequiredModules_test
+
 elab "#my_test" : command => do
   -- functionality of `#redundant_imports`
   let expected := #[`ImportGraph.RequiredModules]

ImportGraphTest/expected.dot

@@ -0,0 +1,13 @@
+digraph "import_graph" {
+  "ImportGraph.Unused" [style=filled, fillcolor=lightgray];
+  "ImportGraph.CurrentModule";
+  "ImportGraph.Lean.Name";
+  "ImportGraph.RequiredModules";
+  "ImportGraph.Imports";
+  "ImportGraph.RequiredModules" -> "ImportGraph.Imports";
+  "ImportGraph.Cli";
+  "ImportGraph.CurrentModule" -> "ImportGraph.Cli";
+  "ImportGraph.Imports" -> "ImportGraph.Cli";
+  "ImportGraph.Lean.Name" -> "ImportGraph.Cli";
+  "ImportGraph.Unused" -> "ImportGraph.Cli";
+}

lakefile.toml

@@ -1,6 +1,6 @@
 name = "importGraph"
 defaultTargets = ["ImportGraph", "graph"]
-testRunner = "test"
+testRunner = "ImportGraphTest"
 
 [[require]]
 name = "Cli"
@@ -24,6 +24,5 @@ root = "Main"
 # Remove this line if you do not need such functionality.
 supportInterpreter = true
 
-[[lean_exe]]
-name = "test"
-srcDir = "scripts"
+[[lean_lib]]
+name = "ImportGraphTest"