refactor(lsp): further refactor of setting (#4705)

Author: benfdking

sqlmesh/lsp/main.py

@@ -88,6 +88,8 @@ def __init__(
     def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
         """Create a new LSPContext instance using the configured context class.
 
+        On success, sets self.lsp_context and returns the created context.
+
         Args:
             paths: List of paths to pass to the context constructor
 
@@ -96,7 +98,9 @@ def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
         """
         try:
             context = self.context_class(paths=paths)
-            return LSPContext(context)
+            lsp_context = LSPContext(context)
+            self.lsp_context = lsp_context
+            return lsp_context
         except Exception as e:
             self.server.log_trace(f"Error creating context: {e}")
             return None
@@ -132,9 +136,7 @@ def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
                         for ext in ("py", "yml", "yaml"):
                             config_path = folder_path / f"config.{ext}"
                             if config_path.exists():
-                                lsp_context = self._create_lsp_context([folder_path])
-                                if lsp_context:
-                                    self.lsp_context = lsp_context
+                                if self._create_lsp_context([folder_path]):
                                     loaded_sqlmesh_message(ls, folder_path)
                                     return  # Exit after successfully loading any config
             except Exception as e:
@@ -299,9 +301,7 @@ def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> Non
 
             # Reload the entire context and create a new LSPContext
             if self.lsp_context is not None:
-                new_lsp_context = self._create_lsp_context(list(self.lsp_context.context.configs))
-                if new_lsp_context:
-                    self.lsp_context = new_lsp_context
+                if self._create_lsp_context(list(self.lsp_context.context.configs)):
                     return
 
             context = self._context_get_or_load(uri)
@@ -672,9 +672,7 @@ def _ensure_context_in_folder(self, folder_uri: Path) -> None:
         for ext in ("py", "yml", "yaml"):
             config_path = folder_uri / f"config.{ext}"
             if config_path.exists():
-                lsp_context = self._create_lsp_context([folder_uri])
-                if lsp_context:
-                    self.lsp_context = lsp_context
+                if self._create_lsp_context([folder_uri]):
                     loaded_sqlmesh_message(self.server, folder_uri)
                     return
 
@@ -683,9 +681,7 @@ def _ensure_context_in_folder(self, folder_uri: Path) -> None:
             for ext in ("py", "yml", "yaml"):
                 config_path = workspace_folder / f"config.{ext}"
                 if config_path.exists():
-                    lsp_context = self._create_lsp_context([workspace_folder])
-                    if lsp_context:
-                        self.lsp_context = lsp_context
+                    if self._create_lsp_context([workspace_folder]):
                         loaded_sqlmesh_message(self.server, workspace_folder)
                         return
 
@@ -714,9 +710,7 @@ def _ensure_context_for_document(
             for ext in ("py", "yml", "yaml"):
                 config_path = path / f"config.{ext}"
                 if config_path.exists():
-                    lsp_context = self._create_lsp_context([path])
-                    if lsp_context:
-                        self.lsp_context = lsp_context
+                    if self._create_lsp_context([path]):
                         loaded = True
                         # Re-check context for the document now that it's loaded
                         return self._ensure_context_for_document(document_uri)
@@ -728,9 +722,7 @@ def _ensure_context_for_document(
                 for ext in ("py", "yml", "yaml"):
                     config_path = workspace_folder / f"config.{ext}"
                     if config_path.exists():
-                        lsp_context = self._create_lsp_context([workspace_folder])
-                        if lsp_context:
-                            self.lsp_context = lsp_context
+                        if self._create_lsp_context([workspace_folder]):
                             loaded_sqlmesh_message(self.server, workspace_folder)
                             return