From 104e0d1ecdba1f43341c8c0823e4ffbd15668084 Mon Sep 17 00:00:00 2001
From: Keijo Kapp <keijo.kapp@gmail.com>
Date: Sun, 12 Jun 2022 12:14:16 +0300
Subject: [PATCH] Added depth limit

---
 README.md   |  1 +
 bin/cli.js  | 14 ++++++++++++
 lib/api.js  |  1 +
 lib/tree.js | 61 +++++++++++++++++++++++++++++++++++++----------------
 4 files changed, 59 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 9cbc5ce2..1ca0b2d9 100644
--- a/README.md
+++ b/README.md
@@ -228,6 +228,7 @@ Property | Type | Default | Description
 `requireConfig` | String | null | RequireJS config for resolving aliased modules
 `webpackConfig` | String | null | Webpack config for resolving aliased modules
 `tsConfig` | String\|Object | null | TypeScript config for resolving aliased modules - Either a path to a tsconfig file or an object containing the config
+`depth` | Number | null | Maximum dependency depth from source files to display
 `layout` | String | dot | Layout to use in the graph
 `rankdir` | String | LR | Sets the [direction](https://graphviz.gitlab.io/_pages/doc/info/attrs.html#d:rankdir) of the graph layout
 `fontName` | String | Arial | Font name to use in the graph
diff --git a/bin/cli.js b/bin/cli.js
index d8f14cd7..d8451d5d 100755
--- a/bin/cli.js
+++ b/bin/cli.js
@@ -31,6 +31,7 @@ program
 	.option('--require-config <file>', 'path to RequireJS config')
 	.option('--webpack-config <file>', 'path to webpack config')
 	.option('--ts-config <file>', 'path to typescript config')
+	.option('--depth <integer>', 'maximum depth from source files to draw')
 	.option('--include-npm', 'include shallow NPM modules', false)
 	.option('--no-color', 'disable color in output and image', false)
 	.option('--no-spinner', 'disable progress spinner', false)
@@ -118,6 +119,19 @@ if (config.tsConfig) {
 	config.tsConfig = obj.raw;
 }
 
+if ('depth' in program) {
+	config.depth = program.depth;
+}
+
+if ('depth' in config) {
+	config.depth = Number(config.depth);
+
+	if (!Number.isInteger(config.depth) || config.depth < 0) {
+		console.log('%s %s', chalk.red('✖'), 'Invalid depth');
+		process.exit(1);
+	}
+}
+
 if (program.includeNpm) {
 	config.includeNpm = program.includeNpm;
 }
diff --git a/lib/api.js b/lib/api.js
index 76a989ee..a7b5d43e 100644
--- a/lib/api.js
+++ b/lib/api.js
@@ -13,6 +13,7 @@ const defaultConfig = {
 	requireConfig: null,
 	webpackConfig: null,
 	tsConfig: null,
+	depth: null,
 	rankdir: 'LR',
 	layout: 'dot',
 	fontName: 'Arial',
diff --git a/lib/tree.js b/lib/tree.js
index 2db6de2a..62c7ec7b 100644
--- a/lib/tree.js
+++ b/lib/tree.js
@@ -114,10 +114,6 @@ class Tree {
 		const pathCache = {};
 
 		files.forEach((file) => {
-			if (visited[file]) {
-				return;
-			}
-
 			Object.assign(depTree, dependencyTree({
 				filename: file,
 				directory: this.baseDir,
@@ -148,7 +144,7 @@ class Tree {
 			}));
 		});
 
-		let tree = this.convertTree(depTree, {}, pathCache, npmPaths);
+		let tree = this.convertTree(visited, depTree, this.config.depth);
 
 		for (const npmKey in npmPaths) {
 			const id = this.processPath(npmKey, pathCache);
@@ -171,27 +167,56 @@ class Tree {
 	/**
 	 * Convert deep tree produced by dependency-tree to a
 	 * shallow (one level deep) tree used by madge.
-	 * @param  {Object} depTree
+	 * @param  {Object} modules
 	 * @param  {Object} tree
-	 * @param  {Object} pathCache
+	 * @param  {number?} depthLimit
 	 * @return {Object}
 	 */
-	convertTree(depTree, tree, pathCache) {
-		for (const key in depTree) {
-			const id = this.processPath(key, pathCache);
-
-			if (!tree[id]) {
-				tree[id] = [];
-
-				for (const dep in depTree[key]) {
-					tree[id].push(this.processPath(dep, pathCache));
+	convertTree(modules, tree, depthLimit) {
+		const self = this;
+		const depths = {};
+		const deepDependencies = {};
+
+		function calculateDepths(tree, depth) {
+			if (depth <= depthLimit) {
+				for (const dependency in tree) {
+					depths[dependency] = true;
+					calculateDepths(modules[dependency], depth + 1);
 				}
+			}
+		}
+
+		function getDeepDependencies(dependency) {
+			if (deepDependencies[dependency] === null) {
+				return [];
+			}
 
-				this.convertTree(depTree[key], tree, pathCache);
+			if (!(dependency in deepDependencies)) {
+				deepDependencies[dependency] = null;
+				deepDependencies[dependency] = [...new Set(Object.keys(modules[dependency]).flatMap(
+					(dependency) => dependency in depths ? [dependency] : getDeepDependencies(dependency)
+				))];
 			}
+
+			return deepDependencies[dependency];
+		}
+
+		const pathCache = {};
+		const result = {};
+
+		if (!Number.isInteger(depthLimit)) {
+			Object.entries(modules).forEach(([module, dependencies]) => {
+				result[self.processPath(module, pathCache)] = Object.keys(dependencies).map((dependency) => self.processPath(dependency, pathCache));
+			});
+		} else {
+			calculateDepths(tree, 0);
+
+			Object.keys(depths).forEach((module) => {
+				result[self.processPath(module, pathCache)] = getDeepDependencies(module).map((dependency) => self.processPath(dependency, pathCache));
+			});
 		}
 
-		return tree;
+		return result;
 	}
 
 	/**