Commits

Sergey Astanin committed 4a31c83

rewrite _SubSequencer to n-way split (not only binary partitioning)

Comments (0)

Files changed (1)

 __all__ = [ "chop", "partition", "split" ]
 __author__ = "Sergey Astanin"
 __license__ = "MIT"
-__version__ = "0.2"
+__version__ = "0.3"
 
 if platform.python_version_tuple()[0] == "2":
     _range = xrange
 
 class _SubSequencer:
     """
-    Lazily process a sequence in single pass and split into two.
-    Computes both output sequences even if only one of them is consumed.
+    Lazily process a sequence in single pass and split into many.
+    Compute all output sequences even if only one of them is consumed.
+    Maintain as many subsequence queues as there are predicate values.
     """
-    def __init__(self, condition, sequence):
-        self.cond = condition
-        self.goods = collections.deque([])
-        self.bads = collections.deque([])
+    def __init__(self, predicate, sequence):
+        self.predicate = predicate
+        self.subseqs = dict()  # mapping of { predicate_value: subseq }
         self.seq = iter(sequence)
-    def getNext(self, getGood=True):
-        if getGood:
-            these, those, cond = self.goods, self.bads, self.cond
+
+    def subqueue(self, predicate_value):
+        "Get or create a subsequence queue."
+        if not predicate_value in self.subseqs:
+            subseq = self.subseqs[predicate_value] = collections.deque([])
         else:
-            these, those, cond = self.bads, self.goods, lambda x: not self.cond(x)
-        if these:
-            return these.popleft()
-        else:
-            while 1: # exit on StopIteration
+            subseq = self.subseqs[predicate_value]
+        return subseq
+
+    def subappend(self, predicate_value, item):
+        "Append a processed item to subsequence queue."
+        subseq = self.subqueue(predicate_value)
+        subseq.append(item)
+
+    def subnext(self, predicate_value):
+        "Next value in the subsequence corresponding to predicate_value."
+        subseq = self.subqueue(predicate_value)
+        if subseq:
+            return subseq.popleft()
+        else:  # subsequence queue is empty
+            while True:  # exit on StopIteration
                 n = next(self.seq)
-                if cond(n):
+                pv = self.predicate(n)
+                if pv == predicate_value:
                     return n
                 else:
-                    those.append(n)
+                    self.subappend(pv, n)
 
 def partition(condition, sequence):
     """
     cond = condition if condition else bool  # eval as bool if condition is None
     ss = _SubSequencer(cond, sequence)
     def condition_holds():
-        while 1:
-            yield ss.getNext(getGood=True)
+        while True:
+            yield ss.subnext(True)
     def condition_doesnt_hold():
         while 1:
-            yield ss.getNext(getGood=False)
+            yield ss.subnext(False)
     return condition_holds(), condition_doesnt_hold()
 
 def _take(n, sequence):