execlib/execlog/handler.py

import re
import logging
from pathlib import Path

from inotify_simple import flags
from starlette.endpoints import WebSocketEndpoint


logger = logging.getLogger(__name__)

#page_re = re.compile(r'https?:\/\/.*?\/(.*?)(?:\?|\.html|$)')
page_re = re.compile(r'https?:\/\/.*?\/(.*?)$')

def client_reload_wrap(reloaded_file):
    rpath = Path(reloaded_file)
    static_extensions = ['.js', '.css']

    if rpath.suffix in static_extensions:
        return lambda _: True
    else:
        return lambda c: Path(c).with_suffix('.html') == rpath
    

class Handler(WebSocketEndpoint):
    '''
    Subclasses WebSocketEndpoint to be attached to live reload endpoints.

    Note: Reload model
        - Served HTML files are generated from templates that include livereload JS and the
          target livereload server (port manually set prior to site build).
        - When pages are visited (be they served from NGINX or via the development
          server), the livereload.js attempts to connect to the known livereload WS
          endpoint.
        - FastAPI routes the request to _this_ endpoint, and `on_connect` is called.
        - Upon successful connection, the livereload JS client sends a "hello" message.
          This is picked up as the first post-acceptance message, and captured by the
          `on_receive` method.
        - `on_receive` subsequently initiates a formal handshake, sending back a "hello"
          command and waiting the "info" command from the client.
        - If the "info" command is received successfully and contains the requesting
          page's URL, the handshake completes and the websocket is added to the class'
          `live_clients` tracker.
        - Later, when a file in a watch path of the server's watcher is _modified_,
          `reload_clients` will be called from within the originating server's event loop,
          and pass in the FS event associated with the file change. `client_reload_wrap`
          is used to wrap a boolean checker method for whether or not to reload clients
          given the FS event.

          TODO: flesh out the reload wrapper to incorporate more complex filters and/or
          transformations when determining when to reload certain clients.
    '''
    encoding = 'json'
    live_clients = {}

    async def on_connect(self, websocket):
        await websocket.accept()

    async def on_receive(self, websocket, data):
        '''
        Note: On page names
            When websockets connect, they simply communicate the exact URL from the origin
            page. The websocket is then indexed to possibly variable page names (often
            without an `.html` suffix, but occasionally with). The `client_reload_wrap` is
            then responsible for taking this client page name and normalizing it to be
            matched with full file names (i.e., suffix always included).
        '''
        url = await self._lr_handshake(websocket, data)

        if url is None:
            logger.warning('Client handshake failed, ignoring')
            return

        origin_m = page_re.search(url)
        if origin_m is not None:
            origin_page = origin_m.group(1) 

            # assume index.html if connected to empty name
            if origin_page == '':
                origin_page = 'index.html'
        else:
            origin_page = '<unidentified>.null'

        self.live_clients[origin_page] = websocket
        logger.info(f'Reloader connected to [{origin_page}] ({len(self.live_clients)} live clients)')

    async def on_disconnect(self, websocket, close_code):
        remove_page = None
        for page, ws in self.live_clients.items():
            if ws == websocket:
                remove_page = page

        if remove_page is not None:
            logger.info(f'Client for [{remove_page}] disconnected, removing')
            self.live_clients.pop(remove_page)

    @classmethod
    async def reload_clients(cls, event):
        '''
        Method targeted as a watcher callback. This async method is scheduled in a
        thread-safe manner by the watcher to be ran in the FastAPI event loop.
        '''
        logger.info(f'> [{event.name}] changed on disk')
        should_reload = client_reload_wrap(event.name)

        for page, ws in cls.live_clients.items():
            if should_reload(page):
                logger.info(f'>> Reloading client for [{page}]')
                await ws.send_json({
                    'command' : 'reload',
                    'path'    : page,
                    'liveCSS' : True,
                    'liveImg' : True,
                })
        
    @staticmethod
    async def _lr_handshake(websocket, hello):
        '''
        Handshake with livereload.js

        1. client send 'hello'
        2. server reply 'hello'
        3. client send 'info'
        '''
        # 1. await client hello after accept
        #hello = await websocket.receive_json()

        if hello.get('command') != 'hello':
            logger.warning('Client handshake failed at "hello" stage')
            return 
        
        # 2. send hello to client
        await websocket.send_json({
            'command': 'hello',
            'protocols': [
                'http://livereload.com/protocols/official-7',
            ],
            'serverName': 'livereload-tornado',
        })

        # 3. await info response
        info = await websocket.receive_json()

        if info.get('command') != 'info':
            logger.warning('Client handshake failed at "info" stage')
            return None
        elif 'url' not in info:
            logger.warning('Info received from client, but no URL provided')
            return None

        return info['url']