From e54e3b3d48f934d3a4d44b9f4ff262f742a4aaf9 Mon Sep 17 00:00:00 2001
From: Josef Engelfrost <josef@josf.se>
Date: Thu, 18 Oct 2018 08:54:23 +0200
Subject: [PATCH] feat(engine-nunjucks): Configurable extension locations; Use
 usePatternlabConfig()

---
 packages/engine-nunjucks/README.md            | 22 ++++-
 .../engine-nunjucks/lib/engine_nunjucks.js    | 97 +++++++++++++------
 2 files changed, 87 insertions(+), 32 deletions(-)

diff --git a/packages/engine-nunjucks/README.md b/packages/engine-nunjucks/README.md
index b85c5e5d7..c5910c5ae 100644
--- a/packages/engine-nunjucks/README.md
+++ b/packages/engine-nunjucks/README.md
@@ -17,16 +17,34 @@ Level of Support is more or less full. Partial calls and lineage hunting are sup
 
 ## Extending the Nunjucks instance
 
-To add custom filters or make customizations to the nunjucks instance, create a file named `patternlab-nunjucks-config.js` in the root of your Pattern Lab project. `patternlab-nunjucks-config.js` should export a function with the Nunjucks environment as parameter.
+To add custom filters or make customizations to the nunjucks instance, add the following to `patternlab-config.json`:
 
+```json
+  {
+    ...
+    "engines": {
+      "nunjucks": {
+        "extensions": [
+          "nunjucks-extensions/*.js"
+        ]
+      }
+    }
+  }
 ```
+
+...or use the default file name: `patternlab-nunjucks-config.js` (in the root of your Pattern Lab project).
+
+Each file providing extensions should export a function with the Nunjucks environment as parameter.
+
+```js
 module.exports = function (env) {
   [YOUR CUSTOM CODE HERE]
 };
 ```
 
 Example: `patternlab-nunjucks-config.js` file that uses lodash and adds three custom filters.
-```
+
+```js
 var _shuffle = require('lodash/shuffle'),
   _take = require('lodash/take');
 
diff --git a/packages/engine-nunjucks/lib/engine_nunjucks.js b/packages/engine-nunjucks/lib/engine_nunjucks.js
index 5cfc4f803..3840cdf9f 100644
--- a/packages/engine-nunjucks/lib/engine_nunjucks.js
+++ b/packages/engine-nunjucks/lib/engine_nunjucks.js
@@ -22,39 +22,10 @@
 
 const fs = require('fs-extra');
 const path = require('path');
-const plPath = process.cwd();
-const plConfig = require(path.join(plPath, 'patternlab-config.json'));
 const nunjucks = require('nunjucks');
 const partialRegistry = [];
 
-// Create Pattern Loader
-// Since Pattern Lab includes are not path based we need a custom loader for Nunjucks.
-function PatternLoader() {}
-
-PatternLoader.prototype.getSource = function(name) {
-  const fullPath = path.resolve(
-    plConfig.paths.source.patterns,
-    partialRegistry[name]
-  );
-  return {
-    src: fs.readFileSync(fullPath, 'utf-8'),
-    path: fullPath,
-    noCache: true
-  };
-};
-
-const env = new nunjucks.Environment(new PatternLoader());
-
-// Load any user Defined configurations
-try {
-  const nunjucksConfig = require(path.join(
-    plPath,
-    'patternlab-nunjucks-config.js'
-  ));
-  if (typeof nunjucksConfig === 'function') {
-    nunjucksConfig(env);
-  }
-} catch (err) {}
+let env;
 
 // Nunjucks Engine
 const engine_nunjucks = {
@@ -153,6 +124,72 @@ const engine_nunjucks = {
     this.spawnFile(config, '_00-head.njk');
     this.spawnFile(config, '_01-foot.njk');
   },
+
+  /**
+   * Accept a Pattern Lab config object from the core and use the settings to
+   * load helpers.
+   *
+   * @param {object} config - the global config object from core
+   */
+  usePatternLabConfig: function(config) {
+    // Create Pattern Loader
+    // Since Pattern Lab includes are not path based we need a custom loader for Nunjucks.
+    function PatternLoader() {}
+
+    PatternLoader.prototype.getSource = function(name) {
+      const fullPath = path.resolve(
+        config.paths.source.patterns,
+        partialRegistry[name]
+      );
+      return {
+        src: fs.readFileSync(fullPath, 'utf-8'),
+        path: fullPath,
+        noCache: true,
+      };
+    };
+
+    env = new nunjucks.Environment(new PatternLoader());
+
+    let extensions;
+
+    try {
+      // Look for helpers in the config
+      extensions = config.engines.nunjucks.helpers;
+
+      if (typeof extensions === 'string') {
+        extensions = [extensions];
+      }
+    } catch (error) {
+      // No defined path(s) found, look in default location
+
+      const configPath = 'patternlab-nunjucks-config.js';
+      if (fs.existsSync(path.join(process.cwd(), configPath))) {
+        extensions = [configPath];
+      }
+    }
+
+    if (extensions) {
+      extensions.forEach(extensionPath => {
+        // Load any user Defined configurations
+        const nunjucksConfigPath = path.join(process.cwd(), extensionPath);
+
+        try {
+          const nunjucksConfig = require(nunjucksConfigPath);
+          if (typeof nunjucksConfig === 'function') {
+            nunjucksConfig(env);
+          } else {
+            console.error(
+              `Failed to load Nunjucks extension: Expected ${extensionPath} to export a function.`
+            );
+          }
+        } catch (err) {
+          console.error(
+            `Failed to load Nunjucks extension ${nunjucksConfigPath}.`
+          );
+        }
+      });
+    }
+  },
 };
 
 module.exports = engine_nunjucks;