From 9bf004e55dfdc40c6733a2a3b195264b7700a5b5 Mon Sep 17 00:00:00 2001
From: andreasgriffin <andreasgriffin@proton.me>
Date: Tue, 17 Mar 2026 20:30:36 +0000
Subject: [PATCH] Refine storage typing with protocol

---
 bitcoin_safe_lib/storage.py | 19 ++++++++++++-------
 pyproject.toml              |  1 -
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/bitcoin_safe_lib/storage.py b/bitcoin_safe_lib/storage.py
index 65cdda7..8f0d36f 100644
--- a/bitcoin_safe_lib/storage.py
+++ b/bitcoin_safe_lib/storage.py
@@ -41,7 +41,7 @@
 from base64 import urlsafe_b64encode as b64e
 from collections.abc import Callable, Iterable
 from pathlib import Path
-from typing import Any, TypeAlias, TypeGuard, TypeVar
+from typing import Any, ClassVar, Protocol, TypeAlias, TypeGuard, TypeVar
 
 import bdkpython as bdk
 from cryptography.fernet import Fernet
@@ -56,10 +56,15 @@
 T = TypeVar("T", bound="BaseSaveableClass")
 ClassArgs: TypeAlias = dict[str, Any]  # noqa: UP040
 ClassKwargs: TypeAlias = dict[str, ClassArgs]  # noqa: UP040
-SaveableClass: TypeAlias = type["BaseSaveableClass"]  # noqa: UP040
-EnumClass: TypeAlias = type[enum.Enum]  # noqa: UP040
-KnownClass: TypeAlias = SaveableClass | EnumClass  # noqa: UP040
-KnownClasses: TypeAlias = dict[str, KnownClass]  # noqa: UP040
+KnownClasses: TypeAlias = dict[str, type[Any]]  # noqa: UP040
+
+
+class SupportsFromDumpClass(Protocol):
+    known_classes: ClassVar[KnownClasses]
+
+    @classmethod
+    def from_dump(cls, dct: dict[str, Any], class_kwargs: ClassKwargs | None = None) -> Any: ...
+
 
 logger = logging.getLogger(__name__)
 
@@ -165,11 +170,11 @@ def load(self, filename: str, password: str | None = None) -> str:
 
 class ClassSerializer:
     @staticmethod
-    def _is_saveable_class(obj_cls: KnownClass) -> TypeGuard[SaveableClass]:
+    def _is_saveable_class(obj_cls: type[Any]) -> TypeGuard[type[SupportsFromDumpClass]]:
         return issubclass(obj_cls, BaseSaveableClass)
 
     @staticmethod
-    def _is_enum_class(obj_cls: KnownClass) -> TypeGuard[EnumClass]:
+    def _is_enum_class(obj_cls: type[Any]) -> TypeGuard[type[enum.Enum]]:
         return issubclass(obj_cls, enum.Enum)
 
     @staticmethod
diff --git a/pyproject.toml b/pyproject.toml
index 8e41901..3584f33 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,4 +52,3 @@ known-first-party = ["bitcoin_safe_lib"]
 
 [tool.ruff.format]
 # Ruff formatter is Black-compatible; keep defaults or tweak here.
-