fix type annotations and context manager signatures in persistent peerstore

Author: Winter-Soren

libp2p/peer/persistent/async_/peerstore.py

@@ -856,7 +856,10 @@ async def __aenter__(self) -> "AsyncPersistentPeerStore":
         return self
 
     async def __aexit__(
-        self, exc_type: type, exc_val: Exception, exc_tb: object
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
     ) -> None:
         """Async context manager exit."""
         await self.close_async()

libp2p/peer/persistent/datastore/leveldb_sync.py

@@ -184,11 +184,15 @@ def close(self) -> None:
                 self.db.close()
                 self.db = None
 
-    def __enter__(self):
+    def __enter__(self) -> "LevelDBDatastoreSync":
         """Context manager entry."""
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
         """Context manager exit."""
         self.close()
-        return False

libp2p/peer/persistent/datastore/rocksdb_sync.py

@@ -194,11 +194,15 @@ def close(self) -> None:
                 self.db.close()
                 self.db = None
 
-    def __enter__(self):
+    def __enter__(self) -> "RocksDBDatastoreSync":
         """Context manager entry."""
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
         """Context manager exit."""
         self.close()
-        return False

libp2p/peer/persistent/datastore/sqlite_sync.py

@@ -217,12 +217,16 @@ def close(self) -> None:
                     self.connection = None
                     self._closed = True
 
-    def __enter__(self):
+    def __enter__(self) -> "SQLiteDatastoreSync":
         """Context manager entry."""
         self._ensure_connection()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
         """Context manager exit."""
         self.close()
-        return False

libp2p/peer/persistent/serialization.py

@@ -13,6 +13,11 @@
 
 from libp2p.crypto.keys import KeyPair
 from libp2p.peer.envelope import Envelope
+from libp2p.peer.pb.crypto_pb2 import (
+    KeyType as PBKeyType,
+    PrivateKey as PBPrivateKey,
+    PublicKey as PBPublicKey,
+)
 from libp2p.peer.peerstore import PeerRecordState
 
 from .pb import (
@@ -151,24 +156,39 @@ def serialize_keypair(keypair: KeyPair) -> bytes:
 
         # Serialize public key
         if keypair.public_key:
-            from libp2p.peer.pb.crypto_pb2 import PublicKey as PBPublicKey
-
             pb_public_key = PBPublicKey()
 
-            # Map key types
             key_type = keypair.public_key.get_type()
-            pb_public_key.Type = key_type.value
+            # Map from libp2p KeyType to protobuf KeyType
+            if key_type.value == 0:  # RSA
+                pb_public_key.Type = PBKeyType.RSA
+            elif key_type.value == 1:  # Ed25519
+                pb_public_key.Type = PBKeyType.Ed25519
+            elif key_type.value == 2:  # Secp256k1
+                pb_public_key.Type = PBKeyType.Secp256k1
+            elif key_type.value == 3:  # ECDSA
+                pb_public_key.Type = PBKeyType.ECDSA
+            else:
+                raise SerializationError(f"Unsupported key type: {key_type}")
             pb_public_key.Data = keypair.public_key.serialize()
             pb_keys.public_key.CopyFrom(pb_public_key)
 
         # Serialize private key
         if keypair.private_key:
-            from libp2p.peer.pb.crypto_pb2 import PrivateKey as PBPrivateKey
-
             pb_private_key = PBPrivateKey()
 
             key_type = keypair.private_key.get_type()
-            pb_private_key.Type = key_type.value
+            # Map from libp2p KeyType to protobuf KeyType
+            if key_type.value == 0:  # RSA
+                pb_private_key.Type = PBKeyType.RSA
+            elif key_type.value == 1:  # Ed25519
+                pb_private_key.Type = PBKeyType.Ed25519
+            elif key_type.value == 2:  # Secp256k1
+                pb_private_key.Type = PBKeyType.Secp256k1
+            elif key_type.value == 3:  # ECDSA
+                pb_private_key.Type = PBKeyType.ECDSA
+            else:
+                raise SerializationError(f"Unsupported key type: {key_type}")
             pb_private_key.Data = keypair.private_key.serialize()
             pb_keys.private_key.CopyFrom(pb_private_key)
 

libp2p/peer/persistent/sync/peerstore.py

@@ -814,6 +814,11 @@ def __enter__(self) -> "SyncPersistentPeerStore":
         """Context manager entry."""
         return self
 
-    def __exit__(self, exc_type: type, exc_val: Exception, exc_tb: object) -> None:
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
         """Context manager exit."""
         self.close()