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 

16from ..util.asyncio import py38args 

17 

18class ActionNoYes(argparse.Action): 

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

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('--'), "ActionNoYes options must start with --" 

25 super().__init__(['--' + option_strings[0][2:], '--no-' + option_strings[0][2:]], dest, nargs=0, const=None, default=default, required=required, help=help) 

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

27 if option_string.startswith('--no-'): 

28 setattr(namespace, self.dest, False) 

29 else: 

30 setattr(namespace, self.dest, True) 

31 

32class AsyncCLIDaemon: 

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

34 

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

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

37 

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

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

40 

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

42 start method. 

43 

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

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

46 

47 Two usage patterns for this are supported: 

48 

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

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

51 

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

53 the loop when SIGINT is received. 

54 

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

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

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

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

59 

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

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

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

63 """ 

64 

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

66 loop = kwargs.pop('loop', None) 

67 if loop is None: 

68 loop = asyncio.get_running_loop() 

69 self.__exitcode = loop.create_future() 

70 self.initializing = loop.create_task( 

71 self.start(*args, **kwargs), 

72 **py38args(name="Initialization of %r" % (self,)) 

73 ) 

74 

75 def stop(self, exitcode): 

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

77 self.__exitcode.set_result(exitcode) 

78 

79 @classmethod 

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

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

82 on keyboard interrupt. 

83 

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

85 that with loops created in sync_main. 

86 """ 

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

88 

89 try: 

90 asyncio.get_running_loop().add_signal_handler( 

91 signal.SIGTERM, 

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

93 ) 

94 except NotImplementedError: 

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

96 pass 

97 

98 try: 

99 await main.initializing 

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

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

102 # management. 

103 logging.info("Application ready.") 

104 # Common options are 143 or 0 

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

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

107 exitcode = await main.__exitcode 

108 except KeyboardInterrupt: 

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

110 sys.exit(3) 

111 else: 

112 sys.exit(exitcode) 

113 finally: 

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

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

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

117 # would be even weirder. 

118 pass 

119 else: 

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

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

122 # result and shut down cleanly again 

123 await main.initializing 

124 

125 # And no matter whether that happened during initialization 

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

127 await main.shutdown() 

128 

129 @classmethod 

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

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

132 on keyboard interrupt.""" 

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