chore: type hints for Connection

Author: dimaqq

juju/client/connection.py

@@ -1,5 +1,6 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 import base64
 import json
@@ -9,12 +10,13 @@
 import warnings
 import weakref
 from http.client import HTTPSConnection
-from typing import Dict, Literal, Optional, Sequence
+from typing import Any, Literal, Sequence
 
 import macaroonbakery.bakery as bakery
 import macaroonbakery.httpbakery as httpbakery
 import websockets
 from dateutil.parser import parse
+from typing_extensions import Self
 
 from juju import errors, jasyncio, tag, utils
 from juju.client import client
@@ -27,19 +29,6 @@
 log = logging.getLogger("juju.client.connection")
 
 
-def facade_versions(name, versions):
-    """facade_versions returns a new object that correctly returns a object in
-    format expected by the connection facades inspection.
-    :param name: name of the facade
-    :param versions: versions to support by the facade
-    """
-    if name.endswith("Facade"):
-        name = name[: -len("Facade")]
-    return {
-        name: {"versions": versions},
-    }
-
-
 class Monitor:
     """Monitor helper class for our Connection class.
 
@@ -59,7 +48,7 @@ class Monitor:
     DISCONNECTING = "disconnecting"
     DISCONNECTED = "disconnected"
 
-    def __init__(self, connection):
+    def __init__(self, connection: Connection):
         self.connection = weakref.ref(connection)
         self.reconnecting = jasyncio.Lock()
         self.close_called = jasyncio.Event()
@@ -117,28 +106,40 @@ class Connection:
 
     MAX_FRAME_SIZE = 2**22
     "Maximum size for a single frame.  Defaults to 4MB."
-    facades: Dict[str, int]
-    _specified_facades: Dict[str, Sequence[int]]
+    facades: dict[str, int]
+    _specified_facades: dict[str, Sequence[int]]
+    bakery_client: Any
+    usertag: str | None
+    password: str | None
+    name: str
+    __request_id__: int
+    endpoints: list[tuple[str, str]] | None  # Set by juju/controller.py
+    is_debug_log_connection: bool
+    monitor: Monitor
+    proxy: Any  # Need to find types for this library
+    max_frame_size: int
+    _retries: int
+    _retry_backoff: float
+    uuid: str | None
 
     @classmethod
     async def connect(
         cls,
         endpoint=None,
-        uuid=None,
-        username=None,
-        password=None,
+        uuid: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
         cacert=None,
         bakery_client=None,
-        max_frame_size=None,
+        max_frame_size: int | None = None,
         retries=3,
         retry_backoff=10,
-        specified_facades: Optional[
-            Dict[str, Dict[Literal["versions"], Sequence[int]]]
-        ] = None,
+        specified_facades: dict[str, dict[Literal["versions"], Sequence[int]]]
+        | None = None,
         proxy=None,
         debug_log_conn=None,
         debug_log_params={},
-    ):
+    ) -> Self:
         """Connect to the websocket.
 
         If uuid is None, the connection will be to the controller. Otherwise it
@@ -270,7 +271,7 @@ def ws(self):
         return self._ws
 
     @property
-    def username(self):
+    def username(self) -> str | None:
         if not self.usertag:
             return None
         return self.usertag[len("user-") :]
@@ -534,7 +535,7 @@ async def _do_ping():
             log.debug("ping failed because of closed connection")
             pass
 
-    async def rpc(self, msg, encoder=None):
+    async def rpc(self, msg: dict, encoder=None) -> dict:
         """Make an RPC to the API. The message is encoded as JSON
         using the given encoder if any.
         :param msg: Parameters for the call (will be encoded as JSON).
@@ -744,8 +745,13 @@ async def _try_endpoint(endpoint, cacert, delay):
             # only executed if inner loop's else did not continue
             # (i.e., inner loop did break due to successful connection)
             break
+        else:
+            # impossible, work around https://github.com/microsoft/pyright/issues/8791
+            assert False  # noqa: B011
+
         for task in tasks:
             task.cancel()
+
         self._ws = result[0]
         self.addr = result[1]
         self.endpoint = result[2]

juju/client/connector.py

@@ -1,8 +1,10 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 import copy
 import logging
+from typing import Any
 
 import macaroonbakery.httpbakery as httpbakery
 from packaging import version
@@ -33,9 +35,9 @@ class Connector:
 
     def __init__(
         self,
-        max_frame_size=None,
-        bakery_client=None,
-        jujudata=None,
+        max_frame_size: int | None = None,
+        bakery_client: Any | None = None,
+        jujudata: Any | None = None,
     ):
         """Initialize a connector that will use the given parameters
         by default when making a new connection
@@ -52,7 +54,7 @@ def is_connected(self):
         """Report whether there is a currently connected controller or not"""
         return self._connection is not None
 
-    def connection(self):
+    def connection(self) -> Connection:
         """Return the current connection; raises an exception if there
         is no current connection.
         """