Use *args instead of an iterable for SigningKeypair.digest

nkms/crypto/keypairs.py

@@ -127,21 +127,18 @@ class SigningKeypair(object):
         s = int.from_bytes(sig[2], byteorder='big')
         return (v, r, s)
 
-    def digest(self, data: Iterable) -> bytes:
+    def digest(self, *args):
         """
         Accepts an iterable containing bytes and digests it.
 
-        :param Iterable data: Data to digest
+        :param bytes *args: Data to hash
 
         :rtype: bytes
         :return: bytestring of digested data
         """
-        if type(data) == bytes:
-            return sha3.keccak_256(data).digest()
-
         hash = sha3.keccak_256()
-        for item in data:
-            hash.update(item)
+        for arg in args:
+            hash.update(arg)
         return hash.digest()
 
     def sign(self, msghash):

tests/crypto/test_keypairs.py

@@ -33,18 +33,8 @@ class TestSigningKeypair(unittest.TestCase):
                                            pubkey=self.keypair_a.pub_key)
         self.assertTrue(verify_sig)
 
-    def test_digest_iterable(self):
-        # Test digest equality
-        digest_a = self.keypair_a.digest([self.msg])
-        digest_b = self.keypair_a.digest((self.msg))
-        digest_c = self.keypair_a.digest(self.msg)
-
-        self.assertEqual(digest_a, digest_b)
-        self.assertEqual(digest_a, digest_c)
-        self.assertEqual(digest_b, digest_c)
-
-        # Test multiple items
-        digest_a = self.keypair_a.digest([b'foo', b'bar'])
+    def test_digest(self):
+        digest_a = self.keypair_a.digest(b'foo', b'bar')
         digest_b = self.keypair_a.digest(b'foobar')
 
         self.assertEqual(digest_a, digest_b)