refactor: deduplicate ipython execution count

Author: shuoweil

packages/bigframes/bigframes/core/utils.py

@@ -249,3 +249,16 @@ def timedelta_to_micros(
         ) * 1_000_000 + timedelta.microseconds
 
     raise TypeError(f"Unrecognized input type: {type(timedelta)}")
+
+
+def get_ipython_execution_count() -> typing.Optional[int]:
+    """Returns the current IPython cell execution count if running in a notebook, else None."""
+    try:
+        import IPython
+
+        ipy = IPython.get_ipython()
+        if ipy is not None and hasattr(ipy, "execution_count"):
+            return ipy.execution_count
+    except (ImportError, NameError):
+        pass
+    return None

packages/bigframes/bigframes/display/anywidget.py

@@ -92,15 +92,9 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
 
         self._dataframe = dataframe
 
-        self._cell_execution_count = None
-        try:
-            import IPython
+        from bigframes.core.utils import get_ipython_execution_count
 
-            ipy = IPython.get_ipython()
-            if ipy is not None and hasattr(ipy, "execution_count"):
-                self._cell_execution_count = ipy.execution_count
-        except (ImportError, NameError):
-            pass
+        self._cell_execution_count = get_ipython_execution_count()
 
         super().__init__()
 

packages/bigframes/bigframes/session/_io/bigquery/__init__.py

@@ -317,14 +317,9 @@ def start_query_with_job(
     Starts query job and waits for results.
     """
     if cell_execution_count is None:
-        try:
-            import IPython
+        from bigframes.core.utils import get_ipython_execution_count
 
-            ipy = IPython.get_ipython()
-            if ipy is not None and hasattr(ipy, "execution_count"):
-                cell_execution_count = ipy.execution_count
-        except (ImportError, NameError):
-            pass
+        cell_execution_count = get_ipython_execution_count()
 
     # Note: Ensure no additional labels are added to job_config after this
     # point, as `add_and_trim_labels` ensures the label count does not
@@ -382,14 +377,9 @@ def start_query_job_optional(
     https://docs.cloud.google.com/bigquery/docs/running-queries#optional-job-creation
     """
     if cell_execution_count is None:
-        try:
-            import IPython
-
-            ipy = IPython.get_ipython()
-            if ipy is not None and hasattr(ipy, "execution_count"):
-                cell_execution_count = ipy.execution_count
-        except (ImportError, NameError):
-            pass
+        from bigframes.core.utils import get_ipython_execution_count
+
+        cell_execution_count = get_ipython_execution_count()
 
     add_and_trim_labels(job_config, session=session)
     try:

packages/bigframes/bigframes/session/execution_spec.py

@@ -81,14 +81,9 @@ def with_compute_options(self, compute_options: ComputeOptions) -> ExecutionSpec
 
         cell_execution_count = self.cell_execution_count
         if cell_execution_count is None:
-            try:
-                import IPython
-
-                ipy = IPython.get_ipython()
-                if ipy is not None and hasattr(ipy, "execution_count"):
-                    cell_execution_count = ipy.execution_count
-            except (ImportError, NameError):
-                pass
+            from bigframes.core.utils import get_ipython_execution_count
+
+            cell_execution_count = get_ipython_execution_count()
 
         return dataclasses.replace(
             self,

packages/bigframes/bigframes/session/metrics.py

@@ -74,14 +74,9 @@ def from_job(
             )
 
         if cell_execution_count is None:
-            try:
-                import IPython
+            from bigframes.core.utils import get_ipython_execution_count
 
-                ipy = IPython.get_ipython()
-                if ipy is not None and hasattr(ipy, "execution_count"):
-                    cell_execution_count = ipy.execution_count
-            except (ImportError, NameError):
-                pass
+            cell_execution_count = get_ipython_execution_count()
 
         metadata = cls(
             job_id=query_job.job_id,
@@ -147,14 +142,9 @@ def from_row_iterator(
             )
 
         if cell_execution_count is None:
-            try:
-                import IPython
+            from bigframes.core.utils import get_ipython_execution_count
 
-                ipy = IPython.get_ipython()
-                if ipy is not None and hasattr(ipy, "execution_count"):
-                    cell_execution_count = ipy.execution_count
-            except (ImportError, NameError):
-                pass
+            cell_execution_count = get_ipython_execution_count()
 
         # fmt: off
         return cls(
@@ -328,14 +318,9 @@ def on_event(self, envelope: Any):
                 self.bytes_processed += bytes_processed
 
                 if cell_execution_count is None:
-                    try:
-                        import IPython
-
-                        ipy = IPython.get_ipython()
-                        if ipy is not None and hasattr(ipy, "execution_count"):
-                            cell_execution_count = ipy.execution_count
-                    except (ImportError, NameError):
-                        pass
+                    from bigframes.core.utils import get_ipython_execution_count
+
+                    cell_execution_count = get_ipython_execution_count()
 
                 metadata = JobMetadata(
                     job_type="polars",

packages/bigframes/tests/unit/display/test_anywidget.py

@@ -185,18 +185,15 @@ def test_cell_execution_count_propagation(mock_df):
     """Test that the captured cell_execution_count is propagated to to_pandas_batches."""
     from bigframes.display.anywidget import TableWidget
 
-    # Mock IPython to return a specific execution count
     mock_ipy = mock.Mock()
     mock_ipy.execution_count = 42
 
     with mock.patch("IPython.get_ipython", return_value=mock_ipy):
         with bigframes.option_context("display.render_mode", "anywidget"):
             widget = TableWidget(mock_df)
 
-    # Verify captured execution count
     assert widget._cell_execution_count == 42
 
-    # to_pandas_batches should be called with cell_execution_count=42
     mock_df.to_pandas_batches.assert_called_with(
         page_size=widget.page_size,
         cell_execution_count=42,

packages/bigframes/tests/unit/session/test_read_gbq_colab.py

@@ -161,17 +161,14 @@ def test_execution_history_filtering():
 
     session = mocks.create_bigquery_session()
 
-    # Add mock jobs to session metrics
     job1 = metrics.JobMetadata(job_id="job_1", job_type="query", query="SELECT 1")
     job2 = metrics.JobMetadata(job_id="job_2", job_type="query", query="SELECT 2")
     session._metrics.jobs.extend([job1, job2])
 
-    # Verify filtering by job_ids isolates the target execution
     history_job1 = session.execution_history(job_ids=["job_1"]).to_dataframe()
     assert len(history_job1) == 1
     assert history_job1.iloc[0]["job_id"] == "job_1"
 
-    # Verify filtering by events isolates the target execution
     event2 = mock.Mock()
     event2.job_id = "job_2"
     history_job2 = session.execution_history(events=[event2]).to_dataframe()