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)
