simplify

Author: SKairinos

codeforlife/apps/__init__.py

@@ -0,0 +1,7 @@
+"""
+© Ocado Group
+Created on 28/03/2025 at 15:37:48(+00:00).
+"""
+
+from .celery import CeleryServer
+from .django import DjangoServer

codeforlife/servers/__init__.py

@@ -8,14 +8,17 @@
 import atexit
 import os
 import subprocess
+import sys
 import typing as t
 
-from celery import Celery as _Celery
+import celery
 
 from ..types import LogLevel
 
+CELERY_MAIN_PATH = os.path.join(celery.__path__[0], "__main__.py")
 
-class CeleryApplication(_Celery):
+
+class CeleryServer(celery.Celery):
     """A server for a Celery app."""
 
     def __init__(
@@ -25,8 +28,18 @@ def __init__(
     ):
         """Initialize a Celery app.
 
+        Examples:
+            ```
+            from codeforlife.servers import CeleryServer, DjangoServer
+
+            # Make sure to set up Django before initializing!
+            DjangoServer.setup()
+
+            celery_app = CeleryServer().app
+            ```
+
         Args:
-            app: The Celery app name.
+            app: The dot-path to the Celery app.
             debug: A flag designating whether to run the app in debug mode.
 
         Raises:
@@ -39,8 +52,8 @@ def __init__(
             )
 
         super().__init__()
-
-        self.app = app
+        self.app = self
+        self._app = app
 
         # Using a string here means the worker doesn't have to serialize
         # the configuration object to child processes.
@@ -59,6 +72,10 @@ def _debug(self, *args, **kwargs):
 
                 print(f"Request: {self.request!r}")
 
+        if os.path.abspath(sys.argv[0]) != CELERY_MAIN_PATH:
+            self.start_background_workers()
+            self.start_background_beat()
+
     def start_background_workers(
         self,
         workers: t.Optional[t.Union[t.Set[str], t.Dict[str, int]]] = None,
@@ -74,8 +91,6 @@ def start_background_workers(
         """
         # pylint: enable=line-too-long
 
-        print("Starting all Celery workers.")
-
         commands: t.Dict[str, t.List[str]] = {}
 
         def build_command(worker: str, concurrency: int = 0):
@@ -84,7 +99,7 @@ def build_command(worker: str, concurrency: int = 0):
                 "multi",
                 "start",
                 worker,
-                f"--app={self.app}",
+                f"--app={self._app}",
                 f"--loglevel={log_level}",
             ]
 
@@ -106,20 +121,11 @@ def build_command(worker: str, concurrency: int = 0):
 
         for worker, command in commands.items():
             try:
-                process = subprocess.run(
-                    command,
-                    capture_output=True,
-                    text=True,
-                    check=True,
-                )
-                print(f"Successfully started Celery worker '{worker}'.")
-                print(process.stdout)
-
+                subprocess.run(command, check=True)
                 successfully_started_workers.add(worker)
 
