-
Notifications
You must be signed in to change notification settings - Fork 1
/
DataBatcher.py
39 lines (27 loc) · 988 Bytes
/
DataBatcher.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
from collections import deque
class DataBatcher:
def __init__(self, keyFunc, valFunc=None):
self.keyFunc = keyFunc
self.valFunc = valFunc
self.prevKey = None
self.batches = deque()
self.batch = None
def addData(self, data):
key = self.keyFunc(data)
if key == None:
raise Exception('Batching Key cannot be None')
if key != self.prevKey:
if self.batch != None:
self.batches.append({ self.prevKey: self.batch })
self.batch = []
self.prevKey = key
self.batch.append(self.valFunc(data) if self.valFunc else data)
def endBatch(self):
self.batches.append({ self.prevKey: self.batch })
self.prevKey = None
self.batch = None
def getBatches(self):
while(self.batches):
yield self.batches.popleft()
def getBatchesCount(self):
return len(self.batches)