Make a temp document method in workspaces

Author: jasongrout

Diff for: pylsp/python_lsp.py

@@ -593,7 +593,6 @@ def _cell_document__definition(self, cellDocument=None, position=None, **_kwargs
         if notebookDocument is None:
             raise ValueError("Invalid notebook document")
 
-        random_uri = str(uuid.uuid4())
         # cell_list helps us map the diagnostics back to the correct cell later.
         cell_list: List[Dict[str, Any]] = []
 
@@ -623,21 +622,8 @@ def _cell_document__definition(self, cellDocument=None, position=None, **_kwargs
 
             offset += num_lines
 
-        # TODO: make a workspace temp document context manager that yields the random uri and cleans up afterwards
-        workspace.put_document(random_uri, total_source)
-        log.info(f'Making new document {random_uri}')
-        try:
+        with workspace.temp_document(total_source) as random_uri:
             definitions = self.definitions(random_uri, position)
-            log.info(f'Got definitions: {definitions}')
-
-            # {
-            #     'uri': uris.uri_with(document.uri, path=str(d.module_path)),
-            #     'range': {
-            #         'start': {'line': d.line - 1, 'character': d.column},
-            #         'end': {'line': d.line - 1, 'character': d.column + len(d.name)},
-            #     }
-            # }
-            print(definitions)
             for definition in definitions:
                 # TODO: a better test for if a definition is the random_uri
                 if random_uri in definition['uri']:
@@ -650,8 +636,6 @@ def _cell_document__definition(self, cellDocument=None, position=None, **_kwargs
                             definition['range']['start']['line'] -= cell['line_start']
                             definition['range']['end']['line'] -= cell['line_start']
             return definitions
-        finally:
-            workspace.rm_document(random_uri)
 
     def m_text_document__definition(self, textDocument=None, position=None, **_kwargs):
         # textDocument here is just a dict with a uri

Diff for: pylsp/workspace.py

@@ -119,6 +119,17 @@ def put_document(self, doc_uri, source, version=None):
     def put_notebook_document(self, doc_uri, notebook_type, cells, version=None, metadata=None):
         self._docs[doc_uri] = self._create_notebook_document(doc_uri, notebook_type, cells, version, metadata)
 
+    @contextmanager
+    def temp_document(self, source):
+        random_uri = str(uuid.uuid4())
+        try:
+            self.put_document(random_uri, source)
+            yield random_uri
+        finally:
+            self.rm_document(random_uri)
+
+
+
     def add_notebook_cells(self, doc_uri, cells, start):
         self._docs[doc_uri].add_cells(cells, start)