Simplify sampler's logic

As storers always expect a binary file to read from, we can simplify
samplers by delegating to the profiler the creation of the in-memory
file and the interaction with the storer.

Author: Giuseppe Lumia

benchmark/record_frames.py

@@ -6,7 +6,6 @@
 """
 
 import pickle
-from io import BytesIO
 
 from handler import handler
 
@@ -50,11 +49,8 @@ def sample(self, frame):
         # Record custom stack
         self.data.append(frames[-1])
 
-    def dump(self, storer):
-        file = BytesIO()
+    def dump(self, file):
         pickle.dump(self.data, file)
-        file.seek(0)
-        storer.store(file)
 
 
 if __name__ == "__main__":

pylaprof/__init__.py

@@ -107,12 +107,12 @@ def sample(self, frame):
         """
         pass  # pragma: no cover
 
-    def dump(self, storer):
+    def dump(self, file):
         """
         Dump sampling data.
 
-        storer (Storer)
-          Storer to use for data's storage.
+        file
+          A file-like object in binary mode where to write data.
         """
         pass  # pragma: no cover
 
@@ -136,13 +136,10 @@ def sample(self, frame):
             frame = frame.f_back
         self._data[tuple(stack)] += 1
 
-    def dump(self, storer):
-        file = BytesIO()
+    def dump(self, file):
         for stack, hits in self._data.items():
             line = f"{';'.join(stack[::-1])} {hits}\n"
             file.write(line.encode())
-        file.seek(0)
-        storer.store(file)
 
 
 class Profiler(threading.Thread):
@@ -241,7 +238,10 @@ def run(self):
             end = time.time()
 
             if end - start >= self.min_time:
-                self.sampler.dump(self.storer)
+                file = BytesIO()
+                self.sampler.dump(file)
+                file.seek(0)
+                self.storer.store(file)
 
             stop_event.clear()
             self.clean_exit = True

tests/test_profiler.py

@@ -1,6 +1,7 @@
 import sys
 import threading
 from inspect import signature
+from io import BytesIO
 from unittest.mock import MagicMock, Mock
 
 from pylaprof import Profiler, profile
@@ -116,6 +117,8 @@ def test_profiler_run(monkeypatch):
     profiler = Profiler(period=period, sampler=Mock(), storer=Mock())
     profiler._test = Mock(side_effect=lambda ident: ident == ident1)
     profiler._stop_event = Mock()
+    report = b"I'm supposed to be a profiling report"
+    profiler.sampler.dump.side_effect = lambda file: file.write(report)
 
     profiler.start()
     profiler.stop()
@@ -126,6 +129,13 @@ def test_profiler_run(monkeypatch):
     sampler_calls = {args for args, _ in profiler.sampler.sample.call_args_list}
     assert (frame1,) in sampler_calls
     assert (frame2,) not in sampler_calls
+    profiler.sampler.dump.assert_called()
+    profiler.storer.store.assert_called_once()
+    store_args = profiler.storer.store.call_args_list[0][0]
+    file = store_args[0]
+    assert isinstance(file, BytesIO)
+    file.seek(0)
+    assert file.read() == report
     profiler._stop_event.wait.assert_called_with(period)
     profiler._stop_event.clear.assert_called()
     assert profiler.clean_exit is True

tests/test_sampler.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+from io import BytesIO
 from unittest.mock import Mock
 
 from pylaprof import StackCollapse
@@ -85,10 +86,9 @@ def test_stack_collapse_dump():
         "three (some_path/some_module.py:42);one (some_path/some_module.py:12) 4\n"
         + "three_v2 (some_path/some_module.py:82);some_func (some_path/some_other_module.py:11) 2\n"  # noqa
     ).encode()
-    storer = Mock()
+    file = BytesIO()
 
-    stack_collapse.dump(storer=storer)
+    stack_collapse.dump(file=file)
 
-    storer.store.assert_called()
-    args, _ = storer.store.call_args_list[0]
-    assert args[0].read() == exp_file
+    file.seek(0)
+    assert file.read() == exp_file