@@ -477,7 +477,7 @@ fn bittensor_wallet(module: Bound<'_, PyModule>) -> PyResult<()> {
477
477
module. add_class :: < PyKeyfile > ( ) ?;
478
478
module. add_class :: < PyKeypair > ( ) ?;
479
479
module. add_class :: < Wallet > ( ) ?;
480
-
480
+
481
481
// Add submodules to the main module
482
482
register_config_module ( & module) ?;
483
483
register_errors_module ( & module) ?;
@@ -1000,6 +1000,15 @@ except argparse.ArgumentError:
1000
1000
Ok ( PyKeypair { inner : keypair } )
1001
1001
}
1002
1002
1003
+ #[ pyo3( signature = ( password=None ) ) ]
1004
+ fn get_hotkeypub ( & self , password : Option < String > ) -> PyResult < PyKeypair > {
1005
+ let keypair = self
1006
+ . inner
1007
+ . get_hotkeypub ( password)
1008
+ . map_err ( |e| PyErr :: new :: < PyKeyFileError , _ > ( e) ) ?;
1009
+ Ok ( PyKeypair { inner : keypair } )
1010
+ }
1011
+
1003
1012
#[ pyo3( signature = ( keypair, encrypt=true , overwrite=false , save_coldkey_to_env=false , coldkey_password=None ) ) ]
1004
1013
fn set_coldkey (
1005
1014
& mut self ,
@@ -1052,6 +1061,18 @@ except argparse.ArgumentError:
1052
1061
. map_err ( |e| PyErr :: new :: < PyKeyFileError , _ > ( e) )
1053
1062
}
1054
1063
1064
+ #[ pyo3( signature = ( keypair, encrypt=false , overwrite=false ) ) ]
1065
+ fn set_hotkeypub (
1066
+ & mut self ,
1067
+ keypair : PyKeypair ,
1068
+ encrypt : bool ,
1069
+ overwrite : bool ,
1070
+ ) -> PyResult < ( ) > {
1071
+ self . inner
1072
+ . set_hotkeypub ( keypair. inner , encrypt, overwrite)
1073
+ . map_err ( |e| PyErr :: new :: < PyKeyFileError , _ > ( e) )
1074
+ }
1075
+
1055
1076
// Getters
1056
1077
#[ getter( coldkey) ]
1057
1078
fn coldkey_py_property ( & self ) -> PyResult < PyKeypair > {
@@ -1077,6 +1098,14 @@ except argparse.ArgumentError:
1077
1098
Ok ( PyKeypair { inner : keypair } )
1078
1099
}
1079
1100
1101
+ #[ getter( hotkeypub) ]
1102
+ fn hotkeypub_py_property ( & self ) -> PyResult < PyKeypair > {
1103
+ let keypair = self . inner . hotkeypub_property ( ) . map_err ( |e| {
1104
+ PyErr :: new :: < PyKeyFileError , _ > ( format ! ( "Failed to get hotkeypub: {:?}" , e) )
1105
+ } ) ?;
1106
+ Ok ( PyKeypair { inner : keypair } )
1107
+ }
1108
+
1080
1109
#[ getter]
1081
1110
fn coldkey_file ( & self ) -> PyResult < PyKeyfile > {
1082
1111
self . inner
@@ -1101,6 +1130,14 @@ except argparse.ArgumentError:
1101
1130
. map_err ( |e| PyErr :: new :: < PyKeyFileError , _ > ( e) )
1102
1131
}
1103
1132
1133
+ #[ getter]
1134
+ fn hotkeypub_file ( & self ) -> PyResult < PyKeyfile > {
1135
+ self . inner
1136
+ . hotkeypub_file ( )
1137
+ . map ( |inner| PyKeyfile { inner } )
1138
+ . map_err ( |e| PyErr :: new :: < PyKeyFileError , _ > ( e) )
1139
+ }
1140
+
1104
1141
#[ getter]
1105
1142
fn name ( & self ) -> String {
1106
1143
self . inner . get_name ( )
@@ -1218,6 +1255,20 @@ except argparse.ArgumentError:
1218
1255
} )
1219
1256
}
1220
1257
1258
+ #[ pyo3( text_signature = "($self)" ) ]
1259
+ fn unlock_hotkeypub ( & mut self ) -> PyResult < PyKeypair > {
1260
+ self . inner
1261
+ . unlock_hotkeypub ( )
1262
+ . map ( |inner| PyKeypair { inner } )
1263
+ . map_err ( |e| match e {
1264
+ KeyFileError :: DecryptionError ( _) => PyErr :: new :: < PyPasswordError , _ > ( format ! (
1265
+ "Decryption failed: {}" ,
1266
+ e. to_string( )
1267
+ ) ) ,
1268
+ _ => PyErr :: new :: < PyKeyFileError , _ > ( format ! ( "Failed to unlock hotkey: {:?}" , e) ) ,
1269
+ } )
1270
+ }
1271
+
1221
1272
#[ pyo3(
1222
1273
name = "create_new_coldkey" ,
1223
1274
signature = ( n_words=Some ( 12 ) , use_password=None , overwrite=None , suppress=None , save_coldkey_to_env=None , coldkey_password=None )
@@ -1377,4 +1428,21 @@ except argparse.ArgumentError:
1377
1428
inner : self . inner . clone ( ) ,
1378
1429
} )
1379
1430
}
1431
+
1432
+ #[ pyo3( signature = ( ss58_address=None , public_key=None , overwrite=None ) ) ]
1433
+ fn regenerate_hotkeypub (
1434
+ & mut self ,
1435
+ ss58_address : Option < String > ,
1436
+ public_key : Option < String > ,
1437
+ overwrite : Option < bool > ,
1438
+ ) -> PyResult < Self > {
1439
+ let new_inner_wallet = self
1440
+ . inner
1441
+ . regenerate_hotkeypub ( ss58_address, public_key, overwrite. unwrap_or ( false ) )
1442
+ . map_err ( |e| PyErr :: new :: < PyKeyFileError , _ > ( e) ) ?;
1443
+ self . inner = new_inner_wallet;
1444
+ Ok ( Wallet {
1445
+ inner : self . inner . clone ( ) ,
1446
+ } )
1447
+ }
1380
1448
}
0 commit comments