-            except subprocess.CalledProcessError as e:
-                print(f"Error starting Celery worker '{worker}': {e}")
-                print(e.stderr)
+            except Exception as ex:  # pylint: disable=broad-exception-caught
+                print(f"Error starting Celery worker '{worker}': {ex}")
 
         if successfully_started_workers:
             atexit.register(
@@ -141,28 +147,21 @@ def stop_background_workers(
             log_level: The log level.
         """
 
-        print("Stopping all Celery workers.")
-
         try:
-            process = subprocess.run(
+            subprocess.run(
                 [
                     "celery",
                     "multi",
                     "stopwait",
                     *workers,
-                    f"--app={self.app}",
+                    f"--app={self._app}",
                     f"--loglevel={log_level}",
                 ],
-                capture_output=True,
-                text=True,
                 check=True,
             )
-            print("Successfully stopped all Celery workers.")
-            print(process.stdout)
 
-        except subprocess.CalledProcessError as error:
-            print(f"Error stopping all Celery workers: {error}")
-            print(error.stderr)
+        except Exception as ex:  # pylint: disable=broad-exception-caught
+            print(f"Error stopping all Celery workers: {ex}")
 
     def start_background_beat(self, log_level: LogLevel = "INFO"):
         """Start Celery beat using the 'celery --app=app beat' command.
@@ -174,64 +173,23 @@ def start_background_beat(self, log_level: LogLevel = "INFO"):
             The background process running Celery beat.
         """
 
-        print("Starting Celery beat.")
-
         try:
             process = subprocess.Popen(  # pylint: disable=consider-using-with
                 [
                     "celery",
-                    f"--app={self.app}",
+                    f"--app={self._app}",
                     "beat",
                     f"--loglevel={log_level}",
                 ],
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.DEVNULL,
             )
-            print("Successfully started Celery beat.")
 
-            atexit.register(self.stop_background_beat, process)
+            atexit.register(process.terminate)
 
             return process
 
-        except subprocess.CalledProcessError as error:
-            print(f"Error starting Celery beat: {error}")
-
-        return None
-
-    def stop_background_beat(self, process: subprocess.Popen):
-        """Stop a Celery beat process.
-
-        Args:
-            process: The process to stop.
-        """
-
-        print("Stopping Celery beat.")
-
-        try:
-            process.terminate()
-            print("Successfully stopped Celery beat.")
-
         except Exception as ex:  # pylint: disable=broad-exception-caught
-            print(f"Error stopping Celery beat: {ex}")
-
-    def handle_startup(self):
-        """Handle the startup procedure of a Celery app.
-
-        Examples:
-            ```
-            from codeforlife.apps import CeleryApplication, DjangoApplication
-
-            # Make sure to set up Django before starting!
-            DjangoApplication.setup()
+            print(f"Error starting Celery beat: {ex}")
 
-            celery_app = CeleryApplication().handle_startup()
-            ```
-
-        Returns:
-            The Celery app instance for convenience.
-        """
-
-        self.start_background_workers()
-        self.start_background_beat()
-
-        return self
+        return None

codeforlife/apps/celery.py renamed to codeforlife/servers/celery.py

@@ -5,6 +5,7 @@
 
 import multiprocessing
 import os
+import sys
 
 from django import setup
 from django.core.asgi import get_asgi_application
@@ -14,20 +15,35 @@
 
 
 # pylint: disable-next=abstract-method
-class DjangoApplication(BaseApplication):
+class DjangoServer(BaseApplication):
     """A server for a Django app.
 
     Based off of:
     https://gist.github.com/Kludex/c98ed6b06f5c0f89fd78dd75ef58b424
     https://docs.gunicorn.org/en/stable/custom.html
     """
 
-    def __init__(self, workers: int = int(os.getenv("WORKERS", "0"))):
+    def __init__(
+        self,
+        app: str = "application",
+        workers: int = int(os.getenv("WORKERS", "0")),
+    ):
         """Initialize a Django app.
 
         Before starting, all migrations will be applied.
 
+        Examples:
+            ```
+            from codeforlife.apps import DjangoServer
+
+            # Make sure to set up Django before initializing!
+            DjangoServer.setup()
+
+            django_app = DjangoServer().wsgi_app
+            ```
+
         Args:
+            app: The dot-path to the Django app.
             workers: The number of Gunicorn workers. 0 will auto-calculate.
         """
 
@@ -40,10 +56,16 @@ def __init__(self, workers: int = int(os.getenv("WORKERS", "0"))):
             "worker_class": "uvicorn.workers.UvicornWorker",
         }
 
-        self.application = get_asgi_application()
+        self.asgi_app = get_asgi_application()
+        self.wsgi_app = get_wsgi_application()
 
         super().__init__()
 
+        # Auto-run if the entry script is the location of the Django app.
+        file_name = os.path.basename(sys.argv[0])
+        if os.path.splitext(file_name)[0] == app:
+            self.run()
+
     def load_config(self):
         config = {
             key: value
@@ -54,7 +76,7 @@ def load_config(self):
             self.cfg.set(key.lower(), value)
 
     def load(self):
-        return self.application
+        return self.asgi_app
 
     @staticmethod
     def setup(settings_module: str = "settings"):
@@ -67,30 +89,3 @@ def setup(settings_module: str = "settings"):
         os.environ.setdefault("DJANGO_SETTINGS_MODULE", settings_module)
 
         setup()
-
-    def handle_startup(self, name: str):
-        """Handle the startup procedure of a Django app.
-
-        Examples:
-            ```
-            from codeforlife.apps import DjangoApplication
-
-            # Make sure to set up Django before starting!
-            DjangoApplication.setup()
-
-            django_app = DjangoApplication().handle_startup(__name__)
-            ```
-
-        Args:
-            name: The name of the file calling this function.
-
-        Returns:
-            An instance of a WSGI app.
-        """
-
-        if name == "__main__":
-            self.run()
-        else:
-            return get_wsgi_application()
-
-        return None