Include websockets by default

core-api · tomchristie · Jan 27, 2016 · Jan 27, 2016 · Jan 28, 2016 · Jan 28, 2016
commit b33695ef0e73581dd73135c41d2143800c0d865d
diff --git a/coreapi/commandline.py b/coreapi/commandline.py
@@ -65,7 +65,8 @@ def get_client():
     credentials = get_credentials()
     headers = get_headers()
     http_transport = coreapi.transports.HTTPTransport(credentials, headers)
-    return coreapi.Client(transports=[http_transport])
+    websocket_transport = coreapi.transports.WebSocketsTransport()
+    return coreapi.Client(transports=[http_transport, websocket_transport])
 
 
 def get_document():
@@ -134,6 +135,28 @@ def get(url):
         set_history(history)
 
 
+@click.command(help='Watch a document at the given URL.')
+@click.argument('url')
+def watch(url):
+    client = get_client()
+    history = get_history()
+    heading = click.style('Watching %s' % url, bold=True)
+    watched = client.get(url)
+    try:
+        for doc in watched:
+            click.clear()
+            click.echo(heading)
+            click.echo(display(doc))
+            if isinstance(doc, coreapi.Document):
+                history = history.add(doc)
+                set_document(doc)
+                set_history(history)
+    except coreapi.exceptions.ErrorMessage as exc:
+        click.echo(display(exc.error))
+        sys.exit(1)
+
+
+
 @click.command(help='Load a document from disk.')
 @click.argument('input_file', type=click.File('rb'))
 @click.option('--format', default='corejson', type=click.Choice(['corejson', 'hal', 'hyperschema']))
@@ -546,6 +569,7 @@ def history_forward():
 
 
 client.add_command(get)
+client.add_command(watch)
 client.add_command(show)
 client.add_command(action)
 client.add_command(reload_document, name='reload')

diff --git a/coreapi/transports/__init__.py b/coreapi/transports/__init__.py
@@ -3,14 +3,15 @@
 from coreapi.exceptions import TransportError
 from coreapi.transports.base import BaseTransport
 from coreapi.transports.http import HTTPTransport
+from coreapi.transports.websockets import WebSocketsTransport
 import itypes
 
 
 __all__ = [
-    'BaseTransport', 'HTTPTransport'
+    'BaseTransport', 'HTTPTransport', 'WebSocketsTransport'
 ]
 
-default_transports = itypes.List([HTTPTransport()])
+default_transports = itypes.List([HTTPTransport(), WebSocketsTransport()])
 
 
 def determine_transport(url, transports=default_transports):

diff --git a/coreapi/transports/websockets.py b/coreapi/transports/websockets.py
@@ -0,0 +1,60 @@
+from coreapi.codecs import negotiate_decoder, default_decoders
+from coreapi.compat import force_bytes
+from coreapi.transports.base import BaseTransport
+from websocket import create_connection
+from websocket._exceptions import WebSocketConnectionClosedException
+import json
+import jsonpatch
+
+
+def _get_headers_and_body(content):
+    head, body = content.split('\n\n', 1)
+    key_value_pairs = [line.split(':', 1) for line in head.splitlines()]
+    headers = dict([
+        (key.strip().lower(), value.strip())
+        for key, value in key_value_pairs
+    ])
+    return (headers, body)
+
+
+def _decode_content(headers, content, decoders=None, base_url=None):
+    content_type = headers.get('content-type')
+    codec = negotiate_decoder(content_type, decoders=decoders)
+    return codec.load(content, base_url=base_url)
+
+
+def _apply_diff(content, diff):
+    # TODO: Negotiate diff scheme.
+    previous = json.loads(content.decode('utf-8'))
+    new = jsonpatch.apply_patch(previous, diff)
+    return force_bytes(json.dumps(new))
+
+
+def _get_request(decoders=None):
+    # TODO: Include User-Agent, X-Accept-Diff
+    if decoders is None:
+        decoders = default_decoders
+
+    accept = ', '.join([decoder.media_type for decoder in decoders])
+    return 'Accept: %s\n\n' % accept
+
+
+class WebSocketsTransport(BaseTransport):
+    schemes = ['ws', 'wss']
+
+    def transition(self, link, params=None, decoders=None, link_ancestors=None):
+        url = link.url
+        ws = create_connection(url)
+        request = _get_request(decoders)
+        ws.send(request)
+        content = ws.recv()
+        headers, body = _get_headers_and_body(content)
+        yield _decode_content(headers, body, decoders=decoders, base_url=url)
+        while True:
+            try:
+                diff = ws.recv()
+            except WebSocketConnectionClosedException:
+                return
+            patch = jsonpatch.JsonPatch.from_string(diff)
+            body = json.dumps(jsonpatch.apply_patch(json.loads(body), patch))
+            yield _decode_content(headers, body, decoders=decoders, base_url=url)