Source Code for Module ClusterShell.MsgTree

  1  # 
  2  # Copyright (C) 2007-2016 CEA/DAM 
  3  # Copyright (C) 2016 Stephane Thiell <sthiell@stanford.edu> 
  4  # 
  5  # This file is part of ClusterShell. 
  6  # 
  7  # ClusterShell is free software; you can redistribute it and/or 
  8  # modify it under the terms of the GNU Lesser General Public 
  9  # License as published by the Free Software Foundation; either 
 10  # version 2.1 of the License, or (at your option) any later version. 
 11  # 
 12  # ClusterShell is distributed in the hope that it will be useful, 
 13  # but WITHOUT ANY WARRANTY; without even the implied warranty of 
 14  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU 
 15  # Lesser General Public License for more details. 
 16  # 
 17  # You should have received a copy of the GNU Lesser General Public 
 18  # License along with ClusterShell; if not, write to the Free Software 
 19  # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA 
 20   
 21  """ 
 22  MsgTree 
 23   
 24  ClusterShell message tree module. The purpose of MsgTree is to provide a 
 25  shared message tree for storing message lines received from ClusterShell 
 26  Workers (for example, from remote cluster commands). It should be 
 27  efficient, in term of algorithm and memory consumption, especially when 
 28  remote messages are the same. 
 29  """ 
 30   
 31  from itertools import ifilterfalse, imap 
 32  from operator import itemgetter 
 33   
 34  # MsgTree behavior modes 
 35  MODE_DEFER = 0 
 36  MODE_SHIFT = 1 
 37  MODE_TRACE = 2 
 38   
 39   
 40 -class MsgTreeElem(object): 
 41      """ 
 42      Class representing an element of the MsgTree and its associated 
 43      message. Object of this class are returned by the various MsgTree 
 44      methods like messages() or walk(). The object can then be used as 
 45      an iterator over the message lines or casted into a string. 
 46      """ 
 47 -    def __init__(self, msgline=None, parent=None, trace=False): 
 48          """ 
 49          Initialize message tree element. 
 50          """ 
 51          # structure 
 52          self.parent = parent 
 53          self.children = {} 
 54          if trace:  # special behavior for trace mode 
 55              self._shift = self._shift_trace 
 56          else: 
 57              self._shift = self._shift_notrace 
 58          # content 
 59          self.msgline = msgline 
 60          self.keys = None 
 61   
 62 -    def __len__(self): 
 63          """Length of whole message string.""" 
 64          return len(str(self)) 
 65   
 66 -    def __eq__(self, other): 
 67          """Comparison method compares whole message strings.""" 
 68          return str(self) == str(other) 
 69   
 70 -    def _add_key(self, key): 
 71          """Add a key to this tree element.""" 
 72          if self.keys is None: 
 73              self.keys = set([key]) 
 74          else: 
 75              self.keys.add(key) 
 76   
 77 -    def _shift_notrace(self, key, target_elem): 
 78          """Shift one of our key to specified target element.""" 
 79          if self.keys and len(self.keys) == 1: 
 80              shifting = self.keys 
 81              self.keys = None 
 82          else: 
 83              shifting = set([key]) 
 84              if self.keys: 
 85                  self.keys.difference_update(shifting) 
 86   
 87          if not target_elem.keys: 
 88              target_elem.keys = shifting 
 89          else: 
 90              target_elem.keys.update(shifting) 
 91   
 92          return target_elem 
 93   
 94 -    def _shift_trace(self, key, target_elem): 
 95          """Shift one of our key to specified target element (trace 
 96          mode: keep backtrace of keys).""" 
 97          if not target_elem.keys: 
 98              target_elem.keys = set([key]) 
 99          else: 
