1 files changed, 20 insertions, 11 deletions
diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index 66101ab92e..66dd915a86 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -9,6 +9,7 @@ from __future__ import print_function
 import collections
 import json
 import logging
+import math
 import os
 
 import numpy
@@ -76,7 +77,8 @@ class StatsManager(object):
   """
 
   # pylint: disable=W0102
-  def __init__(self, smid='', title='', order=[], hide_domains=[]):
+  def __init__(self, smid='', title='', order=[], hide_domains=[],
+               accept_nan=True):
     """Initialize infrastructure for data and their statistics."""
     self._title = title
     self._data = collections.defaultdict(list)
@@ -85,6 +87,7 @@ class StatsManager(object):
     self._order = order
     self._hide_domains = hide_domains
     self._summary = {}
+    self._accept_nan = accept_nan
     self._logger = logging.getLogger('StatsManager')
 
   def AddSample(self, domain, sample):
@@ -93,14 +96,20 @@ class StatsManager(object):
     Args:
       domain: the domain name for the sample.
       sample: one time sample for domain, expect type float.
+
+    Raises:
+      StatsManagerError: if trying to add NaN and |_accept_nan| is false
     """
-    if isinstance(sample, int):
+    try:
       sample = float(sample)
-    if isinstance(sample, float):
-      self._data[domain].append(sample)
-      return
-    self._logger.warn('sample %s for domain %s is not a number, thus ignored.',
-                      sample, domain)
+    except ValueError:
+      # if we don't accept nan this will be caught below
+      self._logger.debug('sample %s for domain %s is not a number. Making NaN',
+                         sample, domain)
+      sample = float('NaN')
+    if not self._accept_nan and math.isnan(sample):
+      raise StatsManagerError('accept_nan is false. Cannot add NaN sample.')
+    self._data[domain].append(sample)
 
   def SetUnit(self, domain, unit):
     """Set the unit for a domain.
@@ -126,10 +135,10 @@ class StatsManager(object):
     for domain, data in self._data.iteritems():
       data_np = numpy.array(data)
       self._summary[domain] = {
-          'mean': data_np.mean(),
-          'min': data_np.min(),
-          'max': data_np.max(),
-          'stddev': data_np.std(),
+          'mean': numpy.nanmean(data_np),
+          'min': numpy.nanmin(data_np),
+          'max': numpy.nanmax(data_np),
+          'stddev': numpy.nanstd(data_np),
           'count': data_np.size,
       }