add from_counter convenience classmethod

ipyflow · Mar 20, 2022 · f62a24b · f62a24b
1 parent ce29abf
commit f62a24b
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/nbsafety/data_model/code_cell.py b/nbsafety/data_model/code_cell.py
@@ -245,12 +245,16 @@ def all_cells_most_recently_run_for_each_id(
     ) -> Generator["ExecutedCodeCell", None, None]:
         yield from cls._current_cell_by_cell_id.values()
 
+    @classmethod
+    def from_counter(cls, ctr: int) -> "ExecutedCodeCell":
+        return cls._cell_by_cell_ctr[ctr]
+
     @classmethod
     def from_timestamp(cls, ts: TimestampOrCounter) -> "ExecutedCodeCell":
         if isinstance(ts, Timestamp):
-            return cls._cell_by_cell_ctr[ts.cell_num]
+            return cls.from_counter(ts.cell_num)
         else:
-            return cls._cell_by_cell_ctr[ts]
+            return cls.from_counter(ts)
 
     @classmethod
     def from_id(cls, cell_id: CellId) -> Optional["ExecutedCodeCell"]: