Skip to content

Commit

Permalink
better
Browse files Browse the repository at this point in the history
  • Loading branch information
altendky committed Jan 28, 2025
1 parent 9d8b2cb commit e5baa6d
Show file tree
Hide file tree
Showing 6 changed files with 76 additions and 30 deletions.
22 changes: 4 additions & 18 deletions crates/chia-datalayer/src/merkle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use chia_protocol::Bytes32;
use chia_py_streamable_macro::{PyJsonDict, PyStreamable};
use chia_sha2::Sha256;
use chia_streamable_macro::Streamable;
use chia_traits::{FromJsonDict, Streamable, ToJsonDict};
use chia_traits::Streamable;
use num_traits::ToBytes;
use std::cmp::Ordering;
use std::collections::{HashMap, HashSet, VecDeque};
Expand Down Expand Up @@ -337,7 +337,7 @@ fn internal_hash(left_hash: &Hash, right_hash: &Hash) -> Hash {

#[cfg_attr(
feature = "py-bindings",
pyclass(eq, eq_int, frozen),
// pyclass(eq, eq_int, frozen),
derive(PyJsonDict, PyStreamable)
)]
#[repr(u8)]
Expand All @@ -347,20 +347,6 @@ pub enum Side {
Right = 1,
}

// #[cfg(feature = "py-bindings")]
// impl FromJsonDict for Side {
// fn from_json_dict(o: &Bound<'_, PyAny>) -> PyResult<Self> {
// o.extract::<Side>()
// }
// }
//
// #[cfg(feature = "py-bindings")]
// impl ToJsonDict for Side {
// fn to_json_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
// Ok((*self as u8).into_pyobject(py)?.unbind().into_any())
// }
// }

