Generate data from another DataFlow, but with a fixed total count.
| 230 | |
| 231 | |
| 232 | class FixedSizeData(ProxyDataFlow): |
| 233 | """ Generate data from another DataFlow, but with a fixed total count. |
| 234 | """ |
| 235 | def __init__(self, ds, size, keep_state=True): |
| 236 | """ |
| 237 | Args: |
| 238 | ds (DataFlow): input dataflow |
| 239 | size (int): size |
| 240 | keep_state (bool): keep the iterator state of ``ds`` |
| 241 | between calls to :meth:`__iter__()`, so that the |
| 242 | next call will continue the previous iteration over ``ds``, |
| 243 | instead of reinitializing an iterator. |
| 244 | |
| 245 | Example: |
| 246 | |
| 247 | .. code-block:: none |
| 248 | |
| 249 | ds produces: 1, 2, 3, 4, 5; 1, 2, 3, 4, 5; ... |
| 250 | FixedSizeData(ds, 3, True): 1, 2, 3; 4, 5, 1; 2, 3, 4; ... |
| 251 | FixedSizeData(ds, 3, False): 1, 2, 3; 1, 2, 3; ... |
| 252 | FixedSizeData(ds, 6, False): 1, 2, 3, 4, 5, 1; 1, 2, 3, 4, 5, 1;... |
| 253 | """ |
| 254 | super(FixedSizeData, self).__init__(ds) |
| 255 | self._size = int(size) |
| 256 | self.itr = None |
| 257 | self._keep = keep_state |
| 258 | |
| 259 | def __len__(self): |
| 260 | return self._size |
| 261 | |
| 262 | def reset_state(self): |
| 263 | super(FixedSizeData, self).reset_state() |
| 264 | self.itr = self.ds.__iter__() |
| 265 | self._guard = DataFlowReentrantGuard() |
| 266 | |
| 267 | def __iter__(self): |
| 268 | with self._guard: |
| 269 | if self.itr is None: |
| 270 | self.itr = self.ds.__iter__() |
| 271 | cnt = 0 |
| 272 | while True: |
| 273 | try: |
| 274 | dp = next(self.itr) |
| 275 | except StopIteration: |
| 276 | self.itr = self.ds.__iter__() |
| 277 | dp = next(self.itr) |
| 278 | |
| 279 | cnt += 1 |
| 280 | yield dp |
| 281 | if cnt == self._size: |
| 282 | if not self._keep: |
| 283 | self.itr = None |
| 284 | return |
| 285 | |
| 286 | |
| 287 | class MapData(ProxyDataFlow): |
no outgoing calls
no test coverage detected