From af2f118bb93d34f75f968b1f603dbc3d85b11f58 Mon Sep 17 00:00:00 2001
From: Andrew Snare <andrew.snare@databricks.com>
Date: Wed, 10 Jul 2024 14:12:12 +0200
Subject: [PATCH] Support for /-style building of paths.

---
 src/databricks/labs/blueprint/paths.py | 31 ++++++++++++++++++++++++++
 tests/unit/test_paths.py               | 25 +++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/src/databricks/labs/blueprint/paths.py b/src/databricks/labs/blueprint/paths.py
index 625a516..bc3cc52 100644
--- a/src/databricks/labs/blueprint/paths.py
+++ b/src/databricks/labs/blueprint/paths.py
@@ -408,6 +408,18 @@ def _str_normcase(self):
         # Compatibility property (python 3.12+), accessed via equality comparison. This can't be avoided.
         return str(self)
 
+    @classmethod
+    def _from_parts(cls, *args) -> NoReturn:
+        # Compatibility method (python <= 3.11), accessed via reverse /-style building. This can't be avoided.
+        # See __rtruediv__ for more information.
+        raise TypeError("trigger NotImplemented")
+
+    @property
+    def _raw_paths(self) -> NoReturn:
+        # Compatibility method (python 3.12+), accessed via reverse /-style building. This can't be avoided.
+        # See __rtruediv__ for more information.
+        raise TypeError("trigger NotImplemented")
+
     def __lt__(self, other):
         if not isinstance(other, type(self)):
             return NotImplemented
@@ -566,6 +578,25 @@ def is_reserved(self):
     def joinpath(self, *pathsegments):
         return self.with_segments(self, *pathsegments)
 
+    def __truediv__(self, other):
+        try:
+            return self.with_segments(*self._parts(), other)
+        except TypeError:
+            return NotImplemented
+
+    def __rtruediv__(self, other):
+        # Note: this is only invoked if __truediv__ has already returned NotImplemented.
+        # For the case of Path / WorkspacePath this means the underlying __truediv__ is invoked.
+        # The base-class implementations all access internals but yield NotImplemented if TypeError is raised. As
+        # such we stub those internals (_from_parts and _raw_path) to trigger the NotImplemented path and ensure that
+        # control ends up here.
+        try:
+            if isinstance(other, PurePath):
+                return type(other)(other, *self._parts())
+            return self.with_segments(other, *self._parts())
+        except TypeError:
+            return NotImplemented
+
     @classmethod
     def _compile_pattern(cls, pattern: str, case_sensitive: bool) -> re.Pattern:
         flags = 0 if case_sensitive else re.IGNORECASE
diff --git a/tests/unit/test_paths.py b/tests/unit/test_paths.py
index 5b9eda2..dd032a6 100644
--- a/tests/unit/test_paths.py
+++ b/tests/unit/test_paths.py
@@ -364,6 +364,31 @@ def test_joinpath() -> None:
     assert WorkspacePath(ws, "home").joinpath("jane") == WorkspacePath(ws, "home/jane")
 
 
+def test_join_dsl() -> None:
+    """Test that the /-based DSL can be used to build new paths."""
+    ws = create_autospec(WorkspaceClient)
+
+    # First the forward style options.
+    assert WorkspacePath(ws, "/home/bob") / "data" == WorkspacePath(ws, "/home/bob/data")
+    assert WorkspacePath(ws, "/home/bob") / "data" / "base" == WorkspacePath(ws, "/home/bob/data/base")
+    assert WorkspacePath(ws, "/home/bob") / "data/base" == WorkspacePath(ws, "/home/bob/data/base")
+    assert WorkspacePath(ws, "home") / "bob" == WorkspacePath(ws, "home/bob")
+    # New root
+    assert WorkspacePath(ws, "whatever") / "/home" == WorkspacePath(ws, "/home")
+    # Mix types: eventual type is less-associative
+    assert WorkspacePath(ws, "/home/bob") / PurePosixPath("data") == WorkspacePath(ws, "/home/bob/data")
+
+    # Building from the other direction; same as above.
+    assert "/home/bob" / WorkspacePath(ws, "data") == WorkspacePath(ws, "/home/bob/data")
+    assert "/home/bob" / WorkspacePath(ws, "data") / "base" == WorkspacePath(ws, "/home/bob/data/base")
+    assert "/home/bob" / WorkspacePath(ws, "data/base") == WorkspacePath(ws, "/home/bob/data/base")
+    assert "home" / WorkspacePath(ws, "bob") == WorkspacePath(ws, "home/bob")
+    # New root
+    assert "whatever" / WorkspacePath(ws, "/home") == WorkspacePath(ws, "/home")
+    # Mix types: eventual type is less-associative
+    assert PurePosixPath("/home/bob") / WorkspacePath(ws, "data") == PurePosixPath("/home/bob/data")
+
+
 def test_match() -> None:
     """Test that glob matching works."""
     ws = create_autospec(WorkspaceClient)