在 Python 中接管键盘中断信号

开发 后端
假设有这样一个需求,你需要从 Redis 中持续不断读取数据,并把这些数据写入到 MongoDB 中。

 [[313841]]

假设有这样一个需求,你需要从 Redis 中持续不断读取数据,并把这些数据写入到 MongoDB 中。你可能会这样写代码:

 

  1. import json 
  2. import redis 
  3. import pymongo 
  4.  
  5.  
  6. client = redis.Redis() 
  7. handler = pymongo.MongoClient().example.col 
  8.  
  9.  
  10. while True
  11.     data_raw = client.blpop('data', timeout=300) 
  12.     if not data_raw: 
  13.         continue 
  14.     data = json.loads(data_raw[1].decode()) 
  15.     handler.insert_one(data) 

但这样写有一个问题,就是每来一条数据都要连接一次 MongoDB,大量时间浪费在了网络 I/O上。

于是大家会把代码改成下面这样:

  1. import json 
  2. import redis 
  3. import pymongo 
  4.  
  5.  
  6. client = redis.Redis() 
  7. handler = pymongo.MongoClient().example.col 
  8.  
  9.  
  10. to_be_insert = [] 
  11. while True
  12.     data_raw = client.blpop('data', timeout=300) 
  13.     if not data_raw: 
  14.         continue 
  15.     data = json.loads(data_raw[1].decode()) 
  16.     to_be_insert.append(data) 
  17.     if len(to_be_insert) >= 1000: 
  18.         handler.insert_many(to_be_insert) 
  19.         to_be_insert = [] 

每凑够1000条数据,批量写入到 MongoDB 中。

现在又面临另外一个问题。假设因为某种原因,我需要更新这个程序,于是我按下了键盘上的Ctrl + C强制关闭了这个程序。而此时to_be_insert列表里面有999条数据将会永久丢失——它们已经被从 Redis 中删除了,但又没有来得及写入 MongoDB 中。

我想实现,当我按下 Ctrl + C 时,程序不再从 Redis 中读取数据,但会先把to_be_insert中的数据(无论有几条)都插入 MongoDB 中。最后再关闭程序。

要实现这个需求,就必须在我们按下Ctrl + C时,程序还能继续运行一段代码。可问题是按下Ctrl + C时,程序就直接结束了,如何还能再运行一段代码?

实际上,当我们按下键盘上的Ctrl + C时,Python 收到一个名为SIGINT的信号。具体规则可以阅读官方文档。收到信号以后,Python 会调用一个信号回调函数。只不过默认的回调函数就是让程序抛出一个 KeyboardInterrupt异常导致程序关闭。现在,我们可以设法让 Python 使用我们自定义的一段函数来作为信号回调函数。

要使用信号,我们需用导入 Python 的signal库。然后自定义一个信号回调函数,当 Python 收到某个信号时,调用这个函数。

所以我们修改一下上面的代码:

 

  1. import signal 
  2. import json 
  3. import redis 
  4. import pymongo 
  5.  
  6.  
  7. client = redis.Redis() 
  8. handler = pymongo.MongoClient().example.col 
  9. stop = False 
  10.  
  11.  
  12. def keyboard_handler(signum, frame): 
  13.     global stop 
  14.     stop = True 
  15.  
  16.  
  17. signal.signal(signal.SIGINT, keyboard_handler) 
  18.  
  19. to_be_insert = [] 
  20. while not stop: 
  21.     data_raw = client.blpop('data', timeout=300) 
  22.     if not data_raw: 
  23.         continue 
  24.     data = json.loads(data_raw[1].decode()) 
  25.     to_be_insert.append(data) 
  26.     if len(to_be_insert) >= 1000: 
  27.         handler.insert_many(to_be_insert) 
  28.         to_be_insert = [] 
  29.  
  30. if to_be_insert: 
  31.     handler.insert_many(to_be_insert) 

我们定义了一个全局变量stop,默认为 False,所以默认情况下,while not stop所在的循环体会持续运行。

我们定义了一个函数keyboard_handler,它的作用是修改全局变量stop为 True。需要注意的是,在函数里面修改全局变量,必须先使用global 变量名声明这个变量为全局变量。否则无法修改。

修改以后,while not stop循环停止,于是程序进入:

 

  1. if to_be_insert: 
  2.     handler.insert_many(to_be_insert) 

只要列表里面有数据,就会批量插入 MongoDB 中。然后程序结束。

整段代码的关键就在signal.signal(signal.SIGINT, keyboard_handler)这里把信号SIGINT与函数keyboard_handler关联上了,于是,在上面这段代码运行的任何时候,只要按下键盘的Ctrl + C,程序就会进入keyboard_handler函数里面,优先执行这个函数里面的代码。执行完成以后,回到之前中断的地方,继续执行之前没有完成的代码。而由于在函数里面我已经修改了stop的值,所以原来的循环不能继续执行,于是进入最后的收尾工作。

需要注意的是,如果你的整个代码全都是使用 Python 写的,那么 signal可以在你程序的任何阶段触发,只要你按下 Ctrl + C,立刻就会进入设置好的信号回调函数中。

但如果你的代码中,有一部分代码是使用 C 语言写的,那么当你按下Ctrl + C以后,可能需要等这段C 语言的代码运行完成以后,才会进入你设置的信号回调函数中。

责任编辑:华轩 来源: 未闻Code
相关推荐

2023-12-05 18:09:04

2020-03-12 13:00:36

Python断点设置编程语言

2011-05-31 09:16:15

JavaScript

2024-01-03 10:17:51

Linux通信

2023-12-11 07:21:36

2024-04-02 08:30:40

RustUnix信号服务器

2021-12-11 19:00:54

Java中断机制

2016-01-08 15:59:08

2020-10-13 19:04:58

Bash信号捕获Shell脚本

2021-12-20 07:51:16

Linux函数应用层

2015-07-08 16:46:05

iOS键盘

2021-08-03 09:33:55

HTTP网络协议TCP

2021-08-03 10:10:16

HTTP连接管理

2021-01-11 12:53:28

线程Java管理

2024-10-21 09:06:15

2022-02-17 20:34:12

Python短路机制开发

2022-03-29 10:56:46

Pythonblinker信号库

2010-06-13 15:37:24

TCP协议

2024-04-07 11:50:53

2023-03-01 09:39:05

点赞
收藏

51CTO技术栈公众号