Merge pull request #1189 from dimaqq/chore-connnection-types

#1189

Pull the Connection type hints from #1097

Only minor cleanup beyond that, no logical changes.

Author: jujubot

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,37 +10,29 @@
 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, TypeAlias, overload
 
 from juju import errors, jasyncio, tag, utils
 from juju.client import client
 from juju.utils import IdQueue
 from juju.version import CLIENT_VERSION
 
+from .facade import TypeEncoder, _Json, _RichJson
 from .facade_versions import client_facade_versions, known_unsupported_facades
 
+SpecifiedFacades: TypeAlias = "dict[str, dict[Literal['versions'], Sequence[int]]]"
+_WebSocket: TypeAlias = "websockets.legacy.client.WebSocketClientProtocol"
+
 LEVELS = ["TRACE", "DEBUG", "INFO", "WARNING", "ERROR"]
 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 +52,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 +110,41 @@ 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
+    messages: IdQueue
+    _ws: _WebSocket | 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: SpecifiedFacades | 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 +276,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-") :]
@@ -299,7 +305,7 @@ def _get_ssl(self, cert=None):
             context.check_hostname = False
         return context
 
-    async def _open(self, endpoint, cacert):
+    async def _open(self, endpoint, cacert) -> tuple[_WebSocket, str, str, str]:
         if self.is_debug_log_connection:
             assert self.uuid
             url = f"wss://user-{self.username}:{self.password}@{endpoint}/model/{self.uuid}/log"
@@ -372,7 +378,7 @@ async def close(self, to_reconnect=False):
         if self.proxy is not None:
             self.proxy.close()
 
-    async def _recv(self, request_id):
+    async def _recv(self, request_id: int) -> dict[str, Any]:
         if not self.is_open:
             raise websockets.exceptions.ConnectionClosed(
                 websockets.frames.Close(
@@ -534,7 +540,19 @@ async def _do_ping():
             log.debug("ping failed because of closed connection")
             pass
 
-    async def rpc(self, msg, encoder=None):
+    @overload
+    async def rpc(
+        self, msg: dict[str, _Json], encoder: None = None
+    ) -> dict[str, _Json]: ...
+
+    @overload
+    async def rpc(
+        self, msg: dict[str, _RichJson], encoder: TypeEncoder
+    ) -> dict[str, _Json]: ...
+
+    async def rpc(
+        self, msg: dict[str, Any], encoder: json.JSONEncoder | None = None
+    ) -> dict[str, _Json]:
         """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).
@@ -710,7 +728,9 @@ async def _connect(self, endpoints):
         if len(endpoints) == 0:
             raise errors.JujuConnectionError("no endpoints to connect to")
 
-        async def _try_endpoint(endpoint, cacert, delay):
+        async def _try_endpoint(
+            endpoint, cacert, delay
+        ) -> tuple[_WebSocket, str, str, str]:
             if delay:
                 await jasyncio.sleep(delay)
             return await self._open(endpoint, cacert)
@@ -722,6 +742,8 @@ async def _try_endpoint(endpoint, cacert, delay):
             jasyncio.ensure_future(_try_endpoint(endpoint, cacert, 0.1 * i))
             for i, (endpoint, cacert) in enumerate(endpoints)
         ]
+        result: tuple[_WebSocket, str, str, str] | None = None
+
         for attempt in range(self._retries + 1):
             for task in jasyncio.as_completed(tasks):
                 try:
@@ -744,8 +766,12 @@ 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
+
         for task in tasks:
             task.cancel()
+
+        assert result  # loop raises or sets the result
+
         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.
         """

juju/client/facade.py

@@ -1,5 +1,6 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 import argparse
 import builtins
@@ -13,13 +14,22 @@
 from collections import defaultdict
 from glob import glob
 from pathlib import Path
-from typing import Any, Dict, List, Mapping, Sequence
+from typing import Any, Mapping, Sequence
 
 import packaging.version
 import typing_inspect
+from typing_extensions import TypeAlias
 
 from . import codegen
 
+# Plain JSON, what is received from Juju
+_JsonLeaf: TypeAlias = "None | bool | int | float | str"
+_Json: TypeAlias = "_JsonLeaf|list[_Json]|dict[str, _Json]"
+
+# Type-enriched JSON, what can be sent to Juju
+_RichLeaf: TypeAlias = "_JsonLeaf|Type"
+_RichJson: TypeAlias = "_RichLeaf|list[_RichJson]|dict[str, _RichJson]"
+
 _marker = object()
 
 JUJU_VERSION = re.compile(r"[0-9]+\.[0-9-]+[\.\-][0-9a-z]+(\.[0-9]+)?")
@@ -634,7 +644,7 @@ class {name}Facade(Type):
 
 
 class TypeEncoder(json.JSONEncoder):
-    def default(self, obj):
+    def default(self, obj: _RichJson) -> _Json:
         if isinstance(obj, Type):
             return obj.serialize()
         return json.JSONEncoder.default(self, obj)
@@ -653,7 +663,7 @@ def __eq__(self, other):
 
         return self.__dict__ == other.__dict__
 
-    async def rpc(self, msg):
+    async def rpc(self, msg: dict[str, _RichJson]) -> _Json:
         result = await self.connection.rpc(msg, encoder=TypeEncoder)
         return result
 
@@ -704,13 +714,13 @@ def _parse_nested_list_entry(expr, result_dict):
             return cls(**d)
         return None
 
-    def serialize(self):
+    def serialize(self) -> dict[str, _Json]:
         d = {}
         for attr, tgt in self._toSchema.items():
             d[tgt] = getattr(self, attr)
         return d
 
-    def to_json(self):
+    def to_json(self) -> str:
         return json.dumps(self.serialize(), cls=TypeEncoder, sort_keys=True)
 
     def __contains__(self, key):
@@ -917,8 +927,8 @@ def generate_definitions(schemas):
 
 
 def generate_facades(
-    schemas: Dict[str, List[Schema]],
-) -> Dict[str, Dict[int, codegen.Capture]]:
+    schemas: dict[str, list[Schema]],
+) -> dict[str, dict[int, codegen.Capture]]:
     captures = defaultdict(codegen.Capture)
 
     # Build the Facade classes

juju/client/overrides.py

@@ -1,8 +1,9 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 import re
-from collections import namedtuple
+from typing import Any, NamedTuple
 
 from . import _client, _definitions
 from .facade import ReturnMapping, Type, TypeEncoder
@@ -22,6 +23,12 @@
 ]
 
 
