From ce28b637abbd7f663bf03bbe1d17232b56a5f61f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gina=20H=C3=A4u=C3=9Fge?= <osd@foosel.net>
Date: Mon, 29 Aug 2016 17:50:14 +0200
Subject: [PATCH] Better solution for access to .md files when running from
 source

The SelectedFileSystemLoader doesn't collect everything and the
kitchen sink from disk...
---
 src/octoprint/server/__init__.py | 26 ++------------------------
 src/octoprint/util/jinja.py      | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/src/octoprint/server/__init__.py b/src/octoprint/server/__init__.py
index 796e2ab4..476000dd 100644
--- a/src/octoprint/server/__init__.py
+++ b/src/octoprint/server/__init__.py
@@ -738,32 +738,10 @@ class Server(object):
 
 		loaders = [app.jinja_loader, filesystem_loader]
 		if octoprint.util.is_running_from_source():
+			from octoprint.util.jinja import SelectedFileSystemLoader
 			root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
 			allowed = ["AUTHORS.md", "CHANGELOG.md", "SUPPORTERS.md", "THIRDPARTYLICENSES.md"]
-
-			class SourceRootFilesystemLoader(jinja2.FileSystemLoader):
-				def __init__(self, template_filter, prefix, *args, **kwargs):
-					jinja2.FileSystemLoader.__init__(self, *args, **kwargs)
-					self._filter = template_filter
-					if not prefix.endswith("/"):
-						prefix += "/"
-					self._prefix = prefix
-
-				def get_source(self, environment, template):
-					if not template.startswith(self._prefix):
-						raise jinja2.TemplateNotFound(template)
-
-					template = template[len(self._prefix):]
-					if not self._filter(template):
-						raise jinja2.TemplateNotFound(template)
-
-					return jinja2.FileSystemLoader.get_source(self, environment, template)
-
-				def list_templates(self):
-					templates = jinja2.FileSystemLoader.list_templates(self)
-					return map(lambda t: self._prefix + t, filter(self._filter, templates))
-
-			loaders.append(SourceRootFilesystemLoader(lambda t: t in allowed, "_data/", root))
+			loaders.append(SelectedFileSystemLoader(root, allowed, prefix="_data/"))
 
 		jinja_loader = jinja2.ChoiceLoader(loaders)
 		app.jinja_loader = jinja_loader
diff --git a/src/octoprint/util/jinja.py b/src/octoprint/util/jinja.py
index 12d70461..650c28cf 100644
--- a/src/octoprint/util/jinja.py
+++ b/src/octoprint/util/jinja.py
@@ -54,6 +54,32 @@ class FilteredFileSystemLoader(FileSystemLoader):
 		return all(filter_results)
 
 
+class SelectedFileSystemLoader(FileSystemLoader):
+	def __init__(self, searchpath, files, prefix=None, **kwargs):
+		FileSystemLoader.__init__(self, searchpath, **kwargs)
+		self.files = files
+
+		if prefix is not None and not prefix.endswith("/"):
+			prefix += "/"
+		self.prefix = prefix
+
+	def get_source(self, environment, template):
+		if not template.startswith(self.prefix):
+			raise TemplateNotFound(template)
+
+		template = template[len(self.prefix):]
+		if not template in self.files:
+			raise TemplateNotFound(template)
+
+		return FileSystemLoader.get_source(self, environment, template)
+
+	def list_templates(self):
+		return [self._prefixed(f) for f in self.files if any(map(lambda folder: os.path.exists(os.path.join(folder, f)), self.searchpath))]
+
+	def _prefixed(self, name):
+		return self.prefix + name if self.prefix else name
+
+
 def get_all_template_paths(loader):
 	def walk_folder(folder):
 		files = []
@@ -71,6 +97,12 @@ def get_all_template_paths(loader):
 				result += walk_folder(folder)
 			return filter(loader.path_filter, result)
 
+		elif isinstance(loader, SelectedFileSystemLoader):
+			result = []
+			for folder in loader.searchpath:
+				result += filter(lambda x: os.path.exists(f), [os.path.join(folder, f) for f in loader.files])
+			return result
+
 		elif isinstance(loader, FileSystemLoader):
 			result = []
 			for folder in loader.searchpath: