From b6203b611ab796270a6ddcdd47dc132c3cb91466 Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Mon, 25 Nov 2024 20:32:11 +0000 Subject: [PATCH] cleanup types --- pyproject.toml | 6 + src/aoa.py | 17 +++ src/config.py | 8 + src/csi.py | 5 +- src/main.py | 16 +- src/preprocess.py | 25 +-- src/visualise/__init__.py | 24 +-- src/visualise/static/plot.js | 2 +- typings/flask_sock/__init__.pyi | 56 +++++++ typings/simple_websocket/__init__.pyi | 16 ++ typings/simple_websocket/aiows.pyi | 163 ++++++++++++++++++++ typings/simple_websocket/asgi.pyi | 23 +++ typings/simple_websocket/errors.pyi | 22 +++ typings/simple_websocket/ws.pyi | 211 ++++++++++++++++++++++++++ 14 files changed, 565 insertions(+), 29 deletions(-) create mode 100644 pyproject.toml create mode 100644 src/aoa.py create mode 100644 src/config.py create mode 100644 typings/flask_sock/__init__.pyi create mode 100644 typings/simple_websocket/__init__.pyi create mode 100644 typings/simple_websocket/aiows.pyi create mode 100644 typings/simple_websocket/asgi.pyi create mode 100644 typings/simple_websocket/errors.pyi create mode 100644 typings/simple_websocket/ws.pyi diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..693b92b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,6 @@ +[tool.ruff] +line-length = 88 + +[tool.pyright] +typeCheckingMode = "strict" +reportMissingTypeStubs = "warning" diff --git a/src/aoa.py b/src/aoa.py new file mode 100644 index 0000000..91340ee --- /dev/null +++ b/src/aoa.py @@ -0,0 +1,17 @@ +import numpy as np +import numpy.typing as npt +import config + + +class AoA: + def __init__(self): + self.historical_data = np.array([]) + pass + + def update(self, data: npt.NDArray[np.complex128]): + self.historical_data = np.concatenate([self.historical_data, data], axis=0) + if self.historical_data.shape[0] > config.AOA_SLIDING_WINDOW_SIZE: + self.historical_data = self.historical_data[ + -config.AOA_SLIDING_WINDOW_SIZE : + ] + pass diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..916f4c7 --- /dev/null +++ b/src/config.py @@ -0,0 +1,8 @@ +import os + +PREPROCESSING_SHORT_TERM_WINDOW_SIZE = 5 +PREPROCESSING_LONG_TERM_ALPHA = 0.1 +AOA_SLIDING_WINDOW_SIZE = 20 + +FEITCSI_IP_ADDRESS = os.getenv("IP_ADDRESS", "10.0.12.62") +FEITCSI_PORT = 8008 diff --git a/src/csi.py b/src/csi.py index 7386528..b984036 100644 --- a/src/csi.py +++ b/src/csi.py @@ -1,5 +1,6 @@ import struct import numpy as np +import numpy.typing as npt RATE_MCS_MOD_TYPE_POS = 8 @@ -93,13 +94,13 @@ class CSIHeader: class CSI: @staticmethod def parseCsiData(data: bytes, header: CSIHeader): - csi_matrix = np.zeros( + csi_matrix: npt.NDArray[np.complex128] = np.zeros( ( header.num_subcarriers, header.num_rx, header.num_tx, ), - dtype=complex, + dtype=np.complex128, ) pos = 0 for j in range(header.num_rx): diff --git a/src/main.py b/src/main.py index 1738c55..b060b82 100644 --- a/src/main.py +++ b/src/main.py @@ -1,10 +1,13 @@ import logging import socket +from typing import Callable from csi import CSI +from aoa import AoA from preprocess import Preprocessor import visualise import threading import struct +import config logging.basicConfig( level=logging.INFO, @@ -21,7 +24,7 @@ class FeitServer: frame_format: str = "VHT", ): self.ip = ip - self.port = 8008 + self.port = config.FEITCSI_PORT self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) self.server.connect((self.ip, self.port)) self.start_string = ( @@ -32,8 +35,11 @@ class FeitServer: self.server.send(b"stop\n") self.server.send(self.start_string.encode()) - def listen(self, callback): + def listen(self, callback: Callable[[CSI], None]): while True: + # This is the max size of a UDP packet. The size of the actual CSI + # packet will depend on the frame format and channel width, which + # changes the number of subcarriers data = self.server.recv(65535) try: csidata = CSI(data) @@ -42,15 +48,17 @@ class FeitServer: logging.error("Failed to parse CSI data") -def process_data(data): +def process_data(data: CSI): logging.debug("Got CSI frame") processed = preprocess.preprocess(data) visualise.add_data(processed) + aoa.update(processed) if __name__ == "__main__": - server = FeitServer("10.0.12.62") + server = FeitServer(config.FEITCSI_IP_ADDRESS) preprocess = Preprocessor() + aoa = AoA() # Start webapp in background thread webapp = threading.Thread(target=visualise.start) diff --git a/src/preprocess.py b/src/preprocess.py index c7807ca..366215a 100644 --- a/src/preprocess.py +++ b/src/preprocess.py @@ -2,19 +2,18 @@ import numpy as np from csi import CSI import logging from queue import Queue +import config +import numpy.typing as npt logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) class Preprocessor: - NUM_SHORT_TERM_ENTRIES = 10 - LONG_TERM_ALPHA = 0.1 - def __init__(self): - self.prev_entries = Queue(maxsize=100) - self.short_term_avg = np.zeros((1,), dtype=complex) - self.long_term_avg = np.zeros((1,), dtype=complex) + self.prev_entries: Queue[npt.NDArray[np.complex128]] = Queue(maxsize=100) + self.short_term_avg = np.zeros((1,), dtype=np.complex128) + self.long_term_avg = np.zeros((1,), dtype=np.complex128) def preprocess(self, csi: CSI): h = csi.matrix @@ -22,18 +21,20 @@ class Preprocessor: # Assume that all csi matrices will have the same shape if self.short_term_avg.shape != h_hat.shape: - self.short_term_avg = np.zeros(h_hat.shape, dtype=complex) - self.long_term_avg = np.zeros(h_hat.shape, dtype=complex) + self.short_term_avg = np.zeros(h_hat.shape, dtype=np.complex128) + self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex128) self.prev_entries = Queue(maxsize=100) while self.prev_entries.full(): old_value = self.prev_entries.get() - self.short_term_avg -= old_value / self.NUM_SHORT_TERM_ENTRIES + self.short_term_avg -= ( + old_value / config.PREPROCESSING_SHORT_TERM_WINDOW_SIZE + ) self.prev_entries.put(h_hat) - self.short_term_avg += h_hat / self.NUM_SHORT_TERM_ENTRIES + self.short_term_avg += h_hat / config.PREPROCESSING_SHORT_TERM_WINDOW_SIZE self.long_term_avg = ( - self.long_term_avg * (1 - self.LONG_TERM_ALPHA) - + self.short_term_avg * self.LONG_TERM_ALPHA + self.long_term_avg * (1 - config.PREPROCESSING_LONG_TERM_ALPHA) + + self.short_term_avg * config.PREPROCESSING_LONG_TERM_ALPHA ) # Remove static components diff --git a/src/visualise/__init__.py b/src/visualise/__init__.py index d7443a8..05732b6 100644 --- a/src/visualise/__init__.py +++ b/src/visualise/__init__.py @@ -1,8 +1,8 @@ -from flask import Flask, render_template, request +from flask import Flask, render_template from flask_sock import Sock import numpy as np import numpy.typing as npt -import json +from simple_websocket import Server app = Flask(__name__) sock = Sock(app) @@ -16,17 +16,21 @@ def preprocessed(): return render_template("preprocessed.html") -subscribers = {} +type Subscriber = Server +subscriber_settings: dict[Subscriber, tuple[int, int, int]] = {} @sock.route("/data") -def get_data(sock): +def get_data(sock: Subscriber): while True: msg = sock.receive() - subscribers[sock] = list(map(int, msg.split())) + if len(msg.split()) != 3: + break + subcarrier, rx, tx = map(int, msg.split()) + subscriber_settings[sock] = (subcarrier, rx, tx) -def add_data(new_data): +def add_data(new_data: npt.NDArray[np.complex128]): global data if data.size == 0: data = np.expand_dims(new_data, axis=0) @@ -36,17 +40,17 @@ def add_data(new_data): # Only keep latest 100 entries if data.shape[0] > 100: data = data[-100:] - to_remove = [] - for subscriber in subscribers: + to_remove: list[Subscriber] = [] + for subscriber in subscriber_settings: try: - subcarrier, rx, tx = subscribers[subscriber] + subcarrier, rx, tx = subscriber_settings[subscriber] subscriber.send(new_data.real[subcarrier, rx, tx]) except Exception as e: to_remove.append(subscriber) print(e) for subscriber in to_remove: - del subscribers[subscriber] + del subscriber_settings[subscriber] def start(): diff --git a/src/visualise/static/plot.js b/src/visualise/static/plot.js index 4040704..a481928 100644 --- a/src/visualise/static/plot.js +++ b/src/visualise/static/plot.js @@ -12,7 +12,7 @@ Plotly.newPlot( { x: [], y: [], - type: "scatter", // Line chart + type: "scatter", mode: "lines+markers", // Line + markers }, ], diff --git a/typings/flask_sock/__init__.pyi b/typings/flask_sock/__init__.pyi new file mode 100644 index 0000000..6e139ce --- /dev/null +++ b/typings/flask_sock/__init__.pyi @@ -0,0 +1,56 @@ +""" +This type stub file was generated by pyright. +""" + +from functools import wraps +from flask import Blueprint, Response, current_app, request +from simple_websocket import ConnectionClosed, Server +from typing import Callable + +class Sock: + """Instantiate the Flask-Sock extension. + + :param app: The Flask application instance. If not provided, it must be + initialized later by calling the :func:`Sock.init_app` method. + """ + def __init__(self, app=...) -> None: ... + def init_app(self, app): # -> None: + """Initialize the Flask-Socket extension. + + + :param app: The Flask application instance. This method only needs to + be called if the application instance was not passed as + an argument to the constructor. + """ + ... + + def route( + self, path: str, bp: Blueprint | None = None + ) -> Callable[ + [Callable[[Server], None]], Callable[[Server], None] + ]: # -> Callable[..., None]: + """Decorator to create a WebSocket route. + + The decorated function will be invoked when a WebSocket client + establishes a connection, with a WebSocket connection object passed + as an argument. Example:: + + @sock.route('/ws') + def websocket_route(ws): + # The ws object has the following methods: + # - ws.send(data) + # - ws.receive(timeout=None) + # - ws.close(reason=None, message=None) + + If the route has variable components, the ``ws`` argument needs to be + included before them. + + :param path: the URL associated with the route. + :param bp: the blueprint on which to register the route. If not given, + the route is attached directly to the Flask application + instance. When a blueprint is used, the application is + responsible for the blueprint's registration. + :param kwargs: additional route options. See the Flask documentation + for the ``app.route`` decorator for details. + """ + ... diff --git a/typings/simple_websocket/__init__.pyi b/typings/simple_websocket/__init__.pyi new file mode 100644 index 0000000..8da6b2c --- /dev/null +++ b/typings/simple_websocket/__init__.pyi @@ -0,0 +1,16 @@ +""" +This type stub file was generated by pyright. +""" + +from .ws import Client, Server +from .aiows import AioClient, AioServer +from .errors import ConnectionClosed, ConnectionError + +__all__ = [ + "Client", + "Server", + "AioClient", + "AioServer", + "ConnectionClosed", + "ConnectionError", +] diff --git a/typings/simple_websocket/aiows.pyi b/typings/simple_websocket/aiows.pyi new file mode 100644 index 0000000..d21fd4d --- /dev/null +++ b/typings/simple_websocket/aiows.pyi @@ -0,0 +1,163 @@ +""" +This type stub file was generated by pyright. +""" + +class AioBase: + def __init__(self, connection_type=..., receive_bytes=..., ping_interval=..., max_message_size=...) -> None: + ... + + async def connect(self): # -> None: + ... + + async def handshake(self): # -> None: + ... + + async def send(self, data): # -> None: + """Send data over the WebSocket connection. + + :param data: The data to send. If ``data`` is of type ``bytes``, then + a binary message is sent. Else, the message is sent in + text format. + """ + ... + + async def receive(self, timeout=...): # -> None: + """Receive data over the WebSocket connection. + + :param timeout: Amount of time to wait for the data, in seconds. Set + to ``None`` (the default) to wait indefinitely. Set + to 0 to read without blocking. + + The data received is returned, as ``bytes`` or ``str``, depending on + the type of the incoming message. + """ + ... + + async def close(self, reason=..., message=...): # -> None: + """Close the WebSocket connection. + + :param reason: A numeric status code indicating the reason of the + closure, as defined by the WebSocket specification. The + default is 1000 (normal closure). + :param message: A text message to be sent to the other side. + """ + ... + + def choose_subprotocol(self, request): # -> None: + ... + + + +class AioServer(AioBase): + """This class implements a WebSocket server. + + Instead of creating an instance of this class directly, use the + ``accept()`` class method to create individual instances of the server, + each bound to a client request. + """ + def __init__(self, request, subprotocols=..., receive_bytes=..., ping_interval=..., max_message_size=...) -> None: + ... + + @classmethod + async def accept(cls, aiohttp=..., asgi=..., sock=..., headers=..., subprotocols=..., receive_bytes=..., ping_interval=..., max_message_size=...): # -> WebSocketASGI | Self: + """Accept a WebSocket connection from a client. + + :param aiohttp: The request object from aiohttp. If this argument is + provided, ``asgi``, ``sock`` and ``headers`` must not + be set. + :param asgi: A (scope, receive, send) tuple from an ASGI request. If + this argument is provided, ``aiohttp``, ``sock`` and + ``headers`` must not be set. + :param sock: A connected socket to use. If this argument is provided, + ``aiohttp`` and ``asgi`` must not be set. The ``headers`` + argument must be set with the incoming request headers. + :param headers: A dictionary with the incoming request headers, when + ``sock`` is used. + :param subprotocols: A list of supported subprotocols, or ``None`` (the + default) to disable subprotocol negotiation. + :param receive_bytes: The size of the receive buffer, in bytes. The + default is 4096. + :param ping_interval: Send ping packets to clients at the requested + interval in seconds. Set to ``None`` (the + default) to disable ping/pong logic. Enable to + prevent disconnections when the line is idle for + a certain amount of time, or to detect + unresponsive clients and disconnect them. A + recommended interval is 25 seconds. + :param max_message_size: The maximum size allowed for a message, in + bytes, or ``None`` for no limit. The default + is ``None``. + """ + ... + + async def handshake(self): # -> None: + ... + + def choose_subprotocol(self, request): # -> None: + """Choose a subprotocol to use for the WebSocket connection. + + The default implementation selects the first protocol requested by the + client that is accepted by the server. Subclasses can override this + method to implement a different subprotocol negotiation algorithm. + + :param request: A ``Request`` object. + + The method should return the subprotocol to use, or ``None`` if no + subprotocol is chosen. + """ + ... + + + +class AioClient(AioBase): + """This class implements a WebSocket client. + + Instead of creating an instance of this class directly, use the + ``connect()`` class method to create an instance that is connected to a + server. + """ + def __init__(self, url, subprotocols=..., headers=..., receive_bytes=..., ping_interval=..., max_message_size=..., ssl_context=...) -> None: + ... + + @classmethod + async def connect(cls, url, subprotocols=..., headers=..., receive_bytes=..., ping_interval=..., max_message_size=..., ssl_context=..., thread_class=..., event_class=...): # -> Self: + """Returns a WebSocket client connection. + + :param url: The connection URL. Both ``ws://`` and ``wss://`` URLs are + accepted. + :param subprotocols: The name of the subprotocol to use, or a list of + subprotocol names in order of preference. Set to + ``None`` (the default) to not use a subprotocol. + :param headers: A dictionary or list of tuples with additional HTTP + headers to send with the connection request. Note that + custom headers are not supported by the WebSocket + protocol, so the use of this parameter is not + recommended. + :param receive_bytes: The size of the receive buffer, in bytes. The + default is 4096. + :param ping_interval: Send ping packets to the server at the requested + interval in seconds. Set to ``None`` (the + default) to disable ping/pong logic. Enable to + prevent disconnections when the line is idle for + a certain amount of time, or to detect an + unresponsive server and disconnect. A recommended + interval is 25 seconds. In general it is + preferred to enable ping/pong on the server, and + let the client respond with pong (which it does + regardless of this setting). + :param max_message_size: The maximum size allowed for a message, in + bytes, or ``None`` for no limit. The default + is ``None``. + :param ssl_context: An ``SSLContext`` instance, if a default SSL + context isn't sufficient. + """ + ... + + async def handshake(self): # -> None: + ... + + async def close(self, reason=..., message=...): # -> None: + ... + + + diff --git a/typings/simple_websocket/asgi.pyi b/typings/simple_websocket/asgi.pyi new file mode 100644 index 0000000..c2624d7 --- /dev/null +++ b/typings/simple_websocket/asgi.pyi @@ -0,0 +1,23 @@ +""" +This type stub file was generated by pyright. +""" + +class WebSocketASGI: + def __init__(self, scope, receive, send, subprotocols=...) -> None: + ... + + @classmethod + async def accept(cls, scope, receive, send, subprotocols=...): # -> WebSocketASGI: + ... + + async def receive(self): + ... + + async def send(self, data): # -> None: + ... + + async def close(self): # -> None: + ... + + + diff --git a/typings/simple_websocket/errors.pyi b/typings/simple_websocket/errors.pyi new file mode 100644 index 0000000..4d2fff1 --- /dev/null +++ b/typings/simple_websocket/errors.pyi @@ -0,0 +1,22 @@ +""" +This type stub file was generated by pyright. +""" + +class SimpleWebsocketError(RuntimeError): + ... + + +class ConnectionError(SimpleWebsocketError): + """Connection error exception class.""" + def __init__(self, status_code=...) -> None: + ... + + + +class ConnectionClosed(SimpleWebsocketError): + """Connection closed exception class.""" + def __init__(self, reason=..., message=...) -> None: + ... + + + diff --git a/typings/simple_websocket/ws.pyi b/typings/simple_websocket/ws.pyi new file mode 100644 index 0000000..ac83896 --- /dev/null +++ b/typings/simple_websocket/ws.pyi @@ -0,0 +1,211 @@ +""" +This type stub file was generated by pyright. +""" + +class Base: + def __init__( + self, + sock=..., + connection_type=..., + receive_bytes=..., + ping_interval=..., + max_message_size=..., + thread_class=..., + event_class=..., + selector_class=..., + ) -> None: ... + def handshake(self): # -> None: + ... + def send(self, data: bytes | str) -> None: # -> None: + """Send data over the WebSocket connection. + + :param data: The data to send. If ``data`` is of type ``bytes``, then + a binary message is sent. Else, the message is sent in + text format. + """ + ... + + def receive(self, timeout: float | None = ...) -> bytes | str: # -> None: + """Receive data over the WebSocket connection. + + :param timeout: Amount of time to wait for the data, in seconds. Set + to ``None`` (the default) to wait indefinitely. Set + to 0 to read without blocking. + + The data received is returned, as ``bytes`` or ``str``, depending on + the type of the incoming message. + """ + ... + + def close(self, reason=..., message=...): # -> None: + """Close the WebSocket connection. + + :param reason: A numeric status code indicating the reason of the + closure, as defined by the WebSocket specification. The + default is 1000 (normal closure). + :param message: A text message to be sent to the other side. + """ + ... + + def choose_subprotocol(self, request): # -> None: + ... + +class Server(Base): + """This class implements a WebSocket server. + + Instead of creating an instance of this class directly, use the + ``accept()`` class method to create individual instances of the server, + each bound to a client request. + """ + def __init__( + self, + environ, + subprotocols=..., + receive_bytes=..., + ping_interval=..., + max_message_size=..., + thread_class=..., + event_class=..., + selector_class=..., + ) -> None: ... + @classmethod + def accept( + cls, + environ, + subprotocols=..., + receive_bytes=..., + ping_interval=..., + max_message_size=..., + thread_class=..., + event_class=..., + selector_class=..., + ): # -> Self: + """Accept a WebSocket connection from a client. + + :param environ: A WSGI ``environ`` dictionary with the request details. + Among other things, this class expects to find the + low-level network socket for the connection somewhere + in this dictionary. Since the WSGI specification does + not cover where or how to store this socket, each web + server does this in its own different way. Werkzeug, + Gunicorn, Eventlet and Gevent are the only web servers + that are currently supported. + :param subprotocols: A list of supported subprotocols, or ``None`` (the + default) to disable subprotocol negotiation. + :param receive_bytes: The size of the receive buffer, in bytes. The + default is 4096. + :param ping_interval: Send ping packets to clients at the requested + interval in seconds. Set to ``None`` (the + default) to disable ping/pong logic. Enable to + prevent disconnections when the line is idle for + a certain amount of time, or to detect + unresponsive clients and disconnect them. A + recommended interval is 25 seconds. + :param max_message_size: The maximum size allowed for a message, in + bytes, or ``None`` for no limit. The default + is ``None``. + :param thread_class: The ``Thread`` class to use when creating + background threads. The default is the + ``threading.Thread`` class from the Python + standard library. + :param event_class: The ``Event`` class to use when creating event + objects. The default is the `threading.Event`` + class from the Python standard library. + :param selector_class: The ``Selector`` class to use when creating + selectors. The default is the + ``selectors.DefaultSelector`` class from the + Python standard library. + """ + ... + + def handshake(self): # -> None: + ... + def choose_subprotocol(self, request): # -> None: + """Choose a subprotocol to use for the WebSocket connection. + + The default implementation selects the first protocol requested by the + client that is accepted by the server. Subclasses can override this + method to implement a different subprotocol negotiation algorithm. + + :param request: A ``Request`` object. + + The method should return the subprotocol to use, or ``None`` if no + subprotocol is chosen. + """ + ... + +class Client(Base): + """This class implements a WebSocket client. + + Instead of creating an instance of this class directly, use the + ``connect()`` class method to create an instance that is connected to a + server. + """ + def __init__( + self, + url, + subprotocols=..., + headers=..., + receive_bytes=..., + ping_interval=..., + max_message_size=..., + ssl_context=..., + thread_class=..., + event_class=..., + ) -> None: ... + @classmethod + def connect( + cls, + url: str, + subprotocols=..., + headers=..., + receive_bytes=..., + ping_interval=..., + max_message_size=..., + ssl_context=..., + thread_class=..., + event_class=..., + ): # -> Self: + """Returns a WebSocket client connection. + + :param url: The connection URL. Both ``ws://`` and ``wss://`` URLs are + accepted. + :param subprotocols: The name of the subprotocol to use, or a list of + subprotocol names in order of preference. Set to + ``None`` (the default) to not use a subprotocol. + :param headers: A dictionary or list of tuples with additional HTTP + headers to send with the connection request. Note that + custom headers are not supported by the WebSocket + protocol, so the use of this parameter is not + recommended. + :param receive_bytes: The size of the receive buffer, in bytes. The + default is 4096. + :param ping_interval: Send ping packets to the server at the requested + interval in seconds. Set to ``None`` (the + default) to disable ping/pong logic. Enable to + prevent disconnections when the line is idle for + a certain amount of time, or to detect an + unresponsive server and disconnect. A recommended + interval is 25 seconds. In general it is + preferred to enable ping/pong on the server, and + let the client respond with pong (which it does + regardless of this setting). + :param max_message_size: The maximum size allowed for a message, in + bytes, or ``None`` for no limit. The default + is ``None``. + :param ssl_context: An ``SSLContext`` instance, if a default SSL + context isn't sufficient. + :param thread_class: The ``Thread`` class to use when creating + background threads. The default is the + ``threading.Thread`` class from the Python + standard library. + :param event_class: The ``Event`` class to use when creating event + objects. The default is the `threading.Event`` + class from the Python standard library. + """ + ... + + def handshake(self): # -> None: + ... + def close(self, reason=..., message=...): # -> None: + ...