¿Cómo controlar el cerrado de tu script en Python?

Me he topado con casos donde mi script no cierra correctamente al cerrarlo con Ctrl+C o Ctrl+D. Algunos ejemplos de esto:

1. Scripts que trabajan con puertos como TCP o Web
2. Scripts que hacen multithreading
3. Scripts que lanzan trabajos en segundo plano.

Cuando uno cierra la aplicación de forma controlada, por ejemplo, al presionar «q» para salir, todo se termina correctamente. Sin embargo, algunas veces tenemos la tendencia de forzar las cosas con Ctrl+C y, en ese momento, es donde todo comienza a salir mal.

En esta entrada, veremos dos alternativas, dependiendo de lo que estás haciendo.

Por destrucción de clases

En ocasiones, encapsulamos objetos que contienen funciones. Por ejemplo, encapsulamos un lanzador de un programa con «Shell». Sin embargo, ocupamos que ese script finalice al salir el programa o al eliminar nuestro objecto.

Para ello, podemos usar el destructor de la clase. Por ejemplo:

import socket
import subprocess
import time

class WebService():
    def __init__(self, location="web-receptor", host='127.0.0.1', port=1337):
        self.process = None
        self.socket = None

        # Initialise node
        self.process = subprocess.Popen(['node', location+'/server.js'], \
            stdout=subprocess.PIPE)
        time.sleep(3)
        # Get socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_address = (host, port)
        sock.connect(server_address)
        self.socket = sock
    
    def send_data(self, data):
        self.socket.sendall(data.encode())

El problema con el código de arriba es que lanza un proceso, que no siempre termina al destruir la clase. Para ello, forzamos el destructor con:

    def __del__(self):
        self.process.terminate()

El destructor __del__ nos permite ejecutar algunas instrucciones antes de destruir la clase, tal como ocurre con las clases en C++.

Con el módulo atexit

Si queremos ejecutar instrucciones justo antes de que el programa termine, podemos usar atexit. Para usarlo, se puede usar la función register o usar el decorador.

Con el decorador:

#!/usr/bin/env python3
import atexit

@atexit.register
def before_exit():
    print("Good bye cruel world")

if __name__ == "__main__":
    input("Waiting for a key")
    print("Exited correctly")

En este caso, el decorador registra la función before_exit como una función que se debe ejecutar al final de la ejecución.

Con la función register:

#!/usr/bin/env python3
import atexit

def before_exit():
    print("Good bye cruel world")

atexit.register(before_exit)

if __name__ == "__main__":
    input("Waiting for a key")
    print("Exited correctly")

Ambas representaciones son equivalentes. La función before_exit se ejecutará justo antes de terminar la ejecución después de haber presionado Ctrl+C.