Added .suffix override for notebooks in WorkspacePath (#1557)

This PR allows for more unified work with workspace notebooks and files

Fix #1455

Author: nfx

src/databricks/labs/ucx/mixins/fixtures.py

@@ -45,7 +45,7 @@
     Query,
     QueryInfo,
 )
-from databricks.sdk.service.workspace import ImportFormat
+from databricks.sdk.service.workspace import ImportFormat, Language
 
 from databricks.labs.ucx.workspace_access.groups import MigratedGroup
 
@@ -556,12 +556,20 @@ def create(*, scope: str, principal: str, permission: workspace.AclPermission):
 
 @pytest.fixture
 def make_notebook(ws, make_random):
-    def create(*, path: str | None = None, content: BinaryIO | None = None, **kwargs):
+    def create(
+        *,
+        path: str | None = None,
+        content: BinaryIO | None = None,
+        language: Language = Language.PYTHON,
+        format: ImportFormat = ImportFormat.SOURCE,  # pylint:  disable=redefined-builtin
+        overwrite: bool = False,
+    ) -> str:
         if path is None:
-            path = f"/Users/{ws.current_user.me().user_name}/sdk-{make_random(4)}.py"
+            path = f"/Users/{ws.current_user.me().user_name}/sdk-{make_random(4)}"
         if content is None:
             content = io.BytesIO(b"print(1)")
-        ws.workspace.upload(path, content, **kwargs)
+        path = str(path)
+        ws.workspace.upload(path, content, language=language, format=format, overwrite=overwrite)
         return path
 
     yield from factory("notebook", create, lambda x: ws.workspace.delete(x))

src/databricks/labs/ucx/mixins/wspath.py

@@ -12,8 +12,7 @@
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
-from databricks.sdk.service.workspace import ObjectInfo, ObjectType, ExportFormat, ImportFormat
-
+from databricks.sdk.service.workspace import ObjectInfo, ObjectType, ExportFormat, ImportFormat, Language
 
 logger = logging.getLogger(__name__)
 
@@ -137,6 +136,8 @@ def __init__(self, ws: WorkspaceClient, path: str):
 class WorkspacePath(Path):
     """Experimental implementation of pathlib.Path for Databricks Workspace."""
 
+    _SUFFIXES = {'.py': Language.PYTHON, '.sql': Language.SQL, '.scala': Language.SCALA, '.R': Language.R}
+
     _ws: WorkspaceClient
     _flavour: _DatabricksFlavour
     _accessor: _DatabricksAccessor
@@ -270,6 +271,24 @@ def open(self, mode="r", buffering=-1, encoding=None, errors=None, newline=None)
             return _TextUploadIO(self._ws, self.as_posix())
         raise ValueError(f"invalid mode: {mode}")
 
+    @property
+    def suffix(self):
+        """Return the file extension. If the file is a notebook, return the suffix based on the language."""
+        suffix = super().suffix
+        if suffix:
+            return suffix
+        if not self.is_notebook():
+            return ""
+        for sfx, lang in self._SUFFIXES.items():
+            if self._object_info.language == lang:
+                return sfx
+        return ""
+
+    def __lt__(self, other: pathlib.PurePath):
+        if not isinstance(other, pathlib.PurePath):
+            return NotImplemented
+        return self.as_posix() < other.as_posix()
+
     @cached_property
     def _object_info(self) -> ObjectInfo:
         # this method is cached because it is used in multiple is_* methods.

tests/integration/mixins/test_wspath.py

@@ -93,3 +93,20 @@ def test_replace(ws, make_random):
 
     assert not hello_txt.exists()
     assert (with_user / "hello2.txt").read_text() == 'Hello, World!'
+
+
+def test_file_and_notebook_in_same_folder_with_different_suffixes(ws, make_notebook, make_directory):
+    folder = WorkspacePath(ws, make_directory())
+
+    txt_file = folder / "a.txt"
+    py_notebook = folder / 'b'  # notebooks have no file extension
+
+    make_notebook(path=py_notebook, content="display(spark.range(10))")
+    txt_file.write_text("Hello, World!")
+
+    files = {_.name: _ for _ in folder.glob("**/*")}
+    assert len(files) == 2
+
+    assert files['a.txt'].suffix == '.txt'
+    assert files['b'].suffix == '.py'  # suffix is determined from ObjectInfo
+    assert files['b'].read_text() == "# Databricks notebook source\ndisplay(spark.range(10))"