stats_manager: remove expliting printing from stats_manager

StatsManager is a library shared across some modules. It shouldn't print
anything itself, but rather use logging & return a formatted summary
that can be printed.

Also take this opportunity to replace majority of explicit print calls
in powerlog with logger module.

BRANCH=None
BUG=chromium:760267
TEST=manual testing, output as expected still for powerlog

Change-Id: I0ed2ccaa45bbd1261119a1446110beadb075c1a2
Signed-off-by: Ruben Rodriguez Buchillon <coconutruben@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1140026
Reviewed-by: Mengqi Guo <mqg@chromium.org>

3 files changed, 64 insertions, 62 deletions

diff --git a/extra/usb_power/powerlog.py b/extra/usb_power/powerlog.py
index 6d7a8bab26..a431b05b5f 100755
--- a/extra/usb_power/powerlog.py
+++ b/extra/usb_power/powerlog.py
@@ -12,12 +12,13 @@ import argparse
 import array
 from distutils import sysconfig
 import json
+import logging
 import os
+import pprint
 import struct
 import sys
 import time
 import traceback
-from pprint import pprint
 
 import usb
 
@@ -32,12 +33,6 @@ LIB_DIR = os.path.join(sysconfig.get_python_lib(standard_lib=False), 'servo',
 CONFIG_LOCATIONS = [os.getcwd(), os.path.dirname(os.path.realpath(__file__)),
                     LIB_DIR]
 
-# This can be overridden by -v.
-debug = False
-def debuglog(msg):
-  if debug:
-    print(msg)
-
 def logoutput(msg):
   print(msg)
   sys.stdout.flush()
@@ -155,6 +150,7 @@ class Spower(object):
 
   def __init__(self, board, vendor=0x18d1,
                product=0x5020, interface=1, serialname=None):
+    self._logger = logging.getLogger(__name__)
     self._board = board
 
     # Find the stm32.
@@ -184,7 +180,7 @@ class Spower(object):
         # Incompatible pyUsb version.
         dev = dev_list.next()
 
-    debuglog("Found USB device: %04x:%04x" % (vendor, product))
+    self._logger.debug("Found USB device: %04x:%04x", vendor, product)
     self._dev = dev
 
     # Get an endpoint instance.
@@ -198,7 +194,7 @@ class Spower(object):
         i.bInterfaceClass==255 and i.bInterfaceSubClass==0x54)
 
     self._intf = intf
-    debuglog("InterfaceNumber: %s" % intf.bInterfaceNumber)
+    self._logger.debug("InterfaceNumber: %s", intf.bInterfaceNumber)
 
     read_ep = usb.util.find_descriptor(
       intf,
@@ -210,7 +206,7 @@ class Spower(object):
     )
 
     self._read_ep = read_ep
-    debuglog("Reader endpoint: 0x%x" % read_ep.bEndpointAddress)
+    self._logger.debug("Reader endpoint: 0x%x", read_ep.bEndpointAddress)
 
     write_ep = usb.util.find_descriptor(
       intf,
@@ -222,11 +218,11 @@ class Spower(object):
     )
 
     self._write_ep = write_ep
-    debuglog("Writer endpoint: 0x%x" % write_ep.bEndpointAddress)
+    self._logger.debug("Writer endpoint: 0x%x", write_ep.bEndpointAddress)
 
     self.clear_ina_struct()
 
-    debuglog("Found power logging USB endpoint.")
+    self._logger.debug("Found power logging USB endpoint.")
 
   def clear_ina_struct(self):
     """ Clear INA description struct."""
