Skip to content

Commit 53953bb

Browse files
committed
ipython_display_ to set fallback mimetypes
1 parent 1b15ef0 commit 53953bb

File tree

2 files changed

+33
-25
lines changed

2 files changed

+33
-25
lines changed

bigframes/dataframe.py

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -789,9 +789,7 @@ def __repr__(self) -> str:
789789

790790
opts = bigframes.options.display
791791
max_results = opts.max_rows
792-
# anywdiget mode uses the same display logic as the "deferred" mode
793-
# for faster execution
794-
if opts.repr_mode in ("deferred", "anywidget"):
792+
if opts.repr_mode == "deferred":
795793
return formatter.repr_query_job(self._compute_dry_run())
796794

797795
# TODO(swast): pass max_columns and get the true column count back. Maybe
@@ -837,9 +835,7 @@ def _repr_html_(self) -> str:
837835
"""
838836
opts = bigframes.options.display
839837
max_results = opts.max_rows
840-
# For anywidget mode, return deferred representation
841-
# The actual widget display is handled by _ipython_display_()
842-
if opts.repr_mode in ("deferred", "anywidget"):
838+
if opts.repr_mode == "deferred":
843839
return formatter.repr_query_job(self._compute_dry_run())
844840

845841
# Process blob columns first for non-deferred modes
@@ -916,7 +912,7 @@ def obj_ref_rt_to_html(obj_ref_rt) -> str:
916912
html_string += f"[{row_count} rows x {column_count} columns in total]"
917913
return html_string
918914

919-
def _ipython_display_(self):
915+
def _repr_mimebundle_(self, include=None, exclude=None):
920916
"""
921917
Custom display method for IPython/Jupyter environments.
922918
This is called by IPython's display system when the object is displayed.
@@ -944,26 +940,26 @@ def _ipython_display_(self):
944940

945941
# Create and display the widget
946942
widget = display.TableWidget(df)
943+
widget_repr = widget._repr_mimebundle_(include=include, exclude=exclude)
947944

948-
# IPython will automatically display the widget
949-
# since we're returning it from _ipython_display_()
950-
from IPython.display import display as ipython_display
951-
952-
ipython_display(widget)
953-
return # Important: return None to signal we handled display
945+
# Use deferred repr for text/plain of anywidget display.
946+
# This avoids kicking off a query when the user is just
947+
# printing the last expression in a cell.
948+
widget_repr["text/plain"] = repr(self)
949+
widget_repr["text/html"] = self._repr_html_()
950+
return widget_repr
954951

955952
except (AttributeError, ValueError, ImportError):
956953
# Fallback: let IPython use _repr_html_() instead
957954
warnings.warn(
958955
"Anywidget mode is not available. "
959956
"Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use interactive tables. "
960-
f"Falling back to deferred mode. Error: {traceback.format_exc()}"
957+
f"Falling back to static HTML. Error: {traceback.format_exc()}"
961958
)
962959
# Don't return anything - let IPython fall back to _repr_html_()
963-
return
960+
pass
964961

965-
# For other modes, don't handle display - let IPython use _repr_html_()
966-
return
962+
return {"text/html": self._repr_html_(), "text/plain": repr(self)}
967963

968964
def __delitem__(self, key: str):
969965
df = self.drop(columns=[key])

tests/system/small/test_anywidget.py

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -686,21 +686,33 @@ def test_widget_with_unknown_row_count_empty_dataframe(
686686
assert widget.page == 0
687687

688688

689-
def test_repr_html_anywidget_fallback(paginated_bf_df: bf.dataframe.DataFrame):
689+
def test_repr_mimebundle_anywidget_fallback(paginated_bf_df: bf.dataframe.DataFrame):
690690
"""
691-
Test that _repr_html_ falls back to deferred mode when anywidget is not available.
691+
Test that _repr_mimebundle_ falls back to static html when anywidget is not available.
692692
"""
693693
with bf.option_context("display.repr_mode", "anywidget"):
694694
# Use a mock to simulate the absence of the 'anywidget' module.
695695
with mock.patch.dict(
696696
"sys.modules", {"anywidget": None, "IPython": mock.MagicMock()}
697697
):
698-
# The warning is now expected inside the _ipython_display_ call, not _repr_html_
699-
# The test setup doesn't easily allow capturing warnings from ipython display hooks.
700-
# Instead we focus on the fallback behavior of _repr_html_
701-
html = paginated_bf_df._repr_html_()
702-
assert "Computation deferred." in html
703-
assert "Computation will process" in html
698+
bundle = paginated_bf_df._repr_mimebundle_()
699+
assert "application/vnd.jupyter.widget-view+json" not in bundle
700+
assert "text/html" in bundle
701+
html = bundle["text/html"]
702+
assert "page_1_row_1" in html
703+
assert "page_1_row_2" in html
704+
assert "page_2_row_1" not in html
705+
706+
707+
def test_repr_mimebundle_anywidget_success(paginated_bf_df: bf.dataframe.DataFrame):
708+
"""
709+
Test that _repr_mimebundle_ returns a widget view when anywidget is available.
710+
"""
711+
with bf.option_context("display.repr_mode", "anywidget"):
712+
bundle = paginated_bf_df._repr_mimebundle_()
713+
assert "application/vnd.jupyter.widget-view+json" in bundle
714+
assert "text/html" in bundle
715+
assert "text/plain" in bundle
704716

705717

706718
# TODO(b/332316283): Add tests for custom index and multiindex

0 commit comments

Comments
 (0)