diff --git a/cid/builder.py b/cid/builder.py index 188d5c3..f55c4c1 100644 --- a/cid/builder.py +++ b/cid/builder.py @@ -1,7 +1,6 @@ """Builder pattern for CID construction.""" from abc import ABC, abstractmethod -import hashlib from typing import TYPE_CHECKING import multihash @@ -59,8 +58,7 @@ def sum(self, data: bytes) -> "CIDv0": """ from .cid import CIDv0 - digest = hashlib.sha256(data).digest() - mhash = multihash.encode(digest, "sha2-256") + mhash = multihash.sum(data, "sha2-256").encode() return CIDv0(mhash) def get_codec(self) -> str: @@ -113,16 +111,8 @@ def sum(self, data: bytes) -> "CIDv1": """ from .cid import CIDv1 - if self.mh_type == "sha2-256": - digest = hashlib.sha256(data).digest() - elif self.mh_type == "sha2-512": - digest = hashlib.sha512(data).digest() - else: - msg = f"Hash type {self.mh_type} not fully implemented" - raise NotImplementedError(msg) - mh_length = None if self.mh_length == -1 else self.mh_length - mhash = multihash.encode(digest, self.mh_type, mh_length) + mhash = multihash.sum(data, self.mh_type, length=mh_length).encode() return CIDv1(self.codec, mhash) def get_codec(self) -> str: diff --git a/cid/prefix.py b/cid/prefix.py index 8ff42a1..837d9b9 100644 --- a/cid/prefix.py +++ b/cid/prefix.py @@ -1,6 +1,5 @@ """CID Prefix operations for creating CIDs from data.""" -import hashlib from typing import TYPE_CHECKING import multicodec @@ -111,22 +110,10 @@ def sum(self, data: bytes) -> "CIDv0 | CIDv1": :rtype: :py:class:`cid.CIDv0` or :py:class:`cid.CIDv1` :raises NotImplementedError: if hash type is not supported """ - # Hash data using mh_type - if self.mh_type == "sha2-256": - digest = hashlib.sha256(data).digest() - elif self.mh_type == "sha2-512": - digest = hashlib.sha512(data).digest() - else: - # Use multihash library for other types - # This is a simplified implementation - in practice, - # you'd want to support more hash types - msg = f"Hash type {self.mh_type} not fully implemented" - raise NotImplementedError(msg) - # Encode as multihash # Pass None if mh_length is -1 (default), otherwise use specified length mh_length = None if self.mh_length == -1 else self.mh_length - mhash = multihash.encode(digest, self.mh_type, mh_length) + mhash = multihash.sum(data, self.mh_type, length=mh_length).encode() # Create CID if self.version == 0: diff --git a/newsfragments/64.bugfix.rst b/newsfragments/64.bugfix.rst new file mode 100644 index 0000000..80e5c0b --- /dev/null +++ b/newsfragments/64.bugfix.rst @@ -0,0 +1 @@ +Support all standard hash types from multihash library in Prefix.sum(). diff --git a/tests/test_prefix.py b/tests/test_prefix.py index 20c92fb..c060338 100644 --- a/tests/test_prefix.py +++ b/tests/test_prefix.py @@ -114,6 +114,26 @@ def test_v1_factory(self): assert prefix.codec == "raw" assert prefix.mh_type == "sha2-512" + @pytest.mark.parametrize( + "mh_type", + [ + "sha2-256", + "sha2-512", + "sha3-256", + "sha3-512", + "blake2b-256", + "identity", + ], + ) + def test_prefix_sum_various_hash_types(self, mh_type): + """Prefix.sum: correctly creates CID from various hash types""" + prefix = Prefix.v1(codec="raw", mh_type=mh_type) + cid = prefix.sum(b"hello world") + assert isinstance(cid, CIDv1) + assert cid.codec == "raw" + expected_type = "id" if mh_type == "identity" else mh_type + assert cid.prefix().mh_type == expected_type + class TestCIDPrefix: """Tests for CID.prefix() method"""