Tidy up

Author: lgarber-akamai

linode_api4/objects/linode.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import string
 import sys
@@ -7,7 +9,10 @@
 from enum import Enum
 from os import urandom
 from random import randint
-from typing import Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+if TYPE_CHECKING:
+    from linode_api4.objects.volume import Volume
 from urllib import parse
 
 from linode_api4.common import load_and_validate_keys
@@ -40,7 +45,11 @@
 from linode_api4.objects.serializable import JSONObject, StrEnum
 from linode_api4.objects.vpc import VPC, VPCSubnet
 from linode_api4.paginated_list import PaginatedList
-from linode_api4.util import drop_null_keys, generate_device_suffixes
+from linode_api4.util import (
+    drop_null_keys,
+    generate_device_suffixes,
+    normalize_as_list,
+)
 
 PASSWORD_CHARS = string.ascii_letters + string.digits + string.punctuation
 MIN_DEVICE_LIMIT = 8
@@ -1249,14 +1258,14 @@ def config_create(
         kernel: Kernel | str | None = None,
         label: str | None = None,
         devices: (
-            Union["Disk", "Volume", dict[str, Any]]
-            | list[Union["Disk", "Volume", dict[str, Any]]]
+            Disk
+            | Volume
+            | dict[str, Any]
+            | list[Disk | Volume | dict[str, Any]]
             | None
         ) = None,
-        disks: Union["Disk", int] | list[Union["Disk", int]] | None = None,
-        volumes: (
-            Union["Volume", int] | list[Union["Volume", int]] | None
-        ) = None,
+        disks: Disk | int | list[Disk | int] | None = None,
+        volumes: Volume | int | list[Volume | int] | None = None,
         interfaces: list[ConfigInterface | dict[str, Any]] | None = None,
         **kwargs,
     ) -> Config:
@@ -1269,12 +1278,13 @@ def config_create(
         :param label: The config label
         :param disks: The list of disks, starting at sda, to map to this config.
         :param volumes: The volumes, starting after the last disk, to map to this
-            config
+                        config.
         :param devices: A list of devices to assign to this config, in device
-            index order, a raw device mapping dict to pass directly to the API
-            (e.g. ``{"sda": {"disk_id": 123}, "sdb": Volume(...)}``), or
-            a single Disk or Volume.
-            If this is given, you may not include disks or volumes.
+                        index order, a raw device mapping dict to pass directly to the API
+                        (e.g. ``{"sda": {"disk_id": 123}, "sdb": Volume(...)}``), or
+                        a single Disk or Volume.
+                        If this is given, you may not include disks or volumes.
+        :param interfaces: A list of ConfigInterface objects or dicts to assign to this config.
         :param **kwargs: Any other arguments accepted by the api.
 
         :returns: A new Linode Config
@@ -1298,7 +1308,7 @@ def config_create(
             for suffix in generate_device_suffixes(device_limit)
         ]
 
-        def _flatten_device(device: Optional[Union[Disk, Volume]]):
+        def _flatten_device(device: Disk | Volume | None):
             if device is None:
                 return None
             elif isinstance(device, Disk):
@@ -1308,15 +1318,21 @@ def _flatten_device(device: Optional[Union[Disk, Volume]]):
 
             raise TypeError("Disk or Volume expected!")
 
-        def _device_entry(device: Union[Disk, Volume, int], key: str):
-            return (
-                {key: device if isinstance(device, int) else device.id}
-                if isinstance(device, int)
-                else _flatten_device(device)
-            )
+        def _device_entry(device: Disk | Volume | int, key: str):
+            if isinstance(device, (Disk, Volume)):
+                return _flatten_device(device)
+
+            try:
+                device_id = int(device)
+            except (TypeError, ValueError):
+                raise TypeError(
+                    "Disk, Volume, or integer ID expected!"
+                ) from None
+
+            return {key: device_id}
 
         def _build_devices():
-            # Devices is a dict, pass through
+            # Devices is a dict, flatten and pass through
             if isinstance(devices, dict):
                 return {
                     k: (
@@ -1329,33 +1345,22 @@ def _build_devices():
 
             device_list = []
 
-            # "Devices" was provided
             if devices:
-                devices_norm = (
-                    devices if isinstance(devices, (list, tuple)) else [devices]
-                )
-
                 device_list += [
-                    _flatten_device(device) for device in devices_norm
+                    _flatten_device(device)
+                    for device in normalize_as_list(disks)
                 ]
 
             if disks:
-                disks_norm = (
-                    disks if isinstance(disks, (list, tuple)) else [disks]
-                )
-
                 device_list += [
-                    _device_entry(disk, "disk_id") for disk in disks_norm
+                    _device_entry(disk, "disk_id")
+                    for disk in normalize_as_list(disks)
                 ]
 
             if volumes:
-                volumes_norm = (
-                    volumes if isinstance(volumes, (list, tuple)) else [volumes]
-                )
-
                 device_list += [
                     _device_entry(volume, "volume_id")
-                    for volume in volumes_norm
+                    for volume in normalize_as_list(disks)
                 ]
 
             return {

linode_api4/util.py

@@ -3,7 +3,7 @@
 """
 
 import string
-from typing import Any, Dict
+from typing import Any, Dict, List, Tuple, Union
 
 
 def drop_null_keys(data: Dict[Any, Any], recursive=True) -> Dict[Any, Any]:
@@ -30,6 +30,13 @@ def recursive_helper(value: Any) -> Any:
     return recursive_helper(data)
 
 
+def normalize_as_list(value: Any) -> Union[List, Tuple]:
+    """
+    Returns the value wrapped in a list if it isn't already a list or tuple.
+    """
+    return value if isinstance(value, (list, tuple)) else [value]
+
+
 def generate_device_suffixes(n: int) -> list[str]:
     """
     Generate n alphabetical suffixes starting with a, b, c, etc.