Feat/rao changes (#9)

camfairchild · web-flow · commit 65e9d2d6bbbe · 2025-01-21T11:04:10.000-05:00
* update lib and pyi for NeuronInfo fields

* bump ver

* change ver to pre-release

* chore: ruff

* chore: ruff

* use tuple instead

* use older version of bittensor for testing

* add more logging to test and attr fix
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "bt_decode"
-version = "0.4.0"
+version = "0.5.0-a0"
 edition = "2021"
 
 [lib]
diff --git a/bt_decode.pyi b/bt_decode.pyi
@@ -60,6 +60,9 @@ class NeuronInfo:
     stake: List[
         Tuple[bytes, int]
     ]  # map of coldkey to stake on this neuron/hotkey (includes delegations)
+    total_stake: int
+    alpha_stake: int
+    tao_stake: int
     rank: int
     emission: int
     incentive: int
@@ -94,6 +97,9 @@ class NeuronInfoLite:
     stake: List[
         Tuple[bytes, int]
     ]  # map of coldkey to stake on this neuron/hotkey (includes delegations)
+    total_stake: int
+    alpha_stake: int
+    tao_stake: int
     rank: int
     emission: int
     incentive: int
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "bt-decode"
-version = "0.4.0"
+version = "0.5.0-a0"
 description = "A wrapper around the scale-codec crate for fast scale-decoding of Bittensor data structures."
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -41,7 +41,7 @@ build-backend = "maturin"
 [project.optional-dependencies]
 dev = ["black==23.7.0","maturin", "ruff==0.4.7"]
 test = [
-    "bittensor==8.2.0",
+    "bittensor==7.3.1",
     "pytest==7.2.0",
     "pytest-asyncio==0.23.7",
     "pytest-mock==3.12.0",
diff --git a/src/lib.rs b/src/lib.rs
@@ -98,6 +98,9 @@ mod bt_decode {
         axon_info: AxonInfo,
         prometheus_info: PrometheusInfo,
         stake: Vec<(AccountId, Compact<u64>)>, // map of coldkey to stake on this neuron/hotkey (includes delegations)
+        total_stake: Compact<u64>,
+        alpha_stake: Compact<u64>,
+        tao_stake: Compact<u64>,
         rank: Compact<u16>,
         emission: Compact<u64>,
         incentive: Compact<u16>,
@@ -127,6 +130,9 @@ mod bt_decode {
         axon_info: AxonInfo,
         prometheus_info: PrometheusInfo,
         stake: Vec<(AccountId, Compact<u64>)>, // map of coldkey to stake on this neuron/hotkey (includes delegations)
+        total_stake: Compact<u64>,
+        alpha_stake: Compact<u64>,
+        tao_stake: Compact<u64>,
         rank: Compact<u16>,
         emission: Compact<u64>,
         incentive: Compact<u16>,
diff --git a/tests/test_decode_by_type_string.py b/tests/test_decode_by_type_string.py
@@ -1,18 +1,8 @@
-from typing import Any, Callable, Dict, Tuple
+from typing import Any, Dict, Tuple
 
-import dataclasses
-import unittest
 import pytest
 
 import bt_decode
-import substrateinterface, scalecodec
-import bittensor
-
-from . import (
-    get_file_bytes,
-    fix_field as fix_field_fixes,
-    py_getattr as py_getattr_fixes,
-)
 
 
 TEST_TYPE_STRING_SCALE_INFO_DECODING: Dict[str, Tuple[str, Any]] = {
diff --git a/tests/test_decode_delegate_info.py b/tests/test_decode_delegate_info.py
@@ -25,18 +25,20 @@
     "take": lambda x: bittensor.U16_NORMALIZED_FLOAT(x),
     "nominators": lambda x: [(bittensor.u8_key_to_ss58(e[0]), e[1]) for e in x],
 }
-fix_field = lambda key, value, parent_key=None: fix_field_fixes(
-    FIELD_FIXES, key, value, parent_key
-)
+
+
+def fix_field(key, value, parent_key=None):
+    return fix_field_fixes(FIELD_FIXES, key, value, parent_key)
+
 
 ATTR_NAME_FIXES: Dict[str, str] = {
     # None
     "delegate_ss58": "hotkey_ss58",
 }
 
-py_getattr = lambda obj, attr, parent_name=None: py_getattr_fixes(
-    ATTR_NAME_FIXES, obj, attr, parent_name
-)
+
+def py_getattr(obj, attr, parent_name=None):
+    return py_getattr_fixes(ATTR_NAME_FIXES, obj, attr, parent_name)
 
 
 class TestDecodeDelegateInfo(unittest.TestCase):
diff --git a/tests/test_decode_neurons.py b/tests/test_decode_neurons.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional
+from typing import Callable, Dict, List
 
 import dataclasses
 import netaddr
@@ -44,18 +44,18 @@
     "weights": lambda x: [[e[0], e[1]] for e in x],
 }
 
-fix_field = lambda key, value, parent_key=None: fix_field_fixes(
-    FIELD_FIXES, key, value, parent_key
-)
+
+def fix_field(key, value, parent_key=None):
+    return fix_field_fixes(FIELD_FIXES, key, value, parent_key)
 
 
 ATTR_NAME_FIXES: Dict[str, str] = {
     # None
 }
 
-py_getattr = lambda obj, attr, parent_name=None: py_getattr_fixes(
-    ATTR_NAME_FIXES, obj, attr, parent_name
-)
+
+def py_getattr(obj, attr, parent_name=None):
+    return py_getattr_fixes(ATTR_NAME_FIXES, obj, attr, parent_name)
 
 
 class TestDecodeNeuronInfoLite(unittest.TestCase):
diff --git a/tests/test_decode_stake_info.py b/tests/test_decode_stake_info.py
@@ -7,7 +7,6 @@
 import bittensor
 
 from . import (
-    get_file_bytes,
     fix_field as fix_field_fixes,
     py_getattr as py_getattr_fixes,
 )
@@ -22,19 +21,21 @@
     "coldkey": lambda x: bittensor.u8_key_to_ss58(x),
     "hotkey": lambda x: bittensor.u8_key_to_ss58(x),
 }
-fix_field = lambda key, value, parent_key=None: fix_field_fixes(
-    FIELD_FIXES, key, value, parent_key
-)
+
+
+def fix_field(key, value, parent_key=None):
+    return fix_field_fixes(FIELD_FIXES, key, value, parent_key)
+
 
 ATTR_NAME_FIXES: Dict[str, str] = {
     # None
     "coldkey": "coldkey_ss58",
     "hotkey": "hotkey_ss58",
 }
 
-py_getattr = lambda obj, attr, parent_name=None: py_getattr_fixes(
-    ATTR_NAME_FIXES, obj, attr, parent_name
-)
+
+def py_getattr(obj, attr, parent_name=None):
+    return py_getattr_fixes(ATTR_NAME_FIXES, obj, attr, parent_name)
 
 
 class TestDecodeStakeInfo(unittest.TestCase):
diff --git a/tests/test_decode_subnet_hyp.py b/tests/test_decode_subnet_hyp.py
@@ -7,7 +7,6 @@
 import bittensor
 
 from . import (
-    get_file_bytes,
     fix_field as fix_field_fixes,
     py_getattr as py_getattr_fixes,
 )
@@ -21,15 +20,17 @@
 FIELD_FIXES: Dict[str, Callable] = {
     # None
 }
-fix_field = lambda key, value, parent_key=None: fix_field_fixes(
-    FIELD_FIXES, key, value, parent_key
-)
+
+
+def fix_field(key, value, parent_key=None):
+    return fix_field_fixes(FIELD_FIXES, key, value, parent_key)
+
 
 ATTR_NAME_FIXES: Dict[str, str] = {"max_weights_limit": "max_weight_limit"}
 
-py_getattr = lambda obj, attr, parent_name=None: py_getattr_fixes(
-    ATTR_NAME_FIXES, obj, attr, parent_name
-)
+
+def py_getattr(obj, attr, parent_name=None):
+    return py_getattr_fixes(ATTR_NAME_FIXES, obj, attr, parent_name)
 
 
 class TestDecodeSubnetHyperparameters(unittest.TestCase):
diff --git a/tests/test_decode_subnet_info.py b/tests/test_decode_subnet_info.py
@@ -26,9 +26,11 @@
     },
     "owner": lambda x: bittensor.u8_key_to_ss58(x),
 }
