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 = ['/'] + list(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 # TODO: openbmc/openbmc#2994 remove python 2 support 32 try: # python 2 33 self.it = d.iteritems() 34 except AttributeError: # python 3 35 self.it = iter(d.items()) 36 37 def __iter__(self): 38 return self 39 40 # TODO: openbmc/openbmc#2994 remove python 2 support 41 # python 2 42 def next(self): 43 key, value = self._next() 44 path = self.subtree[0] + '/'.join(self.subtree[1:] + self.path) 45 return path, value.get('data') 46 47 # python 3 48 import sys 49 if sys.version_info[0] > 2: 50 __next__ = next 51 52 def _next(self): 53 try: 54 while True: 55 x = next(self.it) 56 if self.depth: 57 if len(self.path) + 1 > self.depth: 58 continue 59 self.itlist.append(self.it) 60 self.path.append(x[0]) 61 # TODO: openbmc/openbmc#2994 remove python 2 support 62 try: # python 2 63 self.it = x[1]['children'].iteritems() 64 except AttributeError: # python 3 65 self.it = iter(x[1]['children'].items()) 66 break 67 68 except StopIteration: 69 if not self.itlist: 70 raise StopIteration 71 72 self.it = self.itlist.pop() 73 self.path.pop() 74 x = self._next() 75 76 return x 77 78 79class PathTreeKeyIterator(PathTreeItemIterator): 80 def __init__(self, path_tree, subtree, depth): 81 super(PathTreeKeyIterator, self).__init__(path_tree, subtree, depth) 82 83 # TODO: openbmc/openbmc#2994 remove python 2 support 84 # python 2 85 def next(self): 86 return super(PathTreeKeyIterator, self).next()[0] 87 88 # python 3 89 import sys 90 if sys.version_info[0] > 2: 91 __next__ = next 92 93 94class PathTree: 95 def __init__(self): 96 self.root = {} 97 self.cache = {} 98 99 def _try_delete_parent(self, elements): 100 if len(elements) == 1: 101 return False 102 103 kids = 'children' 104 elements.pop() 105 d = self.root 106 for k in elements[:-1]: 107 d = d[k][kids] 108 109 if 'data' not in d[elements[-1]] and not d[elements[-1]][kids]: 110 del d[elements[-1]] 111 self._try_delete_parent(elements) 112 113 def _get_node(self, key): 114 kids = 'children' 115 elements = ['/'] + list(filter(bool, key.split('/'))) 116 d = self.root 117 for k in elements[:-1]: 118 try: 119 d = d[k][kids] 120 except KeyError: 121 raise KeyError(key) 122 123 return d[elements[-1]] 124 125 def __iter__(self): 126 return PathTreeItemIterator(self, '/', None) 127 128 def __missing__(self, key): 129 for x in self.iterkeys(): 130 if key == x: 131 return False 132 return True 133 134 def __delitem__(self, key): 135 del self.cache[key] 136 kids = 'children' 137 elements = ['/'] + list(filter(bool, key.split('/'))) 138 d = self.root 139 for k in elements[:-1]: 140 try: 141 d = d[k][kids] 142 except KeyError: 143 raise KeyError(key) 144 145 del d[elements[-1]] 146 self._try_delete_parent(elements) 147 148 def __setitem__(self, key, value): 149 self.cache[key] = value 150 kids = 'children' 151 elements = ['/'] + list(filter(bool, key.split('/'))) 152 d = self.root 153 for k in elements[:-1]: 154 d = d.setdefault(k, {kids: {}})[kids] 155 156 children = d.setdefault(elements[-1], {kids: {}})[kids] 157 d[elements[-1]].update({kids: children, 'data': value}) 158 159 def __getitem__(self, key): 160 if key in self.cache: 161 return self.cache[key] 162 163 return self._get_node(key).get('data') 164 165 def setdefault(self, key, default): 166 if not self.get(key): 167 self.__setitem__(key, default) 168 169 return self.__getitem__(key) 170 171 def get(self, key, default=None): 172 try: 173 x = self.__getitem__(key) 174 except KeyError: 175 x = default 176 177 return x 178 179 def get_children(self, key): 180 return [x for x in self._get_node(key)['children'].keys()] 181 182 def demote(self, key): 183 n = self._get_node(key) 184 if 'data' in n: 185 del n['data'] 186 187 def keys(self, subtree='/', depth=None): 188 return [x for x in self.iterkeys(subtree, depth)] 189 190 def values(self, subtree='/', depth=None): 191 return [x[1] for x in self.iteritems(subtree, depth)] 192 193 def items(self, subtree='/', depth=None): 194 return [x for x in self.iteritems(subtree, depth)] 195 196 def dataitems(self, subtree='/', depth=None): 197 # dataitems() must return an iterable object containing all of the 198 # items explicitly inserted into the tree, rooted at subtree with 199 # depth number of path elements from the subtree root. 200 # 201 # Calling dataitems() to request all of the populated entries is 202 # unfortunately common, and it's (also) unfortunately expensive to 203 # generate (done by iterating over the entire tree). However, given we 204 # have a flat dict whose job is to cache the explicitly inserted values 205 # we can leverage that to provide a cheap-to-calculate answer to 206 # requests for the entire set of populated entries 207 if subtree == '/' and not depth: 208 return self.cache.items() 209 210 return [x for x in self.iteritems(subtree, depth) 211 if x[1] is not None] 212 213 def iterkeys(self, subtree='/', depth=None): 214 if not self.root: 215 # TODO: openbmc/openbmc#2994 remove python 2 support 216 try: # python 2 217 return {}.iterkeys() 218 except AttributeError: # python 3 219 return iter({}.keys()) 220 return PathTreeKeyIterator(self, subtree, depth) 221 222 def iteritems(self, subtree='/', depth=None): 223 if not self.root: 224 # TODO: openbmc/openbmc#2994 remove python 2 support 225 try: # python 2 226 return {}.iteritems() 227 except AttributeError: # python 3 228 return iter({}.items()) 229 return PathTreeItemIterator(self, subtree, depth) 230 231 def dumpd(self, subtree='/'): 232 result = {} 233 d = result 234 235 for k, v in self.iteritems(subtree): 236 elements = ['/'] + list(filter(bool, k.split('/'))) 237 d = result 238 for k in elements: 239 d = d.setdefault(k, {}) 240 if v is not None: 241 d.update(v) 242 243 return result 244