From b5552724dfd5e7149fd07f5b26dc52d2bfbc5d55 Mon Sep 17 00:00:00 2001 From: Kyle Altendorf Date: Mon, 21 Oct 2024 09:58:14 -0400 Subject: [PATCH 1/4] add `bytes*.fill()` helper --- tests/merkle_set.py | 2 +- tests/test_merkle_set.py | 2 +- tests/test_sized_bytes.py | 37 ++++++++++++++++++++++++ wheel/python/chia_rs/sized_byte_class.py | 21 ++++++++++++++ 4 files changed, 60 insertions(+), 2 deletions(-) create mode 100644 tests/test_sized_bytes.py diff --git a/tests/merkle_set.py b/tests/merkle_set.py index e71ffad9e..75fa456b2 100644 --- a/tests/merkle_set.py +++ b/tests/merkle_set.py @@ -45,7 +45,7 @@ MIDDLE = bytes([2]) TRUNCATED = bytes([3]) -BLANK = bytes32([0] * 32) +BLANK = bytes32.zeros prehashed: Dict[bytes, _Hash] = {} diff --git a/tests/test_merkle_set.py b/tests/test_merkle_set.py index 08034a285..595abd987 100644 --- a/tests/test_merkle_set.py +++ b/tests/test_merkle_set.py @@ -66,7 +66,7 @@ def check_tree(leafs: List[bytes32]) -> None: ) for i in range(256): - item = bytes32([i] + [2] * 31) + item = bytes32.fill(i.to_bytes(), fill=b"\x02", align="<") py_included, py_proof = py_tree.is_included_already_hashed(item) assert not py_included ru_included, ru_proof = ru_tree.is_included_already_hashed(item) diff --git a/tests/test_sized_bytes.py b/tests/test_sized_bytes.py new file mode 100644 index 000000000..26c2484ea --- /dev/null +++ b/tests/test_sized_bytes.py @@ -0,0 +1,37 @@ +import pytest + +from chia_rs.sized_bytes import bytes8 + + +def test_fill_empty() -> None: + assert bytes8.fill(b"", b"\x01") == bytes8([1, 1, 1, 1, 1, 1, 1, 1]) + + +def test_fill_non_empty_with_single() -> None: + assert bytes8.fill(b"\x02", b"\x01") == bytes8([1, 1, 1, 1, 1, 1, 1, 2]) + + +def test_fill_non_empty_with_double() -> None: + assert bytes8.fill(b"\x02\x02", b"\x01\x01") == bytes8([1, 1, 1, 1, 1, 1, 2, 2]) + + +def test_fill_needed_with_0_length_fill_raises() -> None: + with pytest.raises(ValueError): + bytes8.fill(b"\x00", fill=b"") + + +def test_fill_not_needed_with_0_length_fill_works() -> None: + blob = b"\x00" * 8 + assert bytes8.fill(blob, fill=b"") == bytes8(blob) + + +def test_fill_not_multiple_raises() -> None: + with pytest.raises(ValueError): + bytes8.fill(b"\x00", fill=b"\x01\x01") + +def test_align_left() -> None: + assert bytes8.fill(b"\x01", fill=b"\x02", align="<") == bytes8([1, 2, 2, 2, 2, 2, 2, 2]) + +def test_invalid_alignment() -> None: + with pytest.raises(ValueError): + bytes8.fill(b"", fill=b"\x00", align="|") diff --git a/wheel/python/chia_rs/sized_byte_class.py b/wheel/python/chia_rs/sized_byte_class.py index 6dcbeb11e..8467b88ee 100644 --- a/wheel/python/chia_rs/sized_byte_class.py +++ b/wheel/python/chia_rs/sized_byte_class.py @@ -78,6 +78,27 @@ def random( def secret(cls: Type[_T_SizedBytes]) -> _T_SizedBytes: return cls.random(r=system_random) + @classmethod + def fill(cls: Type[_T_SizedBytes], blob: bytes, fill: bytes, align: Literal["<", ">"] = ">") -> _T_SizedBytes: + if len(blob) == cls._size: + return cls(blob) + + fill_length = len(fill) + if fill_length == 0: + raise ValueError("fill required but length is zero") + + div, mod = divmod(cls._size - len(blob), fill_length) + if mod != 0: + raise ValueError("invalid fill value, range to be filled must be multiple of fil size") + + all_fill = fill * div + if align == "<": + return cls(blob + all_fill) + elif align == ">": + return cls(all_fill + blob) + + raise ValueError(f"invalid alignment: {align!r}") + def __str__(self) -> str: return self.hex() From 072fc1e8a364b8587f55ffd347b6bb33c7b5afdd Mon Sep 17 00:00:00 2001 From: Kyle Altendorf Date: Mon, 21 Oct 2024 10:33:45 -0400 Subject: [PATCH 2/4] literal import --- wheel/python/chia_rs/sized_byte_class.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wheel/python/chia_rs/sized_byte_class.py b/wheel/python/chia_rs/sized_byte_class.py index 8467b88ee..d81d54497 100644 --- a/wheel/python/chia_rs/sized_byte_class.py +++ b/wheel/python/chia_rs/sized_byte_class.py @@ -5,6 +5,7 @@ from typing import ( BinaryIO, Iterable, + Literal, Optional, SupportsBytes, SupportsIndex, From 5452a51ed7eeff143f9d6887a0240c6a25d4d330 Mon Sep 17 00:00:00 2001 From: Kyle Altendorf Date: Mon, 21 Oct 2024 10:55:09 -0400 Subject: [PATCH 3/4] mypy --- tests/test_sized_bytes.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_sized_bytes.py b/tests/test_sized_bytes.py index 26c2484ea..a7b8aef13 100644 --- a/tests/test_sized_bytes.py +++ b/tests/test_sized_bytes.py @@ -34,4 +34,5 @@ def test_align_left() -> None: def test_invalid_alignment() -> None: with pytest.raises(ValueError): - bytes8.fill(b"", fill=b"\x00", align="|") + # type ignore since we are intentionally testing a bad case + bytes8.fill(b"", fill=b"\x00", align="|") # type: ignore[arg-type] From 35237b7c51a1cfc9a852b64cb3a80e8028d79a1d Mon Sep 17 00:00:00 2001 From: Kyle Altendorf Date: Mon, 21 Oct 2024 11:07:07 -0400 Subject: [PATCH 4/4] fixup --- tests/test_merkle_set.py | 2 +- tests/test_sized_bytes.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_merkle_set.py b/tests/test_merkle_set.py index 595abd987..3f42d58c3 100644 --- a/tests/test_merkle_set.py +++ b/tests/test_merkle_set.py @@ -66,7 +66,7 @@ def check_tree(leafs: List[bytes32]) -> None: ) for i in range(256): - item = bytes32.fill(i.to_bytes(), fill=b"\x02", align="<") + item = bytes32.fill(bytes([i]), fill=b"\x02", align="<") py_included, py_proof = py_tree.is_included_already_hashed(item) assert not py_included ru_included, ru_proof = ru_tree.is_included_already_hashed(item) diff --git a/tests/test_sized_bytes.py b/tests/test_sized_bytes.py index a7b8aef13..ec3576003 100644 --- a/tests/test_sized_bytes.py +++ b/tests/test_sized_bytes.py @@ -29,8 +29,12 @@ def test_fill_not_multiple_raises() -> None: with pytest.raises(ValueError): bytes8.fill(b"\x00", fill=b"\x01\x01") + def test_align_left() -> None: - assert bytes8.fill(b"\x01", fill=b"\x02", align="<") == bytes8([1, 2, 2, 2, 2, 2, 2, 2]) + assert bytes8.fill(b"\x01", fill=b"\x02", align="<") == bytes8( + [1, 2, 2, 2, 2, 2, 2, 2] + ) + def test_invalid_alignment() -> None: with pytest.raises(ValueError):