From a57fa986f8c352cb2277bbb23d7fd42b101c4f2c Mon Sep 17 00:00:00 2001
From: Zijian Zhang <doomspec@outlook.com>
Date: Thu, 12 Oct 2023 22:48:34 -0400
Subject: [PATCH] feat: add draw_module_tree()

---
 README.md      |  5 +++++
 moduler/gui.py | 15 ++++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 86860a9..8044fbd 100644
--- a/README.md
+++ b/README.md
@@ -138,6 +138,11 @@ def buggy_foo():
     foo(a=b)
 ```
 
+## Visualize your codebase
+
+You can put the function `draw_module_tree()` in any of your modules to visualize the tree structure of it. Call it in the `__init__.py` of your package is recommended.
+
+
 ## Installation
 
 You can install Moduler by
diff --git a/moduler/gui.py b/moduler/gui.py
index 0a07a9d..07731e5 100644
--- a/moduler/gui.py
+++ b/moduler/gui.py
@@ -1,5 +1,19 @@
 from moduler import Struct
 import plotly.graph_objects as go
+import inspect
+from moduler.core import build_module_tree
+
+def draw_module_tree():
+    # Get current module by inspecting the stack
+    frame = inspect.stack()[1]
+    module = inspect.getmodule(frame[0])
+    struct = build_module_tree(module)
+    draw_treemap(struct)
+
+
+"""
+## Draw treemap
+"""
 
 def draw_treemap(struct: Struct):
     ids, labels, parents, texts = extract_tree_ingredients(struct)
@@ -97,7 +111,6 @@ def hypenate_texts(texts: str, line_width=40):
 
 
 if __name__ == '__main__':
-    from moduler.core import build_module_tree
     import moduler
     struct = build_module_tree(moduler)
     draw_treemap(struct)
\ No newline at end of file