119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
from typing import Any, Callable, Dict, Generic, Optional, Set, TypeVar
|
|
|
|
from .connection import Connection, ConnectionStatus
|
|
from .envelope import Envelope, EnvelopeNotify, EnvelopeRequest
|
|
from .util import make_id
|
|
from .watchable import Watchable
|
|
from . import TransportBase
|
|
|
|
T = TypeVar('T')
|
|
Thunk = Callable[[], None]
|
|
|
|
|
|
class Connection(Generic[T]):
|
|
def __init__(self, transport: TransportBase, device_id: str, description: str, methods: T, send_envelope: Callable[[Connection[T], Envelope], None]) -> None:
|
|
self.status = Watchable(ConnectionStatus.CONNECTING)
|
|
self._close_cbs: Set[Thunk] = set()
|
|
self.description = description
|
|
self._transport = transport
|
|
self._device_id = device_id
|
|
self._other_device_id: Optional[str] = None
|
|
self._methods = methods
|
|
self._send_envelope = send_envelope
|
|
self._deferred_requests: Dict[str, Deferred[Any]] = {}
|
|
self._last_seen = 0
|
|
|
|
@property
|
|
def is_closed(self) -> bool:
|
|
return self.status == ConnectionStatus.CLOSED
|
|
|
|
def on_close(self, func: Thunk) -> Thunk:
|
|
if self.is_closed:
|
|
raise RpcErrorUseAfterClose('the connection is closed')
|
|
|
|
self._close_cbs.add(func)
|
|
|
|
def del_cb():
|
|
self._close_cbs.remove(func)
|
|
|
|
return del_cb
|
|
|
|
def close(self) -> None:
|
|
if self.is_closed:
|
|
return
|
|
|
|
self.status = ConnectionStatus.CLOSED
|
|
|
|
for func in self._close_cbs:
|
|
func()
|
|
|
|
self._close_cbs.clear()
|
|
|
|
def handle_incoming_envelope(env: Envelope[T]) -> None:
|
|
# TODO: maybe this function should be in a lock to ensure it only runs one at a time
|
|
if self.is_closed:
|
|
raise RpcErrorUseAfterClose('the connection is closed')
|
|
|
|
# TODO: throw error if status is ERROR?
|
|
if env.kind == 'NOTIFY':
|
|
if not hasattr(self._methods, env.method):
|
|
logger.warn(f'> error in NOTIFY handler: no notify method called {env.method}')
|
|
else:
|
|
try:
|
|
self._methods[env.method](*env.args)
|
|
except BaseException as exc:
|
|
logger.warn(f'> error when running NOTIFY method: {env} {exc}')
|
|
|
|
elif env.kind == 'REQUEST':
|
|
try:
|
|
if not hasattr(self._methods, env.method):
|
|
raise RpcErrorUnknownMethod(f'unknown method in REQUEST: {env.method}')
|
|
|
|
data = self._methods[env.method](*env.args)
|
|
response_env_data = EnvelopeResponseWithData(kind='RESPONSE', from_device_id=self._device_id, envelope_id=env.envelope_id, data=data)
|
|
self._send_envelope(self, response_env_data)
|
|
except BaseException as exc:
|
|
response_env_error = EnvelopeResponseWithError(kind='RESPONSE', from_device_id=self.device_id, envelope_id=env.envelope_id, error=str(exc))
|
|
|
|
self._send_envelope(self, response_error)
|
|
elif env.kind == 'RESPONSE':
|
|
deferred = self._deferred_requests.get(env.envelope_id)
|
|
|
|
if not deferred:
|
|
return
|
|
|
|
if env.data:
|
|
deferred.resolve(env.data)
|
|
elif env.error:
|
|
deferred.reject(RpcErrorFromMethod(env.error))
|
|
else:
|
|
logger.warn('> RESPONSE has neither data nor error. This should never happen')
|
|
|
|
deferred.reject(RpcError('> RESPONSE has neither data nor error??'))
|
|
|
|
self._deferred_requests.remove(env.envelope_id)
|
|
|
|
def notify(self, method: MethodKey, *args: Parameters[BagType[MethodKey]]) -> None:
|
|
if self.is_closed:
|
|
raise RpcErrorUseAfterClose('the connection is closed')
|
|
|
|
env = EnvelopeNotify(
|
|
kind='NOTIFY',
|
|
from_device_id=self._device_id,
|
|
envelope_id=f'env:{make_id()}',
|
|
method=method,
|
|
args=args
|
|
)
|
|
self._send_envelope(self, env)
|
|
|
|
def request(self, method: str, *args: Parameters[BagType[MethodKey]]) -> ReturnType[BagType[MethodKey]]:
|
|
if self.is_closed:
|
|
raise RpcErrorUseAfterClose('the connection is closed')
|
|
|
|
env = EnvelopeRequest(kind='REQUEST', from_device_id=self._device_id, envelope_id=f'env:{make_id}', method=method, args=args)
|
|
deferred = make_deferred()
|
|
self._deferred_requests[env.envelope_id] = deferred
|
|
self._send_envelope(self, env)
|
|
|
|
return deferred.promise
|