fix: resolve plotly rendering issue by using ipython html for job pro…

bigframes/formatting_helpers.py

@@ -16,6 +16,7 @@
 # TODO(orrbradford): cleanup up typings and documenttion in this file
 
 import datetime
+import random
 from typing import Any, Optional, Union
 
 import google.api_core.exceptions as api_core_exceptions
@@ -57,9 +58,9 @@ def repr_query_job_html(query_job: Optional[bigquery.QueryJob]):
         Pywidget html table.
     """
     if query_job is None:
-        return widgets.HTML("No job information available")
+        return display.HTML("No job information available")
     if query_job.dry_run:
-        return widgets.HTML(
+        return display.HTML(
             f"Computation deferred. Computation will process {get_formatted_bytes(query_job.total_bytes_processed)}"
         )
     table_html = "<style> td {text-align: left;}</style>"
@@ -125,16 +126,20 @@ def wait_for_query_job(
     Returns:
         A row iterator over the query results.
     """
-    loading_bar = widgets.HTML(get_query_job_loading_html(query_job))
     if progress_bar == "auto":
         progress_bar = "notebook" if in_ipython() else "terminal"
 
     try:
         if progress_bar == "notebook":
-            display.display(loading_bar)
+            display_id = str(random.random())
+            loading_bar = display.HTML(get_query_job_loading_html(query_job))
+            display.display(loading_bar, display_id=display_id)
             query_result = query_job.result(max_results=max_results)
             query_job.reload()
-            loading_bar.value = get_query_job_loading_html(query_job)
+            display.update_display(
+                display.HTML(get_query_job_loading_html(query_job)),
+                display_id=display_id,
+            )
         elif progress_bar == "terminal":
             initial_loading_bar = get_query_job_loading_string(query_job)
             print(initial_loading_bar)
