some fixes

Signed-off-by: Philippe Leprince <philippe@ynput.io>

mkdocs_hooks.py

@@ -1,10 +1,71 @@
 import os
-import sys
 from pathlib import Path
 from shutil import rmtree
 import json
+import glob
+import logging
 
 TMP_FILE = "./missing_init_files.json"
+NFILES = []
+
+# -----------------------------------------------------------------------------
+
+
+class ColorFormatter(logging.Formatter):
+    grey = "\x1b[38;20m"
+    green = "\x1b[32;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    fmt = (
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s "  # noqa
+        "(%(filename)s:%(lineno)d)"
+    )
+
+    FORMATS = {
+        logging.DEBUG: grey + fmt + reset,
+        logging.INFO: green + fmt + reset,
+        logging.WARNING: yellow + fmt + reset,
+        logging.ERROR: red + fmt + reset,
+        logging.CRITICAL: bold_red + fmt + reset,
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+ch = logging.StreamHandler()
+ch.setFormatter(ColorFormatter())
+
+logging.basicConfig(
+    level=logging.INFO,
+    handlers=[ch],
+)
+
+
+# -----------------------------------------------------------------------------
+
+
+def create_init_file(dirpath, msg):
+    global NFILES
+    ini_file = f"{dirpath}/__init__.py"
+    Path(ini_file).touch()
+    NFILES.append(ini_file)
+    logging.info(f"{msg}: created '{ini_file}'")
+
+
+def create_parent_init_files(dirpath: str, rootpath: str, msg: str):
+    parent_path = dirpath
+    while parent_path != rootpath:
+        parent_path = os.path.dirname(parent_path)
+        parent_init = os.path.join(parent_path, "__init__.py")
+        if not os.path.exists(parent_init):
+            create_init_file(parent_path, msg)
+        else:
+            break
 
 
 def add_missing_init_files(*roots, msg=""):
@@ -19,28 +80,26 @@ def add_missing_init_files(*roots, msg=""):
     Returns:
         None
     """
-    nfiles = []
 
     for root in roots:
         if not os.path.exists(root):
             continue
-        for dirpath, dirs, files in os.walk(root):
+        rootpath = os.path.abspath(root)
+        for dirpath, dirs, files in os.walk(rootpath):
             if "__init__.py" in files:
                 continue
-            else:
-                Path(f"{dirpath}/__init__.py").touch()
-                nfiles.append(f"{dirpath}/__init__.py")
-                sys.stdout.write(
-                    "\r\x1b[K" + f"{msg}: created {len(nfiles)} "
-                    "temp '__init__.py' files"
-                )
-                sys.stdout.flush()
+
+            if "." in dirpath:
+                continue
+
+            if not glob.glob(os.path.join(dirpath, "*.py")):
+                continue
+
+            create_init_file(dirpath, msg)
+            create_parent_init_files(dirpath, rootpath, msg)
 
     with open(TMP_FILE, "w") as f:
-        json.dump(nfiles, f)
-
-    sys.stdout.write("\n")
-    sys.stdout.flush()
+        json.dump(NFILES, f)
 
 
 def remove_missing_init_files(msg=""):
@@ -55,26 +114,26 @@ def remove_missing_init_files(msg=""):
     Returns:
         None
     """
-    with open(TMP_FILE, "r") as f:
-        nfiles = json.load(f)
+    global NFILES
+    nfiles = []
+    if os.path.exists(TMP_FILE):
+        with open(TMP_FILE, "r") as f:
+            nfiles = json.load(f)
+    else:
+        nfiles = NFILES
 
     for file in nfiles:
         Path(file).unlink()
-        sys.stdout.write(
-            "\r\x1b[K" + f"{msg}: removed {len(nfiles)} temp '__init__.py' files"
-        )
-        sys.stdout.flush()
+        logging.info(f"{msg}: removed {file}")
 
     os.remove(TMP_FILE)
-
-    sys.stdout.write("\n")
-    sys.stdout.flush()
+    NFILES = []
 
 
 def remove_pychache_dirs(msg=""):
     """
-    This function walks the current directory and removes all existing '__pycache__'
-    directories.
+    This function walks the current directory and removes all existing
+    '__pycache__' directories.
 
     Args:
         msg: An optional message to display during the removal process.
@@ -89,19 +148,13 @@ def remove_pychache_dirs(msg=""):
             pydir = Path(f"{dirpath}/__pycache__")
             rmtree(pydir)
             nremoved += 1
-            sys.stdout.write(
-                "\r\x1b[K" + f"{msg}: removed {nremoved} '__pycache__' directories"
-            )
-            sys.stdout.flush()
+            logging.info(f"{msg}: removed '{pydir}'")
 
     if not nremoved:
-        sys.stdout.write(f"{msg}: no __pycache__ dirs found")
-
-    sys.stdout.write("\n")
-    sys.stdout.flush()
+        logging.info(f"{msg}: no __pycache__ dirs found")
 
 
-# mkdocs hooks -----------------------------------------------------------------
+# mkdocs hooks ----------------------------------------------------------------
 
 
 def on_startup(command, dirty):
@@ -114,13 +167,19 @@ def on_pre_build(config):
     temporary `__init__.py` files to directories that do not contain one, to
     make sure mkdocs doesn't ignore them.
     """
-    add_missing_init_files(
-        "client",
-        "server",
-        "services",
-        "tests",
-        msg="HOOK    -  on_pre_build",
-    )
+    try:
+        add_missing_init_files(
+            "client",
+            "server",
+            "services",
+            msg="HOOK    -  on_pre_build",
+        )
+    except BaseException as e:
+        logging.error(e)
+        remove_missing_init_files(
+            msg="HOOK    -  on_post_build: cleaning up on error !"
+        )
+        raise
 
 
 def on_post_build(config):