Merge pull request #85 from j9ac9k/make-console-friendly-again

Replace problematic utf-8 characters with ASCII

pyproject.toml

@@ -88,6 +88,7 @@ module = [
   "pyproj.database",
   "pyproj.enums",
   "pyproj.transformer",
+  "pythonjsonlogger",
   "rasterio",
   "rasterio.coords",
   "rasterio.crs",
@@ -107,5 +108,6 @@ module = [
   "skimage",
   "skimage.measure",
   "trimesh",
+  "websocket"
 ]
 ignore_missing_imports = true

src/codem/lib/console.py

@@ -0,0 +1,14 @@
+from typing import Any
+
+
+class DummyConsole:
+    """
+    Class that acts as a drop in replacement for rich's Console, in case
+    rich is not present on the system, or we are not wanting rich output
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self.level = float("-inf")
+
+    def print(self, *args: Any, **kwargs: Any) -> None:
+        print(*args)

src/codem/lib/log.py

@@ -7,17 +7,63 @@
 """
 import logging
 import os
+from typing import Any
+from typing import Dict
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from codem.preprocessing import CodemParameters
+    from vcd.preprocessing import VCDParameters
+
+try:
+    import websocket
+    from pythonjsonlogger import jsonlogger
+except ImportError:
+    pass
+else:
+
+    class CustomJsonFormatter(jsonlogger.JsonFormatter):
+        def add_fields(
+            self,
+            log_record: Dict[str, Any],
+            record: logging.LogRecord,
+            message_dict: Dict[str, Any],
+        ) -> None:
+            super().add_fields(log_record, record, message_dict)
+            if log_record.get("level"):
+                log_record["level"] = log_record["level"].upper()
+            else:
+                log_record["level"] = record.levelname
+
+            if log_record.get("type") is None:
+                log_record["type"] = "log_message"
+            return None
+
+    class WebSocketHandler(logging.Handler):
+        def __init__(self, level: str, websocket: "websocket.WebSocket") -> None:
+            super().__init__(level)
+            self.ws = websocket
+            # TODO: check if websocket is already connected?
+
+        def emit(self, record: logging.LogRecord) -> None:
+            msg = self.format(record)
+            _ = self.ws.send(msg)
+            return None
+
+        def close(self) -> None:
+            self.ws.close()
+            return super().close()
 
 
 class Log:
-    def __init__(self, config: dict):
+    def __init__(self, config: Dict[str, Any]):
         """
         Creates logging formatting and structure
 
         Parameters
         ----------
-        verbose: bool
-            Verbose mode switch
+        config:
+            Dictionary representing the runtime config
         """
 
         self.logger = logging.getLogger("codem")
@@ -26,14 +72,34 @@ def __init__(self, config: dict):
         # disable loggers
         logging.getLogger("matplotlib.font_manager").disabled = True
 
-        # configure Formatting
+        # File Handler for Logging
         log_format = logging.Formatter(
             "%(asctime)s - %(name)s - %(levelname)s - %(funcName)s:%(lineno)d - %(message)s"
         )
         file_handler = logging.FileHandler(
-            os.path.join(config["OUTPUT_DIR"], "log.txt")
+            os.path.join(config.get("OUTPUT_DIR", "."), "log.txt")
         )
         file_handler.setLevel(logging.DEBUG)
         file_handler.setFormatter(log_format)
         self.logger.addHandler(file_handler)
-        config['log'] = self
+        self.relay = None
+
+        # Supplemental Handler
+        if config["LOG_TYPE"] == "rich":
+            from rich.logging import RichHandler
+
+            log_handler = RichHandler()
+        elif config["LOG_TYPE"] == "websocket":
+            formatter = CustomJsonFormatter()
+            self.relay = websocket.WebSocket()
+            self.relay.connect(f'ws://{config["WEBSOCKET_URL"]}/websocket')
+            log_handler = WebSocketHandler("DEBUG", websocket=self.relay)
+            log_handler.setFormatter(formatter)
+        else:
+            log_handler = logging.StreamHandler()
+        log_handler.setLevel("DEBUG")
+        self.logger.addHandler(log_handler)
+
+    def __del__(self) -> None:
+        if isinstance(self.logger, WebSocketHandler):
+            self.logger.close()

src/codem/lib/progress.py

@@ -0,0 +1,38 @@
+import json
+from contextlib import ContextDecorator
+from typing import Any
+from typing import Dict
+
+
+try:
+    import websocket
+except ImportError:
+    pass
+else:
+
+    class WebSocketProgress(ContextDecorator):
+        def __init__(self, url: str) -> None:
+            super().__init__()
+            self.ws = websocket.WebSocket()
+            self.tasks: Dict[str, int] = {}
+            self.current: Dict[str, int] = {}
+            self.url = url
+
+        def __enter__(self) -> Any:
+            self.ws.connect(f"ws://{self.url}/websocket")
+            return self
+
+        def __exit__(self, *args: Any, **kwargs: Any) -> None:
+            self.ws.close()
+            return None
+
+        def advance(self, name: str, value: int) -> None:
+            self.current[name] += value
+            new_value = self.current[name]
+            self.ws.send(json.dumps({"advance": new_value, "type": "progress"}))
+            return None
+
+        def add_task(self, title: str, total: int) -> str:
+            self.tasks[title] = total
+            self.current[title] = 0
+            return title