@@ -278,8 +274,8 @@ class Spower(object):
     Returns:
       bytes read, or None on failure.
     """
-    debuglog("Spower.wr_command(write_list=[%s] (%d), read_count=%s)" % (
-          list(bytearray(write_list)), len(write_list), read_count))
+    self._logger.debug("Spower.wr_command(write_list=[%s] (%d), read_count=%s)",
+                       list(bytearray(write_list)), len(write_list), read_count)
 
     # Clean up args from python style to correct types.
     write_length = 0
@@ -293,17 +289,17 @@ class Spower(object):
       cmd = write_list
       ret = self._write_ep.write(cmd, wtimeout)
 
-    debuglog("RET: %s " % ret)
+    self._logger.debug("RET: %s ", ret)
 
     # Read back response if necessary.
     if read_count:
       bytesread = self._read_ep.read(512, rtimeout)
-      debuglog("BYTES: [%s]" % bytesread)
+      self._logger.debug("BYTES: [%s]", bytesread)
 
       if len(bytesread) != read_count:
         pass
 
-      debuglog("STATUS: 0x%02x" % int(bytesread[0]))
+      self._logger.debug("STATUS: 0x%02x", int(bytesread[0]))
       if read_count == 1:
         return bytesread[0]
       else:
@@ -316,7 +312,8 @@ class Spower(object):
     try:
       while True:
         ret = self.wr_command("", read_count=512, rtimeout=100, wtimeout=50)
-        debuglog("Try Clear: read %s" % "success" if ret == 0 else "failure")
+        self._logger.debug("Try Clear: read %s",
+                           "success" if ret == 0 else "failure")
     except:
       pass
 
@@ -324,7 +321,8 @@ class Spower(object):
     """Reset the power interface on the stm32"""
     cmd = struct.pack("<H", self.CMD_RESET)
     ret = self.wr_command(cmd, rtimeout=50, wtimeout=50)
-    debuglog("Command RESET: %s" % "success" if ret == 0 else "failure")
+    self._logger.debug("Command RESET: %s",
+                       "success" if ret == 0 else "failure")
 
   def reset(self):
     """Try resetting the USB interface until success.
@@ -343,7 +341,7 @@ class Spower(object):
       except Exception as e:
         self.clear()
         self.clear()
-        debuglog("TRY %d of %d: %s" % (count, max_reset_retry, e))
+        self._logger.debug("TRY %d of %d: %s", count, max_reset_retry, e)
         time.sleep(count * 0.01)
     raise Exception("Power", "Failed to reset")
 
@@ -351,7 +349,8 @@ class Spower(object):
     """Stop any active data acquisition."""
     cmd = struct.pack("<H", self.CMD_STOP)
     ret = self.wr_command(cmd)
-    debuglog("Command STOP: %s" % "success" if ret == 0 else "failure")
+    self._logger.debug("Command STOP: %s",
+                       "success" if ret == 0 else "failure")
 
   def start(self, integration_us):
     """Start data acquisition.
@@ -368,10 +367,10 @@ class Spower(object):
     actual_us = 0
     if len(read) == 5:
       ret, actual_us = struct.unpack("<BI", read)
-      debuglog("Command START: %s %dus" % (
-               "success" if ret == 0 else "failure", actual_us))
+      self._logger.debug("Command START: %s %dus",
+               "success" if ret == 0 else "failure", actual_us)
     else:
-      debuglog("Command START: FAIL")
+      self._logger.debug("Command START: FAIL")
 
     return actual_us
 
@@ -417,7 +416,8 @@ class Spower(object):
     cmd = struct.pack("<HQ", self.CMD_SETTIME, timestamp_us)
     ret = self.wr_command(cmd)
 
-    debuglog("Command SETTIME: %s" % "success" if ret == 0 else "failure")
+    self._logger.debug("Command SETTIME: %s",
+                       "success" if ret == 0 else "failure")
 
   def add_ina(self, bus, ina_type, addr, extra, resistance, data=None):
     """Add an INA to the data acquisition list.
@@ -440,7 +440,8 @@ class Spower(object):
         name = "ina%d_%02x" % (bus, addr)
       self.append_ina_struct(name, resistance, bus, addr,
                              data=data, ina_type=ina_type)
-    debuglog("Command ADD_INA: %s" % "success" if ret == 0 else "failure")
+    self._logger.debug("Command ADD_INA: %s",
+                       "success" if ret == 0 else "failure")
 
   def report_header_size(self):
     """Helper function to calculate power record header size."""
@@ -470,18 +471,18 @@ class Spower(object):
       cmd = struct.pack("<H", self.CMD_NEXT)
       bytesread = self.wr_command(cmd, read_count=expected_bytes)
     except usb.core.USBError as e:
-      print("READ LINE FAILED %s" % e)
+      self._logger.error("READ LINE FAILED %s", e)
       return None
 
     if len(bytesread) == 1:
       if bytesread[0] != 0x6:
