refactor: tidy codes

Signed-off-by: Jack Cherng <[email protected]>

plugin/client.py

@@ -4,7 +4,6 @@
 import json
 import os
 import weakref
-from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
 from functools import wraps
@@ -63,8 +62,8 @@
 
 @dataclass
 class WindowAttr:
-    client_ref: weakref.ReferenceType[CopilotPlugin] | None = None
-    """The weak reference of the LSP client instance for the window."""
+    client: CopilotPlugin | None = None
+    """The LSP client instance for the window."""
 
 
 def _guard_view(*, failed_return: Any = None) -> Callable[[T_Callable], T_Callable]:
@@ -112,7 +111,7 @@ class CopilotPlugin(NpmClientHandler):
     server_version_gh = ""
     """The version of the Github Copilot 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."""
 
     _account_status = AccountStatus(
@@ -125,7 +124,8 @@ class CopilotPlugin(NpmClientHandler):
     def __init__(self, session: weakref.ref[Session]) -> None:
         super().__init__(session)
         if sess := session():
-            self.window_attrs[sess.window.id()].client_ref = weakref.ref(self)
+            self.window_attrs[sess.window] = WindowAttr()
+            self.window_attrs[sess.window].client = self
 
         self._activity_indicator = ActivityIndicator(self.update_status_bar_text)
 
@@ -260,12 +260,7 @@ def set_account_status(
 
     @classmethod
     def from_view(cls, view: sublime.View) -> CopilotPlugin | None:
-        if (
-            (window := view.window())
-            and (self_ref := cls.window_attrs[window.id()].client_ref)
-            and (self := self_ref())
-            and self.is_valid_for_view(view)
-        ):
+        if (window := view.window()) and (self := cls.window_attrs[window].client) and self.is_valid_for_view(view):
             return self
         return None