Tidy up

Author: matthewelwell

src/common/core/main.py

@@ -1,7 +1,6 @@
 import contextlib
 import logging
 import os
-import pathlib
 import sys
 import typing
 from tempfile import gettempdir
@@ -12,7 +11,7 @@
 
 from common.core.cli import healthcheck
 from common.core.constants import DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME
-from common.core.utils import clear_directory
+from common.core.utils import clear_directory, make_writable_directory
 
 logger = logging.getLogger(__name__)
 
@@ -38,25 +37,13 @@ def ensure_cli_env() -> typing.Generator[None, None, None]:
     # TODO @khvn26 Move logging setup to here
 
     # Prometheus multiproc support
-    prom_dir = pathlib.Path(
-        os.environ.setdefault(
-            "PROMETHEUS_MULTIPROC_DIR",
-            os.path.join(gettempdir(), DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME),
-        )
+    prom_dir = os.environ.setdefault(
+        "PROMETHEUS_MULTIPROC_DIR",
+        os.path.join(gettempdir(), DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME),
     )
-
-    if prom_dir.exists():
+    if os.path.exists(prom_dir):
         clear_directory(prom_dir)
-
-    prom_dir.mkdir(parents=True, exist_ok=True)
-
-    # While `mkdir` sets mode=0o777 by default, this can be affected by umask resulting in
-    # lesser permissions for other users. This step ensures the directory is writable for
-    # all users.
-    try:
-        prom_dir.chmod(0o777)
-    except Exception:
-        pass
+    make_writable_directory(prom_dir)
 
     # Currently we don't install Flagsmith modules as a package, so we need to add
     # $CWD to the Python path to be able to import them

src/common/core/utils.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 import pathlib
 import random
 import shutil
@@ -201,16 +202,28 @@ def using_database_replica(
     return manager.db_manager(chosen_replica)
 
 
-def clear_directory(dir_: pathlib.Path) -> None:
+def clear_directory(directory_path: str) -> None:
     """
     Safely clear a directory including all subdirectories and files.
     """
-    for p in dir_.rglob("*"):
+    for p in pathlib.Path(directory_path).rglob("*"):
         try:
             # Ensure that the cleanup doesn't silently fail on
             # files and subdirs created by other users.
             p.chmod(0o777)
-        except Exception:  # pragma: no cover
+        except (PermissionError, FileNotFoundError):  # pragma: no cover
             pass
 
-    shutil.rmtree(dir_, ignore_errors=True)
+    shutil.rmtree(directory_path, ignore_errors=True)
+
+
+def make_writable_directory(directory_path: str) -> None:
+    os.makedirs(directory_path, exist_ok=True)
+
+    try:
+        # While `mkdir` sets mode=0o777 by default, this can be affected by umask
+        # resulting in lesser permissions for other users. This step ensures the
+        # directory is writable for all users.
+        os.chmod(directory_path, 0o777)
+    except PermissionError:
+        pass