Improve Perf interface

babi/main.py

@@ -4,6 +4,7 @@ from typing import Optional
 from typing import Sequence
 
 from babi.file import File
+from babi.perf import perf_log
 from babi.screen import EditResult
 from babi.screen import make_stdscr
 from babi.screen import Screen
@@ -32,22 +33,23 @@ def _edit(screen: Screen) -> EditResult:
 
 
 def c_main(stdscr: 'curses._CursesWindow', args: argparse.Namespace) -> None:
-    screen = Screen(stdscr, [File(f) for f in args.filenames or [None]])
-    with screen.perf.log(args.perf_log), screen.history.save():
-        while screen.files:
-            screen.i = screen.i % len(screen.files)
-            res = _edit(screen)
-            if res == EditResult.EXIT:
-                del screen.files[screen.i]
-                screen.status.clear()
-            elif res == EditResult.NEXT:
-                screen.i += 1
-                screen.status.clear()
-            elif res == EditResult.PREV:
-                screen.i -= 1
-                screen.status.clear()
-            else:
-                raise AssertionError(f'unreachable {res}')
+    with perf_log(args.perf_log) as perf:
+        screen = Screen(stdscr, args.filenames or [None], perf)
+        with screen.history.save():
+            while screen.files:
+                screen.i = screen.i % len(screen.files)
+                res = _edit(screen)
+                if res == EditResult.EXIT:
+                    del screen.files[screen.i]
+                    screen.status.clear()
+                elif res == EditResult.NEXT:
+                    screen.i += 1
+                    screen.status.clear()
+                elif res == EditResult.PREV:
+                    screen.i -= 1
+                    screen.status.clear()
+                else:
+                    raise AssertionError(f'unreachable {res}')
 
 
 def main(argv: Optional[Sequence[str]] = None) -> int:

babi/perf.py

@@ -29,19 +29,28 @@ class Perf:
             self._records.append((self._name, time.monotonic() - self._time))
             self._name = self._time = None
 
-    @contextlib.contextmanager
-    def log(self, filename: Optional[str]) -> Generator[None, None, None]:
-        if filename is None:
-            yield
-        else:
-            self._prof = cProfile.Profile()
-            self.start('startup')
-            try:
-                yield
-            finally:
-                self.end()
-                self._prof.dump_stats(f'{filename}.pstats')
-                with open(filename, 'w') as f:
-                    f.write('μs\tevent\n')
-                    for name, duration in self._records:
-                        f.write(f'{int(duration * 1000 * 1000)}\t{name}\n')
+    def init_profiling(self) -> None:
+        self._prof = cProfile.Profile()
+        self.start('startup')
+
+    def save_profiles(self, filename: str) -> None:
+        assert self._prof is not None
+        self._prof.dump_stats(f'{filename}.pstats')
+        with open(filename, 'w') as f:
+            f.write('μs\tevent\n')
+            for name, duration in self._records:
+                f.write(f'{int(duration * 1000 * 1000)}\t{name}\n')
+
+
+@contextlib.contextmanager
+def perf_log(filename: Optional[str]) -> Generator[Perf, None, None]:
+    perf = Perf()
+    if filename is None:
+        yield perf
+    else:
+        perf.init_profiling()
+        try:
+            yield perf
+        finally:
+            perf.end()
+            perf.save_profiles(filename)

babi/screen.py

@@ -66,13 +66,14 @@ class Screen:
     def __init__(
             self,
             stdscr: 'curses._CursesWindow',
-            files: List[File],
+            filenames: List[Optional[str]],
+            perf: Perf,
     ) -> None:
         self.stdscr = stdscr
-        self.files = files
+        self.files = [File(f) for f in filenames]
         self.i = 0
         self.history = History()
-        self.perf = Perf()
+        self.perf = perf
         self.status = Status()
         self.margin = Margin.from_current_screen()
         self.cut_buffer: Tuple[str, ...] = ()