cros_utils: Make GetLogger have type guarantees

At present, GetLogger could return None, Logger, or MockLogger. No code
seems to be built for when GetLogger returns None, and it would instead
be nice to have that safety.

The original InitLogger relies on modifying global state to ensure
the main_logger is not None. This does not make the type analysis happy,
so instead lets have InitLogger return the logger we initted.

This also does a minor clean up of the code.

BUG=None
TEST=./run_tests_for.py cros_utils/*

Change-Id: I2f5b142e088ceb0299d7ac2513b7c3c6f3800800
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/4674780
Tested-by: Jordan Abrahams-Whitehead <ajordanr@google.com>
Reviewed-by: George Burgess <gbiv@chromium.org>
Commit-Queue: Jordan Abrahams-Whitehead <ajordanr@google.com>

1 files changed, 9 insertions, 9 deletions

diff --git a/cros_utils/logger.py b/cros_utils/logger.py
index e9b9d1ba..7df3f8ff 100644
--- a/cros_utils/logger.py
+++ b/cros_utils/logger.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright 2019 The ChromiumOS Authors
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
@@ -10,6 +9,7 @@
 import os.path
 import sys
 import traceback
+from typing import Union
 
 
 # TODO(yunlian@google.com): Use GetRoot from misc
@@ -19,7 +19,7 @@ def GetRoot(scr_name):
     return (os.path.dirname(abs_path), os.path.basename(abs_path))
 
 
-class Logger(object):
+class Logger:
     """Logging helper class."""
 
     MAX_LOG_FILES = 10
@@ -67,7 +67,7 @@ class Logger(object):
     def _CreateLogFileHandle(self, name):
         fd = None
         try:
-            fd = open(name, "w")
+            fd = open(name, "w", encoding="utf-8")
         except IOError:
             print("Warning: could not open %s for writing." % name)
         return fd
@@ -215,7 +215,7 @@ class Logger(object):
         self.stderr.flush()
 
 
-class MockLogger(object):
+class MockLogger:
     """Logging helper class."""
 
     MAX_LOG_FILES = 10
@@ -363,7 +363,8 @@ def InitLogger(script_name, log_dir, print_console=True, mock=False):
     """Initialize a global logger. To be called only once."""
     # pylint: disable=global-statement
     global main_logger
-    assert not main_logger, "The logger has already been initialized"
+    if main_logger:
+        return main_logger
     rootdir, basefilename = GetRoot(script_name)
     if not log_dir:
         log_dir = rootdir
@@ -371,12 +372,11 @@ def InitLogger(script_name, log_dir, print_console=True, mock=False):
         main_logger = Logger(log_dir, basefilename, print_console)
     else:
         main_logger = MockLogger(log_dir, basefilename, print_console)
+    return main_logger
 
 
-def GetLogger(log_dir="", mock=False):
-    if not main_logger:
-        InitLogger(sys.argv[0], log_dir, mock=mock)
-    return main_logger
+def GetLogger(log_dir="", mock=False) -> Union[Logger, MockLogger]:
+    return InitLogger(sys.argv[0], log_dir, mock=mock)
 
 
 def HandleUncaughtExceptions(fun):