Allow the renaming comm to rename things other that dataframes (#140)

* Allow the renaming comm to rename things other that dataframes -- SQL cell variable results could be scalars nowadays.
MSeal · Dec 21, 2022 · c411c0b · c411c0b
1 parent c725f77
commit c411c0b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 20 deletions.
diff --git a/src/dx/comms/rename.py b/src/dx/comms/rename.py
@@ -1,6 +1,5 @@
 from typing import Optional
 
-import pandas as pd
 import structlog
 from IPython import get_ipython
 from IPython.core.interactiveshell import InteractiveShell
@@ -16,8 +15,11 @@ def _recv(msg):
         handle_renaming_comm(msg)
 
 
+UNSET = object()
+
+
 def handle_renaming_comm(msg: dict, ipython_shell: Optional[InteractiveShell] = None):
-    """Implementation behind renaming a SQL cell dataframe via comms"""
+    """Implementation behind renaming a SQL cell output variable via comms"""
     content = msg.get("content")
     if not content:
         return
@@ -32,16 +34,8 @@ def handle_renaming_comm(msg: dict, ipython_shell: Optional[InteractiveShell] =
             # shell will be passed in from test suite, otherwise go with global shell.
             ipython_shell = get_ipython()
 
-        value_to_rename = ipython_shell.user_ns.get(data["old_name"])
-
-        # Do not rename unless old_name mapped onto exactly a dataframe.
-        #
-        # (Handles case when it maps onto None, indicating that the old name
-        #  hasn't been assigned to at all yet (i.e. user gestured to rename
-        #  SQL cell dataframe name before the SQL cell has even been run the
-        #  first time yet))
-        #
-        if isinstance(value_to_rename, pd.DataFrame):
+        value_to_rename = ipython_shell.user_ns.get(data["old_name"], UNSET)
+        if value_to_rename is not UNSET:
             # New name can be empty string, indicating to drop reference to the var.
             if data["new_name"]:
                 ipython_shell.user_ns[data["new_name"]] = value_to_rename

diff --git a/tests/test_comms.py b/tests/test_comms.py
@@ -96,26 +96,28 @@ def test_rename_skipped_missing_values(
         handle_renaming_comm(msg, ipython_shell=get_ipython)
         assert "new_df" not in get_ipython.user_ns.keys()
 
-    def test_rename_skipped_no_dataframe(
+    def test_rename_non_dataframes(
         self,
         get_ipython: TerminalInteractiveShell,
     ):
         """
-        Test that a rename comm is skipped if the variable matching
-        the `old_name` is not a pandas DataFrame.
+        Test that a rename comm is also works if non-dataframe is referenced.
         """
         msg = {
             "content": {
                 "data": {
-                    "old_name": "old_df",
-                    "new_name": "new_df",
+                    "old_name": "old_var",
+                    "new_name": "new_var",
                 }
             }
         }
-        get_ipython.user_ns["old_df"] = "I am not a dataframe; I am a string."
+        the_string = "I am not a dataframe; I am a string."
+        get_ipython.user_ns["old_var"] = the_string
         handle_renaming_comm(msg, ipython_shell=get_ipython)
-        assert "new_df" not in get_ipython.user_ns.keys()
-        assert "old_df" in get_ipython.user_ns.keys()
+        assert (
+            "new_var" in get_ipython.user_ns.keys() and get_ipython.user_ns["new_var"] == the_string
+        )
+        assert "old_var" not in get_ipython.user_ns.keys()
 
 
 class TestAssignmentComm: