feat: normalize percent-encoded uris used as workspace dict keys

URIs used as dictionary keys in the Workspace class were stored and
looked up without normalization. If a client sent a percent-encoded URI
(e.g. file:///C%3A/foo) on one request and a decoded form (file:///C:/foo)
on another, lookups would fail with KeyError since the raw strings differ.

Apply urllib.parse.unquote to all URIs before using them as keys in
_text_documents, _notebook_documents, _cell_in_notebook, and _folders
dictionaries. This ensures equivalent URIs that differ only in percent
encoding map to the same entry.

Author: edvilme

pygls/workspace/workspace.py

@@ -20,6 +20,7 @@
 import logging
 import os
 from typing import Dict, Optional, Sequence, Union
+from urllib.parse import unquote
 
 from lsprotocol import types
 from lsprotocol.types import (
@@ -90,7 +91,7 @@ def _create_text_document(
         )
 
     def add_folder(self, folder: WorkspaceFolder):
-        self._folders[folder.uri] = folder
+        self._folders[unquote(folder.uri)] = folder
 
     @property
     def notebook_documents(self):
@@ -127,10 +128,10 @@ def get_notebook_document(
            The requested notebook document if found, ``None`` otherwise.
         """
         if notebook_uri is not None:
-            return self._notebook_documents.get(notebook_uri)
+            return self._notebook_documents.get(unquote(notebook_uri))
 
         if cell_uri is not None:
-            notebook_uri = self._cell_in_notebook.get(cell_uri)
+            notebook_uri = self._cell_in_notebook.get(unquote(cell_uri))
             if notebook_uri is None:
                 return None
 
@@ -145,7 +146,9 @@ def get_text_document(self, doc_uri: str) -> TextDocument:
 
         See https://github.com/Microsoft/language-server-protocol/issues/177
         """
-        return self._text_documents.get(doc_uri) or self._create_text_document(doc_uri)
+        return self._text_documents.get(unquote(doc_uri)) or self._create_text_document(
+            doc_uri
+        )
 
     def is_local(self):
 
@@ -161,7 +164,7 @@ def put_notebook_document(self, params: types.DidOpenNotebookDocumentParams):
         notebook = params.notebook_document
 
         # Create a fresh instance to ensure our copy cannot be accidentally modified.
-        self._notebook_documents[notebook.uri] = copy.deepcopy(notebook)
+        self._notebook_documents[unquote(notebook.uri)] = copy.deepcopy(notebook)
 
         for cell_document in params.cell_text_documents:
             self.put_text_document(cell_document, notebook_uri=notebook.uri)
@@ -184,31 +187,31 @@ def put_text_document(
         """
         doc_uri = text_document.uri
 
-        self._text_documents[doc_uri] = self._create_text_document(
+        self._text_documents[unquote(doc_uri)] = self._create_text_document(
             doc_uri,
             source=text_document.text,
             version=text_document.version,
             language_id=text_document.language_id,
         )
 
         if notebook_uri:
-            self._cell_in_notebook[doc_uri] = notebook_uri
+            self._cell_in_notebook[unquote(doc_uri)] = unquote(notebook_uri)
 
     def remove_notebook_document(self, params: types.DidCloseNotebookDocumentParams):
         notebook_uri = params.notebook_document.uri
-        self._notebook_documents.pop(notebook_uri, None)
+        self._notebook_documents.pop(unquote(notebook_uri), None)
 
         for cell_document in params.cell_text_documents:
             self.remove_text_document(cell_document.uri)
 
     def remove_text_document(self, doc_uri: str):
-        self._text_documents.pop(doc_uri, None)
-        self._cell_in_notebook.pop(doc_uri, None)
+        self._text_documents.pop(unquote(doc_uri), None)
+        self._cell_in_notebook.pop(unquote(doc_uri), None)
 
     def remove_folder(self, folder_uri: str):
-        self._folders.pop(folder_uri, None)
+        self._folders.pop(unquote(folder_uri), None)
         try:
-            del self._folders[folder_uri]
+            del self._folders[unquote(folder_uri)]
         except KeyError:
             pass
 
@@ -222,7 +225,7 @@ def root_uri(self):
 
     def update_notebook_document(self, params: types.DidChangeNotebookDocumentParams):
         uri = params.notebook_document.uri
-        notebook = self._notebook_documents[uri]
+        notebook = self._notebook_documents[unquote(uri)]
         notebook.version = params.notebook_document.version
 
         if params.change.metadata:
@@ -274,5 +277,5 @@ def update_text_document(
         change: types.TextDocumentContentChangeEvent,
     ):
         doc_uri = text_doc.uri
-        self._text_documents[doc_uri].apply_change(change)
-        self._text_documents[doc_uri].version = text_doc.version
+        self._text_documents[unquote(doc_uri)].apply_change(change)
+        self._text_documents[unquote(doc_uri)].version = text_doc.version

tests/test_workspace.py

@@ -440,3 +440,107 @@ def test_null_workspace():
 
     assert workspace.root_uri is None
     assert workspace.root_path is None
+
+
+# -- Percent-encoding normalization tests --
+
+ENCODED_DOC_URI = "file:///C%3A/path/to/file.py"
+DECODED_DOC_URI = "file:///C:/path/to/file.py"
+ENCODED_DOC = types.TextDocumentItem(
+    uri=ENCODED_DOC_URI, language_id="python", version=0, text="# encoded"
+)
+
+ENCODED_NB_URI = "file:///C%3A/path/to/notebook.ipynb"
+DECODED_NB_URI = "file:///C:/path/to/notebook.ipynb"
+ENCODED_NOTEBOOK = types.NotebookDocument(
+    uri=ENCODED_NB_URI,
+    notebook_type="jupyter-notebook",
+    version=0,
+    cells=[
+        types.NotebookCell(
+            kind=types.NotebookCellKind.Code,
+            document="nb-cell-scheme://C%3A/path/to/notebook.ipynb#cell1",
+        ),
+    ],
+)
+ENCODED_NB_CELL = types.TextDocumentItem(
+    uri="nb-cell-scheme://C%3A/path/to/notebook.ipynb#cell1",
+    language_id="python",
+    version=0,
+    text="# cell",
+)
+
+
+def test_get_text_document_percent_encoded(workspace):
+    """Looking up a document with a decoded URI after storing with an encoded one."""
+    workspace.put_text_document(ENCODED_DOC)
+    assert workspace.get_text_document(DECODED_DOC_URI).source == "# encoded"
+
+
+def test_get_text_document_percent_decoded(workspace):
+    """Looking up a document with an encoded URI after storing with a decoded one."""
+    decoded_doc = types.TextDocumentItem(
+        uri=DECODED_DOC_URI, language_id="python", version=0, text="# decoded"
+    )
+    workspace.put_text_document(decoded_doc)
+    assert workspace.get_text_document(ENCODED_DOC_URI).source == "# decoded"
+
+
+def test_remove_text_document_percent_encoded(workspace):
+    """Removing a document stored with an encoded URI using a decoded URI."""
+    workspace.put_text_document(ENCODED_DOC)
+    workspace.remove_text_document(DECODED_DOC_URI)
+    assert workspace.get_text_document(DECODED_DOC_URI)._source is None
+
+
+def test_get_notebook_document_percent_encoded(workspace):
+    """Looking up a notebook with a decoded URI after storing with an encoded one."""
+    params = types.DidOpenNotebookDocumentParams(
+        notebook_document=ENCODED_NOTEBOOK,
+        cell_text_documents=[ENCODED_NB_CELL],
+    )
+    workspace.put_notebook_document(params)
+
+    notebook = workspace.get_notebook_document(notebook_uri=DECODED_NB_URI)
+    assert notebook is not None
+    assert notebook.uri == ENCODED_NB_URI
+
+
+def test_update_notebook_document_percent_encoded(workspace):
+    """Updating a notebook stored with an encoded URI using a decoded URI."""
+    params = types.DidOpenNotebookDocumentParams(
+        notebook_document=ENCODED_NOTEBOOK,
+        cell_text_documents=[ENCODED_NB_CELL],
+    )
+    workspace.put_notebook_document(params)
+
+    update_params = types.DidChangeNotebookDocumentParams(
+        notebook_document=types.VersionedNotebookDocumentIdentifier(
+            uri=DECODED_NB_URI, version=5
+        ),
+        change=types.NotebookDocumentChangeEvent(
+            metadata={"updated": True},
+        ),
+    )
+    workspace.update_notebook_document(update_params)
+
+    notebook = workspace.get_notebook_document(notebook_uri=ENCODED_NB_URI)
+    assert notebook.version == 5
+    assert notebook.metadata == {"updated": True}
+
+
+def test_add_folder_percent_encoded(workspace):
+    """Looking up a folder with a decoded URI after adding with an encoded one."""
+    encoded_uri = "file:///C%3A/workspace"
+    decoded_uri = "file:///C:/workspace"
+    workspace.add_folder(types.WorkspaceFolder(uri=encoded_uri, name="ws"))
+    assert decoded_uri in workspace.folders
+
+
+def test_remove_folder_percent_encoded(workspace):
+    """Removing a folder stored with an encoded URI using a decoded URI."""
+    encoded_uri = "file:///C%3A/workspace"
+    decoded_uri = "file:///C:/workspace"
+    workspace.add_folder(types.WorkspaceFolder(uri=encoded_uri, name="ws"))
+    workspace.remove_folder(decoded_uri)
+    assert decoded_uri not in workspace.folders