Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/zmq/sugar/tracker.py: 38%

1"""Tracker for zero-copy messages with 0MQ.""" 

2 

3# Copyright (C) PyZMQ Developers 

4# Distributed under the terms of the Modified BSD License. 

5 

6from __future__ import annotations 

7 

8import time 

9from threading import Event 

10from typing import Final 

11 

12from zmq.backend import Frame 

13from zmq.error import NotDone 

14 

15 

16class MessageTracker: 

17 """A class for tracking if 0MQ is done using one or more messages. 

18 

19 When you send a 0MQ message, it is not sent immediately. The 0MQ IO thread 

20 sends the message at some later time. Often you want to know when 0MQ has 

21 actually sent the message though. This is complicated by the fact that 

22 a single 0MQ message can be sent multiple times using different sockets. 

23 This class allows you to track all of the 0MQ usages of a message. 

24 

25 Parameters 

26 ---------- 

27 towatch : Event, MessageTracker, zmq.Frame 

28 This objects to track. This class can track the low-level 

29 Events used by the Message class, other MessageTrackers or 

30 actual Messages. 

31 """ 

32 

33 events: set[Event] 

34 peers: set[MessageTracker] 

35 

36 def __init__(self, *towatch: MessageTracker | Event | Frame) -> None: 

37 """Create a message tracker to track a set of messages. 

38 

39 Parameters 

40 ---------- 

41 *towatch : tuple of Event, MessageTracker, Message instances. 

42 This list of objects to track. This class can track the low-level 

43 Events used by the Message class, other MessageTrackers or 

44 actual Messages. 

45 """ 

46 self.events = set() 

47 self.peers = set() 

48 for obj in towatch: 

49 if isinstance(obj, Event): 

50 self.events.add(obj) 

51 elif isinstance(obj, MessageTracker): 

52 self.peers.add(obj) 

53 elif isinstance(obj, Frame): 

54 if not obj.tracker: 

55 raise ValueError("Not a tracked message") 

56 self.peers.add(obj.tracker) 

57 else: 

58 raise TypeError(f"Require Events or Message Frames, not {type(obj)}") 

59 

60 @property 

61 def done(self) -> bool: 

62 """Is 0MQ completely done with the message(s) being tracked?""" 

63 for evt in self.events: 

64 if not evt.is_set(): 

65 return False 

66 for pm in self.peers: 

67 if not pm.done: 

68 return False 

69 return True 

70 

71 def wait(self, timeout: float | int = -1) -> None: 

72 """Wait for 0MQ to be done with the message or until `timeout`. 

73 

74 Parameters 

75 ---------- 

76 timeout : float 

77 default: -1, which means wait forever. 

78 Maximum time in (s) to wait before raising NotDone. 

79 

80 Returns 

81 ------- 

82 None 

83 if done before `timeout` 

84 

85 Raises 

86 ------ 

87 NotDone 

88 if `timeout` reached before I am done. 

89 """ 

90 tic = time.time() 

91 remaining: float 

92 if timeout is False or timeout < 0: 

93 remaining = 3600 * 24 * 7 # a week 

94 else: 

95 remaining = timeout 

96 for evt in self.events: 

97 if remaining < 0: 

98 raise NotDone 

99 evt.wait(timeout=remaining) 

100 if not evt.is_set(): 

101 raise NotDone 

102 toc = time.time() 

103 remaining -= toc - tic 

104 tic = toc 

105 

106 for peer in self.peers: 

107 if remaining < 0: 

108 raise NotDone 

109 peer.wait(timeout=remaining) 

110 toc = time.time() 

111 remaining -= toc - tic 

112 tic = toc 

113 

114 

115_FINISHED_TRACKER: Final[MessageTracker] = MessageTracker() 

116 

117__all__ = ['MessageTracker', '_FINISHED_TRACKER']