summaryrefslogtreecommitdiffstats
path: root/obmc/utils/pathtree.py
blob: 221495e4bb036299c9ac892b332ab580121c0a65 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
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)
OpenPOWER on IntegriCloud