clean up class, make true optional import

Author: isabelizimm

pins/databricks/fs.py

@@ -2,8 +2,6 @@
 from io import BytesIO
 from pathlib import Path, PurePath
 
-from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import NotFound
 from fsspec import AbstractFileSystem
 
 from pins.errors import PinsError
@@ -13,18 +11,18 @@ class DatabricksFs(AbstractFileSystem):
     protocol = "dbc"
 
     def ls(self, path, detail=False, **kwargs):
-        return self._databricks_ls(self, path, detail)
+        return self._databricks_ls(path, detail)
 
     def exists(self, path: str, **kwargs):
-        return self._databricks_exists(self, path)
+        return self._databricks_exists(path)
 
     def open(self, path: str, mode: str = "rb", *args, **kwargs):
         if mode != "rb":
             raise NotImplementedError
-        return self._databricks_open(self, path)
+        return self._databricks_open(path)
 
     def get(self, rpath, lpath, recursive=False, **kwargs):
-        self._databricks_get(self, rpath, lpath, recursive, **kwargs)
+        self._databricks_get(rpath, lpath, recursive, **kwargs)
 
     def mkdir(self, path, create_parents=True, **kwargs):
         if not create_parents:
@@ -50,11 +48,18 @@ def rm(self, path, recursive=True, maxdepth=None) -> None:
             raise NotImplementedError
         if maxdepth is not None:
             raise NotImplementedError
-        if self._databricks_exists(self, path):
-            self._databricks_rm_dir(self, path)
+        if self._databricks_exists(path):
+            self._databricks_rm_dir(path)
 
     @staticmethod
     def _databricks_put(lpath, rpath):
+        try:
+            from databricks.sdk import WorkspaceClient
+        except ModuleNotFoundError:
+            raise PinsError(
+                "Install the `databricks-sdk` package for Databricks board support."
+            )
+
         w = WorkspaceClient()
         path = Path(lpath).absolute()
         orig_path = path
@@ -74,15 +79,21 @@ def _upload_files(path):
 
         _upload_files(path)
 
-    @staticmethod
     def _databricks_get(self, board, rpath, lpath, recursive=False, **kwargs):
+        try:
+            from databricks.sdk import WorkspaceClient
+        except ModuleNotFoundError:
+            raise PinsError(
+                "Install the `databricks-sdk` package for Databricks board support."
+            )
+
         w = WorkspaceClient()
         file_type = self._databricks_is_type(rpath)
         if file_type == "file":
             board.fs.get(rpath, lpath, **kwargs)
             return
 
-        def _get_files(self, path, recursive, **kwargs):
+        def _get_files(path, recursive, **kwargs):
             raw_contents = w.files.list_directory_contents(path)
             contents = list(raw_contents)
             details = list(map(self._databricks_content_details, contents))
@@ -96,11 +107,17 @@ def _get_files(self, path, recursive, **kwargs):
                     target_path = PurePath(lpath).joinpath(rel_path)
                     board.fs.get(item_path, str(target_path))
 
-        _get_files(self, rpath, recursive, **kwargs)
+        _get_files(rpath, recursive, **kwargs)
 
-    @staticmethod
     def _databricks_open(self, path):
-        if not self._databricks_exists(self, path):
+        try:
+            from databricks.sdk import WorkspaceClient
+        except ModuleNotFoundError:
+            raise PinsError(
+                "Install the `databricks-sdk` package for Databricks board support."
+            )
+
+        if not self._databricks_exists(path):
             raise PinsError("File or directory does not exist")
         w = WorkspaceClient()
         resp = w.files.download(path)
@@ -109,7 +126,6 @@ def _databricks_open(self, path):
         f.seek(0)
         return f
 
-    @staticmethod
     def _databricks_exists(self, path: str):
         if self._databricks_is_type(path) == "nothing":
             return False
@@ -118,6 +134,14 @@ def _databricks_exists(self, path: str):
 
     @staticmethod
     def _databricks_is_type(path: str):
+        try:
+            from databricks.sdk import WorkspaceClient
+            from databricks.sdk.errors import NotFound
+        except ModuleNotFoundError:
+            raise PinsError(
+                "Install the `databricks-sdk` package for Databricks board support."
+            )
+
         w = WorkspaceClient()
         try:
             w.files.get_metadata(path)
@@ -131,9 +155,15 @@ def _databricks_is_type(path: str):
         else:
             return "file"
 
-    @staticmethod
     def _databricks_ls(self, path, detail):
