further simplfied it

Author: benfdking

sqlmesh/lsp/main.py

@@ -16,12 +16,8 @@
 
 logger = logging.getLogger(__name__)
 
-GLOBAL_CONTEXT: t.Optional[
-    t.Tuple[
-        Context,
-        t.Dict[str, t.Tuple[Context, t.Union[Model, ModelAudit]]],
-    ]
-] = None
+GLOBAL_CONTEXT: t.Optional[Context,] = None
+FILE_MAP: t.Dict[str, t.Union[Model, ModelAudit]] = {}
 
 
 server = LanguageServer("sqlmesh_lsp", __version__)
@@ -30,31 +26,25 @@
 def ensure_context_for_document(document: TextDocument) -> TextDocument:
     """Ensure that a context exists for the given document if applicable by searching for a config.py or config.yml file in the parent directories."""
     # If the context is already loaded, return the document, if it is part of the same context
-    global GLOBAL_CONTEXT
+    global GLOBAL_CONTEXT, FILE_MAP
     if GLOBAL_CONTEXT is not None:
-        GLOBAL_CONTEXT[0].load()
-        if document.uri in GLOBAL_CONTEXT[1]:
+        GLOBAL_CONTEXT.load()
+        if document.uri in FILE_MAP:
             return document
         else:
-            for model in GLOBAL_CONTEXT[0]._models.values():
+            for model in GLOBAL_CONTEXT._models.values():
                 if model._path is None:
                     continue
                 path = model._path.resolve()
                 if path == document.path:
-                    GLOBAL_CONTEXT = (
-                        GLOBAL_CONTEXT[0],
-                        GLOBAL_CONTEXT[1] | {document.uri: (GLOBAL_CONTEXT[0], model)},
-                    )
+                    FILE_MAP[document.uri] = model
                     return document
-            for audit in GLOBAL_CONTEXT[0]._audits.values():
+            for audit in GLOBAL_CONTEXT._audits.values():
                 if audit._path is None:
                     continue
                 path = audit._path.resolve()
                 if path == document.path:
-                    GLOBAL_CONTEXT = (
-                        GLOBAL_CONTEXT[0],
-                        GLOBAL_CONTEXT[1] | {document.uri: (GLOBAL_CONTEXT[0], audit)},
-                    )
+                    FILE_MAP[document.uri] = audit
                     return document
             return document
 
@@ -68,8 +58,7 @@ def ensure_context_for_document(document: TextDocument) -> TextDocument:
             config_path = path / f"config.{ext}"
             if config_path.exists():
                 with suppress(Exception):
-                    handle = Context(paths=[path])
-                    GLOBAL_CONTEXT = (handle, {})
+                    GLOBAL_CONTEXT = Context(paths=[path])
                     server.show_message(f"Context loaded for: {path}")
                     loaded = True
                     return ensure_context_for_document(document)
@@ -88,7 +77,7 @@ def formatting(
         context = GLOBAL_CONTEXT
         if context is None:
             raise Exception(f"No context found for document: {document.path}")
-        context[0].format(paths=(Path(document.path),))
+        context.format(paths=(Path(document.path),))
         with open(document.path, "r+", encoding="utf-8") as file:
             return [
                 types.TextEdit(