From e02840f90bb3c2e4d03ff9439bff90da184fd992 Mon Sep 17 00:00:00 2001
From: Eric Vergnaud <eric.vergnaud@databricks.com>
Date: Fri, 14 Feb 2025 08:50:33 +0100
Subject: [PATCH] Avoid private property access (#1444)

Our current implementation directly accesses internal `lsprotocol` data
structures, which is hacky.
This PR fixes the issue.

Closes #1378
---
 .../labs/remorph/transpiler/lsp/lsp_engine.py | 20 +++++++++++++------
 tests/resources/lsp_transpiler/lsp_server.py  | 20 +++++++++++++------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/src/databricks/labs/remorph/transpiler/lsp/lsp_engine.py b/src/databricks/labs/remorph/transpiler/lsp/lsp_engine.py
index eb808b7ed..083430209 100644
--- a/src/databricks/labs/remorph/transpiler/lsp/lsp_engine.py
+++ b/src/databricks/labs/remorph/transpiler/lsp/lsp_engine.py
@@ -13,8 +13,7 @@
 import attrs
 import yaml
 
-# see https://github.com/databrickslabs/remorph/issues/1378
-# pylint: disable=import-private-name
+from lsprotocol import types as types_module
 from lsprotocol.types import (
     InitializeParams,
     ClientCapabilities,
@@ -32,7 +31,6 @@
     LanguageKind,
     Range as LSPRange,
     Position as LSPPosition,
-    _SPECIAL_PROPERTIES,
     DiagnosticSeverity,
 )
 from pygls.lsp.client import BaseLanguageClient
@@ -157,9 +155,19 @@ class TranspileDocumentResponse:
     jsonrpc: str = attrs.field(default="2.0")
 
 
-_SPECIAL_PROPERTIES.extend(
-    [f"{TranspileDocumentRequest.__name__}.method", f"{TranspileDocumentRequest.__name__}.jsonrpc"]
-)
+def install_special_properties():
+    is_special_property = getattr(types_module, "is_special_property")
+
+    def customized(cls: type, property_name: str) -> bool:
+        if cls is TranspileDocumentRequest and property_name in {"method", "jsonrpc"}:
+            return True
+        return is_special_property(cls, property_name)
+
+    setattr(types_module, "is_special_property", customized)
+
+
+install_special_properties()
+
 METHOD_TO_TYPES[TRANSPILE_TO_DATABRICKS_METHOD] = (
     TranspileDocumentRequest,
     TranspileDocumentResponse,
diff --git a/tests/resources/lsp_transpiler/lsp_server.py b/tests/resources/lsp_transpiler/lsp_server.py
index ee0ae7f1f..1eea89873 100644
--- a/tests/resources/lsp_transpiler/lsp_server.py
+++ b/tests/resources/lsp_transpiler/lsp_server.py
@@ -7,8 +7,7 @@
 
 import attrs
 
-# see https://github.com/databrickslabs/remorph/issues/1378
-# pylint: disable=import-private-name
+from lsprotocol import types as types_module
 from lsprotocol.types import (
     InitializeParams,
     INITIALIZE,
@@ -23,7 +22,6 @@
     Range,
     Position,
     METHOD_TO_TYPES,
-    _SPECIAL_PROPERTIES,
     DiagnosticSeverity,
     LanguageKind,
 )
@@ -75,9 +73,19 @@ class TranspileDocumentResponse:
     jsonrpc: str = attrs.field(default="2.0")
 
 
-_SPECIAL_PROPERTIES.extend(
-    [f"{TranspileDocumentRequest.__name__}.method", f"{TranspileDocumentRequest.__name__}.jsonrpc"]
-)
+def install_special_properties():
+    is_special_property = getattr(types_module, "is_special_property")
+
+    def customized(cls: type, property_name: str) -> bool:
+        if cls is TranspileDocumentRequest and property_name in {"method", "jsonrpc"}:
+            return True
+        return is_special_property(cls, property_name)
+
+    setattr(types_module, "is_special_property", customized)
+
+
+install_special_properties()
+
 METHOD_TO_TYPES[TRANSPILE_TO_DATABRICKS_METHOD] = (
     TranspileDocumentRequest,
     TranspileDocumentResponse,