Coverage for aiocoap/util/cli.py: 87%

1# SPDX-FileCopyrightText: Christian Amsüss and the aiocoap contributors 

2# 

3# SPDX-License-Identifier: MIT 

4 

5"""Helpers for creating server-style applications in aiocoap 

6 

7Note that these are not particular to aiocoap, but are used at different places 

8in aiocoap and thus shared here.""" 

9 

10import argparse 

11import sys 

12import logging 

13import asyncio 

14import signal 

15 

16 

17class ActionNoYes(argparse.Action): 

18 """Simple action that automatically manages --{,no-}something style options""" 

19 

20 # adapted from Omnifarious's code on 

21 # https://stackoverflow.com/questions/9234258/in-python-argparse-is-it-possible-to-have-paired-no-something-something-arg#9236426 

22 def __init__(self, option_strings, dest, default=True, required=False, help=None): 

23 assert len(option_strings) == 1, "ActionNoYes takes only one option name" 

24 assert option_strings[0].startswith( 

25 "--" 

26 ), "ActionNoYes options must start with --" 

27 super().__init__( 

28 ["--" + option_strings[0][2:], "--no-" + option_strings[0][2:]], 

29 dest, 

30 nargs=0, 

31 const=None, 

32 default=default, 

33 required=required, 

34 help=help, 

35 ) 

36 

37 def __call__(self, parser, namespace, values, option_string=None): 

38 if option_string.startswith("--no-"): 

39 setattr(namespace, self.dest, False) 

40 else: 

41 setattr(namespace, self.dest, True) 

42 

43 

44class AsyncCLIDaemon: 

45 """Helper for creating daemon-style CLI prorgrams. 

46 

47 Note that this currently doesn't create a Daemon in the sense of doing a 

48 daemon-fork; that could be added on demand, though. 

49 

50 Subclass this and implement the :meth:`start` method as an async 

51 function; it will be passed all the constructor's arguments. 

52 

53 When all setup is complete and the program is operational, return from the 

54 start method. 

55 

56 Implement the :meth:`shutdown` coroutine and to do cleanup; what actually 

57 runs your program will, if possible, call that and await its return. 

58 

59 Two usage patterns for this are supported: 

60 

61 * Outside of an async context, run run ``MyClass.sync_main()``, typically 

62 in the program's ``if __name__ == "__main__":`` section. 

63 

64 In this mode, the loop that is started is configured to safely shut down 

65 the loop when SIGINT is received. 

66 

67 * To run a subclass of this in an existing loop, start it with 

68 ``MyClass(...)`` (possibly passing in the loop to run it on if not already 

69 in an async context), and then awaiting its ``.initializing`` future. To 

70 stop it, await its ``.shutdown()`` method. 

71 

72 Note that with this usage pattern, the :meth:`.stop()` method has no 

73 effect; servers that ``.stop()`` themselves need to signal their desire 

74 to be shut down through other channels (but that is an atypical case). 

75 """ 

76 

77 def __init__(self, *args, **kwargs): 

78 loop = kwargs.pop("loop", None) 

79 if loop is None: 

80 loop = asyncio.get_running_loop() 

81 self.__exitcode = loop.create_future() 

82 self.initializing = loop.create_task( 

83 self.start(*args, **kwargs), name="Initialization of %r" % (self,) 

84 ) 

85 

86 def stop(self, exitcode): 

87 """Stop the operation (and exit sync_main) at the next convenience.""" 

88 self.__exitcode.set_result(exitcode) 

89 

90 @classmethod 

91 async def _async_main(cls, *args, **kwargs): 

92 """Run the application in an AsyncIO main loop, shutting down cleanly 

93 on keyboard interrupt. 

94 

95 This is not exposed publicly as it messes with the loop, and we only do 

96 that with loops created in sync_main. 

97 """ 

98 main = cls(*args, **kwargs) 

99 

100 try: 

101 asyncio.get_running_loop().add_signal_handler( 

102 signal.SIGTERM, 

103 lambda: main.__exitcode.set_result(143), 

104 ) 

105 except NotImplementedError: 

106 # Impossible on win32 -- just won't make that clean of a shutdown. 

107 pass 

108 

109 try: 

110 await main.initializing 

111 # This is the time when we'd signal setup completion by the parent 

112 # exiting in case of a daemon setup, or to any other process 

113 # management. 

114 logging.info("Application ready.") 

115 # Common options are 143 or 0 

116 # (<https://github.com/go-task/task/issues/75#issuecomment-339466142> and 

117 # <https://unix.stackexchange.com/questions/10231/when-does-the-system-send-a-sigterm-to-a-process>) 

118 exitcode = await main.__exitcode 

119 except KeyboardInterrupt: 

120 logging.info("Keyboard interupt received, shutting down") 

121 sys.exit(3) 

122 else: 

123 sys.exit(exitcode) 

124 finally: 

125 if main.initializing.done() and main.initializing.exception(): 

126 # The exception if initializing is what we are just watching 

127 # fly by. No need to trigger it again, and running shutdown 

128 # would be even weirder. 

129 pass 

130 else: 

131 # May be done, then it's a no-op, or we might have received a 

132 # signal during startup in which case we better fetch the 

133 # result and shut down cleanly again 

134 await main.initializing 

135 

136 # And no matter whether that happened during initialization 

137 # (which now has finished) or due to a regular signal... 

138 await main.shutdown() 

139 

140 @classmethod 

141 def sync_main(cls, *args, **kwargs): 

142 """Run the application in an AsyncIO main loop, shutting down cleanly 

143 on keyboard interrupt.""" 

144 asyncio.run(cls._async_main(*args, **kwargs))