@@ -171,16 +176,19 @@ def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None):
         progress_bar (str, Optional):
             Which progress bar to show.
     """
-    loading_bar = widgets.HTML(get_base_job_loading_html(job))
     if progress_bar == "auto":
         progress_bar = "notebook" if in_ipython() else "terminal"
 
     try:
         if progress_bar == "notebook":
-            display.display(loading_bar)
+            display_id = str(random.random())
+            loading_bar = display.HTML(get_base_job_loading_html(job))
+            display.display(loading_bar, display_id=display_id)
             job.result()
             job.reload()
-            loading_bar.value = get_base_job_loading_html(job)
+            display.update_display(
+                display.HTML(get_base_job_loading_html(job)), display_id=display_id
+            )
         elif progress_bar == "terminal":
             inital_loading_bar = get_base_job_loading_string(job)
             print(inital_loading_bar)

tests/system/small/test_progress_bar.py

@@ -12,101 +12,92 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import re
 import tempfile
 
 import pandas as pd
 
 import bigframes as bf
 import bigframes.formatting_helpers as formatting_helpers
 
+job_load_message_regex = r"\w+ job [\w-]+ is \w+\."
+
 
 def test_progress_bar_dataframe(
     penguins_df_default_index: bf.dataframe.DataFrame, capsys
 ):
-    bf.options.display.progress_bar = "notebook"
+    bf.options.display.progress_bar = "terminal"
+    capsys.readouterr()  # clear output
     penguins_df_default_index.to_pandas()
-    html_check = "HTML(value="
-    open_job_check = "Open Job"
-    lines = capsys.readouterr().out.split("\n")
-    lines = [line for line in lines if len(line) > 0]
-    assert len(lines) > 0
+
+    assert_loading_msg_exist(capsys.readouterr().out)
     assert penguins_df_default_index.query_job is not None
-    for line in lines:
-        assert html_check in line and open_job_check in line
 
 
 def test_progress_bar_series(penguins_df_default_index: bf.dataframe.DataFrame, capsys):
-    bf.options.display.progress_bar = "notebook"
+    bf.options.display.progress_bar = "terminal"
     series = penguins_df_default_index["body_mass_g"].head(10)
+    capsys.readouterr()  # clear output
     series.to_pandas()
-    html_check = "HTML(value="
-    open_job_check = "Open Job"
-    lines = capsys.readouterr().out.split("\n")
-    lines = [line for line in lines if len(line) > 0]
-    assert len(lines) > 0
+
+    assert_loading_msg_exist(capsys.readouterr().out)
     assert series.query_job is not None
-    for line in lines:
-        assert html_check in line and open_job_check in line
 
 
 def test_progress_bar_scalar(penguins_df_default_index: bf.dataframe.DataFrame, capsys):
-    bf.options.display.progress_bar = "notebook"
+    bf.options.display.progress_bar = "terminal"
+    capsys.readouterr()  # clear output
     penguins_df_default_index["body_mass_g"].head(10).mean()
-    html_check = "HTML(value="
-    open_job_check = "Open Job"
-    lines = capsys.readouterr().out.split("\n")
-    lines = [line for line in lines if len(line) > 0]
-    assert len(lines) > 0
-    for line in lines:
-        assert html_check in line and open_job_check in line
+
+    assert_loading_msg_exist(capsys.readouterr().out)
 
 
 def test_progress_bar_read_gbq(session: bf.Session, penguins_table_id: str, capsys):
-    bf.options.display.progress_bar = "notebook"
+    bf.options.display.progress_bar = "terminal"
+    capsys.readouterr()  # clear output
     session.read_gbq(penguins_table_id)
-    html_check = "HTML(value="
-    open_job_check = "Open Job"
-    lines = capsys.readouterr().out.split("\n")
-    lines = [line for line in lines if len(line) > 0]
-    assert len(lines) > 0
-    for line in lines:
-        assert html_check in line and open_job_check in line
+
+    assert_loading_msg_exist(capsys.readouterr().out)
 
 
 def test_progress_bar_extract_jobs(
     penguins_df_default_index: bf.dataframe.DataFrame, gcs_folder, capsys
 ):
-    bf.options.display.progress_bar = "notebook"
+    bf.options.display.progress_bar = "terminal"
     path = gcs_folder + "test_read_csv_progress_bar*.csv"
+    capsys.readouterr()  # clear output
     penguins_df_default_index.to_csv(path)
-    html_check = "HTML(value="
-    open_job_check = "Open Job"
-    lines = capsys.readouterr().out.split("\n")
-    lines = [line for line in lines if len(line) > 0]
-    assert len(lines) > 0
-    for line in lines:
-        assert html_check in line and open_job_check in line
+
+    assert_loading_msg_exist(capsys.readouterr().out)
 
 
 def test_progress_bar_load_jobs(
     session: bf.Session, penguins_pandas_df_default_index: pd.DataFrame, capsys
 ):
-    bf.options.display.progress_bar = "notebook"
+    bf.options.display.progress_bar = "terminal"
     with tempfile.TemporaryDirectory() as dir:
         path = dir + "/test_read_csv_progress_bar*.csv"
         penguins_pandas_df_default_index.to_csv(path, index=False)
+        capsys.readouterr()  # clear output
         session.read_csv(path)
-    html_check = "HTML(value="
-    open_job_check = "Open Job"
-    lines = capsys.readouterr().out.split("\n")
+
+    assert_loading_msg_exist(capsys.readouterr().out)
+
+
+def assert_loading_msg_exist(capystOut: str, pattern=job_load_message_regex):
+    numLoadingMsg = 0
+    lines = capystOut.split("\n")
     lines = [line for line in lines if len(line) > 0]
+
     assert len(lines) > 0
     for line in lines:
-        assert html_check in line and open_job_check in line
+        if re.match(pattern, line) is not None:
+            numLoadingMsg += 1
+    assert numLoadingMsg > 0
 
 
 def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame):
-    bf.options.display.progress_bar = "notebook"
+    bf.options.display.progress_bar = "terminal"
     penguins_df_default_index._block._expr._session.bqclient.default_query_job_config.use_query_cache = (
         False
     )