expephalon/runner.py

#!/usr/bin/env python3

import argparse
import subprocess
import signal

try:
    import bindglobal
    BG = bindglobal.BindGlobal()
except ImportError:
    BG = False

from expephalon.settings import DEBUG

class Expephalon:
    def __init__(self, webserver=None, beat=None, worker=None, verbose=None, reload=None, bind=None):
        self.webserver = webserver
        self.beat = beat
        self.worker = worker
        self.verbose = verbose
        self.reload = reload
        self.bind = bind

        self.webserver_process = None
        self.beat_process = None
        self.worker_process = None

    def start_webserver(self, bind=None):
        if self.verbose:
            print("Starting Django development webserver")
        
        command = "/usr/bin/env python3 manage.py runserver"
        
        if bind:
            command += f" {bind}"

        self.webserver_process = subprocess.Popen(command.split())
        
    def start_beat(self):
        if self.verbose:
            print("Starting Celery beat scheduler")
        self.beat_process = subprocess.Popen("/usr/bin/env celery -A expephalon beat --scheduler django_celery_beat.schedulers:DatabaseScheduler -l debug".split())

    def start_worker(self):  
        if self.verbose:
            print("Starting Celery worker")  
        self.worker_process = subprocess.Popen("/usr/bin/env celery -A expephalon worker -l debug".split())

    def stop_process(self, process):
        process.send_signal(signal.SIGINT)
        try:
            process.wait(10)
        except subprocess.TimeoutExpired:
            process.kill()

    def stop_webserver(self):
        if self.webserver_process:
            if self.verbose:
                print("Stopping Django development webserver") 
            self.stop_process(self.webserver_process)
            self.webserver_process = None
        elif self.verbose:
            print("Django development webserver not running, so not stopping")

    def stop_beat(self):
        if self.beat_process:
            if self.verbose:
                print("Stopping Celery beat scheduler") 
            self.stop_process(self.beat_process)
            self.beat_process = None
        elif self.verbose:
            print("Celery beat scheduler not running, so not stopping")
    
    def stop_worker(self):
        if self.worker_process:
            if self.verbose:
                print("Stopping Celery worker") 
            self.stop_process(self.worker_process)
            self.worker_process = None
        elif self.verbose:
            print("Celery worker not running, so not stopping")

    def start(self):
        if self.webserver:
            self.start_webserver(self.bind)
        if self.beat != False:
            self.start_beat()
        if self.worker != False:
            self.start_worker()
        if BG and self.reload != False:
            BG.gbind("<Triple-Escape>", self.restart)

    def stop(self):
        if BG and self.reload != False:
            BG.gunbind("<Triple-Escape>")
        if self.webserver_process:
            self.stop_webserver()
        if self.beat_process:
            self.stop_beat()
        if self.worker:
            self.stop_worker()
    
    def restart(self, event=None):
        if self.verbose:
            print("Escape was pressed three times, restarting")
        self.stop()
        self.start()

    def loop(self):
        try:
            self.start()
            while True:
                pass
        except KeyboardInterrupt:
            if self.verbose:
                print("Ctrl+C was pressed, stopping")
            EXITCODE = 0
        except Exception as e:
            print("An uncaught exception has occurred:\n" + str(e))
            EXITCODE = 1
        finally:
            self.stop()
            exit(EXITCODE)

if __name__ == "__main__":
    TRUE = "store_true"
    parser = argparse.ArgumentParser()
    wsgroup = parser.add_mutually_exclusive_group()
    wsgroup.add_argument("--webserver", "-W", action=TRUE, help="Enforce start of Django development webserver when debug mode is off")
    wsgroup.add_argument("--no-webserver", "-w", action=TRUE, help="Prevent start of Django development webserver when debug mode if on")
    parser.add_argument("--no-beat", "-b", action=TRUE, help="Prevent start of Celery beat scheduler")
    parser.add_argument("--no-worker", "-c", action=TRUE, help="Prevent start of Celery worker")
    parser.add_argument("--verbose", "-v", action=TRUE, help="Print additional information while running")
    parser.add_argument("--no-reload", "-r", action=TRUE, help="Disable reload on triple-pressing Escape key")
    parser.add_argument("--bind", "-B", type=str, help="Local port/IP:Port combination to listen on (default: 127.0.0.1:8000)")

    args = parser.parse_args()
    
    Expephalon(webserver=False if args.no_webserver or (DEBUG and not args.webserver) else True, beat=not args.no_beat, worker=not args.no_worker, verbose=args.verbose, reload=not args.no_reload, bind=args.bind).loop()