final

earthsnake/syncer/connection.py

@@ -0,0 +1,118 @@
+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