Make KeyWallet support deepcopy

icon-project · Apr 26, 2023 · 25940c4 · 25940c4
1 parent 11f9f7b
commit 25940c4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/iconsdk/wallet/wallet.py b/iconsdk/wallet/wallet.py
@@ -188,6 +188,9 @@ def __eq__(self, other: KeyWallet) -> bool:
     def __ne__(self, other: KeyWallet) -> bool:
         return not self.__eq__(other)
 
+    def __deepcopy__(self, memodict={}) -> KeyWallet:
+        return KeyWallet.load(self.private_key)
+
 
 def public_key_to_address(public_key: bytes) -> str:
     if not (len(public_key) == 65 and public_key[0] == 4):

diff --git a/tests/wallet/test_wallet.py b/tests/wallet/test_wallet.py
@@ -1,3 +1,4 @@
+import copy
 from typing import Union, Dict
 
 import pytest
@@ -61,6 +62,14 @@ def test_to_dict(self):
         wallet2 = KeyWallet.from_dict(jso, password)
         assert wallet2 == wallet
 
+    def test_copy(self):
+        wallet = KeyWallet.create()
+        wallet2 = copy.deepcopy(wallet)
+        assert wallet == wallet2
+
+        wallet3 = copy.copy(wallet)
+        assert wallet == wallet3
+
 
 def test_public_key_to_address():
     wallet: KeyWallet = KeyWallet.create()