让python脚本在合适的位置优雅退出

被bash的语法恶心久了，就越来越不会写了，写python越来越多，其实感觉在脚本方面，python完全可以替代Linux sh。

写了这么一个python脚本，将一个数据库里的数据取出来，调整一下，存到另一个数据中去，数据量比较大， 然后代码跑了很久，中途想ctrl+c结束掉，调整下日志输出，但又怕出现什么中间状态而搞乱数据，这时就想着， 下次写脚本，一定要考虑一下退出的优雅性。

其实经常就会写这样的代码，一个循环里，不停地干某件事，当想退出时，就直接ctrl+c退出，脚本退出， 但同时打印一坨堆栈信息，乱七八糟，也不知道是从哪句代码结束的。

举个例子如下：

import time
import sys 
import signal

loop = True

try:
    while loop:
        print("sleep begin")
        time.sleep(5)
        print("sleep end")
except KeyboardInterrupt:
    pass

print("sleep finished")

上面代码中的time.sleep(5)就相当于在进行逻辑处理，ctrl+c会产生Interrupt信号， 该信号的默认处理就是抛出KeyboardInterrupt异常。上面的代码并没有实现优雅退出的功能。 一旦按下ctrl+c，就打印sleep finished然后退出了。注意，好像在Python里是不能屏蔽或者忽略Error的，比如IOError，ZeroDivisionError之类，python的哲学是所有Error必须要被处理，所以你可以catch这些exception，但KeyboardInterrupt是由信号产生的，非错误类exception，所以可以通过修改信号处理函数来截获这些信号。

接下来想到的思路就是，在循环里，先忽略Interrupt信号，然后进行逻辑处理， 再恢复Interrupt信号的默认处理，如此确保time.sleep(5)不会被打断。代码如下：

import time
import sys 
import signal

loop = True

def int_loop(signum, frame):
    print("set loop false")
    frame.f_globals['loop'] = False
    #loop = False

while loop:
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    print("sleep begin")
    time.sleep(5)
    print("sleep end")
    signal.signal(signal.SIGINT, int_loop)

print("sleep finished")

上面的代码使用signal.signal()来设置一个信号的处理handler，先设置为SIG_IGN，完了再设置回为int_loop函数， 该函数设置全局变量loop为False来结束循环。这里特别注意，信号处理函数不能直接拿到全局变量loop， 必须通过帧对象frame来取得全局变量，或者临时变量。

上面的代码从理论上讲没有问题，但实际却效果不好，一点也不优雅，因为设置了信号处理函数的时间间隔很短， 即使你不停地按ctrl+c，也很难让程序退出。

如果直接简化来写，像如下代码一样：

import time
import sys 
import signal

loop = True

def int_loop(signum, frame):
    print("set loop false")
    frame.f_globals['loop'] = False

signal.signal(signal.SIGINT, int_loop)


while loop:
    print("sleep begin")
    time.sleep(5)
    print("sleep end")

print("sleep finished")

你会发现time.sleep(5)还是会被打断，但是sleep begin和sleep end确实能保证成对出现， 也就是说，上面的代码已经实现了优雅退出功能，因为如果你的代码是逻辑运算代码，一个循环是会被完整执行的， 只是这里代码比较特殊，sleep()会被打断，中断恢复时，会从下一条语句开始执行。 大多数时候这么写，就已经OK了。不过，像文件读写，系统IO类的操作，不确定被中断打断后， 是否会继续IO操作，未作验证，留个疑问在此。

其实还有更完美的方案，就是使用多线程，开个线程进行逻辑处理，主线程用于专门处理SIGINT信号。代码如下：

import time
import sys 
import signal

loop = True


def int_loop(signum, frame):
    print("set loop false")
    frame.f_globals['loop'] = False

signal.signal(signal.SIGINT, int_loop)


from threading import Thread

def noInterrupt():
    while loop:
        print("sleep begin")
        time.sleep(5)
        print("sleep end")

th = Thread(target=noInterrupt)
th.start()

while loop:
    time.sleep(3600)

print("join thread")
th.join()

上面的代码有两点值得注意： 1. 子线程是不会收到ctrl+c触发的SIGINT信号的，所以要考虑子线程退出的条件 2. 当主线程运行到join()时，会阻塞，此时也是接收不到SIGINT信号的，所以必须要循环sleep()

当然上面的代码写复杂了，实际上已经没有必要修改默认的信号处理函数了。直接用try...except捕获异常就可以了，写为：

import time
import sys 
import signal

loop = True

from threading import Thread

def noInterrupt():
    while loop:
        print("sleep begin")
        time.sleep(5)
        print("sleep end")

th = Thread(target=noInterrupt)
th.start()
try:
    while True:
        time.sleep(3600)
except KeyboardInterrupt:
    loop = False

print("join thread")
th.join()

总之，实现python脚本优雅退出的方法好几种，大体来看，还是使用多线程比较完美一些。

这样，不管你什么时候按下ctrl+c中止你的脚本，它都将在一个合适的位置优雅退出。