diff --git a/renumics/spotlight/viewer.py b/renumics/spotlight/viewer.py index 8fb9522e..af8b339a 100644 --- a/renumics/spotlight/viewer.py +++ b/renumics/spotlight/viewer.py @@ -130,6 +130,7 @@ class Viewer: _host: str _requested_port: Union[int, Literal["auto"]] _server: Optional[Server] + _df: Optional[pd.DataFrame] def __init__( self, @@ -139,6 +140,7 @@ def __init__( self._host = host self._requested_port = port self._server = None + self._df = None def show( self, @@ -177,6 +179,7 @@ def show( embedders (disabled by default). """ + self._df = None if is_pathtype(dataset): dataset = Path(dataset).absolute() if dataset.is_dir(): @@ -259,6 +262,7 @@ def close(self, wait: Union[bool, Literal["forever"]] = False) -> None: else: self._server.wait_for_frontend_disconnect() + self._df = self._server.get_df() _VIEWERS.remove(self) self._server.stop() self._server = None @@ -292,8 +296,7 @@ def df(self) -> Optional[pd.DataFrame]: """ if self._server: return self._server.get_df() - - return None + return self._df @property def host(self) -> str: diff --git a/tests/integration/pandas/test_pandas.py b/tests/integration/pandas/test_pandas.py new file mode 100644 index 00000000..a5ff5c6b --- /dev/null +++ b/tests/integration/pandas/test_pandas.py @@ -0,0 +1,21 @@ +import pandas as pd + +from renumics import spotlight + + +def test_read_df() -> None: + df = pd.DataFrame({"foo": range(4)}) + viewer = spotlight.show(df, wait=False, no_browser=True) + df_after_close = viewer.df + viewer.close() + assert df_after_close is not None + assert "foo" in df_after_close + + +def test_read_df_after_close() -> None: + df = pd.DataFrame({"foo": range(4)}) + viewer = spotlight.show(df, wait=False, no_browser=True) + viewer.close() + df_after_close = viewer.df + assert df_after_close is not None + assert "foo" in df_after_close