From eac85914a69e8359aca3dd737ff1d6026419f79b Mon Sep 17 00:00:00 2001
From: Patrik Henningsson <patrik.henningsson@gmail.com>
Date: Sun, 5 Feb 2017 21:45:27 +0100
Subject: [PATCH] Add support for dependencyFilter function

---
 README.md   |  1 +
 lib/api.js  |  3 ++-
 lib/tree.js | 19 +++++++++++++-----
 test/api.js | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 17a7c3b8..643b2ffb 100644
--- a/README.md
+++ b/README.md
@@ -169,6 +169,7 @@ Property | Type | Default | Description
 `graphVizOptions` | Object | false | Custom GraphViz [options](http://www.graphviz.org/content/attrs)
 `graphVizPath` | String | null | Custom GraphViz path
 `detectiveOptions` | Object | false | Custom `detective` options for [dependency-tree](https://github.com/dependents/node-dependency-tree)
+`dependencyFilter` | Function | false | Function called with a dependency filepath (exclude substree by returning false)
 
 > Note that when running the CLI it's possible to use a runtime configuration file. The config should placed in `.madgerc` in your project or home folder. Look [here](https://github.com/dominictarr/rc#standards) for alternative locations for the file. Here's an example:
 
diff --git a/lib/api.js b/lib/api.js
index 64c5f67e..fc52405b 100644
--- a/lib/api.js
+++ b/lib/api.js
@@ -22,7 +22,8 @@ const defaultConfig = {
 	cyclicNodeColor: '#ff6c60',
 	edgeColor: '#757575',
 	graphVizOptions: false,
-	graphVizPath: false
+	graphVizPath: false,
+	dependencyFilter: false
 };
 
 class Madge {
diff --git a/lib/tree.js b/lib/tree.js
index b5af3512..245031c3 100644
--- a/lib/tree.js
+++ b/lib/tree.js
@@ -81,7 +81,7 @@ class Tree {
 						}
 
 						walk.sync(srcPath, (filePath, stat) => {
-							if (!this.filterPath(filePath) || !stat.isFile()) {
+							if (this.isNpmPath(filePath) || !stat.isFile()) {
 								return;
 							}
 
@@ -117,7 +117,16 @@ class Tree {
 				requireConfig: this.config.requireConfig,
 				webpackConfig: this.config.webpackConfig,
 				visited: visited,
-				filter: this.filterPath.bind(this),
+				filter: (dependencyFilePath, traversedFilePath) => {
+					let dependencyFilterRes = true;
+					const isNpmPath = this.isNpmPath(dependencyFilePath);
+
+					if (this.config.dependencyFilter) {
+						dependencyFilterRes = this.config.dependencyFilter(dependencyFilePath, traversedFilePath, this.baseDir);
+					}
+
+					return !isNpmPath && (dependencyFilterRes || dependencyFilterRes === undefined);
+				},
 				detective: this.config.detectiveOptions,
 				nonExistent: nonExistent
 			}));
@@ -184,12 +193,12 @@ class Tree {
 	}
 
 	/**
-	 * Filter out some paths from found files
+	 * Check if path is from NPM folder
 	 * @param  {String} path
 	 * @return {Boolean}
 	 */
-	filterPath(path) {
-		return this.config.includeNpm || path.indexOf('node_modules') < 0;
+	isNpmPath(path) {
+		return path.indexOf('node_modules') >= 0;
 	}
 
 	/**
diff --git a/test/api.js b/test/api.js
index 442b208e..5d35db75 100644
--- a/test/api.js
+++ b/test/api.js
@@ -100,6 +100,64 @@ describe('API', () => {
 		}).catch(done);
 	});
 
+	describe('dependencyFilter', () => {
+		it('will stop traversing when returning false', (done) => {
+			madge(__dirname + '/cjs/a.js', {
+				dependencyFilter: () => {
+					return false;
+				}
+			}).then((res) => {
+				res.obj().should.eql({
+					a: []
+				});
+				done();
+			}).catch(done);
+		});
+
+		it('will not stop traversing when not returning anything', (done) => {
+			madge(__dirname + '/cjs/a.js', {
+				dependencyFilter: () => {}
+			}).then((res) => {
+				res.obj().should.eql({
+					a: ['b', 'c'],
+					b: ['c'],
+					c: []
+				});
+				done();
+			}).catch(done);
+		});
+
+		it('will pass arguments to the function', (done) => {
+			let counter = 0;
+
+			madge(__dirname + '/cjs/a.js', {
+				dependencyFilter: (dependencyFilePath, traversedFilePath, baseDir) => {
+					if (counter === 0) {
+						dependencyFilePath.should.match(/test\/cjs\/b\.js$/);
+						traversedFilePath.should.match(/test\/cjs\/a\.js$/);
+						baseDir.should.match(/test\/cjs$/);
+					}
+
+					if (counter === 1) {
+						dependencyFilePath.should.match(/test\/cjs\/c\.js$/);
+						traversedFilePath.should.match(/test\/cjs\/a\.js$/);
+						baseDir.should.match(/test\/cjs$/);
+					}
+
+					if (counter === 2) {
+						dependencyFilePath.should.match(/test\/cjs\/c\.js$/);
+						traversedFilePath.should.match(/test\/cjs\/b\.js$/);
+						baseDir.should.match(/test\/cjs$/);
+					}
+
+					counter++;
+				}
+			}).then(() => {
+				done();
+			}).catch(done);
+		});
+	});
+
 	describe('obj()', () => {
 		it('returns dependency object', (done) => {
 			madge(__dirname + '/cjs/a.js').then((res) => {