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 return self.cache[key] 161 162 def setdefault(self, key, default): 163 if not self.get(key): 164 self.__setitem__(key, default) 165 166 return self.__getitem__(key) 167 168 def get(self, key, default=None): 169 try: 170 x = self.__getitem__(key) 171 except KeyError: 172 x = default 173 174 return x 175 176 def get_children(self, key): 177 return [x for x in self._get_node(key)['children'].keys()] 178 179 def demote(self, key): 180 n = self._get_node(key) 181 if 'data' in n: 182 del n['data'] 183 184 def keys(self, subtree='/', depth=None): 185 return [x for x in self.iterkeys(subtree, depth)] 186 187 def values(self, subtree='/', depth=None): 188 return [x[1] for x in self.iteritems(subtree, depth)] 189 190 def items(self, subtree='/', depth=None): 191 return [x for x in self.iteritems(subtree, depth)] 192 193 def dataitems(self, subtree='/', depth=None): 194 return [x for x in self.iteritems(subtree, depth) 195 if x[1] is not None] 196 197 def iterkeys(self, subtree='/', depth=None): 198 if not self.root: 199 # TODO: openbmc/openbmc#2994 remove python 2 support 200 try: # python 2 201 return {}.iterkeys() 202 except AttributeError: # python 3 203 return iter({}.keys()) 204 return PathTreeKeyIterator(self, subtree, depth) 205 206 def iteritems(self, subtree='/', depth=None): 207 if not self.root: 208 # TODO: openbmc/openbmc#2994 remove python 2 support 209 try: # python 2 210 return {}.iteritems() 211 except AttributeError: # python 3 212 return iter({}.items()) 213 return PathTreeItemIterator(self, subtree, depth) 214 215 def dumpd(self, subtree='/'): 216 result = {} 217 d = result 218 219 for k, v in self.iteritems(subtree): 220 elements = ['/'] + list(filter(bool, k.split('/'))) 221 d = result 222 for k in elements: 223 d = d.setdefault(k, {}) 224 if v is not None: 225 d.update(v) 226 227 return result 228