refactor: better type hints for vault (#20)

Author: NTFSvolume

pyproject.toml

@@ -34,7 +34,7 @@ license = "Apache-2.0"
 license-files = ["LICENSE"]
 readme = "README.md"
 requires-python = ">=3.11"
-version = "2.3.0"
+version = "2.3.1"
 
 [project.optional-dependencies]
 cli = [

src/mega/data_structures.py

@@ -29,7 +29,6 @@
 
 NodeKey = tuple[int, int, int, int]
 ComposedFileKey = tuple[int, int, int, int, int, int, int, int]
-ComposedFolderKey = tuple[int, int, int, int]
 
 
 class ByteSize(int):
@@ -182,7 +181,7 @@ class Crypto(_DictDumper):
     iv: tuple[int, int]
     meta_mac: tuple[int, int]
 
-    full_key: ComposedFileKey | ComposedFolderKey
+    full_key: ComposedFileKey | NodeKey
     share_key: NodeKey | None
 
     @classmethod
@@ -206,19 +205,18 @@ def compose(
                 raise RuntimeError(f"Invalid key, expected key len for files is 8, got {len(composed_key) = }")
 
         else:
-            if len(key) != 4:
-                raise RuntimeError(f"Invalid key, expected key len for folders is 4, got {len(key) = }")
-
             iv = ()
             meta_mac = ()
             composed_key = key
+            if len(composed_key) != 4:
+                raise RuntimeError(f"Invalid key, expected key len for folders is 4, got {len(composed_key) = }")
 
         return Crypto(key, iv, meta_mac, composed_key, None)  # pyright: ignore[reportArgumentType]
 
     @classmethod
     def decompose(
         cls,
-        composed_key: ComposedFileKey | ComposedFolderKey,
+        composed_key: ComposedFileKey | NodeKey,
         node_type: NodeType = NodeType.FILE,
         share_key: NodeKey | None = None,
     ) -> Crypto:

src/mega/vault.py

@@ -1,17 +1,17 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING, Final, cast
+from typing import TYPE_CHECKING, Final, NamedTuple, cast
 
 from mega.crypto import b64_to_a32, decrypt_key
 
 if TYPE_CHECKING:
     from mega.data_structures import (
         ComposedFileKey,
-        ComposedFolderKey,
         GetNodesResponse,
         Node,
         NodeID,
+        NodeKey,
         SharedKeys,
         UserID,
     )
@@ -21,6 +21,11 @@
 # An special owner used for exported (AKA public) file/folders
 
 
+class KeyLookup(NamedTuple):
+    full_key: ComposedFileKey | NodeKey
+    share_key: NodeKey | None
+
+
 @dataclasses.dataclass(slots=True, frozen=True)
 class MegaVault:
     master_key: tuple[int, ...] = ()
@@ -43,7 +48,7 @@ def update(self, api_resp: GetNodesResponse) -> None:
             if key := new_keys.get(node_id):
                 self._shared_keys.setdefault(owner, {})[node_id] = key
 
-    def __getitem__(self, node: Node) -> tuple[ComposedFileKey | ComposedFolderKey, tuple[int, ...] | None]:
+    def __getitem__(self, node: Node) -> KeyLookup:
         share_key: tuple[int, ...] | None = None
 
         # my files/folders
@@ -74,7 +79,7 @@ def __getitem__(self, node: Node) -> tuple[ComposedFileKey | ComposedFolderKey,
         if len(full_key) not in (4, 8):
             raise RuntimeError(f"Invalid key len found for node {node.id} ({len(full_key)})")
 
-        return cast("ComposedFileKey | ComposedFolderKey", full_key), share_key
+        return KeyLookup(cast("ComposedFileKey | NodeKey", full_key), share_key)
 
     def save_public_key(self, node_id: NodeID, share_key: tuple[int, ...]) -> None:
         self._shared_keys[_EXPORTED][node_id] = share_key

tests/test_deserialization.py

@@ -91,7 +91,9 @@ def test_attributes_serialization(attrs: AttributesSerialized, extras: tuple[str
 
 def test_storage_metrics() -> None:
     metrics = StorageMetrics.parse([35342640609, 2088, 33, 0, 0])
-    assert metrics
+    assert metrics.bytes_used == 35342640609
+    assert metrics.files == 2088
+    assert metrics.folders == 33
 
 
 def test_storage_quota() -> None: