Merge pull request #9 from AlecThomson/logging

Logging

fixms/fix_ms.py

@@ -10,8 +10,8 @@
 
 from fixms.fix_ms_corrs import fix_ms_corrs
 from fixms.fix_ms_dir import fix_ms_dir
+from fixms.logger import logger
 
-logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 

fixms/fix_ms_corrs.py

@@ -18,7 +18,9 @@
 from casacore.tables import makecoldesc, table
 from tqdm import tqdm
 
-logger = logging.getLogger(__name__)
+from fixms.logger import TqdmToLogger, logger
+
+TQDM_OUT = TqdmToLogger(logger, level=logging.INFO)
 logger.setLevel(logging.INFO)
 
 
@@ -350,7 +352,9 @@ def fix_ms_corrs(
         else:
             # Only perform this correction if the data column was
             # successfully renamed.
-            for data_chunk in tqdm(data_chunks, total=nchunks):
+            for data_chunk in tqdm(
+                data_chunks, total=nchunks, file=TQDM_OUT, desc="Rotating correlations"
+            ):
                 data_chunk_cor = convert_correlations(
                     data_chunk,
                     pol_axis,
@@ -364,6 +368,10 @@ def fix_ms_corrs(
                 tab.flush()
                 start_row += len(data_chunk_cor)
 
+    logger.info(
+        f"Finished correcting {data_column} of {str(ms)}. Written to {corrected_data_column} column."
+    )
+
 
 def cli():
     import argparse

fixms/fix_ms_dir.py

@@ -12,8 +12,11 @@
 
 import numpy as np
 from casacore.tables import table, tablecopy, tableexists, taql
+from tqdm import trange
 
-logger = logging.getLogger(__name__)
+from fixms.logger import TqdmToLogger, logger
+
+TQDM_OUT = TqdmToLogger(logger, level=logging.INFO)
 logger.setLevel(logging.INFO)
 
 RAD2DEG = 180.0 / math.pi
@@ -226,6 +229,7 @@ def decs_rad(dec_string):
 
 
 def fix_ms_dir(ms):
+    logger.info("Fixing FEED directions in %s" % (ms))
     # Check that the observation wasn't in pol_fixed mode
     with table("%s/ANTENNA" % (ms), readonly=True, ack=False) as ta:
         ant_mount = ta.getcol("MOUNT", 0, 1)
@@ -274,7 +278,7 @@ def fix_ms_dir(ms):
         offset_times = t1.getcol("TIME")
         offset_intervals = t1.getcol("INTERVAL")
         logger.info("Found %d offsets in FEED table for beam %d" % (n_offsets, beam))
-        for offset_index in range(n_offsets):
+        for offset_index in trange(n_offsets, desc="Fixing offsets", file=TQDM_OUT):
             offset = t1.getcol("BEAM_OFFSET")[offset_index]
             logger.info(
                 "Offset %d : t=%f-%f : (%fd,%fd)"
@@ -288,8 +292,8 @@ def fix_ms_dir(ms):
             )
 
         # Update the beam position for each field
-        for field in range(n_fields):
-            with table(ms, readonly=True, ack=False) as t:
+        for field in trange(n_fields, desc="Fixing fields", file=TQDM_OUT):
+            with table(ms, readonly=True, ack=False) as _:
                 # Get times for the specified field
                 tfdata = taql(
                     "select from $t where FIELD_ID==$field and FEED1==$beam and ANTENNA1==0 and ANTENNA2==0"
@@ -347,6 +351,8 @@ def fix_ms_dir(ms):
             tp.putcol("PHASE_DIR", ms_phase)
             tp.putcol("REFERENCE_DIR", ms_phase)
 
+    logger.info("Finished fixed FEED directions")
+
 
 def cli():
     import argparse

fixms/logger.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Logging module for FixMS"""
+
+import io
+import logging
+
+
+class TqdmToLogger(io.StringIO):
+    """
+    Output stream for TQDM which will output to logger module instead of
+    the StdOut.
+    """
+
+    logger = None
+    level = None
+    buf = ""
+
+    def __init__(self, logger, level=None):
+        super(TqdmToLogger, self).__init__()
+        self.logger = logger
+        self.level = level or logging.INFO
+
+    def write(self, buf):
+        self.buf = buf.strip("\r\n\t ")
+
+    def flush(self):
+        self.logger.log(self.level, self.buf)
+
+
+# Create formatter
+# formatter = logging.Formatter(
+#     "SPICE: %(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s"
+# )
+class CustomFormatter(logging.Formatter):
+    grey = "\x1b[38;20m"
+    blue = "\x1b[34;20m"
+    green = "\x1b[32;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    format_str = "%(asctime)s.%(msecs)03d %(module)s - %(funcName)s: %(message)s"
+
+    FORMATS = {
+        logging.DEBUG: f"{blue}FixMS-%(levelname)s{reset} {format_str}",
+        logging.INFO: f"{green}FixMS-%(levelname)s{reset} {format_str}",
+        logging.WARNING: f"{yellow}FixMS-%(levelname)s{reset} {format_str}",
+        logging.ERROR: f"{red}FixMS-%(levelname)s{reset} {format_str}",
+        logging.CRITICAL: f"{bold_red}FixMS-%(levelname)s{reset} {format_str}",
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt, "%Y-%m-%d %H:%M:%S")
+        return formatter.format(record)
+
+
+def get_fixms_logger(
+    name: str = "fixms", attach_handler: bool = True
+) -> logging.Logger:
+    """Will construct a logger object.
+
+    Args:
+        name (str, optional): Name of the logger to attempt to use. This is ignored if in a prefect flowrun. Defaults to 'arrakis'.
+        attach_handler (bool, optional): Attacjes a custom StreamHandler. Defaults to True.
+
+    Returns:
+        logging.Logger: The appropriate logger
+    """
+    logging.captureWarnings(True)
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.WARNING)
+
+    if attach_handler:
+        # Create console handler and set level to debug
+        ch = logging.StreamHandler()
+        ch.setLevel(logging.DEBUG)
+
+        # Add formatter to ch
+        ch.setFormatter(CustomFormatter())
+        logger.addHandler(ch)
+
+    return logger
+
+
+logger = get_fixms_logger()