chore: better types for Type.rpc and Connection.rpc

Author: dimaqq

juju/client/connection.py

@@ -16,13 +16,14 @@
 import macaroonbakery.httpbakery as httpbakery
 import websockets
 from dateutil.parser import parse
-from typing_extensions import Self
+from typing_extensions import Self, 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 _JSON, _RICH_JSON, TypeEncoder
 from .facade_versions import client_facade_versions, known_unsupported_facades
 
 LEVELS = ["TRACE", "DEBUG", "INFO", "WARNING", "ERROR"]
@@ -536,7 +537,19 @@ async def _do_ping():
             log.debug("ping failed because of closed connection")
             pass
 
-    async def rpc(self, msg: dict, encoder=None) -> dict:
+    @overload
+    async def rpc(
+        self, msg: dict[str, _JSON], encoder: None = None
+    ) -> dict[str, _JSON]: ...
+
+    @overload
+    async def rpc(
+        self, msg: dict[str, _RICH_JSON], encoder: TypeEncoder
+    ) -> dict[str, _JSON]: ...
+
+    async def rpc(
+        self, msg: dict[str, Any], encoder: json.JSONEncoder | None = None
+    ) -> dict[str, Any]:
         """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).

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
+_JSON_LEAF: TypeAlias = None | bool | int | float | str
+_JSON: TypeAlias = "_JSON_LEAF|list[_JSON]|dict[str, _JSON]"
+
+# Type-enriched JSON, what can be sent to Juju
+_RICH_LEAF: TypeAlias = "_JSON_LEAF|Type"
+_RICH_JSON: TypeAlias = "_RICH_LEAF|list[_RICH_JSON]|dict[str, _RICH_JSON]"
+
 _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: _RICH_JSON) -> _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, _RICH_JSON]) -> _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

pyproject.toml

@@ -212,6 +212,7 @@ ignore = [
 [tool.pyright]
 # These are tentative
 # include = ["**/*.py"]
-pythonVersion = "3.8" # check no python > 3.8 features are used
-pythonPlatform = "All"
+pythonVersion = "3.10"
 typeCheckingMode = "strict"
+useLibraryCodeForTypes = true
+reportGeneralTypeIssues = true