refactor: split repr_mimebundle logic, handle deferred mode in html, and fix mypy errors

Author: shuoweil

bigframes/display/html.py

@@ -29,6 +29,7 @@
 import bigframes
 from bigframes._config import display_options, options
 from bigframes.display import plaintext
+import bigframes.formatting_helpers as formatter
 
 if typing.TYPE_CHECKING:
     import bigframes.dataframe
@@ -114,16 +115,16 @@ def create_html_representation(
     from bigframes.series import Series
 
     if isinstance(obj, Series):
+        # Fallback to pandas string representation if the object is not a Series.
+        # This protects against cases where obj might be something else unexpectedly,
+        # or if the pandas Series implementation changes.
         pd_series = pandas_df.iloc[:, 0]
         try:
             html_string = pd_series._repr_html_()
         except AttributeError:
             html_string = f"<pre>{pd_series.to_string()}</pre>"
 
-        # Series doesn't typically show total rows/cols like DF in HTML repr here?
-        # But let's check what it was doing.
-        # Original code just returned _repr_html_ or wrapped to_string.
-        # It didn't append row/col count string for Series (wait, Series usually has length in repr).
+        html_string += f"[{total_rows} rows]"
         return html_string
     else:
         # It's a DataFrame
@@ -225,32 +226,21 @@ def get_anywidget_bundle(
     return widget_repr, widget_metadata
 
 
-def repr_mimebundle(
+def _repr_mimebundle_deferred(
     obj: Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
-    include=None,
-    exclude=None,
-):
-    """
-    Custom display method for IPython/Jupyter environments.
-    """
-    # TODO(b/467647693): Anywidget integration has been tested in Jupyter, VS Code, and
-    # BQ Studio, but there is a known compatibility issue with Marimo that needs to be addressed.
+) -> dict[str, str]:
+    return {
+        "text/plain": formatter.repr_query_job(obj._compute_dry_run()),
+        "text/html": formatter.repr_query_job_html(obj._compute_dry_run()),
+    }
+
+
+def _repr_mimebundle_head(
+    obj: Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
+) -> dict[str, str]:
     from bigframes.series import Series
 
     opts = options.display
-    if opts.repr_mode == "anywidget":
-        try:
-            return get_anywidget_bundle(obj, include=include, exclude=exclude)
-        except ImportError:
-            # Anywidget is an optional dependency, so warn rather than fail.
-            # TODO(shuowei): When Anywidget becomes the default for all repr modes,
-            # remove this warning.
-            warnings.warn(
-                "Anywidget mode is not available. "
-                "Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use interactive tables. "
-                f"Falling back to static HTML. Error: {traceback.format_exc()}"
-            )
-
     blob_cols: list[str]
     if isinstance(obj, Series):
         pandas_df, row_count, query_job = obj._block.retrieve_repr_request_results(
@@ -275,3 +265,34 @@ def repr_mimebundle(
     )
 
     return {"text/html": html_string, "text/plain": text_representation}
+
+
+def repr_mimebundle(
+    obj: Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
+    include=None,
+    exclude=None,
+):
+    """
+    Custom display method for IPython/Jupyter environments.
+    """
+    # TODO(b/467647693): Anywidget integration has been tested in Jupyter, VS Code, and
+    # BQ Studio, but there is a known compatibility issue with Marimo that needs to be addressed.
+
+    opts = options.display
+    if opts.repr_mode == "deferred":
+        return _repr_mimebundle_deferred(obj)
+
+    if opts.repr_mode == "anywidget":
+        try:
+            return get_anywidget_bundle(obj, include=include, exclude=exclude)
+        except ImportError:
+            # Anywidget is an optional dependency, so warn rather than fail.
+            # TODO(shuowei): When Anywidget becomes the default for all repr modes,
+            # remove this warning.
+            warnings.warn(
+                "Anywidget mode is not available. "
+                "Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use interactive tables. "
+                f"Falling back to static HTML. Error: {traceback.format_exc()}"
+            )
+
+    return _repr_mimebundle_head(obj)

bigframes/formatting_helpers.py

@@ -68,7 +68,7 @@ def repr_query_job(query_job: Optional[bigquery.QueryJob]):
         query_job:
             The job representing the execution of the query on the server.
     Returns:
-        Pywidget html table.
+        Formatted string.
     """
     if query_job is None:
         return "No job information available"
@@ -94,6 +94,46 @@ def repr_query_job(query_job: Optional[bigquery.QueryJob]):
     return res
 
 
+def repr_query_job_html(query_job: Optional[bigquery.QueryJob]):
+    """Return query job as a formatted html string.
+    Args:
+        query_job:
+            The job representing the execution of the query on the server.
+    Returns:
+        Html string.
+    """
+    if query_job is None:
+        return "No job information available"
+    if query_job.dry_run:
+        return f"Computation deferred. Computation will process {get_formatted_bytes(query_job.total_bytes_processed)}"
+
+    # We can reuse the plaintext repr for now or make a nicer table.
+    # For deferred mode consistency, let's just wrap the text in a pre block or similar,
+    # but the request implies we want a distinct HTML representation if possible.
+    # However, existing repr_query_job returns a simple string.
+    # Let's format it as a simple table or list.
+
+    res = "<h3>Query Job Info</h3><ul>"
+    for key, value in query_job_prop_pairs.items():
+        job_val = getattr(query_job, value)
+        if job_val is not None:
+            if key == "Job Id":  # add link to job
+                url = get_job_url(
+                    project_id=query_job.project,
+                    location=query_job.location,
+                    job_id=query_job.job_id,
+                )
+                res += f'<li>Job: <a target="_blank" href="{url}">{query_job.job_id}</a></li>'
+            elif key == "Slot Time":
+                res += f"<li>{key}: {get_formatted_time(job_val)}</li>"
+            elif key == "Bytes Processed":
+                res += f"<li>{key}: {get_formatted_bytes(job_val)}</li>"
+            else:
+                res += f"<li>{key}: {job_val}</li>"
+    res += "</ul>"
+    return res
+
+
 current_display: Optional[display.HTML] = None
 current_display_id: Optional[str] = None
 previous_display_html: str = ""
@@ -296,7 +336,7 @@ def get_job_url(
     """
     if project_id is None or location is None or job_id is None:
         return None
-    return f"""https://console.cloud.google.com/bigquery?project={project_id}&j=bq:{location}:{job_id}&page=queryresults"""
+    return f"""https://console.cloud. google.com/bigquery?project={project_id}&j=bq:{location}:{job_id}&page=queryresults"""
 
 
 def render_bqquery_sent_event_html(
@@ -508,7 +548,7 @@ def get_base_job_loading_html(job: GenericJob):
     Returns:
         Html string.
     """
-    return f"""{job.job_type.capitalize()} job {job.job_id} is {job.state}. <a target="_blank" href="{get_job_url(
+    return f"""{job.job_type.capitalize()} job {job.job_id} is {job.state}. <a target=\"_blank\" href="{get_job_url(
         project_id=job.job_id,
         location=job.location,
         job_id=job.job_id,