diff options
| -rw-r--r-- | pysite/base_route.py | 52 | ||||
| -rw-r--r-- | pysite/mixins.py | 63 | ||||
| -rw-r--r-- | pysite/route_manager.py | 35 | ||||
| -rw-r--r-- | pysite/views/api/bot/tag.py | 5 | ||||
| -rw-r--r-- | pysite/views/api/bot/user.py | 6 | ||||
| -rw-r--r-- | pysite/views/ws/__init__.py | 0 | ||||
| -rw-r--r-- | pysite/views/ws/echo.py | 18 | ||||
| -rw-r--r-- | pysite/websockets.py | 97 | ||||
| -rw-r--r-- | requirements.txt | 1 | ||||
| -rw-r--r-- | templates/ws_test.html | 2 | ||||
| -rw-r--r-- | ws_app.py | 28 |
11 files changed, 214 insertions, 93 deletions
diff --git a/pysite/base_route.py b/pysite/base_route.py index 8e6648ee..36b7fef5 100644 --- a/pysite/base_route.py +++ b/pysite/base_route.py @@ -3,15 +3,10 @@ import os import random import string -from _weakref import ref - from flask import Blueprint, jsonify, render_template from flask.views import MethodView -from rethinkdb.ast import Table - from pysite.constants import ErrorCodes -from pysite.database import RethinkDB class BaseView(MethodView): @@ -124,53 +119,6 @@ class APIView(RouteView): return response -class DBViewMixin: - """ - Mixin for views that make use of RethinkDB. It can automatically create a table with the specified primary - key using the attributes set at class-level. - - This class is intended to be mixed in alongside one of the other view classes. For example: - - >>> class MyView(APIView, DBViewMixin): - ... name = "my_view" # Flask internal name for this route - ... path = "/my_view" # Actual URL path to reach this route - ... table_name = "my_table" # Name of the table to create - ... table_primary_key = "username" # Primary key to set for this table - - You may omit `table_primary_key` and it will be defaulted to RethinkDB's default column - "id". - """ - - table_name = "" # type: str - table_primary_key = "id" # type: str - - @classmethod - def setup(cls: "DBViewMixin", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint): - """ - Set up the view by creating the table specified by the class attributes - this will also deal with multiple - inheritance by calling `super().setup()` as appropriate. - - :param manager: Instance of the current RouteManager (used to get a handle for the database object) - :param blueprint: Current Flask blueprint - """ - - if hasattr(super(), "setup"): - super().setup(manager, blueprint) - - if not cls.table_name: - raise RuntimeError("Routes using DBViewMixin must define `table_name`") - - cls._db = ref(manager.db) - manager.db.create_table(cls.table_name, primary_key=cls.table_primary_key) - - @property - def table(self) -> Table: - return self.db.query(self.table_name) - - @property - def db(self) -> RethinkDB: - return self._db() - - class ErrorView(BaseView): """ Error view, shown for a specific HTTP status code, as defined in the class attributes. diff --git a/pysite/mixins.py b/pysite/mixins.py new file mode 100644 index 00000000..891aa185 --- /dev/null +++ b/pysite/mixins.py @@ -0,0 +1,63 @@ +# coding=utf-8 +from _weakref import ref + +from flask import Blueprint + +from rethinkdb.ast import Table + +from pysite.database import RethinkDB + + +class DBMixin: + """ + Mixin for classes that make use of RethinkDB. It can automatically create a table with the specified primary + key using the attributes set at class-level. + + This class is intended to be mixed in alongside one of the other view classes. For example: + + >>> class MyView(APIView, DBMixin): + ... name = "my_view" # Flask internal name for this route + ... path = "/my_view" # Actual URL path to reach this route + ... table_name = "my_table" # Name of the table to create + ... table_primary_key = "username" # Primary key to set for this table + + This class will also work with Websockets: + + >>> class MyWeboscket(Websocket, DBMixin): + ... name = "my_websocket" + ... path = "/my_websocket" + ... table_name = "my_table" + ... table_primary_key = "username" + + You may omit `table_primary_key` and it will be defaulted to RethinkDB's default column - "id". + """ + + table_name = "" # type: str + table_primary_key = "id" # type: str + + @classmethod + def setup(cls: "DBMixin", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint): + """ + Set up the view by creating the table specified by the class attributes - this will also deal with multiple + inheritance by calling `super().setup()` as appropriate. + + :param manager: Instance of the current RouteManager (used to get a handle for the database object) + :param blueprint: Current Flask blueprint + """ + + if hasattr(super(), "setup"): + super().setup(manager, blueprint) + + if not cls.table_name: + raise RuntimeError("Routes using DBViewMixin must define `table_name`") + + cls._db = ref(manager.db) + manager.db.create_table(cls.table_name, primary_key=cls.table_primary_key) + + @property + def table(self) -> Table: + return self.db.query(self.table_name) + + @property + def db(self) -> RethinkDB: + return self._db() diff --git a/pysite/route_manager.py b/pysite/route_manager.py index 3a23619e..b3f71643 100644 --- a/pysite/route_manager.py +++ b/pysite/route_manager.py @@ -5,8 +5,11 @@ import os from flask import Blueprint, Flask +from flask_sockets import Sockets + from pysite.base_route import APIView, BaseView, ErrorView, RouteView from pysite.database import RethinkDB +from pysite.websockets import Websocket TEMPLATES_PATH = "../templates" STATIC_PATH = "../static" @@ -19,6 +22,8 @@ class RouteManager: self.app = Flask( __name__, template_folder=TEMPLATES_PATH, static_folder=STATIC_PATH, static_url_path="/static", ) + self.sockets = Sockets(self.app) + self.db = RethinkDB() self.app.secret_key = os.environ.get("WEBPAGE_SECRET_KEY", "super_secret") self.app.config["SERVER_NAME"] = os.environ.get("SERVER_NAME", "pythondiscord.com:8080") @@ -36,17 +41,29 @@ class RouteManager: self.subdomains = ['api', 'staff'] for sub in self.subdomains: - self.sub_blueprint = Blueprint(sub, __name__, subdomain=sub) + sub_blueprint = Blueprint(sub, __name__, subdomain=sub) - print(f"Loading Blueprint: {self.sub_blueprint.name}") - self.load_views(self.sub_blueprint, f"pysite/views/{sub}") - self.app.register_blueprint(self.sub_blueprint) + print(f"Loading Blueprint: {sub_blueprint.name}") + self.load_views(sub_blueprint, f"pysite/views/{sub}") + self.app.register_blueprint(sub_blueprint) print("") + # Load the websockets + self.ws_blueprint = Blueprint("ws", __name__) + + print("Loading websocket routes...") + self.load_views(self.ws_blueprint, "pysite/views/ws") + self.sockets.register_blueprint(self.ws_blueprint, url_prefix="/ws") + def run(self): - self.app.run( - port=int(os.environ.get("WEBPAGE_PORT", 8080)), debug="FLASK_DEBUG" in os.environ + from gevent.pywsgi import WSGIServer + from geventwebsocket.handler import WebSocketHandler + + server = WSGIServer( + ("", int(os.environ.get("WEBPAGE_PORT", 8080))), + self.app, handler_class=WebSocketHandler ) + server.serve_forever() def load_views(self, blueprint, location="pysite/views"): for filename in os.listdir(location): @@ -65,7 +82,11 @@ class RouteManager: cls is not ErrorView and cls is not RouteView and cls is not APIView and - BaseView in cls.__mro__ + cls is not Websocket and + ( + BaseView in cls.__mro__ or + Websocket in cls.__mro__ + ) ): cls.setup(self, blueprint) print(f">> View loaded: {cls.name: <15} ({module.__name__}.{cls_name})") diff --git a/pysite/views/api/bot/tag.py b/pysite/views/api/bot/tag.py index 363f98fe..2117d948 100644 --- a/pysite/views/api/bot/tag.py +++ b/pysite/views/api/bot/tag.py @@ -2,12 +2,13 @@ from flask import jsonify, request -from pysite.base_route import APIView, DBViewMixin +from pysite.base_route import APIView from pysite.constants import ErrorCodes from pysite.decorators import api_key +from pysite.mixins import DBMixin -class TagView(APIView, DBViewMixin): +class TagView(APIView, DBMixin): path = "/tag" name = "tag" table_name = "tag" diff --git a/pysite/views/api/bot/user.py b/pysite/views/api/bot/user.py index 8c2d8149..174407b8 100644 --- a/pysite/views/api/bot/user.py +++ b/pysite/views/api/bot/user.py @@ -4,10 +4,10 @@ from flask import jsonify from schema import Schema -from pysite.base_route import APIView, DBViewMixin +from pysite.base_route import APIView from pysite.constants import ValidationTypes from pysite.decorators import api_key, api_params - +from pysite.mixins import DBMixin SCHEMA = Schema([ { @@ -22,7 +22,7 @@ REQUIRED_KEYS = [ ] -class UserView(APIView, DBViewMixin): +class UserView(APIView, DBMixin): path = "/user" name = "user" table_name = "users" diff --git a/pysite/views/ws/__init__.py b/pysite/views/ws/__init__.py new file mode 100644 index 00000000..e69de29b --- /dev/null +++ b/pysite/views/ws/__init__.py diff --git a/pysite/views/ws/echo.py b/pysite/views/ws/echo.py new file mode 100644 index 00000000..135adfcf --- /dev/null +++ b/pysite/views/ws/echo.py @@ -0,0 +1,18 @@ +# coding=utf-8 +from pysite.websockets import Websocket + + +class EchoWebsocket(Websocket): + path = "/echo" + name = "ws_echo" + + def on_open(self): + print("Echo | WS opened.") + self.send("Hey, welcome!") + + def on_message(self, message): + print(f"Echo | Message: {message}") + self.send(message) + + def on_close(self): + print("Echo | WS closed.") diff --git a/pysite/websockets.py b/pysite/websockets.py new file mode 100644 index 00000000..f82217ef --- /dev/null +++ b/pysite/websockets.py @@ -0,0 +1,97 @@ +# coding=utf-8 +from flask import Blueprint + +from geventwebsocket.websocket import WebSocket + + +class Websocket: + """ + Base class for representing a Websocket. + + At minimum, you must implement the `on_message(self, message)` function. Without it, you won't be able to handle + any messages, and an error will be thrown! + + If you need access to the database, you can mix-in DBMixin, just like any view class: + + >>> class DBWebsocket(Websocket, DBMixin): + ... name = "db_websocket" + ... path = "/db_websocket" # This will be prefixed with "/ws" by the blueprint + ... table = "ws" + ... + ... def on_message(self, message): + ... self.send( + ... json.loads(self.db.get(self.table_name, message)) + ... ) + + Please note that an instance of this class is created for every websocket connected to the path. This does, however, + mean that you can store any state required by your websocket. + """ + + path = "" # type: str + name = "" # type: str + + def __init__(self, socket: WebSocket): + self.socket = socket + + def on_open(self): + """ + Called once when the websocket is opened. Optional. + """ + + def on_message(self, message: str): + """ + Called when a message is received by the websocket. + """ + + raise NotImplementedError() + + def on_close(self): + """ + Called once when the websocket is closed. Optional. + """ + + def send(self, message, binary=None): + """ + Send a message to the currently-connected websocket, if it's open. + + Nothing will happen if the websocket is closed. + """ + + if not self.socket.closed: + self.socket.send(message, binary=binary) + + @classmethod + def setup(cls: "type(Websocket)", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint): + """ + Set up the websocket object, calling `setup()` on any superclasses as necessary (for example, on the DB + mixin). + + This function will set up a websocket handler so that it behaves in a class-oriented way. It's up to you to + deal with message handling yourself, however. + """ + + if hasattr(super(), "setup"): + super().setup(manager, blueprint) + + if not cls.path or not cls.name: + raise RuntimeError("Websockets must have both `path` and `name` defined") + + def handle(socket: WebSocket): + """ + Wrap the current Websocket class, dispatching events to it as necessary. We're using gevent, so there's + no need to worry about blocking here. + """ + + ws = cls(socket) # Instantiate the current class, passing it the WS object + + ws.on_open() # Call the "on_open" handler + + while not socket.closed: # As long as the socket is open... + message = socket.receive() # Wait for a message + + if not socket.closed: # If the socket didn't just close (there's always a None message on closing) + ws.on_message(message) # Call the "on_message" handler + + ws.on_close() # The socket just closed, call the "on_close" handler + + blueprint.route(cls.path)(handle) # Register the handling function to the WS blueprint diff --git a/requirements.txt b/requirements.txt index 205154bb..cc2dc0e9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ gevent-websocket wsaccel ujson schema +flask_sockets diff --git a/templates/ws_test.html b/templates/ws_test.html index 921226c0..36be61d6 100644 --- a/templates/ws_test.html +++ b/templates/ws_test.html @@ -5,7 +5,7 @@ <h1>Open your JS console to test</h1> <script type="application/javascript"> - let ws = new WebSocket("wss://api.{{ server_name }}/ws/echo"); + let ws = new WebSocket("ws://{{ server_name }}/ws/echo"); ws.onopen = function(event) { console.log("WS opened! Use send() to send a message."); diff --git a/ws_app.py b/ws_app.py deleted file mode 100644 index cab43ac0..00000000 --- a/ws_app.py +++ /dev/null @@ -1,28 +0,0 @@ -# coding=utf-8 -import os - -from geventwebsocket import Resource, WebSocketApplication, WebSocketServer - - -class EchoApplication(WebSocketApplication): - def on_open(self): - print("Connection opened") - - def on_message(self, message, **kwargs): - print(f"<- {message}") - self.ws.send(message) - print(f"-> {message}") - - def on_close(self, reason): - print(reason) - - -app = WebSocketServer( - ('', os.environ.get("WS_PORT", 8080)), - Resource({ - "/ws/echo": EchoApplication # Dicts are ordered in Python 3.6 - }) -) - -if __name__ == "__main__": - app.serve_forever() |