def test_make_repr(self): class Foo: class Bar: x = 1 y = 2 obj = Foo.Bar() prefix = '%s.%s %#x' % (Foo.Bar.__module__, Foo.Bar.__qualname__, id(obj)) r = classes.make_repr() self.assertEqual(r(obj), f'<{prefix}>') r = classes.make_repr('x={self.x} y={self.y}') self.assertEqual(r(obj), f'<{prefix} x=1 y=2>') r = classes.make_repr('sum={sum}', sum=lambda self: self.x + self.y) self.assertEqual(r(obj), f'<{prefix} sum=3>') with self.assertRaises(AssertionError): classes.make_repr('{}', self=None) with self.assertRaises(AssertionError): classes.make_repr('{}', __self_id=None) with self.assertRaises(AssertionError): classes.make_repr(x=lambda self: self.x)
class ConstSchema(Schema): _raw_type = _capnp.ConstSchema __repr__ = classes.make_repr('type={self.type!r}') type = bases.def_mp('type', _to_type, _raw_type.getType)
class Event: def __init__(self): self._flag = False __repr__ = classes.make_repr( '{state}', state=lambda self: 'set' if self._flag else 'unset', ) def is_set(self): return self._flag def set(self): self._flag = True # Let's make a special case when no task is waiting for this # event object (with this, you may call ``Event.set`` out of a # kernel context). This is useful when you want to initialize # events before a kernel context is initialized. try: contexts.get_kernel().unblock(self) except LookupError: pass def clear(self): self._flag = False async def wait(self): while not self._flag: await traps.block(self) return self._flag
class ListSchema(bases.Base): _raw_type = _capnp.ListSchema __repr__ = classes.make_repr('element_type={self.element_type!r}') element_type = bases.def_mp('type', _to_type, _raw_type.getElementType)
class Type(bases.Base): _raw_type = _capnp.Type __repr__ = classes.make_repr('which={self.which}') which = bases.def_p(_raw_type.which) as_struct = bases.def_f0(StructSchema, _raw_type.asStruct) as_enum = bases.def_f0(EnumSchema, _raw_type.asEnum) as_interface = bases.def_f0(InterfaceSchema, _raw_type.asInterface) as_list = bases.def_f0(ListSchema, _raw_type.asList) is_void = bases.def_f0(_raw_type.isVoid) is_bool = bases.def_f0(_raw_type.isBool) is_int8 = bases.def_f0(_raw_type.isInt8) is_int16 = bases.def_f0(_raw_type.isInt16) is_int32 = bases.def_f0(_raw_type.isInt32) is_int64 = bases.def_f0(_raw_type.isInt64) is_uint8 = bases.def_f0(_raw_type.isUInt8) is_uint16 = bases.def_f0(_raw_type.isUInt16) is_uint32 = bases.def_f0(_raw_type.isUInt32) is_uint64 = bases.def_f0(_raw_type.isUInt64) is_float32 = bases.def_f0(_raw_type.isFloat32) is_float64 = bases.def_f0(_raw_type.isFloat64) is_text = bases.def_f0(_raw_type.isText) is_data = bases.def_f0(_raw_type.isData) is_list = bases.def_f0(_raw_type.isList) is_enum = bases.def_f0(_raw_type.isEnum) is_struct = bases.def_f0(_raw_type.isStruct) is_interface = bases.def_f0(_raw_type.isInterface) is_any_pointer = bases.def_f0(_raw_type.isAnyPointer)
class AbstractBundleDir: deploy_instruction_type = classes.abstract_property post_init = classes.abstract_method def __init__(self, path): self.path = path self.deploy_instruction = jsons.load_dataobject( self.deploy_instruction_type, ASSERT.predicate(self.deploy_instruction_path, Path.is_file)) self.post_init() __repr__ = classes.make_repr('path={self.path}') def __eq__(self, other): return self.path == other.path def __hash__(self): return hash(self.path) @property def deploy_instruction_path(self): return self.path / models.BUNDLE_DEPLOY_INSTRUCTION_FILENAME @property def label(self): return self.deploy_instruction.label @property def version(self): return self.deploy_instruction.version
class Value(bases.Base): _raw_type = _capnp.schema.Value __repr__ = classes.make_repr('which={self.which}') which = bases.def_p(_raw_type.which) is_void = bases.def_f0(_raw_type.isVoid) is_bool = bases.def_f0(_raw_type.isBool) is_int8 = bases.def_f0(_raw_type.isInt8) is_int16 = bases.def_f0(_raw_type.isInt16) is_int32 = bases.def_f0(_raw_type.isInt32) is_int64 = bases.def_f0(_raw_type.isInt64) is_uint8 = bases.def_f0(_raw_type.isUint8) is_uint16 = bases.def_f0(_raw_type.isUint16) is_uint32 = bases.def_f0(_raw_type.isUint32) is_uint64 = bases.def_f0(_raw_type.isUint64) is_float32 = bases.def_f0(_raw_type.isFloat32) is_float64 = bases.def_f0(_raw_type.isFloat64) is_text = bases.def_f0(_raw_type.isText) is_data = bases.def_f0(_raw_type.isData) is_list = bases.def_f0(_raw_type.isList) is_enum = bases.def_f0(_raw_type.isEnum) is_struct = bases.def_f0(_raw_type.isStruct) is_interface = bases.def_f0(_raw_type.isInterface) is_any_pointer = bases.def_f0(_raw_type.isAnyPointer) @classes.memorizing_property def text(self): ASSERT.true(self._raw.isText()) return str(self._raw.getText(), 'utf-8')
class Annotation(bases.Base): _raw_type = _capnp.schema.Annotation __repr__ = classes.make_repr('id={self.id} value={self.value}') id = bases.def_p(_raw_type.getId) value = bases.def_mp('value', _to_value, _raw_type.getValue)
class Publisher: def __init__(self, queue, wiredata, *, drop_when_full=True): self._queue = queue self._wiredata = wiredata self._drop_when_full = drop_when_full # For convenience, create socket before ``__enter__``. self.socket = nng.asyncs.Socket(nng.Protocols.PUB0) __repr__ = classes.make_repr('{self.socket!r}') def __enter__(self): self.socket.__enter__() return self def __exit__(self, *args): messages = self._queue.close(graceful=False) if messages: LOG.warning('drop %d messages', len(messages)) return self.socket.__exit__(*args) async def serve(self): LOG.debug('start publisher: %r', self) try: while True: message = await self._queue.get() try: raw_message = self._wiredata.to_lower(message) except Exception: LOG.exception('to_lower error: %r', message) continue try: # For now we publish with no topic. await self.socket.send(raw_message) except nng.Errors.ETIMEDOUT: LOG.warning('send timeout; drop message: %r', message) continue except (queues.Closed, nng.Errors.ECLOSED): pass self._queue.close() LOG.debug('stop publisher: %r', self) def shutdown(self): self._queue.close() async def publish(self, message): await self._queue.put(message) def publish_nonblocking(self, message): try: self._queue.put_nonblocking(message) except queues.Full: if self._drop_when_full: LOG.warning('queue full; drop message: %r', message) else: raise
class Subscriber: def __init__(self, message_type, queue, wiredata, *, drop_when_full=True): self._message_type = message_type self._queue = queue self._wiredata = wiredata self._drop_when_full = drop_when_full # For convenience, create socket before ``__enter__``. self.socket = nng.asyncs.Socket(nng.Protocols.SUB0) # For now we subscribe to empty topic. self.socket.subscribe(b'') __repr__ = classes.make_repr('{self.socket!r}') def __enter__(self): self.socket.__enter__() return self def __exit__(self, exc_type, *args): messages = self._queue.close(graceful=not exc_type) if messages: LOG.warning('drop %d messages', len(messages)) return self.socket.__exit__(exc_type, *args) async def serve(self): LOG.debug('start subscriber: %r', self) try: while True: try: raw_message = await self.socket.recv() except nng.Errors.ETIMEDOUT: LOG.warning('recv timeout') continue try: message = self._wiredata.to_upper(self._message_type, raw_message) except Exception: LOG.warning('to_upper error: %r', raw_message, exc_info=True) continue if self._drop_when_full: try: self._queue.put_nonblocking(message) except queues.Full: LOG.warning('queue full; drop message: %r', message) else: await self._queue.put(message) except (queues.Closed, nng.Errors.ECLOSED): pass self._queue.close() LOG.debug('stop subscriber: %r', self) def shutdown(self): self.socket.close()
class Enumerant(bases.Base): _raw_type = _capnp.schema.Enumerant __repr__ = classes.make_repr( 'name={self.name!r} code_order={self.code_order}' ) name = bases.def_mp('name', bases.to_str, _raw_type.getName) code_order = bases.def_p(_raw_type.getCodeOrder)
class Message: def __init__(self, data=b'', *, msg_p=None): # In case ``__init__`` raises. self._msg_p = None ASSERT.isinstance(data, bytes) if msg_p is None: msg_p = _nng.nng_msg_p() errors.check(_nng.F.nng_msg_alloc(ctypes.byref(msg_p), len(data))) if data: ctypes.memmove(_nng.F.nng_msg_body(msg_p), data, len(data)) else: # We are taking ownership of ``msg_p`` and should not take # any initial data. ASSERT.false(data) self._msg_p = msg_p self.header = Header(self._get) self.body = Body(self._get) __repr__ = classes.make_repr('{self._msg_p}') def __enter__(self): return self def __exit__(self, *_): self._reset() def disown(self): msg_p, self._msg_p = self._msg_p, None return msg_p def copy(self): msg_p = _nng.nng_msg_p() errors.check(_nng.F.nng_msg_dup(ctypes.byref(msg_p), self._get())) return type(self)(msg_p=msg_p) def _get(self): return ASSERT.not_none(self._msg_p) def _reset(self): # You have to check whether ``__init__`` raises. if self._msg_p is not None: _nng.F.nng_msg_free(self._msg_p) self._msg_p = None def __del__(self): # In case you forget to use Message within a context. self._reset()
class Enumerant(bases.Base): _raw_type = _capnp.EnumSchema.Enumerant __repr__ = classes.make_repr( 'proto={self.proto!r} ordinal={self.ordinal} index={self.index}' ) proto = bases.def_mp('proto', _Schema.Enumerant, _raw_type.getProto) ordinal = bases.def_p(_raw_type.getOrdinal) index = bases.def_p(_raw_type.getIndex)
class Field(bases.Base): _raw_type = _capnp.StructSchema.Field __repr__ = classes.make_repr( 'proto={self.proto!r} index={self.index} type={self.type!r}' ) proto = bases.def_mp('proto', _Schema.Field, _raw_type.getProto) index = bases.def_p(_raw_type.getIndex) type = bases.def_mp('type', _to_type, _raw_type.getType)
class AbstractOpsDir: metadata_type = classes.abstract_property check_invariants = classes.abstract_method install = classes.abstract_method start = classes.abstract_method stop = classes.abstract_method stop_all = classes.abstract_method uninstall = classes.abstract_method def __init__(self, path): self.path = path __repr__ = classes.make_repr('path={self.path}') def __eq__(self, other): return self.path == other.path def __hash__(self): return hash(self.path) @property def metadata_path(self): return self.path / models.OPS_DIR_METADATA_FILENAME # XXX: This annotation works around pylint no-member false errors. metadata: object @classes.memorizing_property def metadata(self): # pylint: disable=function-redefined return jsons.load_dataobject( self.metadata_type, ASSERT.predicate(self.metadata_path, Path.is_file), ) @property def label(self): return self.metadata.label @property def version(self): return self.metadata.version @property def refs_dir_path(self): return self.path / models.OPS_DIR_REFS_DIR_NAME @property def volumes_dir_path(self): return self.path / models.OPS_DIR_VOLUMES_DIR_NAME
class Stub: """Stub for interacting with an actor.""" def __init__( self, *, actor, method_names=(), queue=None, name=None, daemon=None, ): self.future = futures.Future() self.queue = queue if queue is not None else queues.Queue() # Create method senders for convenience. if method_names: self.m = make_senders(method_names, self.queue) self._thread = threading.Thread( target=futures.wrap_thread_target(actor, self.future), name=name, args=(self.queue, ), daemon=daemon, ) self._thread.start() __repr__ = classes.make_repr('{self._thread!r}') def __enter__(self): return self def __exit__(self, exc_type, *_): graceful = not exc_type self.shutdown(graceful) try: self.join(None if graceful else NON_GRACE_PERIOD) except futures.Timeout: LOG.warning('actor join timeout: %r', self) def shutdown(self, graceful=True): items = self.queue.close(graceful) if items: LOG.warning('drop %d messages', len(items)) return items def join(self, timeout=None): exc = self.future.get_exception(timeout) if exc: LOG.error('actor crash: %r', self, exc_info=exc)
class Lock: def __init__(self): self._locked = False __repr__ = classes.make_repr( '{state}', state=lambda self: 'locked' if self._locked else 'unlocked', ) async def __aenter__(self): # Unlike ``threading.Lock``, here we return the object. await self.acquire() return self async def __aexit__(self, *_): self.release() def is_owner(self): """Return true if the current task is the owner of this lock. Only an owner may release a lock. This check is mostly useful internally. """ # NOTE: For a ``Lock``, any task owns a locked lock, but for an # ``RLock``, only the task that has locked it is its owner. return self._locked async def acquire(self, blocking=True): """Acquire the lock and return true when locked is acquired.""" if not blocking: return self.acquire_nonblocking() while self._locked: await traps.block(self) self._locked = True return True def acquire_nonblocking(self): """Non-blocking version of ``acquire``.""" if self._locked: return False else: self._locked = True return True def release(self): ASSERT.true(self.is_owner()) self._locked = False contexts.get_kernel().unblock(self)
class Base(bases.Base): _schema_type = type(None) # Sub-class must override this. def __init__(self, message, schema, raw): ASSERT.isinstance(schema, self._schema_type) super().__init__(raw) # Keep a strong reference to the root message to ensure that it # is not garbage-collected before us. self._message = message self.schema = schema __repr__ = classes.make_repr('schema={self.schema} {self!s}') def __str__(self): raise NotImplementedError
class Request: def __init__(self, method, url, **kwargs): self.method = method self.url = url self._kwargs = kwargs __repr__ = classes.make_repr( '{method} {self.url} kwargs={self._kwargs!r}', method=lambda self: self.method.upper(), ) @property def headers(self): return self._kwargs.setdefault('headers', {}) def copy(self): return Request(self.method, self.url, **self._kwargs)
class Node(bases.Base): class Struct(bases.Base): _raw_type = _capnp.schema.Node.Struct is_group = bases.def_p(_raw_type.getIsGroup) _raw_type = _capnp.schema.Node __repr__ = classes.make_repr( 'id={self.id} ' 'scope_id={self.scope_id} ' 'display_name={self.display_name!r} ' 'which={self.which}' ) id = bases.def_p(_raw_type.getId) display_name = bases.def_mp( 'display_name', bases.to_str, _raw_type.getDisplayName, ) display_name_prefix_length = bases.def_p( _raw_type.getDisplayNamePrefixLength ) scope_id = bases.def_p(_raw_type.getScopeId) @classes.memorizing_property def annotations(self): return tuple(map(_Schema.Annotation, self._raw.getAnnotations())) which = bases.def_p(_raw_type.which) is_file = bases.def_f0(_raw_type.isFile) is_struct = bases.def_f0(_raw_type.isStruct) is_enum = bases.def_f0(_raw_type.isEnum) is_interface = bases.def_f0(_raw_type.isInterface) is_const = bases.def_f0(_raw_type.isConst) is_annotation = bases.def_f0(_raw_type.isAnnotation) struct = bases.def_mp('struct', Struct, _raw_type.getStruct)
class Condition: def __init__(self, lock=None): self._lock = lock or Lock() self._waiters = set() # Re-export these methods. self.acquire = self._lock.acquire self.acquire_nonblocking = self._lock.acquire_nonblocking self.release = self._lock.release __repr__ = classes.make_repr('{self._lock!r}') async def __aenter__(self): # Unlike ``threading.Condition``, here we return the object. await self._lock.__aenter__() return self async def __aexit__(self, *args): return await self._lock.__aexit__(*args) async def wait(self): """Wait until notified. To be somehow compatible with ``threading.Condition.wait``, this always return true (since it never times out). """ ASSERT.true(self._lock.is_owner()) waiter = Gate() self._waiters.add(waiter) # NOTE: We have not implemented ``RLock`` yet, but when we do, # be careful **NOT** to call ``release`` here, since you cannot # unlock the lock acquired recursively. self._lock.release() try: await waiter.wait() finally: await self._lock.acquire() return True def notify(self, n=1): ASSERT.true(self._lock.is_owner()) for _ in range(min(n, len(self._waiters))): self._waiters.pop().unblock() def notify_all(self): self.notify(len(self._waiters))
class AdapterBase: def __init__(self, target, fields): self.__target = target self.__fields = ASSERT.not_contains(fields, 'target') __repr__ = classes.make_repr('{self._AdapterBase__target!r}') def __getattr__(self, name): if name == 'target': return self.__target if name in self.__fields: return getattr(self.__target, name) raise AttributeError('disallow accessing field: %s' % name) def disown(self): target, self.__target = self.__target, None return target
class Client: def __init__(self, request_type, response_type, wiredata): self.socket = nng.asyncs.Socket(nng.Protocols.REQ0) self.transceive = Transceiver(self.socket, response_type, wiredata) self.m = collections.Namespace( **{ name: Method(name, request_type, self.transceive) for name in request_type.m } ) __repr__ = classes.make_repr('{self.socket!r}') def __enter__(self): self.socket.__enter__() return self def __exit__(self, *args): return self.socket.__exit__(*args)
class ContextBase(ContextOptions): _name = 'ctx' send = classes.abstract_method recv = classes.abstract_method sendmsg = classes.abstract_method recvmsg = classes.abstract_method def __init__(self, socket): # In case ``__init__`` raises. self._handle = None handle = _nng.nng_ctx() errors.check(_nng.F.nng_ctx_open(ctypes.byref(handle), socket._handle)) self.socket = socket self._handle = handle __repr__ = classes.make_repr('id={self.id} {self.socket}') def __enter__(self): return self def __exit__(self, *_): self.close() def __del__(self): # You have to check whether ``__init__`` raises. if self._handle is not None: self.close() @property def id(self): return _nng.F.nng_ctx_id(self._handle) def close(self): try: errors.check(_nng.F.nng_ctx_close(self._handle)) except errors.Errors.ECLOSED: pass
class Schema(bases.Base): _raw_type = _capnp.Schema __repr__ = classes.make_repr('proto={self.proto!r}') proto = bases.def_mp('proto', _Schema.Node, _raw_type.getProto) is_branded = bases.def_f0(_raw_type.isBranded) # Use explicit functional form to work around cyclic reference in # the ``asX`` methods below. def as_struct(self): return StructSchema(self._raw.asStruct()) def as_enum(self): return EnumSchema(self._raw.asEnum()) def as_interface(self): return InterfaceSchema(self._raw.asInterface()) def as_const(self): return ConstSchema(self._raw.asConst()) short_display_name = bases.def_mp( 'short_display_name', bases.to_str, _raw_type.getShortDisplayName, ) @classes.memorizing_property def name(self): for annotation in self.proto.annotations: # pylint: disable=no-member if annotation.id == CXX_NAME: name = annotation.value.text break else: name = self.short_display_name return ASSERT.not_none(name)
class Field(bases.Base): class Slot(bases.Base): _raw_type = _capnp.schema.Field.Slot had_explicit_default = bases.def_p(_raw_type.getHadExplicitDefault) _raw_type = _capnp.schema.Field __repr__ = classes.make_repr( 'name={self.name!r} code_order={self.code_order}' ) name = bases.def_mp('name', bases.to_str, _raw_type.getName) code_order = bases.def_p(_raw_type.getCodeOrder) which = bases.def_p(_raw_type.which) is_slot = bases.def_f0(_raw_type.isSlot) is_group = bases.def_f0(_raw_type.isGroup) slot = bases.def_mp('slot', Slot, _raw_type.getSlot)
class Server: """Expose an (asynchronous) application object on a socket. This is a fairly simple server for providing remote method calls. If application defines context management (i.e., ``__enter__``), it will be called when server's context management is called. This provides some sorts of server start/stop callbacks to application. """ def __init__( self, application, request_type, response_type, wiredata, *, warning_level_exc_types=(), invalid_request_error=None, internal_server_error=None, ): self._application = application self._request_type = request_type self._response_type = response_type self._wiredata = wiredata self._warning_level_exc_types = frozenset(warning_level_exc_types) self._declared_error_types = utils.get_declared_error_types( self._response_type) # For convenience, create socket before ``__enter__``. self.socket = nng.asyncs.Socket(nng.Protocols.REP0) # Prepared errors. self._invalid_request_error_wire = self._lower_error_or_none( invalid_request_error) self._internal_server_error_wire = self._lower_error_or_none( internal_server_error) def _lower_error_or_none(self, error): if error is None: return None ASSERT.isinstance(error, Exception) error_name = ASSERT(self._match_error_type(error), 'unknown error type: {!r}', error) return self._wiredata.to_lower( self._response_type(error=self._response_type.Error( **{error_name: error}))) def _match_error_type(self, error): # NOTE: We match the exact type rather than calling isinstance # because error types could form a hierarchy, and isinstance # might match a parent error type rather than a child type. return self._declared_error_types.get(type(error)) __repr__ = classes.make_repr('{self.socket!r}') def __enter__(self): self.socket.__enter__() return self def __exit__(self, *args): return self.socket.__exit__(*args) async def serve(self): """Serve requests sequentially. To serve requests concurrently, just spawn multiple tasks running this. """ LOG.debug('start server: %r', self) try: with nng.asyncs.Context(ASSERT.not_none(self.socket)) as context: while True: response = await self._serve(await context.recv()) if response is not None: await context.send(response) except nng.Errors.ECLOSED: pass LOG.debug('stop server: %r', self) def shutdown(self): self.socket.close() async def _serve(self, request): LOG.debug('wire request: %r', request) try: request = self._wiredata.to_upper(self._request_type, request) except Exception: LOG.warning('to_upper error: %r', request, exc_info=True) return self._invalid_request_error_wire try: method_name, method_args = utils.select(request.args) except Exception: LOG.warning('invalid request: %r', request, exc_info=True) return self._invalid_request_error_wire try: method = getattr(self._application, method_name) except AttributeError: LOG.warning('unknown method: %s: %r', method_name, request) return self._invalid_request_error_wire try: result = await method( **{ field.name: getattr(method_args, field.name) for field in dataclasses.fields(method_args) }) except Exception as exc: if type(exc) in self._warning_level_exc_types: # pylint: disable=unidiomatic-typecheck log = LOG.warning exc_info = False else: log = LOG.error exc_info = True log('server error: %r -> %r', request, exc, exc_info=exc_info) response = self._make_error_response(exc) if response is None: return self._internal_server_error_wire else: response = self._response_type(result=self._response_type.Result( **{method_name: result})) try: response = self._wiredata.to_lower(response) except Exception: # It should be an error when a response object that is fully # under our control cannot be lowered correctly. LOG.exception('to_lower error: %r, %r', request, response) return self._internal_server_error_wire LOG.debug('wire response: %r', response) return response def _make_error_response(self, error): error_name = self._match_error_type(error) if error_name is None: return None return self._response_type(error=self._response_type.Error( **{error_name: error}))
class Future: """Future object. The interface is divided into consumer-side and producer-side. Generally you make a ``Future`` object and pass it to both consumer and producer. On the producer side, you usually do: >>> with future.catching_exception(reraise=False): ... future.set_result(42) Then on the consumer side, you get the result: >>> future.get_result() 42 """ def __init__(self): self._condition = threading.Condition() self._completed = False self._result = None self._exception = None self._callbacks = [] self._consumed = False self._finalizer = None def __del__(self): if not (self._completed and self._exception is None and not self._consumed): return if self._finalizer is not None: try: self._finalizer(self._result) except BaseException: LOG.exception('finalizer error') return # Make a special case for None. if self._result is None: return LOG.warning( 'future is garbage-collected but result is never consumed: %s', # Call repr to format self here to avoid resurrecting self. repr(self), ) __repr__ = classes.make_repr( '{state} {self._result!r} {self._exception!r}', state=lambda self: 'completed' if self._completed else 'uncompleted', ) # # Consumer-side interface. # def is_completed(self): return self._completed def get_result(self, timeout=None): with self._condition: self._wait_for_completion(timeout) self._consumed = True if self._exception: raise self._exception # pylint: disable=raising-bad-type return self._result def get_exception(self, timeout=None): with self._condition: self._wait_for_completion(timeout) self._consumed = True return self._exception def _wait_for_completion(self, timeout): if not self._completed: self._condition.wait(timeout) if not self._completed: raise Timeout def add_callback(self, callback): """Add a callback that is called on completion. There are a few caveats of ``add_callback``: * If a callback is added when the future has completed, the callback is executed on the caller thread; otherwise it is executed on the producer thread. * Exceptions raised from the callbacks are logged and then swallowed. And these caveats are the reason that you normally should not use ``add_callback``. """ with self._condition: if not self._completed: self._callbacks.append(callback) return self._call_callback(callback) def _call_callback(self, callback): try: callback(self) except Exception: LOG.exception('callback err: %r, %r', self, callback) def set_finalizer(self, finalizer): """Set finalizer. The finalizer is called when future's result is set but is never consumed. You may use finalizer to release the result object. """ self._finalizer = finalizer # # Producer-side interface. # @contextlib.contextmanager def catching_exception(self, *, reraise): """Catch exception automatically. NOTE: It catches ``BaseException``, not the usual ``Exception``. As a result, when the producer thread raises ``SystemExit``, it is caught and re-raised in the consumer thread; thus, even if the producer thread is not the main thread, it may still call ``sys.exit``, and the ``SystemExit`` might reach the main thread through the future object. """ try: yield self except BaseException as exc: self.set_exception(exc) if reraise: raise def set_result(self, result): """Set future's result and complete the future. Once the future completes, further calling ``set_result`` or ``set_exception`` will be ignored. """ self._set_result_or_exception(result, None) def set_exception(self, exception): """Set future's result with an exception. Otherwise this is the same as ``set_result``. """ self._set_result_or_exception(None, exception) def _set_result_or_exception(self, result, exception): with self._condition: if self._completed: if exception: LOG.error('ignore exception: %r', self, exc_info=exception) else: LOG.error('ignore result: %r, %r', self, result) return self._result = result self._exception = exception self._completed = True callbacks, self._callbacks = self._callbacks, None self._condition.notify_all() for callback in callbacks: self._call_callback(callback)
class CompletionQueue: """Closable queue for waiting future completion. NOTE: Unlike other "regular" queues, since ``CompletionQueue`` only return completed futures, sometimes even when queue is not empty, ``get`` might not return anything if no future is completed yet. """ def __init__(self, iterable=()): self._lock = threading.Lock() self._uncompleted = Multiset(iterable) self._completed = queues.Queue() self._closed = False # Make a copy so that we may modify ``self._uncompleted`` while # iterating over its items. for f in tuple(self._uncompleted): f.add_callback(self._on_completion) # The two ``len(...)`` calls are not thread-safe, but probably does # not matter since this is just ``__repr__``. __repr__ = classes.make_repr( '{state} uncompleted={uncompleted} completed={completed}', state=lambda self: 'closed' if self._closed else 'open', uncompleted=lambda self: len(self._uncompleted), completed=lambda self: len(self._completed), ) def __bool__(self): with self._lock: return bool(self._uncompleted) or bool(self._completed) def __len__(self): with self._lock: return len(self._uncompleted) + len(self._completed) def is_closed(self): with self._lock: return self._closed def close(self, graceful=True): with self._lock: if graceful: items = [] else: items = self._completed.close(False) items.extend(self._uncompleted) self._uncompleted = () self._closed = True return items def get(self, timeout=None): """Return a completed future.""" with self._lock: if self._closed and not self._uncompleted and not self._completed: raise queues.Closed return self._completed.get(timeout) def put(self, future): with self._lock: if self._closed: raise queues.Closed self._uncompleted.add(future) future.add_callback(self._on_completion) def __iter__(self): """Iterate over completed futures. Unlike ``as_completed``, this does not accept ``timeout``. """ return self.as_completed() def as_completed(self, timeout=None): """Iterate over completed futures.""" timer = timers.make(timeout) while True: timer.start() try: future = self.get(timer.get_timeout()) except (queues.Empty, queues.Closed): break timer.stop() yield future def _on_completion(self, future): """Move future from uncompleted set to completed queue.""" with self._lock: if self._completed.is_closed(): return # This queue has been closed non-gracefully. # Since ``self._uncompleted`` is a ``Multiset``, ``remove`` # should not raise ``KeyError``. self._uncompleted.remove(future) self._completed.put(future)
class StreamBase: """In-memory stream base class. The semantics that this class implements is similar to a pipe, not a regular file (and ``close`` only closes the write-end of stream). Compared to a pipe, this class employs an unbounded buffer, and thus a writer is never blocked. This class provides both blocking and non-blocking interface. """ def __init__(self, buffer_type, data_type, newline): self._buffer_type = buffer_type self._data_type = data_type self._newline = newline self._buffer = self._buffer_type() self._closed = False self._gate = locks.Gate() def _make_buffer(self, data): if data: buffer = self._buffer_type(data) buffer.seek(len(data)) else: buffer = self._buffer_type() return buffer __repr__ = classes.make_repr( '{state}', state=lambda self: 'closed' if self._closed else 'open', ) def __aiter__(self): return self async def __anext__(self): line = await self.readline() if not line: raise StopAsyncIteration return line async def read(self, size=-1): while True: data = self.read_nonblocking(size) if data is None: await self._gate.wait() else: return data async def readline(self, size=-1): while True: line = self.readline_nonblocking(size) if line is None: await self._gate.wait() else: return line async def readlines(self, hint=None): if hint is None or hint <= 0: hint = float('+inf') lines = [] num_read = 0 async for line in self: lines.append(line) num_read += len(line) if num_read >= hint: break return lines async def write(self, data): return self.write_nonblocking(data) # # Non-blocking counterparts. # # There is no implementation for ``__iter__`` and ``readlines`` # because their interface is not (easily?) compatible with # non-blocking semantics. # NonblockingMethods = collections.namedtuple( 'NonblockingMethods', ( 'close', 'read', 'readline', 'write', ), ) @property def nonblocking(self): """Expose non-blocking interface via a file-like interface.""" return self.NonblockingMethods( close=self.close, read=self.read_nonblocking, readline=self.readline_nonblocking, write=self.write_nonblocking, ) def close(self): self._closed = True self._gate.unblock() def read_nonblocking(self, size=-1): data = self._buffer.getvalue() if not data: if self._closed: return data else: return None if size < 0: size = len(data) if size == 0: data = self._data_type() elif size >= len(data): self._buffer = self._buffer_type() else: self._buffer = self._make_buffer(data[size:]) data = data[:size] return data def readline_nonblocking(self, size=-1): data = self._buffer.getvalue() if not data: if self._closed: return data else: return None pos = data.find(self._newline) if pos < 0 and size < 0: if self._closed: size = len(data) else: return None elif size < 0 <= pos: size = pos + len(self._newline) elif pos < 0 <= size: pass # Nothing here. else: # pos >= 0 and size >= 0. size = min(size, pos + len(self._newline)) if size == 0: data = self._data_type() elif size >= len(data): self._buffer = self._buffer_type() else: self._buffer = self._make_buffer(data[size:]) data = data[:size] return data def write_nonblocking(self, data): ASSERT.false(self._closed) self._gate.unblock() return self._buffer.write(data)