refactor: update imports to use rapyuta_io_sdk_v2.Client and streamline method calls

Author: guptadev21

riocli/apply/parse.py

@@ -47,10 +47,10 @@
     from collections.abc import Callable, Iterable, Mapping
 
     from rapyuta_io import Client
+    from rapyuta_io_sdk_v2 import Client as v2Client
     from yaspin.api import Yaspin
 
     from riocli.config import Configuration
-    from riocli.v2client import Client as v2Client
 
 DEFAULT_MAX_WORKERS = 6
 DELETE_POLICY_LABEL = "rapyuta.io/deletionPolicy"
@@ -169,6 +169,9 @@ def _apply_manifest(
         assert v2_client is not None
         assert spinner is not None
 
+        if obj_key not in self.objects:
+            return
+
         obj = self.objects[obj_key]
         obj_key = click.style(obj_key, bold=True)
 
@@ -225,6 +228,9 @@ def _delete_manifest(
         assert v2_client is not None
         assert spinner is not None
 
+        if obj_key not in self.objects:
+            return
+
         obj = self.objects[obj_key]
         obj_key = click.style(obj_key, bold=True)
 
@@ -343,7 +349,9 @@ def _load_values_and_secrets(
 
         if value_files is not None:
             for v in value_files:
-                values.merge(self._load_value(v))
+                value = self._load_value(v)
+                if value is not None:
+                    values.merge(value)
 
         # The "rio" namespace exposes the commonly used fields from rio's
         # Configuration file.
@@ -358,7 +366,9 @@ def _load_values_and_secrets(
 
         if secret_files is not None:
             for s in secret_files:
-                secrets.merge(self._load_secret(s))
+                secret = self._load_secret(s)
+                if secret is not None:
+                    secrets.merge(secret)
 
         if "secrets" in values:
             benedict(values["secrets"]).merge(secrets.dict())

riocli/auth/util.py

@@ -20,15 +20,14 @@
 from rapyuta_io.clients.rip_client import AuthTokenLevel
 from rapyuta_io.utils import UnauthorizedError
 from rapyuta_io.utils.rest_client import HttpMethod, RestClient
+from rapyuta_io_sdk_v2 import Client as v2Client
+from rapyuta_io_sdk_v2.utils import handle_server_errors
 
 from riocli.config import Configuration
 from riocli.constants import Colors, Symbols
 from riocli.exceptions import OrganizationNotFound, ProjectNotFound
 from riocli.utils.selector import show_selection
 from riocli.utils.spinner import with_spinner
-from rapyuta_io_sdk_v2 import Client as v2Client
-from rapyuta_io_sdk_v2.utils import handle_server_errors
-from riocli.exceptions import ProjectNotFound, OrganizationNotFound
 
 TOKEN_LEVELS = {0: AuthTokenLevel.LOW, 1: AuthTokenLevel.MED, 2: AuthTokenLevel.HIGH}
 

riocli/configtree/revision.py

@@ -36,9 +36,8 @@
 from riocli.utils.state import StateFile
 
 if TYPE_CHECKING:
-    from yaspin.core import Yaspin
-
     from rapyuta_io_sdk_v2 import Client
+    from yaspin.core import Yaspin
 
 
 class Revision:

riocli/deployment/model.py

@@ -14,11 +14,11 @@
 import time
 
 from munch import Munch, unmunchify
+from rapyuta_io_sdk_v2 import Client
 from typing_extensions import override
 
 from riocli.constants import Status
 from riocli.model import Model
-from rapyuta_io_sdk_v2 import Client
 from riocli.utils.error import (
     RetriesExhausted,
 )

riocli/deployment/util.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import re
 import time
+from typing import Any
 
 import munch
 from rapyuta_io_sdk_v2 import Client
@@ -34,10 +35,11 @@ def fetch_deployments(
     client: Client,
     deployment_name_or_regex: str,
     include_all: bool,
-) -> list:
+) -> list[Any]:
     deployments = client.list_deployments(phases=DEFAULT_PHASES)
     result = []
-    for deployment in deployments.items:
+    items = deployments.items if deployments.items is not None else []
+    for deployment in items:
         if (
             include_all
             or deployment_name_or_regex == deployment.metadata.name
@@ -75,7 +77,7 @@ def poll_deployment(
     name: str,
     retry_count: int = 50,
     sleep_interval: int = 6,
-    ready_phases: list[str] = None,
+    ready_phases: list[str] | None = None,
 ):
     if ready_phases is None:
         ready_phases = []
@@ -84,18 +86,27 @@ def poll_deployment(
     status = deployment.status
 
     for _ in range(retry_count):
-        if status.phase in ready_phases:
+        if status is not None and status.phase in ready_phases:
             return deployment
 
-        if status.phase == DeploymentPhaseConstants.DeploymentPhaseProvisioning.value:
+        if (
+            status is not None
+            and status.phase == DeploymentPhaseConstants.DeploymentPhaseProvisioning.value
+        ):
             errors = status.get("error_codes", [])
             if "DEP_E153" in errors:
                 raise ImagePullError(
                     f"Deployment not running. Phase: Provisioning Status: {status.phase}"
                 )
-        elif status.phase == DeploymentPhaseConstants.DeploymentPhaseSucceeded.value:
+        elif (
+            status is not None
+            and status.phase == DeploymentPhaseConstants.DeploymentPhaseSucceeded.value
+        ):
             return deployment
-        elif status.phase == DeploymentPhaseConstants.DeploymentPhaseStopped.value:
+        elif (
+            status is not None
+            and status.phase == DeploymentPhaseConstants.DeploymentPhaseStopped.value
+        ):
             raise DeploymentNotRunning(
                 f"Deployment not running. Phase: Stopped  Status: {status.phase}"
             )
@@ -104,9 +115,12 @@ def poll_deployment(
         deployment = client.get_deployment(name=name)
         status = deployment.status
 
+    error_codes = []
+    if status is not None and hasattr(status, "error_codes"):
+        error_codes = status.error_codes or []
     msg = (
         f"Retries exhausted: Tried {retry_count} times with {sleep_interval}s interval. "
         f"Deployment: phase={status.phase} status={status.status} \n"
-        f"{process_errors(status.get('error_codes', []))}"
+        f"{process_errors(error_codes) if callable(process_errors) else process_errors}"
     )
     raise RetriesExhausted(msg)

riocli/device/util.py

@@ -314,7 +314,7 @@ def delete_hwil_device(spec: dict, metadata: dict) -> None:
     )
 
     client = new_hwil_client()
-    devices = client.list_devices(query={"name": device_name})
+    devices = client.list_devices(name=device_name)
     if not devices:
         return
 

riocli/disk/model.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 
 from munch import Munch, unmunchify
+from rapyuta_io_sdk_v2 import Client as v2Client
 from typing_extensions import override
 
+from riocli.disk.util import poll_disk
 from riocli.model import Model
-from rapyuta_io_sdk_v2 import Client as v2Client
 
 
 class Disk(Model):
@@ -29,8 +30,9 @@ def create_object(
         self, v2_client: v2Client, retry_count: int, retry_interval: int, *args, **kwargs
     ) -> Munch | None:
         created = v2_client.create_disk(unmunchify(self))  # pyright:ignore[reportArgumentType]
-        _ = v2_client.poll_disk(
-            created.metadata.name,
+        _ = poll_disk(
+            client=v2_client,
+            name=created.metadata.name,
             retry_count=retry_count,
             sleep_interval=retry_interval,
         )

riocli/model/base.py

@@ -17,13 +17,13 @@
 
 from munch import Munch
 from rapyuta_io import Client
+from rapyuta_io_sdk_v2 import Client as v2Client
+from rapyuta_io_sdk_v2.exceptions import HttpAlreadyExistsError, HttpNotFoundError
 
 from riocli.config.config import Configuration
 from riocli.constants import ApplyResult
 from riocli.exceptions import ResourceNotFound
 from riocli.jsonschema.validate import load_schema
-from riocli.v2client.client import Client as v2Client
-from riocli.v2client.error import HttpAlreadyExistsError, HttpNotFoundError
 
 DELETE_POLICY_LABEL = "rapyuta.io/deletionPolicy"
 

riocli/network/model.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 from munch import Munch, unmunchify
+from rapyuta_io_sdk_v2 import Client
 from typing_extensions import override
 
 from riocli.model import Model
-from rapyuta_io_sdk_v2 import Client
+
+from .util import poll_network
 
 
 class Network(Model):
@@ -29,8 +31,11 @@ def create_object(
         self, v2_client: Client, retry_count: int, retry_interval: int, *args, **kwargs
     ) -> Munch | None:
         network = v2_client.create_network(unmunchify(self))  # pyright:ignore[reportArgumentType]
-        _ = v2_client.poll_network(
-            network.metadata.name, retry_count=retry_count, sleep_interval=retry_interval
+        _ = poll_network(
+            client=v2_client,
+            name=network.metadata.name,  # pyright: ignore[reportOptionalMemberAccess]
+            retry_count=retry_count,
+            sleep_interval=retry_interval,
         )
 
     @override

riocli/organization/model.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 
 from munch import Munch, unmunchify
+from rapyuta_io_sdk_v2 import Client
 from typing_extensions import override
 
 from riocli.constants import ApplyResult
 from riocli.model import Model
-from riocli.v2client.client import Client
 
 
 class Organization(Model):