From 364d692db2de650e6b9603d4803639b75edc1188 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gina=20H=C3=A4u=C3=9Fge?= <gina@octoprint.org>
Date: Fri, 24 Nov 2017 13:07:16 +0100
Subject: [PATCH] Support leaf merging for file extension tree

This allows us to add new extensions to existing entries (e.g. a new
extension for gcode files)
---
 src/octoprint/filemanager/__init__.py | 44 +++++++++++++++++++++++++--
 src/octoprint/util/__init__.py        | 31 +++++++++++++++++--
 2 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/src/octoprint/filemanager/__init__.py b/src/octoprint/filemanager/__init__.py
index 513e91e8..2cacb459 100644
--- a/src/octoprint/filemanager/__init__.py
+++ b/src/octoprint/filemanager/__init__.py
@@ -28,7 +28,12 @@ ContentTypeDetector = namedtuple("ContentTypeDetector", "extensions, detector")
 extensions = dict(
 )
 
-def full_extension_tree():
+_cached_tree = None
+def full_extension_tree(force=False):
+	global _cached_tree
+	if _cached_tree is not None and not force:
+		return _cached_tree
+
 	result = dict(
 		# extensions for 3d model files
 		model=dict(
@@ -40,16 +45,51 @@ def full_extension_tree():
 		)
 	)
 
+	def leaf_merger(a, b):
+		supported_leaf_types = (ContentTypeMapping, ContentTypeDetector, list)
+		if not isinstance(a, supported_leaf_types) or not isinstance(b, supported_leaf_types):
+			raise ValueError()
+
+		if isinstance(a, ContentTypeDetector) and isinstance(b, ContentTypeMapping):
+			raise ValueError()
+
+		if isinstance(a, ContentTypeMapping) and isinstance(b, ContentTypeDetector):
+			raise ValueError()
+
+		a_list = a if isinstance(a, list) else a.extensions
+		b_list = b if isinstance(b, list) else b.extensions
+		merged = a_list + b_list
+
+		content_type = None
+		if isinstance(b, ContentTypeMapping):
+			content_type = b.content_type
+		elif isinstance(a, ContentTypeMapping):
+			content_type = a.content_type
+
+		detector = None
+		if isinstance(b, ContentTypeDetector):
+			detector = b.detector
+		elif isinstance(a, ContentTypeDetector):
+			detector = a.detector
+
+		if content_type is not None:
+			return ContentTypeMapping(merged, content_type)
+		elif detector is not None:
+			return ContentTypeDetector(merged, detector)
+		else:
+			return merged
+
 	extension_tree_hooks = octoprint.plugin.plugin_manager().get_hooks("octoprint.filemanager.extension_tree")
 	for name, hook in extension_tree_hooks.items():
 		try:
 			hook_result = hook()
 			if hook_result is None or not isinstance(hook_result, dict):
 				continue
-			result = octoprint.util.dict_merge(result, hook_result)
+			result = octoprint.util.dict_merge(result, hook_result, leaf_merger=leaf_merger)
 		except:
 			logging.getLogger(__name__).exception("Exception while retrieving additional extension tree entries from hook {name}".format(name=name))
 
+	_cached_tree = result
 	return result
 
 def get_extensions(type, subtree=None):
diff --git a/src/octoprint/util/__init__.py b/src/octoprint/util/__init__.py
index a06dffcc..e4b7932c 100644
--- a/src/octoprint/util/__init__.py
+++ b/src/octoprint/util/__init__.py
@@ -459,7 +459,7 @@ def is_running_from_source():
 	return os.path.isdir(os.path.join(root, "src")) and os.path.isfile(os.path.join(root, "setup.py"))
 
 
-def dict_merge(a, b):
+def dict_merge(a, b, leaf_merger=None):
 	"""
 	Recursively deep-merges two dictionaries.
 
@@ -478,10 +478,24 @@ def dict_merge(a, b):
 	    True
 	    >>> dict_merge(None, None) == dict()
 	    True
+	    >>> def leaf_merger(a, b):
+	    ...     if isinstance(a, list) and isinstance(b, list):
+	    ...         return a + b
+	    ...     raise ValueError()
+	    >>> result = dict_merge(dict(l1=[3, 4], l2=[1], a="a"), dict(l1=[1, 2], l2="foo", b="b"), leaf_merger=leaf_merger)
+	    >>> result.get("l1") == [3, 4, 1, 2]
+	    True
+	    >>> result.get("l2") == "foo"
+	    True
+	    >>> result.get("a") == "a"
+	    True
+	    >>> result.get("b") == "b"
+	    True
 
 	Arguments:
 	    a (dict): The dictionary to merge ``b`` into
 	    b (dict): The dictionary to merge into ``a``
+	    leaf_merger (callable): An optional callable to use to merge leaves (non-dict values)
 
 	Returns:
 	    dict: ``b`` deep-merged into ``a``
@@ -499,9 +513,20 @@ def dict_merge(a, b):
 	result = deepcopy(a)
 	for k, v in b.items():
 		if k in result and isinstance(result[k], dict):
-			result[k] = dict_merge(result[k], v)
+			result[k] = dict_merge(result[k], v, leaf_merger=leaf_merger)
 		else:
-			result[k] = deepcopy(v)
+			merged = None
+			if k in result and callable(leaf_merger):
+				try:
+					merged = leaf_merger(result[k], v)
+				except ValueError:
+					# can't be merged by leaf merger
+					pass
+
+			if merged is None:
+				merged = deepcopy(v)
+
+			result[k] = merged
 	return result