1 files changed, 183 insertions, 0 deletions
diff --git a/obmc/utils/pathtree.py b/obmc/utils/pathtree.py
new file mode 100644
index 0000000..221495e
--- /dev/null
+++ b/obmc/utils/pathtree.py
@@ -0,0 +1,183 @@
+# Contributors Listed Below - COPYRIGHT 2016
+# [+] International Business Machines Corp.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
+
+class PathTreeItemIterator(object):
+    def __init__(self, path_tree, subtree, depth):
+        self.path_tree = path_tree
+        self.path = []
+        self.itlist = []
+        self.subtree = ['/'] + filter(bool, subtree.split('/'))
+        self.depth = depth
+        d = path_tree.root
+        for k in self.subtree:
+            try:
+                d = d[k]['children']
+            except KeyError:
+                raise KeyError(subtree)
+        self.it = d.iteritems()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return super(PathTreeItemIterator, self).next()
+
+    def next(self):
+        key, value = self._next()
+        path = self.subtree[0] + '/'.join(self.subtree[1:] + self.path)
+        return path, value.get('data')
+
+    def _next(self):
+        try:
+            while True:
+                x = self.it.next()
+                depth_exceeded = len(self.path) + 1 > self.depth
+                if self.depth and depth_exceeded:
+                    continue
+                self.itlist.append(self.it)
+                self.path.append(x[0])
+                self.it = x[1]['children'].iteritems()
+                break
+
+        except StopIteration:
+            if not self.itlist:
+                raise StopIteration
+
+            self.it = self.itlist.pop()
+            self.path.pop()
+            x = self._next()
+
+        return x
+
+
+class PathTreeKeyIterator(PathTreeItemIterator):
+    def __init__(self, path_tree, subtree, depth):
+        super(PathTreeKeyIterator, self).__init__(path_tree, subtree, depth)
+
+    def next(self):
+        return super(PathTreeKeyIterator, self).next()[0]
+
+
+class PathTree:
+    def __init__(self):
+        self.root = {}
+
+    def _try_delete_parent(self, elements):
+        if len(elements) == 1:
+            return False
+
+        kids = 'children'
+        elements.pop()
+        d = self.root
+        for k in elements[:-1]:
+            d = d[k][kids]
+
+        if 'data' not in d[elements[-1]] and not d[elements[-1]][kids]:
+            del d[elements[-1]]
+            self._try_delete_parent(elements)
+
+    def _get_node(self, key):
+        kids = 'children'
+        elements = ['/'] + filter(bool, key.split('/'))
+        d = self.root
+        for k in elements[:-1]:
+            try:
+                d = d[k][kids]
+            except KeyError:
+                raise KeyError(key)
+
+        return d[elements[-1]]
+
+    def __iter__(self):
+        return self
+
+    def __missing__(self, key):
+        for x in self.iterkeys():
+            if key == x:
+                return False
+        return True
+
+    def __delitem__(self, key):
+        kids = 'children'
+        elements = ['/'] + filter(bool, key.split('/'))
+        d = self.root
+        for k in elements[:-1]:
+            try:
+                d = d[k][kids]
+            except KeyError:
+                raise KeyError(key)
+
+        del d[elements[-1]]
+        self._try_delete_parent(elements)
+
+    def __setitem__(self, key, value):
+        kids = 'children'
+        elements = ['/'] + filter(bool, key.split('/'))
+        d = self.root
+        for k in elements[:-1]:
+            d = d.setdefault(k, {kids: {}})[kids]
+
+        children = d.setdefault(elements[-1], {kids: {}})[kids]
+        d[elements[-1]].update({kids: children, 'data': value})
+
+    def __getitem__(self, key):
+        return self._get_node(key).get('data')
+
+    def setdefault(self, key, default):
+        if not self.get(key):
+            self.__setitem__(key, default)
+
+        return self.__getitem__(key)
+
+    def get(self, key, default=None):
+        try:
+            x = self.__getitem__(key)
+        except KeyError:
+            x = default
+
+        return x
+
+    def get_children(self, key):
+        return [x for x in self._get_node(key)['children'].iterkeys()]
+
+    def demote(self, key):
+        n = self._get_node(key)
+        if 'data' in n:
+            del n['data']
+
+    def keys(self, subtree='/', depth=None):
+        return [x for x in self.iterkeys(subtree, depth)]
+
+    def values(self, subtree='/', depth=None):
+        return [x[1] for x in self.iteritems(subtree, depth)]
+
+    def items(self, subtree='/', depth=None):
+        return [x for x in self.iteritems(subtree, depth)]
+
+    def dataitems(self, subtree='/', depth=None):
+        return [x for x in self.iteritems(subtree, depth)
+                if x[1] is not None]
+
+    def iterkeys(self, subtree='/', depth=None):
+        if not self.root:
+            return {}.iterkeys()
+        return PathTreeKeyIterator(self, subtree, depth)
+
+    def iteritems(self, subtree='/', depth=None):
+        if not self.root:
+            return {}.iteritems()
+        return PathTreeItemIterator(self, subtree, depth)