-        debuglog("READ LINE FAILED bytes: %d ret: %02x" % (
-            len(bytesread), bytesread[0]))
+        self._logger.debug("READ LINE FAILED bytes: %d ret: %02x",
+                           len(bytesread), bytesread[0])
       return None
 
     if len(bytesread) % expected_bytes != 0:
-      debuglog("READ LINE WARNING: expected %d, got %d" % (
-          expected_bytes, len(bytesread)))
+      self._logger.debug("READ LINE WARNING: expected %d, got %d",
+          expected_bytes, len(bytesread))
 
     packet_count = len(bytesread) / expected_bytes
 
@@ -508,13 +509,14 @@ class Spower(object):
     """
     status, size = struct.unpack("<BB", data[0:2])
     if len(data) != self.report_size(size):
-      print("READ LINE FAILED st:%d size:%d expected:%d len:%d" % (
-          status, size, self.report_size(size), len(data)))
+      self._logger.error("READ LINE FAILED st:%d size:%d expected:%d len:%d",
+                         status, size, self.report_size(size), len(data))
     else:
       pass
 
     timestamp = struct.unpack("<Q", data[2:10])[0]
-    debuglog("READ LINE: st:%d size:%d time:%dus" % (status, size, timestamp))
+    self._logger.debug("READ LINE: st:%d size:%d time:%dus", status, size,
+                       timestamp)
     ftimestamp = float(timestamp) / 1000000.
 
     record = {"ts": ftimestamp, "status": status, "berry":self._board}
@@ -535,8 +537,8 @@ class Spower(object):
       elif self._inas[i]['type'] == Spower.INA_SHUNTV:
         val = raw_val * self._inas[i]['uVscale']
 
-      debuglog("READ %d %s: %fs: 0x%04x %f" % (i,
-          name, ftimestamp, raw_val, val))
+      self._logger.debug("READ %d %s: %fs: 0x%04x %f", i, name, ftimestamp,
+                         raw_val, val)
       record[name_tuple] = val
 
     return record
@@ -552,8 +554,7 @@ class Spower(object):
 
     #TODO: validate this.
     self._brdcfg = data;
-    if debug:
-      pprint(data)
+    self._logger.debug(pprint.pformat(data))
 
 
 class powerlog(object):
@@ -592,6 +593,7 @@ class powerlog(object):
       raw_data_dir: directory to save sweetberry readings raw data; if None then
                     do not save the raw data.
     """
+    self._logger = logging.getLogger(__name__)
     self._data = StatsManager()
     self._pwr = {}
     self._use_ms = use_ms
@@ -739,8 +741,8 @@ class powerlog(object):
               aggregate_record[rkey] = record[rkey]
             aggregate_record["boards"].add(record["berry"])
           else:
-            print("break %s, %s" % (record["berry"],
-                                    aggregate_record["boards"]))
+            self._logger.info("break %s, %s", record["berry"],
+                              aggregate_record["boards"])
             break
 
           if aggregate_record["boards"] == set(self._pwr.keys()):
@@ -764,14 +766,14 @@ class powerlog(object):
               pending_records.pop(0)
 
     except KeyboardInterrupt:
-      print('\nCTRL+C caught.')
+      self._logger.info('\nCTRL+C caught.')
 
     finally:
       for key in self._pwr:
         self._pwr[key].stop()
       self._data.CalculateStats()
       if self._print_stats:
-        self._data.PrintSummary()
+        print(self._data.SummaryToString())
       save_dir = 'sweetberry%s' % time.time()
       if self._stats_dir:
         stats_dir = os.path.join(self._stats_dir, save_dir)
@@ -846,9 +848,15 @@ def main(argv=None):
 
   args = parser.parse_args(argv)
 
-  global debug
+  root_logger = logging.getLogger()
   if args.verbose:
-    debug = True
+    root_logger.setLevel(logging.DEBUG)
+  else:
+    root_logger.setLevel(logging.INFO)
+  # if powerlog is used through main log to sys.stdout
+  stdout_handler = logging.StreamHandler(sys.stdout)
+  stdout_handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))
+  root_logger.addHandler(stdout_handler)
 
   integration_us_request = args.integration_us
   if not args.board:
diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index a53f555832..43ff6eee2d 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -8,6 +8,7 @@ from __future__ import print_function
 
 import collections
 import json
+import logging
 import os
 
 import numpy
