blob: a13c6cb8a4cbc5345ed8a5cf95ee7f80444e262e [file] [log] [blame]
Brad Bishop8ffe1e42016-02-11 16:15:40 -05001# Contributors Listed Below - COPYRIGHT 2016
2# [+] International Business Machines Corp.
3#
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
14# implied. See the License for the specific language governing
15# permissions and limitations under the License.
16
17
18class PathTreeItemIterator(object):
19 def __init__(self, path_tree, subtree, depth):
20 self.path_tree = path_tree
21 self.path = []
22 self.itlist = []
23 self.subtree = ['/'] + filter(bool, subtree.split('/'))
24 self.depth = depth
25 d = path_tree.root
26 for k in self.subtree:
27 try:
28 d = d[k]['children']
29 except KeyError:
30 raise KeyError(subtree)
31 self.it = d.iteritems()
32
33 def __iter__(self):
34 return self
35
36 def __next__(self):
37 return super(PathTreeItemIterator, self).next()
38
39 def next(self):
40 key, value = self._next()
41 path = self.subtree[0] + '/'.join(self.subtree[1:] + self.path)
42 return path, value.get('data')
43
44 def _next(self):
45 try:
46 while True:
47 x = self.it.next()
48 depth_exceeded = len(self.path) + 1 > self.depth
49 if self.depth and depth_exceeded:
50 continue
51 self.itlist.append(self.it)
52 self.path.append(x[0])
53 self.it = x[1]['children'].iteritems()
54 break
55
56 except StopIteration:
57 if not self.itlist:
58 raise StopIteration
59
60 self.it = self.itlist.pop()
61 self.path.pop()
62 x = self._next()
63
64 return x
65
66
67class PathTreeKeyIterator(PathTreeItemIterator):
68 def __init__(self, path_tree, subtree, depth):
69 super(PathTreeKeyIterator, self).__init__(path_tree, subtree, depth)
70
71 def next(self):
72 return super(PathTreeKeyIterator, self).next()[0]
73
74
75class PathTree:
76 def __init__(self):
77 self.root = {}
78
79 def _try_delete_parent(self, elements):
80 if len(elements) == 1:
81 return False
82
83 kids = 'children'
84 elements.pop()
85 d = self.root
86 for k in elements[:-1]:
87 d = d[k][kids]
88
89 if 'data' not in d[elements[-1]] and not d[elements[-1]][kids]:
90 del d[elements[-1]]
91 self._try_delete_parent(elements)
92
93 def _get_node(self, key):
94 kids = 'children'
95 elements = ['/'] + filter(bool, key.split('/'))
96 d = self.root
97 for k in elements[:-1]:
98 try:
99 d = d[k][kids]
100 except KeyError:
101 raise KeyError(key)
102
103 return d[elements[-1]]
104
105 def __iter__(self):
106 return self
107
108 def __missing__(self, key):
109 for x in self.iterkeys():
110 if key == x:
111 return False
112 return True
113
114 def __delitem__(self, key):
115 kids = 'children'
116 elements = ['/'] + filter(bool, key.split('/'))
117 d = self.root
118 for k in elements[:-1]:
119 try:
120 d = d[k][kids]
121 except KeyError:
122 raise KeyError(key)
123
124 del d[elements[-1]]
125 self._try_delete_parent(elements)
126
127 def __setitem__(self, key, value):
128 kids = 'children'
129 elements = ['/'] + filter(bool, key.split('/'))
130 d = self.root
131 for k in elements[:-1]:
132 d = d.setdefault(k, {kids: {}})[kids]
133
134 children = d.setdefault(elements[-1], {kids: {}})[kids]
135 d[elements[-1]].update({kids: children, 'data': value})
136
137 def __getitem__(self, key):
138 return self._get_node(key).get('data')
139
140 def setdefault(self, key, default):
141 if not self.get(key):
142 self.__setitem__(key, default)
143
144 return self.__getitem__(key)
145
146 def get(self, key, default=None):
147 try:
148 x = self.__getitem__(key)
149 except KeyError:
150 x = default
151
152 return x
153
154 def get_children(self, key):
155 return [x for x in self._get_node(key)['children'].iterkeys()]
156
157 def demote(self, key):
158 n = self._get_node(key)
159 if 'data' in n:
160 del n['data']
161
162 def keys(self, subtree='/', depth=None):
163 return [x for x in self.iterkeys(subtree, depth)]
164
165 def values(self, subtree='/', depth=None):
166 return [x[1] for x in self.iteritems(subtree, depth)]
167
168 def items(self, subtree='/', depth=None):
169 return [x for x in self.iteritems(subtree, depth)]
170
171 def dataitems(self, subtree='/', depth=None):
172 return [x for x in self.iteritems(subtree, depth)
173 if x[1] is not None]
174
175 def iterkeys(self, subtree='/', depth=None):
176 if not self.root:
177 return {}.iterkeys()
178 return PathTreeKeyIterator(self, subtree, depth)
179
180 def iteritems(self, subtree='/', depth=None):
181 if not self.root:
182 return {}.iteritems()
183 return PathTreeItemIterator(self, subtree, depth)
Brad Bishop4e601a02016-04-01 14:29:46 -0400184
185 def dumpd(self, subtree='/'):
186 result = {}
187 d = result
188
189 for k, v in self.iteritems(subtree):
190 elements = ['/'] + filter(bool, k.split('/'))
191 d = result
192 for k in elements:
193 d = d.setdefault(k, {})
194 if v is not None:
195 d.update(v)
196
197 return result