diff --git a/src/dx/comms/rename.py b/src/dx/comms/rename.py index 011c43a3..3409913f 100644 --- 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 index 4d25d0ac..8d4b01c8 100644 --- 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: