servod: add data trimming to TimelinedStatsManager

This CL introduces the ability to trim to data inside a
TimelinedStatsManager to only the samples where tstart <= t <= tend for
the timeline. This trimming process is in place and operates on the raw
data of the TimelinedStatsmanager.
The logic for this is from: go/power-status-numpy-magic

BUG=chromium:806146
TEST=added unit tests and ran them, seems to work fine.

Change-Id: I372df1b04163f6704ffdc27b363dcabc7c78fb2a
Signed-off-by: Ruben Rodriguez Buchillon <coconutruben@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1163595
Reviewed-by: Mengqi Guo <mqg@chromium.org>
diff --git a/servo/timelined_stats_manager.py b/servo/timelined_stats_manager.py
index 8595932..f4f58cf 100644
--- a/servo/timelined_stats_manager.py
+++ b/servo/timelined_stats_manager.py
@@ -8,6 +8,7 @@
 import copy
 import time
 
+import numpy
 import stats_manager
 
 TIME_KEY = 'time'
@@ -60,6 +61,7 @@
     super(TimelinedStatsManager, self).CalculateStats()
 
   def AddSample(self, domain, sample):
+    # pylint: disable=C6113
     """NotImplemented.
 
     In order to preserve the balanced timeline adding invidual samples is
@@ -100,3 +102,26 @@
     samples.extend(known_domains_missing_nans)
     for domain, sample in samples:
       super(TimelinedStatsManager, self).AddSample(domain, sample)
+
+  def TrimSamples(self, tstart=None, tend=None):
+    """Trim raw data to [tstart, tend].
+
+    Args:
+      tstart: first timestamp to include. Seconds since epoch
+      tend: last timestamp to include. Seconds since epoch
+    """
+    if tstart is None and tend is None:
+      # Avoid doing any work if there will be no trimming.
+      return
+
+    timeline = numpy.array(self._data[self._tkey])
+    if tstart is None:
+      tstart = timeline[0]
+    if tend is None:
+      tend = timeline[-1]
+    # pylint: disable=W0212
+    for domain, samples in self._data.iteritems():
+      sample_arr = numpy.array(samples)
+      trimmed_samples = sample_arr[numpy.bitwise_and(tstart <= timeline,
+                                                     timeline <= tend)]
+      self._data[domain] = trimmed_samples.tolist()
diff --git a/servo/timelined_stats_manager_unittest.py b/servo/timelined_stats_manager_unittest.py
index d76f562..1fe15fc 100644
--- a/servo/timelined_stats_manager_unittest.py
+++ b/servo/timelined_stats_manager_unittest.py
@@ -75,5 +75,35 @@
     with self.assertRaises(stats_manager.StatsManagerError):
       self.data.AddSamples(samples)
 
+  def test_TrimSamples(self):
+    """Ensure that trimming works as expected."""
+    self.data.AddSamples([('A', 10)])
+    tstart = time.time()
+    time.sleep(0.01)
+    self.data.AddSamples([('A', 23)])
+    self.data.AddSamples([('A', 20)])
+    tend = time.time()
+    time.sleep(0.01)
+    self.data.AddSamples([('A', 10)])
+    self.data.TrimSamples(tstart=tstart, tend=tend)
+    self.data.CalculateStats()
+    # Verify that only the samples between the timestamps are left
+    self.assertEqual([23, 20], self.data._data['A'])
+    for samples in self.data._data.itervalues():
+      # Verify that all domains were trimmed to size 2
+      self.assertEqual(2, len(samples))
+
+  def test_TrimSamplesNoStartNoEnd(self):
+    """Ensure that the trimming encompasses the whole dataset."""
+    orig_samples = [10, 23, 20, 10]
+    for sample in orig_samples:
+      self.data.AddSamples([('A', sample)])
+    self.data.CalculateStats()
+    self.data.TrimSamples()
+    self.assertEqual(orig_samples, self.data._data['A'])
+    for samples in self.data._data.itervalues():
+      # Verify that all domains were not trimmed
+      self.assertEqual(len(orig_samples), len(samples))
+
 if __name__ == '__main__':
   unittest.main()