You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

228 lines
7.5 KiB

# -*- coding: utf-8 -*-
# Copyright (c) 2018 gevent
# cython: auto_pickle=False,embedsignature=True,always_allow_keywords=False,infer_types=True
"""
Iterators across greenlets or AsyncResult objects.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from gevent import _semaphore
from gevent import queue
__all__ = [
'IMapUnordered',
'IMap',
]
locals()['Greenlet'] = __import__('gevent').Greenlet
locals()['Semaphore'] = _semaphore.Semaphore
locals()['UnboundQueue'] = queue.UnboundQueue
class Failure(object):
__slots__ = ('exc', 'raise_exception')
def __init__(self, exc, raise_exception=None):
self.exc = exc
self.raise_exception = raise_exception
def _raise_exc(failure):
# For cython.
if failure.raise_exception:
failure.raise_exception()
else:
raise failure.exc
class IMapUnordered(Greenlet): # pylint:disable=undefined-variable
"""
At iterator of map results.
"""
def __init__(self, func, iterable, spawn, maxsize=None, _zipped=False):
"""
An iterator that.
:param callable spawn: The function we use to create new greenlets.
:keyword int maxsize: If given and not-None, specifies the maximum number of
finished results that will be allowed to accumulated awaiting the reader;
more than that number of results will cause map function greenlets to begin
to block. This is most useful is there is a great disparity in the speed of
the mapping code and the consumer and the results consume a great deal of resources.
Using a bound is more computationally expensive than not using a bound.
.. versionchanged:: 1.1b3
Added the *maxsize* parameter.
"""
Greenlet.__init__(self) # pylint:disable=undefined-variable
self.spawn = spawn
self._zipped = _zipped
self.func = func
self.iterable = iterable
self.queue = UnboundQueue() # pylint:disable=undefined-variable
if maxsize:
# Bounding the queue is not enough if we want to keep from
# accumulating objects; the result value will be around as
# the greenlet's result, blocked on self.queue.put(), and
# we'll go on to spawn another greenlet, which in turn can
# create the result. So we need a semaphore to prevent a
# greenlet from exiting while the queue is full so that we
# don't spawn the next greenlet (assuming that self.spawn
# is of course bounded). (Alternatively we could have the
# greenlet itself do the insert into the pool, but that
# takes some rework).
#
# Given the use of a semaphore at this level, sizing the queue becomes
# redundant, and that lets us avoid having to use self.link() instead
# of self.rawlink() to avoid having blocking methods called in the
# hub greenlet.
self._result_semaphore = Semaphore(maxsize) # pylint:disable=undefined-variable
else:
self._result_semaphore = None
self._outstanding_tasks = 0
# The index (zero based) of the maximum number of
# results we will have.
self._max_index = -1
self.finished = False
# We're iterating in a different greenlet than we're running.
def __iter__(self):
return self
def __next__(self):
if self._result_semaphore is not None:
self._result_semaphore.release()
value = self._inext()
if isinstance(value, Failure):
_raise_exc(value)
return value
next = __next__ # Py2
def _inext(self):
return self.queue.get()
def _ispawn(self, func, item, item_index):
if self._result_semaphore is not None:
self._result_semaphore.acquire()
self._outstanding_tasks += 1
g = self.spawn(func, item) if not self._zipped else self.spawn(func, *item)
g._imap_task_index = item_index
g.rawlink(self._on_result)
return g
def _run(self): # pylint:disable=method-hidden
try:
func = self.func
for item in self.iterable:
self._max_index += 1
self._ispawn(func, item, self._max_index)
self._on_finish(None)
except BaseException as e:
self._on_finish(e)
raise
finally:
self.spawn = None
self.func = None
self.iterable = None
self._result_semaphore = None
def _on_result(self, greenlet):
# This method will be called in the hub greenlet (we rawlink)
self._outstanding_tasks -= 1
count = self._outstanding_tasks
finished = self.finished
ready = self.ready()
put_finished = False
if ready and count <= 0 and not finished:
finished = self.finished = True
put_finished = True
if greenlet.successful():
self.queue.put(self._iqueue_value_for_success(greenlet))
else:
self.queue.put(self._iqueue_value_for_failure(greenlet))
if put_finished:
self.queue.put(self._iqueue_value_for_self_finished())
def _on_finish(self, exception):
# Called in this greenlet.
if self.finished:
return
if exception is not None:
self.finished = True
self.queue.put(self._iqueue_value_for_self_failure(exception))
return
if self._outstanding_tasks <= 0:
self.finished = True
self.queue.put(self._iqueue_value_for_self_finished())
def _iqueue_value_for_success(self, greenlet):
return greenlet.value
def _iqueue_value_for_failure(self, greenlet):
return Failure(greenlet.exception, getattr(greenlet, '_raise_exception'))
def _iqueue_value_for_self_finished(self):
return Failure(StopIteration())
def _iqueue_value_for_self_failure(self, exception):
return Failure(exception, self._raise_exception)
class IMap(IMapUnordered):
# A specialization of IMapUnordered that returns items
# in the order in which they were generated, not
# the order in which they finish.
def __init__(self, *args, **kwargs):
# The result dictionary: {index: value}
self._results = {}
# The index of the result to return next.
self.index = 0
IMapUnordered.__init__(self, *args, **kwargs)
def _inext(self):
try:
value = self._results.pop(self.index)
except KeyError:
# Wait for our index to finish.
while 1:
index, value = self.queue.get()
if index == self.index:
break
else:
self._results[index] = value
self.index += 1
return value
def _iqueue_value_for_success(self, greenlet):
return (greenlet._imap_task_index, IMapUnordered._iqueue_value_for_success(self, greenlet))
def _iqueue_value_for_failure(self, greenlet):
return (greenlet._imap_task_index, IMapUnordered._iqueue_value_for_failure(self, greenlet))
def _iqueue_value_for_self_finished(self):
return (self._max_index + 1, IMapUnordered._iqueue_value_for_self_finished(self))
def _iqueue_value_for_self_failure(self, exception):
return (self._max_index + 1, IMapUnordered._iqueue_value_for_self_failure(self, exception))
from gevent._util import import_c_accel
import_c_accel(globals(), 'gevent.__imap')