chore: pull types from wip PR for upcoming Model.wait_for_idle

Author: dimaqq

juju/model.py

@@ -19,7 +19,7 @@
 from datetime import datetime, timedelta
 from functools import partial
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any, Literal, Mapping, overload
 
 import websockets
 import yaml
@@ -58,6 +58,13 @@
 from .url import URL, Schema
 from .version import DEFAULT_ARCHITECTURE
 
+if TYPE_CHECKING:
+    from .application import Application
+    from .machine import Machine
+    from .relation import Relation
+    from .remoteapplication import ApplicationOffer, RemoteApplication
+    from .unit import Unit
+
 log = logging.getLogger(__name__)
 
 
@@ -134,7 +141,35 @@ def __init__(self, model):
         self.model = model
         self.state = dict()
 
-    def _live_entity_map(self, entity_type):
+    @overload
+    def _live_entity_map(
+        self, entity_type: Literal["application"]
+    ) -> dict[str, Application]: ...
+
+    @overload
+    def _live_entity_map(
+        self, entity_type: Literal["applicationOffer"]
+    ) -> dict[str, ApplicationOffer]: ...
+
+    @overload
+    def _live_entity_map(
+        self, entity_type: Literal["machine"]
+    ) -> dict[str, Machine]: ...
+
+    @overload
+    def _live_entity_map(
+        self, entity_type: Literal["relation"]
+    ) -> dict[str, Relation]: ...
+
+    @overload
+    def _live_entity_map(
+        self, entity_type: Literal["remoteApplication"]
+    ) -> dict[str, RemoteApplication]: ...
+
+    @overload
+    def _live_entity_map(self, entity_type: Literal["unit"]) -> dict[str, Unit]: ...
+
+    def _live_entity_map(self, entity_type: str) -> Mapping[str, ModelEntity]:
         """Return an id:Entity map of all the living entities of
         type ``entity_type``.
 
@@ -146,51 +181,51 @@ def _live_entity_map(self, entity_type):
         }
 
     @property
-    def applications(self):
+    def applications(self) -> dict[str, Application]:
         """Return a map of application-name:Application for all applications
         currently in the model.
 
         """
         return self._live_entity_map("application")
 
     @property
-    def remote_applications(self):
+    def remote_applications(self) -> dict[str, RemoteApplication]:
         """Return a map of application-name:Application for all remote
         applications currently in the model.
 
         """
         return self._live_entity_map("remoteApplication")
 
     @property
-    def application_offers(self):
+    def application_offers(self) -> dict[str, ApplicationOffer]:
         """Return a map of application-name:Application for all applications
         offers currently in the model.
         """
         return self._live_entity_map("applicationOffer")
 
     @property
-    def machines(self):
+    def machines(self) -> dict[str, Machine]:
         """Return a map of machine-id:Machine for all machines currently in
         the model.
 
         """
         return self._live_entity_map("machine")
 
     @property
-    def units(self):
+    def units(self) -> dict[str, Unit]:
         """Return a map of unit-id:Unit for all units currently in
         the model.
 
         """
         return self._live_entity_map("unit")
 
     @property
-    def subordinate_units(self):
+    def subordinate_units(self) -> dict[str, Unit]:
         """Return a map of unit-id:Unit for all subordinate units"""
         return {u_name: u for u_name, u in self.units.items() if u.is_subordinate}
 
     @property
-    def relations(self):
+    def relations(self) -> dict[str, Relation]:
         """Return a map of relation-id:Relation for all relations currently in
         the model.
 
@@ -1110,54 +1145,54 @@ def tag(self):
         return tag.model(self.uuid)
 
     @property
-    def applications(self):
+    def applications(self) -> dict[str, Application]:
         """Return a map of application-name:Application for all applications
         currently in the model.
 
         """
         return self.state.applications
 
     @property
-    def remote_applications(self):
+    def remote_applications(self) -> dict[str, RemoteApplication]:
         """Return a map of application-name:Application for all remote
         applications currently in the model.
 
         """
         return self.state.remote_applications
 
     @property
-    def application_offers(self):
+    def application_offers(self) -> dict[str, ApplicationOffer]:
         """Return a map of application-name:Application for all applications
         offers currently in the model.
         """
         return self.state.application_offers
 
     @property
-    def machines(self):
+    def machines(self) -> dict[str, Machine]:
         """Return a map of machine-id:Machine for all machines currently in
         the model.
 
         """
         return self.state.machines
 
     @property
