From 3a2d2e60b9780796accc88cbebbba44e9fecf375 Mon Sep 17 00:00:00 2001 From: Kartik Gupta <88345179+kartikgupta-db@users.noreply.github.com> Date: Thu, 6 Jul 2023 04:06:22 +0200 Subject: [PATCH] [DECO-1115] Add local implementation for `dbutils.widgets` (#93) ## Changes * Added a new install group (`pip install 'databricks-sdk[notebook]'`). This allows us to safely pin ipywidgets for local installs. DBR can safely continue using `pip install databricks-sdk` or directly using the default build from master without conflicting dependencies. * OSS implementation of widgets is imported only on first use (possible only through OSS implementation of dbutils - `RemoteDbutils`). * Add a wrapper for ipywidgets to enable interactive widgets when in interactive **IPython** notebooks. https://user-images.githubusercontent.com/88345179/236443693-1c804107-ba21-4296-ba40-2b1e8e062d16.mov * Add default widgets implementation that returns a default value, when not in an interactive environment. https://user-images.githubusercontent.com/88345179/236443729-51185404-4d28-49c6-ade0-a665e154e092.mov ## Tests - [x] `make test` run locally - [x] `make fmt` applied - [ ] relevant integration tests applied --- databricks/sdk/_widgets/__init__.py | 72 +++++++++++++++ .../sdk/_widgets/default_widgets_utils.py | 42 +++++++++ databricks/sdk/_widgets/ipywidgets_utils.py | 87 +++++++++++++++++++ databricks/sdk/dbutils.py | 13 +++ setup.py | 8 +- tests/test_dbutils.py | 8 +- 6 files changed, 224 insertions(+), 6 deletions(-) create mode 100644 databricks/sdk/_widgets/__init__.py create mode 100644 databricks/sdk/_widgets/default_widgets_utils.py create mode 100644 databricks/sdk/_widgets/ipywidgets_utils.py diff --git a/databricks/sdk/_widgets/__init__.py b/databricks/sdk/_widgets/__init__.py new file mode 100644 index 000000000..4fef42696 --- /dev/null +++ b/databricks/sdk/_widgets/__init__.py @@ -0,0 +1,72 @@ +import logging +import typing +import warnings +from abc import ABC, abstractmethod + + +class WidgetUtils(ABC): + + def get(self, name: str): + return self._get(name) + + @abstractmethod + def _get(self, name: str) -> str: + pass + + def getArgument(self, name: str, default_value: typing.Optional[str] = None): + try: + return self.get(name) + except Exception: + return default_value + + def remove(self, name: str): + self._remove(name) + + @abstractmethod + def _remove(self, name: str): + pass + + def removeAll(self): + self._remove_all() + + @abstractmethod + def _remove_all(self): + pass + + +try: + # We only use ipywidgets if we are in a notebook interactive shell otherwise we raise error, + # to fallback to using default_widgets. Also, users WILL have IPython in their notebooks (jupyter), + # because we DO NOT SUPPORT any other notebook backends, and hence fallback to default_widgets. + from IPython.core.getipython import get_ipython + + # Detect if we are in an interactive notebook by iterating over the mro of the current ipython instance, + # to find ZMQInteractiveShell (jupyter). When used from REPL or file, this check will fail, since the + # mro only contains TerminalInteractiveShell. + if len(list(filter(lambda i: i.__name__ == 'ZMQInteractiveShell', get_ipython().__class__.__mro__))) == 0: + logging.debug("Not in an interactive notebook. Skipping ipywidgets implementation for dbutils.") + raise EnvironmentError("Not in an interactive notebook.") + + # For import errors in IPyWidgetUtil, we provide a warning message, prompting users to install the + # correct installation group of the sdk. + try: + from .ipywidgets_utils import IPyWidgetUtil + + widget_impl = IPyWidgetUtil + logging.debug("Using ipywidgets implementation for dbutils.") + + except ImportError as e: + # Since we are certain that we are in an interactive notebook, we can make assumptions about + # formatting and make the warning nicer for the user. + warnings.warn( + "\nTo use databricks widgets interactively in your notebook, please install databricks sdk using:\n" + "\tpip install 'databricks-sdk[notebook]'\n" + "Falling back to default_value_only implementation for databricks widgets.") + logging.debug(f"{e.msg}. Skipping ipywidgets implementation for dbutils.") + raise e + +except: + from .default_widgets_utils import DefaultValueOnlyWidgetUtils + + widget_impl = DefaultValueOnlyWidgetUtils + logging.debug("Using default_value_only implementation for dbutils.") diff --git a/databricks/sdk/_widgets/default_widgets_utils.py b/databricks/sdk/_widgets/default_widgets_utils.py new file mode 100644 index 000000000..9b61a75f6 --- /dev/null +++ b/databricks/sdk/_widgets/default_widgets_utils.py @@ -0,0 +1,42 @@ +import typing + +from . import WidgetUtils + + +class DefaultValueOnlyWidgetUtils(WidgetUtils): + + def __init__(self) -> None: + self._widgets: typing.Dict[str, str] = {} + + def text(self, name: str, defaultValue: str, label: typing.Optional[str] = None): + self._widgets[name] = defaultValue + + def dropdown(self, + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None): + self._widgets[name] = defaultValue + + def combobox(self, + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None): + self._widgets[name] = defaultValue + + def multiselect(self, + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None): + self._widgets[name] = defaultValue + + def _get(self, name: str) -> str: + return self._widgets[name] + + def _remove(self, name: str): + del self._widgets[name] + + def _remove_all(self): + self._widgets = {} diff --git a/databricks/sdk/_widgets/ipywidgets_utils.py b/databricks/sdk/_widgets/ipywidgets_utils.py new file mode 100644 index 000000000..6f27df438 --- /dev/null +++ b/databricks/sdk/_widgets/ipywidgets_utils.py @@ -0,0 +1,87 @@ +import typing + +from IPython.core.display_functions import display +from ipywidgets.widgets import (ValueWidget, Widget, widget_box, + widget_selection, widget_string) + +from .default_widgets_utils import WidgetUtils + + +class DbUtilsWidget: + + def __init__(self, label: str, value_widget: ValueWidget) -> None: + self.label_widget = widget_string.Label(label) + self.value_widget = value_widget + self.box = widget_box.Box([self.label_widget, self.value_widget]) + + def display(self): + display(self.box) + + def close(self): + self.label_widget.close() + self.value_widget.close() + self.box.close() + + @property + def value(self): + value = self.value_widget.value + if type(value) == str or value is None: + return value + if type(value) == list or type(value) == tuple: + return ','.join(value) + + raise ValueError("The returned value has invalid type (" + type(value) + ").") + + +class IPyWidgetUtil(WidgetUtils): + + def __init__(self) -> None: + self._widgets: typing.Dict[str, DbUtilsWidget] = {} + + def _register(self, name: str, widget: ValueWidget, label: typing.Optional[str] = None): + label = label if label is not None else name + w = DbUtilsWidget(label, widget) + + if name in self._widgets: + self.remove(name) + + self._widgets[name] = w + w.display() + + def text(self, name: str, defaultValue: str, label: typing.Optional[str] = None): + self._register(name, widget_string.Text(defaultValue), label) + + def dropdown(self, + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None): + self._register(name, widget_selection.Dropdown(value=defaultValue, options=choices), label) + + def combobox(self, + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None): + self._register(name, widget_string.Combobox(value=defaultValue, options=choices), label) + + def multiselect(self, + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None): + self._register( + name, + widget_selection.SelectMultiple(value=(defaultValue, ), + options=[("__EMPTY__", ""), *list(zip(choices, choices))]), label) + + def _get(self, name: str) -> str: + return self._widgets[name].value + + def _remove(self, name: str): + self._widgets[name].close() + del self._widgets[name] + + def _remove_all(self): + Widget.close_all() + self._widgets = {} diff --git a/databricks/sdk/dbutils.py b/databricks/sdk/dbutils.py index 60d2ccd69..cc36e532e 100644 --- a/databricks/sdk/dbutils.py +++ b/databricks/sdk/dbutils.py @@ -175,6 +175,19 @@ def __init__(self, config: 'Config' = None): self.fs = _FsUtil(dbfs_ext.DbfsExt(self._client), self.__getattr__) self.secrets = _SecretsUtil(workspace.SecretsAPI(self._client)) + self._widgets = None + + # When we import widget_impl, the init file checks whether user has the + # correct dependencies required for running on notebook or not (ipywidgets etc). + # We only want these checks (and the subsequent errors and warnings), to + # happen when the user actually uses widgets. + @property + def widgets(self): + if self._widgets is None: + from ._widgets import widget_impl + self._widgets = widget_impl() + + return self._widgets @property def _cluster_id(self) -> str: diff --git a/setup.py b/setup.py index 35d592328..7f7106784 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,6 @@ -import io, pathlib +import io +import pathlib + from setuptools import setup, find_packages version_data = {} @@ -12,7 +14,9 @@ python_requires=">=3.7", install_requires=["requests>=2.28.1,<3"], extras_require={"dev": ["pytest", "pytest-cov", "pytest-xdist", "pytest-mock", - "yapf", "pycodestyle", "autoflake", "isort", "wheel"]}, + "yapf", "pycodestyle", "autoflake", "isort", "wheel", + "ipython", "ipywidgets"], + "notebook": ["ipython>=8,<9", "ipywidgets>=8,<9"]}, author="Serge Smertin", author_email="serge.smertin@databricks.com", description="Databricks SDK for Python (Beta)", diff --git a/tests/test_dbutils.py b/tests/test_dbutils.py index ac7e3102b..1361b6c07 100644 --- a/tests/test_dbutils.py +++ b/tests/test_dbutils.py @@ -184,14 +184,14 @@ def test_any_proxy(dbutils_proxy): command = ('\n' ' import json\n' ' (args, kwargs) = json.loads(\'[["a"], {}]\')\n' - ' result = dbutils.widgets.getParameter(*args, **kwargs)\n' + ' result = dbutils.notebook.exit(*args, **kwargs)\n' ' dbutils.notebook.exit(json.dumps(result))\n' ' ') - dbutils, assertions = dbutils_proxy('b', command) + dbutils, assertions = dbutils_proxy('a', command) - param = dbutils.widgets.getParameter('a') + param = dbutils.notebook.exit("a") - assert param == 'b' + assert param == 'a' assertions()