1# 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) 184 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 198