Coverage for src / graphable / cli / commands / serve.py: 100%

1from asyncio import get_event_loop 

2from html import escape 

3from logging import getLogger 

4from pathlib import Path 

5 

6from starlette.applications import Starlette 

7from starlette.responses import HTMLResponse 

8from starlette.routing import Route, WebSocketRoute 

9from starlette.websockets import WebSocket 

10from uvicorn import Config 

11from uvicorn import Server as UvicornServer 

12from watchfiles import awatch 

13 

14from ...views.html import create_topology_html 

15from .core import load_graph 

16 

17logger = getLogger(__name__) 

18 

19 

20class Server: 

21 def __init__(self, path: Path, tag: str | None = None): 

22 self.path = path 

23 self.tag = tag 

24 self.connections: set[WebSocket] = set() 

25 self.app = Starlette( 

26 routes=[ 

27 Route("/", self.index), 

28 WebSocketRoute("/ws", self.websocket_endpoint), 

29 ] 

30 ) 

31 

32 async def index(self, request): 

33 try: 

34 g = load_graph(self.path, tag=self.tag) 

35 html = create_topology_html(g) 

36 return HTMLResponse(html) 

37 except Exception as e: 

38 return HTMLResponse( 

39 f"<h1>Error loading graph</h1><pre>{escape(str(e))}</pre>", 

40 status_code=500, 

41 ) 

42 

43 async def websocket_endpoint(self, websocket: WebSocket): 

44 await websocket.accept() 

45 self.connections.add(websocket) 

46 try: 

47 while True: 

48 await websocket.receive_text() 

49 except Exception: 

50 self.connections.remove(websocket) 

51 

52 async def watch_file(self): 

53 async for changes in awatch(self.path): 

54 logger.info(f"File {self.path} changed, reloading...") 

55 for ws in self.connections: 

56 await ws.send_text("reload") 

57 

58 

59def serve_command(path: Path, port: int = 8000, tag: str | None = None): 

60 server = Server(path, tag=tag) 

61 

62 config = Config(server.app, host="127.0.0.1", port=port, log_level="info") 

63 uv_server = UvicornServer(config) 

64 

65 loop = get_event_loop() 

66 loop.create_task(server.watch_file()) 

67 loop.run_until_complete(uv_server.serve())