@@ -50,7 +51,7 @@ class StatsManager(object):
     >>> stats.AddValue('foobar', 11111.0)
     >>> stats.AddValue('foobar', 22222.0)
     >>> stats.CalculateStats()
-    >>> stats.PrintSummary()
+    >>> print(stats.SummaryToString())
     @@            NAME  COUNT      MEAN   STDDEV       MAX       MIN
     @@   sample_msecs      4     31.25    15.16     50.00     10.00
     @@         foobar      2  16666.50  5555.50  22222.00  11111.00
@@ -60,6 +61,7 @@ class StatsManager(object):
     _data: dict of list of readings for each domain(key)
     _unit: dict of unit for each domain(key)
     _summary: dict of stats per domain (key): min, max, count, mean, stddev
+    _logger = StatsManager logger
 
   Note:
     _summary is empty until CalculateStats() is called, and is updated when
@@ -71,6 +73,7 @@ class StatsManager(object):
     self._data = collections.defaultdict(list)
     self._unit = collections.defaultdict(str)
     self._summary = {}
+    self._logger = logging.getLogger('StatsManager')
 
   def AddValue(self, domain, value):
     """Add one value for a domain.
@@ -84,8 +87,8 @@ class StatsManager(object):
     if isinstance(value, float):
       self._data[domain].append(value)
       return
-    print('Warning: value %s for domain %s is not a number, thus ignored.' %
-          (value, domain))
+    self._logger.warn('value %s for domain %s is not a number, thus ignored.',
+                      value, domain)
 
   def SetUnit(self, domain, unit):
     """Set the unit for a domain.
@@ -98,8 +101,8 @@ class StatsManager(object):
       unit: unit of the domain.
     """
     if domain in self._unit:
-      print('Warning: overwriting the unit of %s, old unit is %s, new unit is '
-            '%s.' % (domain, self._unit[domain], unit))
+      self._logger.warn('overwriting the unit of %s, old unit is %s, new unit '
+                        'is %s.', domain, self._unit[domain], unit)
     self._unit[domain] = unit
 
   def CalculateStats(self):
@@ -118,7 +121,7 @@ class StatsManager(object):
           'count': data_np.size,
       }
 
-  def _SummaryToString(self, prefix=STATS_PREFIX):
+  def SummaryToString(self, prefix=STATS_PREFIX):
     """Format summary into a string, ready for pretty print.
 
     See class description for format example.
@@ -156,15 +159,6 @@ class StatsManager(object):
       formatted_table.append(formatted_row)
     return '\n'.join(formatted_table)
 
-  def PrintSummary(self, prefix=STATS_PREFIX):
-    """Print the formatted summary.
-
-    Args:
-      prefix: start every row in summary string with prefix, for easier reading.
-    """
-    summary_str = self._SummaryToString(prefix=prefix)
-    print(summary_str)
-
   def GetSummary(self):
     """Getter for summary."""
     return self._summary
@@ -177,7 +171,7 @@ class StatsManager(object):
       fname: filename to save summary under.
       prefix: start every row in summary string with prefix, for easier reading.
     """
-    summary_str = self._SummaryToString(prefix=prefix) + '\n'
+    summary_str = self.SummaryToString(prefix=prefix) + '\n'
 
     if not os.path.exists(directory):
       os.makedirs(directory)
diff --git a/extra/usb_power/stats_manager_unittest.py b/extra/usb_power/stats_manager_unittest.py
index 362db3fa77..ed317eca02 100644
--- a/extra/usb_power/stats_manager_unittest.py
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -76,7 +76,7 @@ class TestStatsManager(unittest.TestCase):
     self.data.AddValue('test', 250)
     self.data.SetUnit('test', 'mw')
     self.data.CalculateStats()
-    summary_str = self.data._SummaryToString()
+    summary_str = self.data.SummaryToString()
     self.assertIn('test_mw', summary_str)
 
   def test_DoubleUnitSuffix(self):
@@ -84,7 +84,7 @@ class TestStatsManager(unittest.TestCase):
     self.data.AddValue('test_mw', 250)
     self.data.SetUnit('test_mw', 'mw')
     self.data.CalculateStats()
-    summary_str = self.data._SummaryToString()
+    summary_str = self.data.SummaryToString()
     self.assertIn('test_mw', summary_str)
     self.assertNotIn('test_mw_mw', summary_str)