-        if not self._databricks_exists(self, path):
+        try:
+            from databricks.sdk import WorkspaceClient
+        except ModuleNotFoundError:
+            raise PinsError(
+                "Install the `databricks-sdk` package for Databricks board support."
+            )
+
+        if not self._databricks_exists(path):
             raise PinsError("File or directory does not exist")
         w = WorkspaceClient()
         if self._databricks_is_type(path) == "file":
@@ -159,22 +189,35 @@ def _databricks_ls(self, path, detail):
                 items.append(item_path)
         return items
 
-    @staticmethod
     def _databricks_rm_dir(self, path):
+        try:
+            from databricks.sdk import WorkspaceClient
+        except ModuleNotFoundError:
+            raise PinsError(
+                "Install the `databricks-sdk` package for Databricks board support."
+            )
+
         w = WorkspaceClient()
         raw_contents = w.files.list_directory_contents(path)
         contents = list(raw_contents)
         details = list(map(self._databricks_content_details, contents))
         for item in details:
             item_path = item.get("path")
             if item.get("is_directory"):
-                self._databricks_rm_dir(self, item_path)
+                self._databricks_rm_dir(item_path)
             else:
                 w.files.delete(item_path)
         w.files.delete_directory(path)
 
     @staticmethod
     def _databricks_mkdir(path):
+        try:
+            from databricks.sdk import WorkspaceClient
+        except ModuleNotFoundError:
+            raise PinsError(
+                "Install the `databricks-sdk` package for Databricks board support."
+            )
+
         w = WorkspaceClient()
         w.files.create_directory(path)
 

pins/tests/conftest.py

@@ -28,7 +28,7 @@
     pytest.param(lambda: BoardBuilder("s3"), id="s3", marks=m.fs_s3),
     pytest.param(lambda: BoardBuilder("gcs"), id="gcs", marks=m.fs_gcs),
     pytest.param(lambda: BoardBuilder("abfs"), id="abfs", marks=m.fs_abfs),
-    pytest.param(lambda: DbcBoardBuilder("dbc"), id="dbc", marks=m.fs_dbc)
+    pytest.param(lambda: DbcBoardBuilder("dbc"), id="dbc", marks=m.fs_dbc),
 ]
 
 # rsc should only be used once, because users are created at docker setup time

pins/tests/helpers.py

@@ -40,9 +40,11 @@
 
 def skip_if_dbc(func):
     """Decorator to skip test if board protocol is 'dbc'"""
+
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         import inspect
+
         board = None
 
         # Get function signature to map args to parameter names.
@@ -52,21 +54,22 @@ def wrapper(*args, **kwargs):
         bound_args = sig.bind_partial(*args, **kwargs)
         all_args = {**bound_args.arguments, **kwargs}
 
-        if 'board' in all_args:
-            board = all_args['board']
-        elif 'board_with_cache' in all_args:
-            board = all_args['board_with_cache']
+        if "board" in all_args:
+            board = all_args["board"]
+        elif "board_with_cache" in all_args:
+            board = all_args["board_with_cache"]
         else:
             # Check all arguments for something that looks like a board
             for arg_value in all_args.values():
-                if hasattr(arg_value, 'fs') and hasattr(arg_value.fs, 'protocol'):
+                if hasattr(arg_value, "fs") and hasattr(arg_value.fs, "protocol"):
                     board = arg_value
                     break
 
         if board and board.fs.protocol == "dbc":
             pytest.skip("All Databricks tests must be read only")
 
         return func(*args, **kwargs)
+
     return wrapper
 
 

pins/tests/test_compat.py

@@ -117,7 +117,7 @@ def test_compat_pin_meta(board):
     elif board.fs.protocol == "dbc":
         assert meta.title == "df_csv: a pinned 3 x 2 DataFrame"
         assert meta.description is None
-        assert meta.created == "20220214T163720Z"
+        assert meta.created == "20250410T083026Z"
         assert meta.file == "df_csv.csv"
         assert meta.file_size == 16
         assert meta.pin_hash == "a173cd6a53908980"
@@ -186,7 +186,7 @@ def test_compat_pin_read(board):
 
     # TODO: update when dbc boards are not read-only
     if board.fs.protocol == "dbc":
-        dst_df = pd.DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
+        dst_df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
     else:
         dst_df = pd.read_csv(p_data)
 

pins/tests/test_constructors.py

@@ -192,7 +192,7 @@ def test_constructor_boards(board, df_csv, tmp_cache):
     # check data
     # TODO: update when dbc boards are not read-only
     if board.fs.protocol == "dbc":
-        df_csv = pd.DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
+        df_csv = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
 
     assert_frame_equal(df, df_csv)