1# -*- coding: utf-8 -*- 2# Copyright 2012 Google Inc. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15"""Iterator wrapper for checking wrapped iterator's emptiness or plurality.""" 16 17# TODO: Here and elsewhere (wildcard_iterator, name_expansion), do not reference 18# __iter__ directly because it causes the first element to be instantiated. 19# Instead, implement __iter__ as a return self and implement the next() function 20# which returns (not yields) the values. This necessitates that in the case 21# of the iterator classes, the iterator is used once per class instantiation 22# so that next() calls do not collide, but this semantic has been long-assumed 23# by the iterator classes for the use of __iter__ anyway. 24 25from __future__ import absolute_import 26 27import sys 28 29 30class PluralityCheckableIterator(object): 31 """Iterator wrapper class. 32 33 Allows you to check whether the wrapped iterator is empty and 34 whether it has more than 1 element. This iterator accepts three types of 35 values from the iterator it wraps: 36 1. A yielded element (this is the normal case). 37 2. A raised exception, which will be buffered and re-raised when it 38 is reached in this iterator. 39 3. A yielded tuple of (exception, stack trace), which will be buffered 40 and raised with it is reached in this iterator. 41 """ 42 43 def __init__(self, it): 44 # Need to get the iterator function here so that we don't immediately 45 # instantiate the first element (which could raise an exception). 46 self.orig_iterator = it 47 self.base_iterator = None 48 self.head = [] 49 self.underlying_iter_empty = False 50 # Populate first 2 elems into head so we can check whether iterator has 51 # more than 1 item. 52 for _ in range(0, 2): 53 self._PopulateHead() 54 55 def _PopulateHead(self): 56 if not self.underlying_iter_empty: 57 try: 58 if not self.base_iterator: 59 self.base_iterator = iter(self.orig_iterator) 60 e = self.base_iterator.next() 61 self.underlying_iter_empty = False 62 if isinstance(e, tuple) and isinstance(e[0], Exception): 63 self.head.append(('exception', e[0], e[1])) 64 else: 65 self.head.append(('element', e)) 66 except StopIteration: 67 # Indicates we can no longer call next() on underlying iterator, but 68 # there could still be elements left to iterate in head. 69 self.underlying_iter_empty = True 70 except Exception, e: 71 # Buffer the exception and raise it when the element is accessed. 72 # Also, preserve the original stack trace, as the stack trace from 73 # within plurality_checkable_iterator.next is not very useful. 74 self.head.append(('exception', e, sys.exc_info()[2])) 75 76 def __iter__(self): 77 return self 78 79 def next(self): 80 # Backfill into head each time we pop an element so we can always check 81 # for emptiness and for HasPlurality(). 82 while self.head: 83 self._PopulateHead() 84 item_tuple = self.head.pop(0) 85 if item_tuple[0] == 'element': 86 return item_tuple[1] 87 else: # buffered exception 88 raise item_tuple[1].__class__, item_tuple[1], item_tuple[2] 89 raise StopIteration() 90 91 def IsEmpty(self): 92 return not self.head 93 94 def HasPlurality(self): 95 return len(self.head) > 1 96