allow str pk in get_eth_addr (#12)

kigawas · Jun 2, 2020 · e4d2145 · e4d2145
1 parent d2c67ea
commit e4d2145
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/bip44/utils.py b/bip44/utils.py
@@ -1,3 +1,5 @@
+from typing import Union
+
 from coincurve import PublicKey
 from sha3 import keccak_256 as _keccak_256
 
@@ -11,8 +13,12 @@ def keccak_256(b: bytes) -> bytes:
     return h.digest()
 
 
-def get_eth_addr(pk: bytes) -> str:
+def get_eth_addr(pk: Union[str, bytes]) -> str:
     """Get ETH address from a public key."""
-    if len(pk) != 64:
-        pk = PublicKey(pk).format(False)[1:]
-    return f"0x{keccak_256(pk)[-20:].hex()}"
+
+    pk_bytes = bytes.fromhex(pk) if isinstance(pk, str) else pk
+
+    if len(pk_bytes) != 64:
+        pk_bytes = PublicKey(pk_bytes).format(False)[1:]
+
+    return f"0x{keccak_256(pk_bytes)[-20:].hex()}"
diff --git a/poetry.lock b/poetry.lock
diff --git a/tests/test_bip44.py b/tests/test_bip44.py
@@ -7,6 +7,12 @@
 MNEMONIC_JA = "なみだ　むろん　しひょう　こうつう　はかい　たいうん　さほう　ことり　げんき　おでかけ　ひこく　ざんしょ"
 
 
+def test_utils():
+    pk = "02d9ed78008e7b6c4bdc2beea13230fb3ccb8072728c0986894a3d544485e9b727"
+    assert get_eth_addr(pk) == get_eth_addr(bytes.fromhex(pk))
+    assert get_eth_addr(pk) == "0x7aD23D6eD9a1D98E240988BED0d78e8C81Ec296C".lower()
+
+
 def test_eth_wallet():
     w = Wallet(MNEMONIC)