100              target_elem.keys.add(key) 
101          return target_elem 
102   
103 -    def __getitem__(self, i): 
104          return list(self.lines())[i] 
105   
106 -    def __iter__(self): 
107          """Iterate over message lines up to this element.""" 
108          bottomtop = [] 
109          if self.msgline is not None: 
110              bottomtop.append(self.msgline) 
111              parent = self.parent 
112              while parent.msgline is not None: 
113                  bottomtop.append(parent.msgline) 
114                  parent = parent.parent 
115          return reversed(bottomtop) 
116   
117 -    def lines(self): 
118          """Get an iterator over all message lines up to this element.""" 
119          return iter(self) 
120   
121      splitlines = lines 
122   
123 -    def message(self): 
124          """ 
125          Get the whole message buffer from this tree element. 
126          """ 
127          # concat buffers 
128          return '\n'.join(self.lines()) 
129   
130      __str__ = message 
131   
132 -    def append(self, msgline, key=None): 
133          """ 
134          A new message is coming, append it to the tree element with 
135          optional associated source key. Called by MsgTree.add(). 
136          Return corresponding MsgTreeElem (possibly newly created). 
137          """ 
138          # get/create child element 
139          elem = self.children.get(msgline) 
140          if elem is None: 
141              elem = self.__class__(msgline, self, 
142                                    self._shift == self._shift_trace) 
143              self.children[msgline] = elem 
144   
145          # if no key is given, MsgTree is in MODE_DEFER 
146          # shift down the given key otherwise 
147          # Note: replace with ternary operator in py2.5+ 
148          if key is None: 
149              return elem 
150          else: 
151              return self._shift(key, elem) 
152   
153   
154 -class MsgTree(object): 
155      """ 
156      MsgTree maps key objects to multi-lines messages. 
157   
158      MsgTree is a mutable object. Keys are almost arbitrary values (must 
159      be hashable). Message lines are organized as a tree internally. 
160      MsgTree provides low memory consumption especially on a cluster when 
161      all nodes return similar messages. Also, the gathering of messages is 
162      done automatically. 
163      """ 
164   
165 -    def __init__(self, mode=MODE_DEFER): 
166          """MsgTree initializer 
167   
168          The `mode' parameter should be set to one of the following constant: 
169   
170          MODE_DEFER: all messages are processed immediately, saving memory from 
171          duplicate message lines, but keys are associated to tree elements 
172          usually later when tree is first "walked", saving useless state 
173          updates and CPU time. Once the tree is "walked" for the first time, its 
174          mode changes to MODE_SHIFT to keep track of further tree updates. 
175          This is the default mode. 
176   
177          MODE_SHIFT: all keys and messages are processed immediately, it is more 
178          CPU time consuming as MsgTree full state is updated at each add() call. 
179   
180          MODE_TRACE: all keys and messages and processed immediately, and keys 
181          are kept for each message element of the tree. The special method 
182          walk_trace() is then available to walk all elements of the tree. 
183          """ 
184          self.mode = mode 
185          # root element of MsgTree 
186          self._root = MsgTreeElem(trace=(mode == MODE_TRACE)) 
187          # dict of keys to MsgTreeElem 
188          self._keys = {} 
189   
190 -    def clear(self): 
191          """Remove all items from the MsgTree.""" 
192          self._root = MsgTreeElem(trace=(self.mode == MODE_TRACE)) 
193          self._keys.clear() 
194   
195 -    def __len__(self): 
196          """Return the number of keys contained in the MsgTree.""" 
197          return len(self._keys) 
198   
199 -    def __getitem__(self, key): 
200          """Return the message of MsgTree with specified key. Raises a 
201          KeyError if key is not in the MsgTree.""" 
202          return self._keys[key] 
203   
204 -    def get(self, key, default=None): 
205          """ 
206          Return the message for key if key is in the MsgTree, else default. 
207          If default is not given, it defaults to None, so that this method 
208          never raises a KeyError. 
209          """ 
210          return self._keys.get(key, default) 
211   
212 -    def add(self, key, msgline): 
213          """ 
214          Add a message line associated with the given key to the MsgTree. 
215          """ 
216          # try to get current element in MsgTree for the given key, 
217          # defaulting to the root element 
218          e_msg = self._keys.get(key, self._root) 
219          if self.mode >= MODE_SHIFT: 
220              key_shift = key 
221          else: 
222              key_shift = None 
223          # add child msg and update keys dict 
224          self._keys[key] = e_msg.append(msgline, key_shift) 
225   
226 -    def _update_keys(self): 
227          """Update keys associated to tree elements (MODE_DEFER).""" 
228          for key, e_msg in self._keys.iteritems(): 
229              assert key is not None and e_msg is not None 
230              e_msg._add_key(key) 
231          # MODE_DEFER is no longer valid as keys are now assigned to MsgTreeElems 
232          self.mode = MODE_SHIFT 
233   
234 -    def keys(self): 
235          """Return an iterator over MsgTree's keys.""" 
236          return self._keys.iterkeys() 
237   
238      __iter__ = keys 
239   
240 -    def messages(self, match=None): 
241          """Return an iterator over MsgTree's messages.""" 
242          return imap(itemgetter(0), self.walk(match)) 
243   
244 -    def items(self, match=None, mapper=None): 
245          """ 
246          Return (key, message) for each key of the MsgTree. 
247          """ 
248          if mapper is None: 
249              mapper = lambda k: k 
250          for key, elem in self._keys.iteritems(): 
251              if match is None or match(key): 
252                  yield mapper(key), elem 
253   
254 -    def _depth(self): 
255          """ 
256          Return the depth of the MsgTree, ie. the max number of lines 
257          per message. Added for debugging. 
258          """ 
259          depth = 0 
260          # stack of (element, depth) tuples used to walk the tree 
261          estack = [(self._root, depth)] 
262   
263          while estack: 
264              elem, edepth = estack.pop() 
265              if len(elem.children) > 0: 
266                  estack += [(v, edepth + 1) for v in elem.children.values()] 
267              depth = max(depth, edepth) 
268   
269          return depth 
270   
271 -    def walk(self, match=None, mapper=None): 
272          """ 
273          Walk the tree. Optionally filter keys on match parameter, 
274          and optionally map resulting keys with mapper function. 
275          Return an iterator over (message, keys) tuples for each 
276          different message in the tree. 
277          """ 
278          if self.mode == MODE_DEFER: 
279              self._update_keys() 
280          # stack of elements used to walk the tree (depth-first) 
281          estack = [self._root] 
282          while estack: 
283              elem = estack.pop() 
284              children = elem.children 
285              if len(children) > 0: 
286                  estack += children.values() 
287              if elem.keys: # has some keys 
288                  mkeys = filter(match, elem.keys) 
289                  if len(mkeys): 
290                      yield elem, map(mapper, mkeys) 
291   
292 -    def walk_trace(self, match=None, mapper=None): 
293          """ 
294          Walk the tree in trace mode. Optionally filter keys on match 
295          parameter, and optionally map resulting keys with mapper 
296          function. 
297          Return an iterator over 4-length tuples (msgline, keys, depth, 
298          num_children). 
299          """ 
300          assert self.mode == MODE_TRACE, \ 
301              "walk_trace() is only callable in trace mode" 
302          # stack of (element, depth) tuples used to walk the tree 
303          estack = [(self._root, 0)] 
304          while estack: 
305              elem, edepth = estack.pop() 
306              children = elem.children 
307              nchildren = len(children) 
308              if nchildren > 0: 
309                  estack += [(v, edepth + 1) for v in children.values()] 
310              if elem.keys: 
311                  mkeys = filter(match, elem.keys) 
312                  if len(mkeys): 
313                      yield elem.msgline, map(mapper, mkeys), edepth, nchildren 
314   
315 -    def remove(self, match=None): 
316          """ 
317          Modify the tree by removing any matching key references from the 
318          messages tree. 
319   
320          Example of use: 
321              >>> msgtree.remove(lambda k: k > 3) 
322          """ 
323          # do not walk tree in MODE_DEFER as no key is associated 
324          if self.mode != MODE_DEFER: 
325              estack = [self._root] 
326              # walk the tree to keep only matching keys 
327              while estack: 
328                  elem = estack.pop() 
329                  if len(elem.children) > 0: 
330                      estack += elem.children.values() 
331                  if elem.keys: # has some keys 
332                      elem.keys = set(ifilterfalse(match, elem.keys)) 
333   
334          # remove key(s) from known keys dict 
335          for key in filter(match, self._keys.keys()): 
336              del self._keys[key] 
337