-fix_field = lambda key, value, parent_key=None: fix_field_fixes(
-    FIELD_FIXES, key, value, parent_key
-)
+
+
+def fix_field(key, value, parent_key=None):
+    return fix_field_fixes(FIELD_FIXES, key, value, parent_key)
+
 
 ATTR_NAME_FIXES: Dict[str, str] = {
     "emission_values": "emission_value",
@@ -37,11 +39,12 @@
     "network_connect": "connection_requirements",
     "network_modality": "modality",
     "owner": "owner_ss58",
+    "blocks_since_last_step": "blocks_since_epoch",
 }
 
-py_getattr = lambda obj, attr, parent_name=None: py_getattr_fixes(
-    ATTR_NAME_FIXES, obj, attr, parent_name
-)
+
+def py_getattr(obj, attr, parent_name=None):
+    return py_getattr_fixes(ATTR_NAME_FIXES, obj, attr, parent_name)
 
 
 class TestDecodeSubnetInfo(unittest.TestCase):
@@ -62,7 +65,13 @@ def test_decode_matches_python_impl(self):
             if not attr.startswith("__") and not callable(getattr(subnet_info, attr)):
                 attr_count += 1
 
-                attr_py = py_getattr(subnet_info_py, attr)
+                try:
+                    attr_py = py_getattr(subnet_info_py, attr)
+                except AttributeError as e:
+                    print(f"Error getting attribute {attr}: {e}")
+                    print(subnet_info_py)
+                    raise e
+
                 if dataclasses.is_dataclass(attr_py):
                     attr_rs = getattr(subnet_info, attr)
 