-    def units(self):
+    def units(self) -> dict[str, Unit]:
         """Return a map of unit-id:Unit for all units currently in
         the model.
 
         """
         return self.state.units
 
     @property
-    def subordinate_units(self):
+    def subordinate_units(self) -> dict[str, Unit]:
         """Return a map of unit-id:Unit for all subordiante units currently in
         the model.
 
         """
         return self.state.subordinate_units
 
     @property
-    def relations(self):
+    def relations(self) -> list[Relation]:
         """Return a list of all Relations currently in the model."""
         return list(self.state.relations.values())
 

juju/status.py

@@ -1,5 +1,6 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 import logging
 import warnings
@@ -8,26 +9,26 @@
 
 log = logging.getLogger(__name__)
 
-""" derive_status is used to determine the application status from a set of unit
-status values.
 
-:param statues: list of known unit workload statues
-
-"""
+def derive_status(statuses: list[str]):
+    """Derive status from a set.
 
+    derive_status is used to determine the application status from a set of unit
+    status values.
 
-def derive_status(statues):
+    :param statuses: list of known unit workload statuses
+    """
     current = "unknown"
-    for status in statues:
-        if status in severities and severities[status] > severities[current]:
+    for status in statuses:
+        if status in severity_map and severity_map[status] > severity_map[current]:
             current = status
     return current
 
 
-""" severities holds status values with a severity measure.
+""" severity_map holds status values with a severity measure.
 Status values with higher severity are used in preference to others.
 """
-severities = {
+severity_map = {
     "error": 100,
     "blocked": 90,
     "waiting": 80,

juju/unit.py

@@ -15,6 +15,8 @@
 
 
 class Unit(model.ModelEntity):
+    name: str
+
     @property
     def agent_status(self):
         """Returns the current agent status string."""

juju/url.py

@@ -1,5 +1,6 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 from enum import Enum
 from urllib.parse import urlparse
@@ -8,6 +9,8 @@
 
 
 class Schema(Enum):
+    """Charm URL schema kinds."""
+
     LOCAL = "local"
     CHARM_STORE = "cs"
     CHARM_HUB = "ch"
@@ -20,18 +23,23 @@ def __str__(self):
 
 
 class URL:
+    """Private URL class for this library internals only."""
+
+    name: str
+
     def __init__(
         self,
         schema,
         user=None,
-        name=None,
+        name: str | None = None,
         revision=None,
         series=None,
         architecture=None,
     ):
         self.schema = schema
         self.user = user
-        self.name = name
+        # the parse method will set the correct value later
+        self.name = name  # type: ignore
         self.series = series
 
         # 0 can be a valid revision, hence the more verbose check.
@@ -41,7 +49,7 @@ def __init__(
         self.architecture = architecture
 
     @staticmethod
-    def parse(s, default_store=Schema.CHARM_HUB):
+    def parse(s: str, default_store=Schema.CHARM_HUB) -> URL:
         """Parse parses the provided charm URL string into its respective
         structure.
 
@@ -103,7 +111,7 @@ def __str__(self):
         return f"{self.schema!s}:{self.path()}"
 
 
-def parse_v1_url(schema, u, s):
+def parse_v1_url(schema, u, s) -> URL:
     c = URL(schema)
 
     parts = u.path.split("/")
@@ -135,7 +143,7 @@ def parse_v1_url(schema, u, s):
     return c
 
 
-def parse_v2_url(u, s, default_store):
+def parse_v2_url(u, s, default_store) -> URL:
     if not u.scheme:
         c = URL(default_store)
     elif Schema.CHARM_HUB.matches(u.scheme):

pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
     "macaroonbakery>=1.1,<2.0",
     "pyRFC3339>=1.0,<2.0",
     "pyyaml>=5.1.2",
-    "websockets>=8.1,<14.0",
+    "websockets>=13.0.1,<14.0",
     "paramiko>=2.4.0",
     "pyasn1>=0.4.4",
     "toposort>=1.5,<2",

setup.py

@@ -23,7 +23,7 @@
         "macaroonbakery>=1.1,<2.0",
         "pyRFC3339>=1.0,<2.0",
         "pyyaml>=5.1.2",
-        "websockets>=8.1,<14.0",
+        "websockets>=13.0.1,<14.0",
         "paramiko>=2.4.0",
         "pyasn1>=0.4.4",
         "toposort>=1.5,<2",