Merge pull request #417 from Renumics/feature/read-df-after-close

Feature/read df after close

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:

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