@@ -112,7 +121,7 @@ def test_decode_vec_option_matches_python_impl(self):
 
         subnet_info_list_py = (
             bittensor.SubnetInfo.list_from_vec_u8(  # Option specified internally
-                list(TEST_SUBNET_INFO_HEX["vec option normal"]())
+                TEST_SUBNET_INFO_HEX["vec option normal"]()
             )
         )
 
@@ -128,7 +137,13 @@ def test_decode_vec_option_matches_python_impl(self):
                 ):
                     attr_count += 1
 
-                    attr_py = py_getattr(subnet_info_py, attr)
+                    try:
+                        attr_py = py_getattr(subnet_info_py, attr)
+                    except AttributeError as e:
+                        print(f"Error getting attribute {attr}: {e}")
+                        print(subnet_info_py)
+                        raise e
+
                     if dataclasses.is_dataclass(attr_py):
                         attr_rs = getattr(subnet_info, attr)
 
diff --git a/tests/test_encode_by_type_string.py b/tests/test_encode_by_type_string.py
@@ -162,39 +162,39 @@
 }
 
 
-TEST_TYPE_STRING_PLAIN_DECODING: Dict[str, Tuple[str, Any]] = {
-    "bool": ("01", True),
-    "bool ": ("00", False),
-    "u8": ("01", 1),
-    "u16": ("0100", 1),
-    "u32": ("01000000", 1),
-    "u64": ("0100000000000000", 1),
-    "u128": ("01000000000000000000000000000000", 1),
-    "Compact<u8>": ("00", 0),
-    "Compact<u8> ": ("fd03", 2**8 - 1),
-    "Compact<u16>": ("feff0300", 2**16 - 1),
-    "Compact<u32>": ("03ffffffff", 2**32 - 1),
-    "Compact<u64>": ("13ffffffffffffffff", 2**64 - 1),
-    "Option<u8>": ("010c", 12),
-    "Option<u8>": ("00", None),
-    "Option<u32>": ("00", None),
-    "Option<u32> ": ("0101000000", 1), 
-    "()": ("", ()),
-    "[u8; 4]": ("62616265", (98, 97, 98, 101)),
-    "[u8; 4]": ("62616265", [98, 97, 98, 101]),
-    "Vec<u8>": ("0c010203", (1, 2, 3)),
-    "Vec<u8> ": ("00", []),
-    "Vec<u8> ": ("00", ()),
-    "(u8, u16) ": ("7bffff", (123, 2**16-1)),
-    "str": ("0c666f6f", "foo"),
-}
+TEST_TYPE_STRING_PLAIN_DECODING: list[tuple[str, tuple[str, Any]]] = [
+    ("bool", ("01", True)),
+    ("bool ", ("00", False)),
+    ("u8", ("01", 1)),
+    ("u16", ("0100", 1)),
+    ("u32", ("01000000", 1)),
+    ("u64", ("0100000000000000", 1)),
+    ("u128", ("01000000000000000000000000000000", 1)),
+    ("Compact<u8>", ("00", 0)),
+    ("Compact<u8> ", ("fd03", 2**8 - 1)),
+    ("Compact<u16>", ("feff0300", 2**16 - 1)),
+    ("Compact<u32>", ("03ffffffff", 2**32 - 1)),
+    ("Compact<u64>", ("13ffffffffffffffff", 2**64 - 1)),
+    ("Option<u8>", ("010c", 12)),
+    ("Option<u8>", ("00", None)),
+    ("Option<u32>", ("00", None)),
+    ("Option<u32> ", ("0101000000", 1)),
+    ("()", ("", ())),
+    ("[u8; 4]", ("62616265", (98, 97, 98, 101))),
+    ("[u8; 4]", ("62616265", [98, 97, 98, 101])),
+    ("Vec<u8>", ("0c010203", (1, 2, 3))),
+    ("Vec<u8> ", ("00", [])),
+    ("Vec<u8> ", ("00", ())),
+    ("(u8, u16) ", ("7bffff", (123, 2**16 - 1))),
+    ("str", ("0c666f6f", "foo")),
+]
 
 TEST_TYPES_JSON = "tests/test_types.json"
 
 
 @pytest.mark.parametrize(
     "type_string,test_hex,test_value",
-    [(x, y, z) for x, (y, z) in TEST_TYPE_STRING_PLAIN_DECODING.items()],
+    [(x, y, z) for x, (y, z) in TEST_TYPE_STRING_PLAIN_DECODING],
 )
 class TestEncodeByPlainTypeString:
     # Test combinations of human-readable type strings and hex-encoded values
