Merge pull request #280 from Chia-Network/json-binary

Make BLS types support bytes as input in from_json_dict()
Chia-Network · Oct 16, 2023 · 063857a · 063857a
2 parents 012fb16 + c4aa38d
commit 063857a
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 22 deletions.
diff --git a/chia-bls/src/public_key.rs b/chia-bls/src/public_key.rs
@@ -305,28 +305,45 @@ impl ToJsonDict for PublicKey {
 
 #[cfg(feature = "py-bindings")]
 pub fn parse_hex_string(o: &PyAny, len: usize, name: &str) -> PyResult<Vec<u8>> {
-    use pyo3::exceptions::PyValueError;
-    let s: String = o.extract()?;
-    let s = if let Some(st) = s.strip_prefix("0x") {
-        st
-    } else {
-        &s[..]
-    };
-    let buf = match hex::decode(s) {
-        Err(_) => {
-            return Err(PyValueError::new_err("invalid hex"));
+    use pyo3::exceptions::{PyTypeError, PyValueError};
+    if let Ok(s) = o.extract::<String>() {
+        let s = if let Some(st) = s.strip_prefix("0x") {
+            st
+        } else {
+            &s[..]
+        };
+        let buf = match hex::decode(s) {
+            Err(_) => {
+                return Err(PyValueError::new_err("invalid hex"));
+            }
+            Ok(v) => v,
+        };
+        if buf.len() != len {
+            Err(PyValueError::new_err(format!(
+                "{}, invalid length {} expected {}",
+                name,
+                buf.len(),
+                len
+            )))
+        } else {
+            Ok(buf)
+        }
+    } else if let Ok(buf) = o.extract::<Vec<u8>>() {
+        if buf.len() != len {
+            Err(PyValueError::new_err(format!(
+                "{}, invalid length {} expected {}",
+                name,
+                buf.len(),
+                len
+            )))
+        } else {
+            Ok(buf)
         }
-        Ok(v) => v,
-    };
-    if buf.len() != len {
-        Err(PyValueError::new_err(format!(
-            "{}, invalid length {} expected {}",
-            name,
-            buf.len(),
-            len
-        )))
     } else {
-        Ok(buf)
+        Err(PyTypeError::new_err(format!(
+            "invalid input type for {}",
+            name
+        )))
     }
 }
 

diff --git a/tests/test_streamable.py b/tests/test_streamable.py
@@ -367,6 +367,9 @@ def test_g1_element() -> None:
     d = G1Element.from_json_dict("0xa24d88ce995cea579675377728938eeb3956d5da608414efc9064774dc9653764edeb4823fc8da22c810917bf389c127")
     assert d == a
 
+    d = G1Element.from_json_dict(bytes.fromhex("a24d88ce995cea579675377728938eeb3956d5da608414efc9064774dc9653764edeb4823fc8da22c810917bf389c127"))
+    assert d == a
+
 def test_g2_element() -> None:
 
     a = G2Element.from_bytes(bytes.fromhex("a566b4d972db20765c668ce7fdcd76a4a5a8201dc2d5b1e747e2993fcdd99c8c96c1ca0503ade72809ae6d19c5e8400e10900a24ae56b7c9c84231ed5b7dd4c0790dd1aef56e0820e86994aa02c33bd409d3f17ace74c7fa40b00fe5022cc6d6"))
@@ -380,6 +383,9 @@ def test_g2_element() -> None:
     d = G2Element.from_json_dict("0xa566b4d972db20765c668ce7fdcd76a4a5a8201dc2d5b1e747e2993fcdd99c8c96c1ca0503ade72809ae6d19c5e8400e10900a24ae56b7c9c84231ed5b7dd4c0790dd1aef56e0820e86994aa02c33bd409d3f17ace74c7fa40b00fe5022cc6d6")
     assert a == d
 
+    d = G2Element.from_json_dict(bytes.fromhex("a566b4d972db20765c668ce7fdcd76a4a5a8201dc2d5b1e747e2993fcdd99c8c96c1ca0503ade72809ae6d19c5e8400e10900a24ae56b7c9c84231ed5b7dd4c0790dd1aef56e0820e86994aa02c33bd409d3f17ace74c7fa40b00fe5022cc6d6"))
+    assert a == d
+
 def test_program() -> None:
     p = Program.from_json_dict("0xff8080")
     assert str(p) == "Program(ff8080)"

diff --git a/wheel/chia_rs.pyi b/wheel/chia_rs.pyi
@@ -75,7 +75,7 @@ class AugSchemeMPL:
     @staticmethod
     def derive_child_sk_unhardened(pk: PrivateKey, index: int) -> PrivateKey: ...
     @staticmethod
-    def derive_child_pk_unhardened(pk: PublicKey, index: int) -> PublicKey: ...
+    def derive_child_pk_unhardened(pk: G1Element, index: int) -> G1Element: ...
 
 class G1Element:
     SIZE: ClassVar[int] = ...

diff --git a/wheel/generate_type_stubs.py b/wheel/generate_type_stubs.py
@@ -226,7 +226,7 @@ def derive_child_sk(pk: PrivateKey, index: int) -> PrivateKey: ...
     @staticmethod
     def derive_child_sk_unhardened(pk: PrivateKey, index: int) -> PrivateKey: ...
     @staticmethod
-    def derive_child_pk_unhardened(pk: PublicKey, index: int) -> PublicKey: ...
+    def derive_child_pk_unhardened(pk: G1Element, index: int) -> G1Element: ...
 """
     )