From 6c26af294eca7a294c58cf031e806fb6c9305443 Mon Sep 17 00:00:00 2001
From: Jack Cherng <jfcherng@gmail.com>
Date: Fri, 12 Jul 2024 23:22:02 +0800
Subject: [PATCH] refactor: tidy codes

Signed-off-by: Jack Cherng <jfcherng@gmail.com>
---
 plugin/client.py | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/plugin/client.py b/plugin/client.py
index 28c530a..5e105a2 100644
--- a/plugin/client.py
+++ b/plugin/client.py
@@ -5,7 +5,7 @@
 import re
 import shutil
 import sys
-from collections import defaultdict
+import weakref
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, cast
@@ -47,7 +47,7 @@ class LspBasedpyrightPlugin(NpmClientHandler):
     server_version = ""
     """The version of the language server."""
 
-    window_attrs: defaultdict[WindowId, WindowAttr] = defaultdict(WindowAttr)
+    window_attrs: weakref.WeakKeyDictionary[sublime.Window, WindowAttr] = weakref.WeakKeyDictionary()
     """Per-window attributes. I.e., per-session attributes."""
 
     @classmethod
@@ -65,6 +65,20 @@ def should_ignore(cls, view: sublime.View) -> bool:
             return True
         return False
 
+    @classmethod
+    def can_start(
+        cls,
+        window: sublime.Window,
+        initiating_view: sublime.View,
+        workspace_folders: list[WorkspaceFolder],
+        configuration: ClientConfig,
+    ) -> str | None:
+        if message := super().can_start(window, initiating_view, workspace_folders, configuration):
+            return message
+
+        cls.window_attrs.setdefault(window, WindowAttr())
+        return None
+
     def on_settings_changed(self, settings: DottedDict) -> None:
         super().on_settings_changed(settings)
 
@@ -92,7 +106,7 @@ def on_pre_start(
 
         cls.server_version = cls.parse_server_version()
         cls.update_venv_info(configuration.settings, workspace_folders, window=window)
-        if venv_info := cls.window_attrs[window.id()].venv_info:
+        if venv_info := cls.window_attrs[window].venv_info:
             log_info(f"Using python executable: {venv_info.python_executable}")
             configuration.settings.set("python.pythonPath", str(venv_info.python_executable))
         return None
@@ -151,13 +165,12 @@ def copy_overwrite_dirs(cls) -> None:
     def update_status_bar_text(self, extra_variables: dict[str, Any] | None = None) -> None:
         if not (session := self.weaksession()):
             return
-        window_id = session.window.id()
 
         variables: dict[str, Any] = {
             "server_version": self.server_version,
         }
 
-        if venv_info := self.window_attrs[window_id].venv_info:
+        if venv_info := self.window_attrs[session.window].venv_info:
             variables["venv"] = {
                 "finder_name": venv_info.meta.finder_name,
                 "python_version": venv_info.python_version,
@@ -251,8 +264,7 @@ def update_venv_info(
         *,
         window: sublime.Window,
     ) -> None:
-        window_id = window.id()
-        window_attr = cls.window_attrs[window_id]
+        window_attr = cls.window_attrs[window]
 
         def _update_venv_info() -> None:
             window_attr.venv_info = None