Preserve filesystem prefixes embedded in paths

databricks/sdk/mixins/files.py

@@ -45,6 +45,12 @@
 
 _LOG = logging.getLogger(__name__)
 
+_FILESYSTEM_PREFIXES = re.compile(r"^(?:dbfs|file):")
+
+
+def _strip_filesystem_prefix(path: Union[str, pathlib.Path]) -> str:
+    return _FILESYSTEM_PREFIXES.sub("", str(path), count=1)
+
 
 class _DbfsIO(BinaryIO):
     MAX_CHUNK_SIZE = 1024 * 1024
@@ -364,10 +370,10 @@ def as_string(self) -> str:
 
 class _LocalPath(_Path):
     def __init__(self, path: str):
+        path = str(path)
         if platform.system() == "Windows":
-            self._path = pathlib.Path(str(path).replace("file:///", "").replace("file:", ""))
-        else:
-            self._path = pathlib.Path(str(path).replace("file:", ""))
+            path = path.removeprefix("file:///")
+        self._path = pathlib.Path(path.removeprefix("file:"))
 
     def _is_local(self) -> bool:
         return True
@@ -433,7 +439,7 @@ def __repr__(self) -> str:
 
 class _VolumesPath(_Path):
     def __init__(self, api: files.FilesAPI, src: Union[str, pathlib.Path]):
-        self._path = pathlib.PurePosixPath(str(src).replace("dbfs:", "").replace("file:", ""))
+        self._path = pathlib.PurePosixPath(_strip_filesystem_prefix(src))
         self._api = api
 
     def _is_local(self) -> bool:
@@ -509,7 +515,7 @@ def __repr__(self) -> str:
 
 class _DbfsPath(_Path):
     def __init__(self, api: files.DbfsAPI, src: str):
-        self._path = pathlib.PurePosixPath(str(src).replace("dbfs:", "").replace("file:", ""))
+        self._path = pathlib.PurePosixPath(_strip_filesystem_prefix(src))
         self._api = api
 
     def _is_local(self) -> bool:

tests/test_dbfs_mixins.py

@@ -1,3 +1,5 @@
+import pathlib
+
 import pytest
 
 from databricks.sdk.errors import NotFound
@@ -77,6 +79,34 @@ def test_fs_path(config, path, expected_type):
     assert isinstance(dbfs_ext._path(path), expected_type)
 
 
+@pytest.mark.parametrize("path_type", [_DbfsPath, _VolumesPath])
+@pytest.mark.parametrize(
+    "path,expected",
+    [
+        ("dbfs:/path/to/file", "/path/to/file"),
+        ("file:/path/to/file", "/path/to/file"),
+        ("/dbfs:", "/dbfs:"),
+        ("/path/dbfs:file", "/path/dbfs:file"),
+        ("/path/file:dbfs:", "/path/file:dbfs:"),
+    ],
+)
+def test_remote_path_only_strips_leading_filesystem_prefix(mocker, path_type, path, expected):
+    assert path_type(mocker.Mock(), path).as_string == expected
+
+
+def test_local_path_only_strips_leading_file_prefix(tmp_path):
+    path = tmp_path / "file:name"
+
+    assert _LocalPath(f"file:{path}").as_string == str(path)
+
+
+def test_windows_local_path_preserves_file_uri_behavior(mocker):
+    mocker.patch("databricks.sdk.mixins.files.platform.system", return_value="Windows")
+
+    assert _LocalPath("file:///C:/Temp/file:name").as_string == str(pathlib.Path("C:/Temp/file:name"))
+    assert _LocalPath("file://server/share/file:name").as_string == str(pathlib.Path("//server/share/file:name"))
+
+
 def test_fs_path_invalid(config):
     dbfs_ext = DbfsExt(config)
     with pytest.raises(ValueError) as e:
@@ -92,7 +122,8 @@ def test_dbfs_local_path_mkdir(config, tmp_path):
     assert w.dbfs.exists(f"file:{tmp_path}/test_dir")
 
 
-def test_dbfs_exists(config, mocker):
+@pytest.mark.parametrize("path", ["/abc/def/ghi", "/dbfs:", "/abc/file:name"])
+def test_dbfs_exists(config, mocker, path):
     from databricks.sdk import WorkspaceClient
 
     get_status = mocker.patch(
@@ -101,17 +132,88 @@ def test_dbfs_exists(config, mocker):
     )
 
     client = WorkspaceClient(config=config)
-    client.dbfs.exists("/abc/def/ghi")
+    client.dbfs.exists(path)
 
-    get_status.assert_called_with("/abc/def/ghi")
+    get_status.assert_called_with(path)
 
 
-def test_volume_exists(config, mocker):
+@pytest.mark.parametrize(
+    "path",
+    [
+        "/Volumes/abc/def/ghi",
+        "/Volumes/abc/def/dbfs:",
+        "/Volumes/abc/def/file:name",
+    ],
+)
+def test_volume_exists(config, mocker, path):
     from databricks.sdk import WorkspaceClient
 
     get_metadata = mocker.patch("databricks.sdk.service.files.FilesAPI.get_metadata")
 
     client = WorkspaceClient(config=config)
-    client.dbfs.exists("/Volumes/abc/def/ghi")
+    client.dbfs.exists(path)
+
+    get_metadata.assert_called_with(path)
+
+
+def test_dbfs_recursive_list_preserves_dbfs_directory_name(config, mocker):
+    from databricks.sdk import WorkspaceClient
+    from databricks.sdk.service.files import FileInfo
 
-    get_metadata.assert_called_with("/Volumes/abc/def/ghi")
+    mocker.patch(
+        "databricks.sdk.service.files.DbfsAPI.get_status",
+        return_value=FileInfo(path="/", is_dir=True),
+    )
+    requested_paths = []
+
+    def list_directory(path):
+        requested_paths.append(path)
+        if len(requested_paths) > 2:
+            pytest.fail("recursive listing revisited the root directory")
+        if path == "/":
+            return [FileInfo(path="/dbfs:", is_dir=True)]
+        if path == "/dbfs:":
+            return [FileInfo(path="/dbfs:/file", is_dir=False)]
+        raise AssertionError(f"unexpected path: {path}")
+
+    mocker.patch("databricks.sdk.service.files.DbfsAPI.list", side_effect=list_directory)
+
+    client = WorkspaceClient(config=config)
+    result = list(client.dbfs.list("/", recursive=True))
+
+    assert [entry.path for entry in result] == ["/dbfs:/file"]
+    assert requested_paths == ["/", "/dbfs:"]
+
+
+def test_volume_recursive_list_preserves_file_directory_name(mocker):
+    from databricks.sdk.service.files import DirectoryEntry
+
+    root = "/Volumes/catalog/schema/volume"
+    directory = f"{root}/file:"
+    api = mocker.Mock()
+    requested_paths = []
+
+    def list_directory(path):
+        requested_paths.append(path)
+        if len(requested_paths) > 2:
+            pytest.fail("recursive listing revisited the root directory")
+        if path == root:
+            return [DirectoryEntry(name="file:", path=directory, is_directory=True)]
+        if path == directory:
+            return [
+                DirectoryEntry(
+                    name="data",
+                    path=f"{directory}/data",
+                    is_directory=False,
+                    file_size=1,
+                    last_modified=123,
+                )
+            ]
+        raise AssertionError(f"unexpected path: {path}")
+
+    api.list_directory_contents.side_effect = list_directory
+
+    result = list(_VolumesPath(api, root).list(recursive=True))
+
+    assert [entry.path for entry in result] == [f"{directory}/data"]
+    assert requested_paths == [root, directory]