diff --git a/src/dx/__init__.py b/src/dx/__init__.py index cff41526..03d45c61 100644 --- a/src/dx/__init__.py +++ b/src/dx/__init__.py @@ -1,4 +1,4 @@ -from dx.comms import handle_assignment_comm, handle_resample_comm +from dx.comms import handle_resample_comm from dx.datatypes import * from dx.dx import display, show_docs from dx.formatters import * diff --git a/src/dx/comms/__init__.py b/src/dx/comms/__init__.py index b9640f95..bd87ad2a 100644 --- a/src/dx/comms/__init__.py +++ b/src/dx/comms/__init__.py @@ -1,2 +1 @@ -from dx.comms.assignment import handle_assignment_comm from dx.comms.resample import handle_resample_comm diff --git a/src/dx/comms/assignment.py b/src/dx/comms/assignment.py deleted file mode 100644 index 46e63565..00000000 --- a/src/dx/comms/assignment.py +++ /dev/null @@ -1,57 +0,0 @@ -from typing import Optional - -import structlog -from IPython import get_ipython -from IPython.core.interactiveshell import InteractiveShell - -from dx.filtering import resample_from_db -from dx.types.filters import DEXFilterSettings -from dx.utils.formatting import incrementing_label - -logger = structlog.get_logger(__name__) - - -# ref: https://jupyter-notebook.readthedocs.io/en/stable/comms.html#opening-a-comm-from-the-frontend -def dataframe_assignment(comm, open_msg): - """ - Datalink resample request. - """ - - @comm.on_msg - def _recv(msg): - # Is separate function to make testing easier. - handle_assignment_comm(msg) - - comm.send({"status": "connected", "source": "dataframe_assignment"}) - - -def handle_assignment_comm(msg: dict, ipython_shell: Optional[InteractiveShell] = None): - data = msg.get("content", {}).get("data", {}) - if not data: - return - - if "display_id" in data and "variable_name" in data: - filters = data["filters"] - sample_size = data["sample_size"] - - sql_filter = f"SELECT * FROM {{table_name}} LIMIT {sample_size}" - if filters: - dex_filters = DEXFilterSettings(filters=filters) - sql_filter_str = dex_filters.to_sql_query() - sql_filter = f"SELECT * FROM {{table_name}} WHERE {sql_filter_str} LIMIT {sample_size}" - - sampled_df = resample_from_db( - display_id=data["display_id"], - sql_filter=sql_filter, - filters=filters, - assign_subset=False, - ) - - ipython = ipython_shell or get_ipython() - variable_name = data["variable_name"] - - # if the variable already exists in the user namespace, add a suffix so the previous value isn't overwritten - if variable_name in ipython.user_ns: - variable_name = incrementing_label(variable_name, ipython.user_ns) - logger.debug(f"assigning {len(sampled_df)}-row dataframe to `{variable_name}` in {ipython}") - ipython.user_ns[variable_name] = sampled_df diff --git a/src/dx/settings.py b/src/dx/settings.py index a2886bf8..f514d09b 100644 --- a/src/dx/settings.py +++ b/src/dx/settings.py @@ -5,7 +5,6 @@ import pandas as pd import structlog -from IPython import get_ipython from IPython.core.interactiveshell import InteractiveShell from pandas import set_option as pandas_set_option from pydantic import BaseSettings, validator @@ -69,7 +68,6 @@ class Settings(BaseSettings): # controls dataframe variable tracking, hashing, and storing in sqlite ENABLE_DATALINK: bool = True - ENABLE_ASSIGNMENT: bool = True NUM_PAST_SAMPLES_TRACKED: int = 3 DB_LOCATION: str = ":memory:" @@ -221,49 +219,10 @@ def set_option( if key == "LOG_LEVEL": set_log_level(value) - # allow enabling/disabling comms based on settings - enable_disable_comms( - setting_name=key, - enabled=value, - ipython_shell=ipython_shell, - ) - return raise ValueError(f"`{key}` is not a valid setting") -def enable_disable_comms( - setting_name: str, - enabled: bool, - ipython_shell: Optional[InteractiveShell] = None, -) -> None: - """ - Registers/unregisters a target based on its associated name within Settings. - For example, the following will unregister the "datalink_resample" comm: - >>> enable_disable_comms("ENABLE_DATALINK", False) - And to re-register it: - >>> enable_disable_comms("ENABLE_DATALINK", True) - """ - from dx import comms - - comm_setting_targets = { - "ENABLE_DATALINK": ("datalink_resample", comms.resample.resampler), - "ENABLE_ASSIGNMENT": ("datalink_assignment", comms.assignment.dataframe_assignment), - } - if setting_name not in comm_setting_targets: - return - - ipython_shell = ipython_shell or get_ipython() - if getattr(ipython_shell, "kernel", None) is None: - return - - comm_target, comm_callback = comm_setting_targets[setting_name] - if enabled: - ipython_shell.kernel.comm_manager.register_target(comm_target, comm_callback) - else: - ipython_shell.kernel.comm_manager.unregister_target(comm_target, comm_callback) - - @contextmanager def settings_context(ipython_shell: Optional[InteractiveShell] = None, **option_kwargs): settings = get_settings() diff --git a/tests/test_comms.py b/tests/test_comms.py index 69866722..7b4fbfc9 100644 --- a/tests/test_comms.py +++ b/tests/test_comms.py @@ -1,11 +1,5 @@ -import uuid - -import pandas as pd -from IPython.terminal.interactiveshell import TerminalInteractiveShell - -from dx.comms.assignment import handle_assignment_comm from dx.comms.resample import handle_resample_comm -from dx.types.filters import DEXFilterSettings, DEXResampleMessage +from dx.types.filters import DEXResampleMessage class TestResampleComm: @@ -34,164 +28,3 @@ def test_resample_handled(self, mocker): handle_resample_comm(msg) resample_msg = DEXResampleMessage.parse_obj(msg["content"]["data"]) mock_handle_resample.assert_called_once_with(resample_msg) - - -class TestAssignmentComm: - def test_assignment_handled( - self, - mocker, - get_ipython: TerminalInteractiveShell, - sample_dataframe: pd.DataFrame, - ): - """ - Test that a valid message handled through the assignment comm - will call resample_from_db() with the provided display_id, filters, - and sample size, and will assign a valid pandas DataFrame in the - kernel namespace with the provided variable name. - """ - display_id = str(uuid.uuid4()) - sample_size = 50 - msg = { - "content": { - "data": { - "cell_id": "cell1", - "display_id": display_id, - "filters": [], - "sample_size": sample_size, - "variable_name": "new_df", - } - } - } - mock_resample = mocker.patch( - "dx.comms.assignment.resample_from_db", return_value=sample_dataframe - ) - handle_assignment_comm(msg, ipython_shell=get_ipython) - resample_params = { - "display_id": display_id, - "sql_filter": f"SELECT * FROM {{table_name}} LIMIT {sample_size}", - "filters": [], - "assign_subset": False, - } - mock_resample.assert_called_once_with(**resample_params) - assert "new_df" in get_ipython.user_ns - assert isinstance(get_ipython.user_ns["new_df"], pd.DataFrame) - assert get_ipython.user_ns["new_df"].equals(sample_dataframe) - - def test_assignment_handled_with_filters( - self, - mocker, - get_ipython: TerminalInteractiveShell, - sample_dataframe: pd.DataFrame, - sample_dex_metric_filter: dict, - ): - """ - Test that a valid message handled through the assignment comm - will call resample_from_db() with the provided display_id, filters, - and sample size, and will assign a valid pandas DataFrame in the - kernel namespace with the provided variable name. - """ - display_id = str(uuid.uuid4()) - sample_size = 50 - - filters = [sample_dex_metric_filter] - dex_filters = DEXFilterSettings(filters=filters) - sql_filter = ( - f"SELECT * FROM {{table_name}} WHERE {dex_filters.to_sql_query()} LIMIT {sample_size}" - ) - - msg = { - "content": { - "data": { - "cell_id": "cell1", - "display_id": display_id, - "filters": filters, - "sample_size": sample_size, - "variable_name": "new_df", - } - } - } - mock_resample = mocker.patch( - "dx.comms.assignment.resample_from_db", return_value=sample_dataframe - ) - handle_assignment_comm(msg, ipython_shell=get_ipython) - resample_params = { - "display_id": display_id, - "sql_filter": sql_filter, - "filters": filters, - "assign_subset": False, - } - mock_resample.assert_called_once_with(**resample_params) - assert "new_df" in get_ipython.user_ns - assert isinstance(get_ipython.user_ns["new_df"], pd.DataFrame) - - def test_assignment_handled_with_existing_variable( - self, - mocker, - get_ipython: TerminalInteractiveShell, - sample_dataframe: pd.DataFrame, - ): - """ - Test that a valid assignment message attempting to assign to an existing - variable will be appended with a numeric suffix and will not - affect the original variable. - """ - existing_dataframe_variable = pd.DataFrame({"test": [1, 2, 3]}) - get_ipython.user_ns["df"] = existing_dataframe_variable - - display_id = str(uuid.uuid4()) - sample_size = 50 - msg = { - "content": { - "data": { - "cell_id": "cell1", - "display_id": display_id, - "filters": [], - "sample_size": sample_size, - "variable_name": "df", - } - } - } - mock_resample = mocker.patch( - "dx.comms.assignment.resample_from_db", return_value=sample_dataframe - ) - handle_assignment_comm(msg, ipython_shell=get_ipython) - resample_params = { - "display_id": display_id, - "sql_filter": f"SELECT * FROM {{table_name}} LIMIT {sample_size}", - "filters": [], - "assign_subset": False, - } - mock_resample.assert_called_once_with(**resample_params) - - # new variable should be assigned with a numeric suffix - assert "df_1" in get_ipython.user_ns - assert isinstance(get_ipython.user_ns["df_1"], pd.DataFrame) - assert get_ipython.user_ns["df_1"].equals(sample_dataframe) - - # old variable should still exist - assert "df" in get_ipython.user_ns - assert get_ipython.user_ns["df"].equals(existing_dataframe_variable) - - def test_assignment_skipped( - self, - mocker, - get_ipython: TerminalInteractiveShell, - ): - """ - Test that variable assignment is skipped if `display_id` and - `variable_name` are not provided in the comm message. - """ - msg = { - "content": { - "data": { - "cell_id": "cell1", - "filters": [], - "sample_size": 50, - "variable_name": "new_df", - } - } - } - mock_resample = mocker.patch("dx.comms.assignment.resample_from_db") - handle_assignment_comm(msg, ipython_shell=get_ipython) - mock_resample.assert_not_called() - assert "new_df" not in get_ipython.user_ns