From fde8a7ffa4871267cc4de1ac5a5b1a43fca62b2b Mon Sep 17 00:00:00 2001
From: Andrew Snare <andrew.snare@databricks.com>
Date: Wed, 10 Jul 2024 17:17:04 +0200
Subject: [PATCH] Revert to original .suffix implementation.

---
 src/databricks/labs/blueprint/paths.py | 31 ++++++++++++++------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/databricks/labs/blueprint/paths.py b/src/databricks/labs/blueprint/paths.py
index 1303cf4..78ba3cb 100644
--- a/src/databricks/labs/blueprint/paths.py
+++ b/src/databricks/labs/blueprint/paths.py
@@ -255,7 +255,7 @@ class WorkspacePath(Path):  # pylint: disable=too-many-public-methods
     )
     _cached_object_info: ObjectInfo
 
-    _SUFFIXES = {Language.PYTHON: ".py", Language.SQL: ".sql", Language.SCALA: ".scala", Language.R: ".R"}
+    _SUFFIXES = {".py": Language.PYTHON, ".sql": Language.SQL, ".scala": Language.SCALA, ".R": Language.R}
 
     # Path semantics are posix-like.
     parser = posixpath
@@ -460,19 +460,6 @@ def parts(self):
             return self.drive + self.root, *self._path_parts
         return self._path_parts
 
-    @property
-    def suffix(self):
-        # Super implementations are mostly fine...
-        suffix = super().suffix
-        # ...but if there is no suffix and this path is for a notebook then infer the extension based on the notebook
-        # language.
-        if not suffix and self.is_notebook():
-            try:
-                suffix = self._SUFFIXES.get(self._object_info.language, "")
-            except DatabricksError:
-                pass
-        return suffix
-
     def with_name(self, name):
         parser = self.parser
         if not name or parser.sep in name or name == ".":
@@ -664,6 +651,22 @@ 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():
+            try:
+                if self._object_info.language == lang:
+                    return sfx
+            except DatabricksError:
+                return ""
+        return ""
+
     @property
     def _object_info(self) -> ObjectInfo:
         # this method is cached because it is used in multiple is_* methods.