Fix #2031 - Add pyscript.WebSocket to the mix (#2042)

* Fix #2031 - Add pyscript.WebSocket to the mix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Working on a test case anyone can run

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

pyscript.core/src/stdlib/pyscript/__init__.py

@@ -41,6 +41,7 @@ from pyscript.magic_js import (
     sync,
     window,
 )
+from pyscript.websocket import WebSocket
 
 try:
     from pyscript.event_handling import when

pyscript.core/src/stdlib/pyscript/fetch.py

@@ -1,6 +1,7 @@
 import json
 
 import js
+from pyscript.util import as_bytearray
 
 
 ### wrap the response to grant Pythonic results
@@ -12,14 +13,6 @@ class _Response:
     def __getattr__(self, attr):
         return getattr(self._response, attr)
 
-    def _as_bytearray(self, buffer):
-        ui8a = js.Uint8Array.new(buffer)
-        size = ui8a.length
-        ba = bytearray(size)
-        for i in range(0, size):
-            ba[i] = ui8a[i]
-        return ba
-
     # exposed methods with Pythonic results
     async def arrayBuffer(self):
         buffer = await self._response.arrayBuffer()
@@ -27,14 +20,14 @@ class _Response:
         if hasattr(buffer, "to_py"):
             return buffer.to_py()
         # shims in MicroPython
-        return memoryview(self._as_bytearray(buffer))
+        return memoryview(as_bytearray(buffer))
 
     async def blob(self):
         return await self._response.blob()
 
     async def bytearray(self):
         buffer = await self._response.arrayBuffer()
-        return self._as_bytearray(buffer)
+        return as_bytearray(buffer)
 
     async def json(self):
         return json.loads(await self.text())

pyscript.core/src/stdlib/pyscript/util.py

@@ -1,3 +1,15 @@
+import js
+
+
+def as_bytearray(buffer):
+    ui8a = js.Uint8Array.new(buffer)
+    size = ui8a.length
+    ba = bytearray(size)
+    for i in range(0, size):
+        ba[i] = ui8a[i]
+    return ba
+
+
 class NotSupported:
     """
     Small helper that raises exceptions if you try to get/set any attribute on

pyscript.core/src/stdlib/pyscript/websocket.py

@@ -0,0 +1,64 @@
+import js
+from pyscript.util import as_bytearray
+
+code = "code"
+protocols = "protocols"
+reason = "reason"
+
+
+class EventMessage:
+    def __init__(self, event):
+        self._event = event
+
+    def __getattr__(self, attr):
+        value = getattr(self._event, attr)
+
+        if attr == "data" and not isinstance(value, str):
+            if hasattr(value, "to_py"):
+                return value.to_py()
+            # shims in MicroPython
+            return memoryview(as_bytearray(value))
+
+        return value
+
+
+class WebSocket(object):
+    CONNECTING = 0
+    OPEN = 1
+    CLOSING = 2
+    CLOSED = 3
+
+    def __init__(self, **kw):
+        url = kw["url"]
+        socket = None
+        if protocols in kw:
+            socket = js.WebSocket.new(url, kw[protocols])
+        else:
+            socket = js.WebSocket.new(url)
+        object.__setattr__(self, "_ws", socket)
+
+    def __getattr__(self, attr):
+        return getattr(self._ws, attr)
+
+    def __setattr__(self, attr, value):
+        if attr == "onmessage":
+            setattr(self._ws, attr, lambda e: value(EventMessage(e)))
+        else:
+            setattr(self._ws, attr, value)
+
+    def close(self, **kw):
+        if code in kw and reason in kw:
+            self._ws.close(kw[code], kw[reason])
+        elif code in kw:
+            self._ws.close(kw[code])
+        else:
+            self._ws.close()
+
+    def send(self, data):
+        if isinstance(data, str):
+            self._ws.send(data)
+        else:
+            buffer = js.Uint8Array.new(len(data))
+            for pos, b in enumerate(data):
+                buffer[pos] = b
+            self._ws.send(buffer)