obmc/utils/pathtree.py

# 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)