feat(ImportGraph): add option to exclude imports from other packages (#6151)

Add an option `lake exe graph --to MyProject.XY --noDep` which removes all nodes that do not start in `MyProject`.

Author: joneugster

ImportGraph/Main.lean

@@ -4,7 +4,9 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
 import Mathlib.Util.Imports
+import Mathlib.Lean.Data.NameMap
 import Mathlib.Lean.IO.Process
+import Mathlib.Lean.Name
 import Cli
 
 /-!
@@ -47,20 +49,24 @@ instance : ParseableType Name where
     else
       String.toName s
 
-open IO.FS IO.Process in
+open IO.FS IO.Process Name in
 /-- Implementation of the import graph command line program. -/
 def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
   let to := match args.flag? "to" with
   | some to => to.as! Name
   | none => `Mathlib -- autodetect the main module from the `lakefile.lean`?
   let from? := match args.flag? "from" with
-  | some to => some <| to.as! Name
+  | some fr => some <| fr.as! Name
   | none => none
   searchPathRef.set compileTimeSearchPath
   let dotFile ← unsafe withImportModules [{module := to}] {} (trustLevel := 1024) fun env => do
     let mut graph := env.importGraph
     if let .some f := from? then
       graph := graph.dependenciesOf (NameSet.empty.insert f)
+    if ¬(args.hasFlag "include-deps") then
+      let p := getModule to
+      graph := graph.filterMap (fun n i =>
+        if p.isPrefixOf n then (i.filter (isPrefixOf p)) else none)
     if args.hasFlag "reduce" then
       graph := graph.transitiveReduction
     return asDotGraph graph
@@ -83,6 +89,7 @@ def graph : Cmd := `[Cli|
     reduce;         "Remove transitively redundant edges."
     to : Name;      "Only show the upstream imports of the specified module."
     "from" : Name;  "Only show the downstream dependencies of the specified module."
+    "include-deps"; "Include used files from other projects (e.g. lake packages)"
 
   ARGS:
     ...outputs : String;  "Filename(s) for the output. " ++

Mathlib.lean

@@ -2089,6 +2089,7 @@ import Mathlib.Init.Quot
 import Mathlib.Init.Set
 import Mathlib.Init.ZeroOne
 import Mathlib.Lean.CoreM
+import Mathlib.Lean.Data.NameMap
 import Mathlib.Lean.Elab.Tactic.Basic
 import Mathlib.Lean.EnvExtension
 import Mathlib.Lean.Exception

Mathlib/Lean/Data/NameMap.lean

@@ -0,0 +1,38 @@
+/-
+Copyright (c) 2023 Jon Eugster. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jon Eugster
+-/
+import Lean.Data.NameMap
+
+/-!
+# Additional functions on `Lean.NameMap`.
+
+We provide `NameMap.filter` and `NameMap.filterMap`.
+-/
+
+namespace Lean.NameMap
+
+/--
+`filter f m` returns the `NameMap` consisting of all
+"`key`/`val`"-pairs in `m` where `f key val` returns `true`.
+-/
+def filter (f : Name → α → Bool) (m : NameMap α) : NameMap α :=
+  m.fold process {}
+where
+  process (r : NameMap α) (n : Name) (i : α) :=
+    if f n i then r.insert n i else r
+
+/--
+`filterMap f m` allows to filter a `NameMap` and simultaneously modify the filtered values.
+
+It takes a function `f : Name → α → Option β` and applies `f name` to the value with key `name`.
+The resulting entries with non-`none` value are collected to form the output `NameMap`.
+-/
+def filterMap (f : Name → α → Option β) (m : NameMap α) : NameMap β :=
+  m.fold process {}
+where
+  process (r : NameMap β) (n : Name) (i : α) :=
+    match f n i with
+    | none => r
+    | some b => r.insert n b

Mathlib/Lean/Name.lean

@@ -10,7 +10,7 @@ import Mathlib.Lean.Expr.Basic
 /-!
 # Additional functions on `Lean.Name`.
 
-We provide `Name.getModule : Name → CoreM (Option Name)`,
+We provide `Name.getModule`,
 and `allNames` and `allNamesByModule`.
 -/
 
@@ -48,3 +48,10 @@ def allNamesByModule (p : Name → Bool) : CoreM (Std.HashMap Name (Array Name))
       | none => return names.insert m #[n]
     else
       return names
+
+/-- Returns the very first part of a name: for `Mathlib.Data.Set.Basic` it returns `Mathlib`. -/
+def getModule (name : Name) (s := "") : Name :=
+  match name with
+    | .anonymous => s
+    | .num _ _ => panic s!"panic in `getModule`: did not expect numerical name: {name}."
+    | .str pre s => getModule pre s