I have a simple Python service, where there is a loop that performs some action infinitely. On various signals, sys.exit(0) is called, which causes SystemExit to be raised and then some cleanup should happen if it can.
In a test, i.e. standard unittest.TestCase, I would like to test that this cleanup happens and the loop exits. However, I'm stuck on even getting the signal to be triggered / SystemExit to be raised.
# service.py
import signal
import sys
import time
def main():
def signal_handler(signalnum, _):
# How to get this to block to run in a test?
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
while True:
try:
print("Some action here")
time.sleep(10)
except SystemExit:
# How to get this to block to run in a test?
print("Some cleanup")
break
if __name__ == '__main__':
main()
How can the code enter the SystemExit handler / signal handler in the test environment? An alternative pattern would also be welcome.
You can trigger a SIGINT (or any signal) from another thread after some delay, which is received in the main thread. You can then assert on its effects just as in any other test, as below.
import os
import signal
import time
import threading
import unittest
from unittest.mock import (
Mock,
patch,
)
import service
class TestService(unittest.TestCase):
#patch('service.print')
def test_signal_handling(self, mock_print):
pid = os.getpid()
def trigger_signal():
while len(mock_print.mock_calls) < 1:
time.sleep(0.2)
os.kill(pid, signal.SIGINT)
thread = threading.Thread(target=trigger_signal)
thread.daemon = True
thread.start()
service.main()
self.assertEqual(mock_print.mock_calls[1][1][0], 'Some cleanup')
if __name__ == '__main__':
unittest.main()
Let's refactor that to make it easier to test:
def loop():
try:
print("Some action here")
except:
# clean up and re-raise
print("Some cleanup")
raise
def main():
def signal_handler(signalnum, _):
# How to get this to block to run in a test?
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
while True:
try:
loop_body()
time.sleep(10)
except SystemExit:
break
if __name__ == '__main__':
main()
This doesn't allow easy testing of the signal handling code though. However, that amount is so small, rarely changed and strongly depends on the environment, that it is possible and perhaps even better to test manually.
For clarity, it could be useful to use a context handler, which is usually a good idea when you have setup/shutdown code. You don't mention the setup code, but my Crystall Ball (tm) tells me it exists. It could then be called like this:
try:
with my_service() as service:
while True:
service.run()
sleep(10)
except SystemExit:
# perform graceful shutdown on signal
pass
I'll leave the implementation of that context manager to you, but check out contextlib, which makes it easy and fun.
Related
I have a program that has quite a few functions, each running on a separate thread.
When the user presses Ctrl+C, only 1 thread crashes with an exception, but because of this, the whole program may not work correctly.
Of course, I can write this construction in each function:
try:
do_something()
except KeyboardInterrupt as e:
pass
but, as I said, there are many functions, perhaps there is an option not to prescribe this construction in each function?
Or is it possible to disable Ctrl+C interrupt in cmd settings?
For example, in the registry. The program creates its own registry key in HKEY_CURRENT_USER\Console\MyProgrammKey
UPD 1
signal.signal(signal.SIGINT, signal.SIG_IGN)
It helped in almost all cases except one: a thread that has an infinite loop with the input() function anyway interrupts.
UPD 2
Here is a sample code
import signal, time
from threading import Thread
def one():
while True:
inp = input("INPUT: ")
def two():
while True:
print("I just printing...")
time.sleep(1)
if __name__ == '__main__':
signal.signal(signal.SIGINT, signal.SIG_IGN)
Thread(target=one).start()
Thread(target=two).start()
UPD 3
Screenshot of exception.
Ctrl+C will send SIGINT signal to program, so you could define a global signal handler to ignore that SIGINT, something like next:
test.py:
import signal, os, time
def handler(signum, frame):
pass
signal.signal(signal.SIGINT, handler)
time.sleep(10)
print("done")
During the program run, if you input Ctrl+c, the program will ignore it, and continue to run, finally print done:
$ python3 test.py
^Cdone
Basically, I'm curious if it's possible to execute a block of python code "atomically" without being interrupted by a signal.
For instance, I want to perform operations in a loop, let's say:
for i in range(100):
do_stuff(1)
do_stuff(2)
do_stuff(3)
But I want to finish all of three do_stuff(1), do_stuff(2), do_stuff(3) if do_stuff(1) managed to start. Script should ignore CTRL+C, finish these three instructions and then terminate if SIGINT happened. All of 100 iterations does not have to be executed.
I believe it could be done with a custom signal handler
import signal
def handler(signum, frame):
# wait for the loop iteration finish and exit
signal.signal(signal.SIGINT, handler)
threads and synchronization but I have no idea how to implement it.
Is it possible?
If it is, can it be done nicely? With some kind of a context manager, for example?
for i in range(100):
with atomic_execution():
do_stuff(1)
do_stuff(2)
do_stuff(3)
Edit: in the meantime I created this:
import threading
import sys
import signal
class atomic_execution:
started = 0
execution_in_progress = threading.Lock()
def __enter__(self):
atomic_execution.execution_in_progress.acquire()
def __exit__(self, type, value, traceback):
atomic_execution.execution_in_progress.release()
def handler(signum, frame):
atomic_execution.execution_in_progress.acquire()
sys.exit(0)
signal.signal(signal.SIGINT, handler)
while True:
with atomic_execution():
print(1)
print(2)
print(3)
I am not sure if it's good, though.
This is the basic idea:
import signal
import time
stop = False
def sighandler(*unused):
global stop
print('signal caught')
stop = True
def main():
for i in range(10):
print('a')
time.sleep(0.5)
print('b')
time.sleep(0.5)
print('c')
time.sleep(0.5)
print()
if stop:
print('STOP')
break
if __name__ == '__main__':
signal.signal(signal.SIGINT, sighandler)
main()
I think it is not difficult to make an context manager for this purpose:
on enter:
save the current signal handler
install own handler setting a flag like in the example above
on exit:
restore the original signal handler
exit if the flag was set
But I do not like the idea, because you want to install the handler once before the loop and test the flag many times at each iteration.
I've read a lot of questions on SO and elsewhere on this topic but can't get it working. Perhaps it's because I'm using Windows, I don't know.
What I'm trying to do is download a bunch of files (whose URLs are read from a CSV file) in parallel. I've tried using multiprocessing and concurrent.futures for this with no success.
The main problem is that I can't stop the program on Ctrl-C - it just keeps running. This is especially bad in the case of processes instead of threads (I used multiprocessing for that) because I have to kill each process manually every time.
Here is my current code:
import concurrent.futures
import signal
import sys
import urllib.request
class Download(object):
def __init__(self, url, filename):
self.url = url
self.filename = filename
def perform_download(download):
print('Downloading {} to {}'.format(download.url, download.filename))
return urllib.request.urlretrieve(download.url, filename=download.filename)
def main(argv):
args = parse_args(argv)
queue = []
with open(args.results_file, 'r', encoding='utf8') as results_file:
# Irrelevant CSV parsing...
queue.append(Download(url, filename))
def handle_interrupt():
print('CAUGHT SIGINT!!!!!!!!!!!!!!!!!!!11111111')
sys.exit(1)
signal.signal(signal.SIGINT, handle_interrupt)
with concurrent.futures.ThreadPoolExecutor(max_workers=args.num_jobs) as executor:
futures = {executor.submit(perform_download, d): d for d in queue}
try:
concurrent.futures.wait(futures)
except KeyboardInterrupt:
print('Interrupted')
sys.exit(1)
I'm trying to catch Ctrl-C in two different ways here but none of them works. The latter one (except KeyboardInterrupt) actually gets run but the process won't exit after calling sys.exit.
Before this I used the multiprocessing module like this:
try:
pool = multiprocessing.Pool(processes=args.num_jobs)
pool.map_async(perform_download, queue).get(1000000)
except Exception as e:
pool.close()
pool.terminate()
sys.exit(0)
So what is the proper way to add ability to terminate all worker threads or processes once you hit Ctrl-C in the terminal?
System information:
Python version: 3.6.1 32-bit
OS: Windows 10
You are catching the SIGINT signal in a signal handler and re-routing it as a SystemExit exception. This prevents the KeyboardInterrupt exception to ever reach your main loop.
Moreover, if the SystemExit is not raised in the main thread, it will just kill the child thread where it is raised.
Jesse Noller, the author of the multiprocessing library, explains how to deal with CTRL+C in a old blog post.
import signal
from multiprocessing import Pool
def initializer():
"""Ignore CTRL+C in the worker process."""
signal.signal(SIGINT, SIG_IGN)
pool = Pool(initializer=initializer)
try:
pool.map(perform_download, dowloads)
except KeyboardInterrupt:
pool.terminate()
pool.join()
I don't believe the accepted answer works under Windows, certainly not under current versions of Python (I am running 3.8.5). In fact, it won't run at all since SIGINT and SIG_IGN will be undefined (what is needed is signal.SIGINT and signal.SIG_IGN).
This is a know problem under Windows. A solution I have come up with is essentially the reverse of the accepted solution: The main process must ignore keyboard interrupts and we initialize the process pool to initially set a global flag ctrl_c_entered to False and to set this flag to True if Ctrl-C is entered. Then any multiprocessing worker function (or method) is decorated with a special decorator, handle_ctrl_c, that firsts tests the ctrl_c_entered flag and only if False does it run the worker function after re-enabling keyboard interrupts and establishing a try/catch handler for keyboard interrups. If the ctrl_c_entered flag was True or if a keyboard interrupt occurs during the execution of the worker function, the value returned is an instance of KeyboardInterrupt, which the main process can check to determine whether a Ctrl-C was entered.
Thus all submitted tasks will be allowed to start but will immediately terminate with a return value of a KeyBoardInterrupt exception and the actual worker function will never be called by the decorator function once a Ctrl-C has been entered.
import signal
from multiprocessing import Pool
from functools import wraps
import time
def handle_ctrl_c(func):
"""
Decorator function.
"""
#wraps(func)
def wrapper(*args, **kwargs):
global ctrl_c_entered
if not ctrl_c_entered:
# re-enable keyboard interrups:
signal.signal(signal.SIGINT, default_sigint_handler)
try:
return func(*args, **kwargs)
except KeyboardInterrupt:
ctrl_c_entered = True
return KeyboardInterrupt()
finally:
signal.signal(signal.SIGINT, pool_ctrl_c_handler)
else:
return KeyboardInterrupt()
return wrapper
def pool_ctrl_c_handler(*args, **kwargs):
global ctrl_c_entered
ctrl_c_entered = True
def init_pool():
# set global variable for each process in the pool:
global ctrl_c_entered
global default_sigint_handler
ctrl_c_entered = False
default_sigint_handler = signal.signal(signal.SIGINT, pool_ctrl_c_handler)
#handle_ctrl_c
def perform_download(download):
print('begin')
time.sleep(2)
print('end')
return True
if __name__ == '__main__':
signal.signal(signal.SIGINT, signal.SIG_IGN)
pool = Pool(initializer=init_pool)
results = pool.map(perform_download, range(20))
if any(map(lambda x: isinstance(x, KeyboardInterrupt), results)):
print('Ctrl-C was entered.')
print(results)
I would like the user to use control-c to close a script, but when control-c is pressed it shows the error and reason for close (which make sense). Is there a way to have my own custom output to the screen rather than what is shown? Not sure how to handle that specific error.
You could use try..except to catch KeyboardInterrupt:
import time
def main():
time.sleep(10)
if __name__ == '__main__':
try:
main()
except KeyboardInterrupt:
print('bye')
use the signal module to define a handler for the SIGINT signal:
import signal
import sys
def sigint_handler(signal_number, stack_frame):
print('caught SIGINT, exiting')
sys.exit(-1)
signal.signal(signal.SIGINT, sigint_handler)
raw_input('waiting...')
For general purpose code, handling the KeyboardInterrupt should suffice. For advanced code, such as threading, it is a whole different story. Here's a simple example.
http://docs.python.org/2/library/exceptions.html#exceptions.KeyboardInterrupt
try:
while 1:
x = raw_input("Type something or press CTRL+C to end: ")
print repr(x)
except KeyboardInterrupt:
print "\nWe're done here."
I want to catch KeyboardInterrupt globally, and deal with it nicely. I don't want to encase my entire script in a huge try/except statement. Is there any way to do this?
You could change sys.excepthook if you really don't want to use a try/except.
import sys
def my_except_hook(exctype, value, traceback):
if exctype == KeyboardInterrupt:
print "Handler code goes here"
else:
sys.__excepthook__(exctype, value, traceback)
sys.excepthook = my_except_hook
If this is a script for execution on the command line, you can encapsulate your run-time logic in main(), call it in an if __name__ == '__main__' and wrap that.
if __name__ == '__main__':
try:
main()
except KeyboardInterrupt:
print 'Killed by user'
sys.exit(0)
You can also use signal like this:
import signal, time
def handler(signum, frame):
print 'I just clicked on CTRL-C '
signal.signal(signal.SIGINT, handler)
print "waiting for 10 s"
time.sleep(10)
Output:
waiting for 10 s
^CI just clicked on CTRL-C
N.B: Don't mix the use of signal with threads.
Does your script have a function you call to start it?
main()
then just do:
try:
main()
except:
...
If you don't have a main but just a huge script that runs line-by-line, then you should put it in a main.
There's no other way to do this, apart from encasing your entire script in a main() function and surrounding that with a try..except block - which is pretty much the same:
def main():
# Your script goes here
pass
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
# cleanup code here
pass