Coverage for src/lexigram/web/websocket/connection_id.py: 62%

1"""WebSocket connection ID management. 

2 

3Provides stable, unique identifiers for WebSocket connections using UUID 

4instead of memory addresses. Prevents state leakage when Python reuses 

5memory addresses after garbage collection. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING 

11import uuid 

12import weakref 

13 

14from lexigram.contracts.core.identity import IdGeneratorProtocol 

15 

16if TYPE_CHECKING: 

17 from starlette.websockets import WebSocket 

18 

19 

20class ConnectionIDManager: 

21 """Manages stable UUID-based connection IDs for WebSocket instances. 

22 

23 Each WebSocket is assigned a unique UUID when first registered. 

24 The mapping uses a WeakKeyDictionary so the UUID is automatically 

25 cleaned up when the WebSocket is garbage collected. 

26 """ 

27 

28 def __init__(self, id_generator: IdGeneratorProtocol | None = None) -> None: 

29 """Initialize the connection ID manager. 

30 

31 Args: 

32 id_generator: Optional ID generator (defaults to uuid.uuid4). 

33 """ 

34 self._ids = id_generator 

35 self._ws_to_id: weakref.WeakKeyDictionary[WebSocket, str] = ( 

36 weakref.WeakKeyDictionary() 

37 ) 

38 

39 def get_or_create(self, websocket: WebSocket) -> str: 

40 """Get or create a stable UUID for the WebSocket. 

41 

42 Args: 

43 websocket: The WebSocket connection. 

44 

45 Returns: 

46 A unique string identifier (UUID) for this connection. 

47 """ 

48 if websocket not in self._ws_to_id: 

49 self._ws_to_id[websocket] = ( 

50 self._ids.generate() if self._ids else str(uuid.uuid4()) 

51 ) 

52 return self._ws_to_id[websocket] 

53 

54 def get(self, websocket: WebSocket) -> str | None: 

55 """Get the UUID for a WebSocket, or None if not registered. 

56 

57 Args: 

58 websocket: The WebSocket connection. 

59 

60 Returns: 

61 The UUID if registered, None otherwise. 

62 """ 

63 return self._ws_to_id.get(websocket) 

64 

65 

66__all__ = ["ConnectionIDManager"]