#[cfg_attr(feature = "py-bindings", pyclass)]
#[derive(Clone, Debug, Hash, Eq, PartialEq)]
pub enum InsertLocation {
Expand Down Expand Up @@ -1459,7 +1445,7 @@ impl MerkleBlob {
value: ValueId,
hash: Hash,
reference_kid: Option<KeyId>,
side: Option<Side>,
side: Option<u8>,
) -> PyResult<()> {
let insert_location = match (reference_kid, side) {
(None, None) => InsertLocation::Auto {},
Expand All @@ -1471,7 +1457,7 @@ impl MerkleBlob {
.ok_or(PyValueError::new_err(format!(
"unknown key id passed as insert location reference: {key}"
)))?,
side,
side: Side::from_bytes(&[side])?,
},
_ => {
// TODO: use a specific error
Expand Down
5 changes: 1 addition & 4 deletions crates/chia_py_streamable_macro/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ fn maybe_upper_fields(py_uppercase: bool, fnames: Vec<Ident>) -> Vec<Ident> {
}
}

#[proc_macro_derive(PyStreamable, attributes(py_uppercase, py_pickle, py_enum))]
#[proc_macro_derive(PyStreamable, attributes(py_uppercase, py_pickle))]
pub fn py_streamable_macro(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
let found_crate = crate_name("chia-traits").expect("chia-traits is present in `Cargo.toml`");

Expand All @@ -34,14 +34,11 @@ pub fn py_streamable_macro(input: proc_macro::TokenStream) -> proc_macro::TokenS

let mut py_uppercase = false;
let mut py_pickle = false;
let mut py_enum = false;
for attr in &attrs {
if attr.path().is_ident("py_uppercase") {
py_uppercase = true;
} else if attr.path().is_ident("py_pickle") {
py_pickle = true;
} else if attr.path().is_ident("py_enum") {
py_enum = true;
}
}

Expand Down
4 changes: 2 additions & 2 deletions tests/test_datalayer.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import pytest

from chia_rs.datalayer import InvalidBlobLengthError, LeafNode, MerkleBlob, KeyId, ValueId, Side
from chia_rs.datalayer import InvalidBlobLengthError, LeafNode, MerkleBlob, KeyId, ValueId
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import int64, uint8

Expand Down Expand Up @@ -45,7 +45,7 @@ def test_checking_coverage() -> None:
merkle_blob.insert(KeyId(int64(i)), ValueId(int64(i)), bytes32.zeros)
else:
merkle_blob.insert(
KeyId(int64(i)), ValueId(int64(i)), bytes32.zeros, KeyId(int64(i - 1)), Side.Left
KeyId(int64(i)), ValueId(int64(i)), bytes32.zeros, KeyId(int64(i - 1)), uint8(0)
)

keys = {
Expand Down
71 changes: 66 additions & 5 deletions wheel/python/chia_rs/datalayer.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@ from typing_extensions import Self
from chia.types.blockchain_format.program import Program as ChiaProgram


# TODO: don't duplicate
ReadableBuffer = Union[bytes, bytearray, memoryview]

DATA_SIZE: int
BLOCK_SIZE: int
METADATA_SIZE: int
Expand All @@ -30,29 +33,87 @@ class IndexIsNotAChildError(Exception): ...
class CycleFoundError(Exception): ...
class BlockIndexOutOfBoundsError(Exception): ...

@final
class Side(Enum):
Left: int = ...
Right: int = ...
# TODO: not quite yet
# @final
# class Side(Enum):
# Left: int = ...
# Right: int = ...

@final
class KeyId:
raw: int64

def __init__(self, raw: int64) -> None: ...

# TODO: generate
def __hash__(self) -> int: ...
def __repr__(self) -> str: ...
def __deepcopy__(self, memo: object) -> Self: ...
def __copy__(self) -> Self: ...
@classmethod
def from_bytes(cls, blob: bytes) -> Self: ...
@classmethod
def from_bytes_unchecked(cls, blob: bytes) -> Self: ...
@classmethod
def parse_rust(cls, blob: ReadableBuffer, trusted: bool = False) -> tuple[Self, int]: ...
def to_bytes(self) -> bytes: ...
def __bytes__(self) -> bytes: ...
def stream_to_bytes(self) -> bytes: ...
def get_hash(self) -> bytes32: ...
def to_json_dict(self) -> int64: ...
@classmethod
def from_json_dict(cls, json_dict: int64) -> Self: ...

@final
class ValueId:
raw: int64

def __init__(self, raw: int64) -> None: ...

# TODO: generate
def __hash__(self) -> int: ...
def __repr__(self) -> str: ...
def __deepcopy__(self, memo: object) -> Self: ...
def __copy__(self) -> Self: ...
@classmethod
def from_bytes(cls, blob: bytes) -> Self: ...
@classmethod
def from_bytes_unchecked(cls, blob: bytes) -> Self: ...
@classmethod
def parse_rust(cls, blob: ReadableBuffer, trusted: bool = False) -> tuple[Self, int]: ...
def to_bytes(self) -> bytes: ...
def __bytes__(self) -> bytes: ...
def stream_to_bytes(self) -> bytes: ...
def get_hash(self) -> bytes32: ...
def to_json_dict(self) -> int64: ...
@classmethod
def from_json_dict(cls, json_dict: int64) -> Self: ...

@final
class TreeIndex:
raw: uint32

def __init__(self, raw: uint32) -> None: ...

# TODO: generate
def __hash__(self) -> int: ...
def __repr__(self) -> str: ...
def __deepcopy__(self, memo: object) -> Self: ...
def __copy__(self) -> Self: ...
@classmethod
def from_bytes(cls, blob: bytes) -> Self: ...
@classmethod
def from_bytes_unchecked(cls, blob: bytes) -> Self: ...
@classmethod
def parse_rust(cls, blob: ReadableBuffer, trusted: bool = False) -> tuple[Self, int]: ...
def to_bytes(self) -> bytes: ...
def __bytes__(self) -> bytes: ...
def stream_to_bytes(self) -> bytes: ...
def get_hash(self) -> bytes32: ...
def to_json_dict(self) -> uint32: ...
@classmethod
def from_json_dict(cls, json_dict: uint32) -> Self: ...

@final
class InternalNode:
def __init__(self, parent: Optional[TreeIndex], hash: bytes32, left: TreeIndex, right: TreeIndex) -> None: ...
Expand Down Expand Up @@ -99,7 +160,7 @@ class MerkleBlob:
blob: bytes,
) -> None: ...

def insert(self, key: KeyId, value: ValueId, hash: bytes32, reference_kid: Optional[KeyId] = None, side: Optional[Side] = None) -> None: ...
def insert(self, key: KeyId, value: ValueId, hash: bytes32, reference_kid: Optional[KeyId] = None, side: Optional[uint8] = None) -> None: ...
def upsert(self, key: KeyId, value: ValueId, new_hash: bytes32) -> None: ...
def delete(self, key: KeyId) -> None: ...
def get_raw_node(self, index: TreeIndex) -> Union[InternalNode, LeafNode]: ...
Expand Down
3 changes: 2 additions & 1 deletion wheel/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -648,7 +648,8 @@ pub fn add_datalayer_submodule(py: Python<'_>, parent: &Bound<'_, PyModule>) ->
datalayer.add_class::<LeafNode>()?;
datalayer.add_class::<KeyId>()?;
datalayer.add_class::<ValueId>()?;
datalayer.add_class::<Side>()?;
// TODO: not quite yet
// datalayer.add_class::<Side>()?;
datalayer.add_class::<TreeIndex>()?;

datalayer.add("BLOCK_SIZE", BLOCK_SIZE)?;
Expand Down
1 change: 1 addition & 0 deletions wheel/stubtest.allowlist
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ chia_rs\.chia_rs\..*
# this is offered to help with hinting only and is not intended to be
# runtime accessible. is there a better option for handling this?
chia_rs\.ReadableBuffer
chia_rs\.datalayer\.ReadableBuffer

# TODO: perhaps these should be private as _*
chia_rs\.BlockRecord\.ip_iters_impl
Expand Down

0 comments on commit e5baa6d

Please sign in to comment.