+class _Change(NamedTuple):
+    entity: str
+    type: str
+    data: dict[str, Any]
+
+
 class Delta(Type):
     """A single websocket delta.
 
@@ -42,12 +49,11 @@ class Delta(Type):
     _toSchema = {"deltas": "deltas"}
     _toPy = {"deltas": "deltas"}
 
-    def __init__(self, deltas=None):
+    def __init__(self, deltas: tuple[str, str, dict[str, Any]]):
         """:param deltas: [str, str, object]"""
         self.deltas = deltas
 
-        Change = namedtuple("Change", "entity type data")
-        change = Change(*self.deltas)
+        change = _Change(*self.deltas)
 
         self.entity = change.entity
         self.type = change.type

juju/delta.py

@@ -1,10 +1,12 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
-from .client import client
+from . import model
+from .client import client, overrides
 
 
-def get_entity_delta(d):
+def get_entity_delta(d: overrides.Delta):
     return _delta_types[d.entity](d.deltas)
 
 
@@ -13,12 +15,14 @@ def get_entity_class(entity_type):
 
 
 class EntityDelta(client.Delta):
-    def get_id(self):
+    data: dict[str, str]
+
+    def get_id(self) -> str:
         return self.data["id"]
 
     @classmethod
-    def get_entity_class(cls):
-        return None
+    def get_entity_class(cls) -> type[model.ModelEntity]:
+        raise NotImplementedError()
 
 
 class ActionDelta(EntityDelta):

juju/jasyncio.py

@@ -21,6 +21,7 @@
 )
 from asyncio import (
     CancelledError,
+    Task,
     create_task,
     wait,
 )
@@ -84,7 +85,7 @@
 ROOT_LOGGER = logging.getLogger()
 
 
-def create_task_with_handler(coro, task_name, logger=ROOT_LOGGER):
+def create_task_with_handler(coro, task_name, logger=ROOT_LOGGER) -> Task:
     """Wrapper around "asyncio.create_task" to make sure the task
     exceptions are handled properly.