From 85517ff7901854b229a10914a5234ea58072ecaa Mon Sep 17 00:00:00 2001
From: bozowski <bozowski@google.com>
Date: Fri, 20 Oct 2023 13:07:40 -0700
Subject: [PATCH] Fix prober, color config

---
 src/tools/interop/idt/capture/config.py       |  1 +
 .../capture/ecosystem/play_services/prober.py |  4 +++-
 src/tools/interop/idt/capture/utils/log.py    | 20 ++++++++++++++-----
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/tools/interop/idt/capture/config.py b/src/tools/interop/idt/capture/config.py
index ebcd13f8bbe2b5..6ca0b1bf4fd407 100644
--- a/src/tools/interop/idt/capture/config.py
+++ b/src/tools/interop/idt/capture/config.py
@@ -19,3 +19,4 @@
 
 log_level = logging.INFO
 async_timeout = 45.0
+enable_color = True
diff --git a/src/tools/interop/idt/capture/ecosystem/play_services/prober.py b/src/tools/interop/idt/capture/ecosystem/play_services/prober.py
index 37722597715610..bd461990114f33 100644
--- a/src/tools/interop/idt/capture/ecosystem/play_services/prober.py
+++ b/src/tools/interop/idt/capture/ecosystem/play_services/prober.py
@@ -18,6 +18,8 @@
 import os
 
 from capture.utils.shell import Bash
+from capture.utils import log
+logger = log.get_logger(__file__)
 
 
 class PlayServicesProber:
@@ -25,7 +27,7 @@ class PlayServicesProber:
     def __init__(self, platform):
         self.platform = platform
         self.artifact_dir = self.platform.artifact_dir
-        self.logger = log_format.get_logger(__file__)
+        self.logger = logger
 
     async def _probe_foyer(self) -> None:
         self.logger.info("probing remote services")
diff --git a/src/tools/interop/idt/capture/utils/log.py b/src/tools/interop/idt/capture/utils/log.py
index aae9b5fa56d74d..963694f8ed3d14 100644
--- a/src/tools/interop/idt/capture/utils/log.py
+++ b/src/tools/interop/idt/capture/utils/log.py
@@ -18,7 +18,7 @@
 import logging
 from typing import TextIO
 
-from capture import config
+from .. import config
 
 _CONFIG_LEVEL = config.log_level
 
@@ -33,8 +33,10 @@
 
 reset = "\x1b[0m"
 
-format_pre = cyan + "%(asctime)s %(levelname)s {%(module)s} [%(funcName)s] " + reset
+format_pre = "%(asctime)s %(levelname)s {%(module)s} [%(funcName)s] " + reset
+format_pre = cyan + format_pre if config.enable_color else format_pre
 format_post = "%(message)s" + reset
+format_no_color = format_pre+format_post
 
 FORMATS = {
     logging.DEBUG: format_pre + blue + format_post,
@@ -48,7 +50,7 @@
 class LoggingFormatter(logging.Formatter):
 
     def format(self, record):
-        log_fmt = FORMATS.get(record.levelno)
+        log_fmt = FORMATS.get(record.levelno) if config.enable_color else format_no_color
         formatter = logging.Formatter(log_fmt)
         return formatter.format(record)
 
@@ -68,11 +70,19 @@ def border_print(to_print: str, important: bool = False) -> None:
     len_borders = len(to_print)
     border = f"\n{'_' * len_borders}\n"
     i_border = f"\n{'!' * len_borders}\n" if important else ""
-    print(f"\x1b[35;1m{border}{i_border}{to_print}{i_border}{border}\x1b[0m")
+    if config.enable_color:
+        print("\x1b[35;1m")
+    print(f"{border}{i_border}{to_print}{i_border}{border}")
+    if config.enable_color:
+        print("\x1b[0m")
 
 
 def print_and_write(to_print: str, file: TextIO) -> None:
-    print(f"\x1b[32;1m{to_print}\x1b[0m")
+    if config.enable_color:
+        print("\x1b[32;1m")
+    print(f"{to_print}")
+    if config.enable_color:
+        print("\x1b[0m")
     file.write(to_print)