Merge pull request #156 from opentensor/feat/roman/add-hotkeypub-functionality

Add `hotkeypub` to bittensor-wallet

Author: basfroman

Cargo.lock

@@ -52,6 +52,7 @@ class Wallet:
     def get_coldkey(self, password: Optional[str] = None) -> "Keypair": ...
     def get_coldkeypub(self, password: Optional[str] = None) -> "Keypair": ...
     def get_hotkey(self, password: Optional[str] = None) -> "Keypair": ...
+    def get_hotkeypub(self, password: Optional[str] = None) -> "Keypair": ...
     def set_coldkey(
         self,
         keypair: "Keypair",
@@ -74,19 +75,29 @@ class Wallet:
         save_hotkey_to_env: bool = False,
         hotkey_password: Optional[str] = None,
     ) -> None: ...
+    def set_hotkeypub(
+        self,
+        keypair: "Keypair",
+        encrypt: bool = False,
+        overwrite: bool = False,
+    ) -> None: ...
     @property
     def coldkey(self) -> "Keypair": ...
     @property
     def coldkeypub(self) -> "Keypair": ...
     @property
     def hotkey(self) -> "Keypair": ...
     @property
+    def hotkeypub(self) -> "Keypair": ...
+    @property
     def coldkey_file(self) -> "Keyfile": ...
     @property
     def coldkeypub_file(self) -> "Keyfile": ...
     @property
     def hotkey_file(self) -> "Keyfile": ...
     @property
+    def hotkeypub_file(self) -> "Keyfile": ...
+    @property
     def name(self) -> str: ...
     @property
     def path(self) -> str: ...
@@ -113,6 +124,7 @@ class Wallet:
     def unlock_coldkey(self) -> "Keypair": ...
     def unlock_coldkeypub(self) -> "Keypair": ...
     def unlock_hotkey(self) -> "Keypair": ...
+    def unlock_hotkeypub(self) -> "Keypair": ...
     def new_coldkey(
         self,
         n_words: Optional[int] = 12,
@@ -177,3 +189,9 @@ class Wallet:
         save_hotkey_to_env: Optional[bool] = False,
         hotkey_password: Optional[str] = None,
     ) -> "Wallet": ...
+    def regenerate_hotkeypub(
+        self,
+        ss58_address: Optional[str] = None,
+        public_key: Optional[bytes] = None,
+        overwrite: Optional[bool] = False,
+    ) -> "Wallet": ...

bittensor_wallet/wallet/__init__.pyi

@@ -477,7 +477,7 @@ fn bittensor_wallet(module: Bound<'_, PyModule>) -> PyResult<()> {
     module.add_class::<PyKeyfile>()?;
     module.add_class::<PyKeypair>()?;
     module.add_class::<Wallet>()?;
-    
+
     // Add submodules to the main module
     register_config_module(&module)?;
     register_errors_module(&module)?;
@@ -1000,6 +1000,15 @@ except argparse.ArgumentError:
         Ok(PyKeypair { inner: keypair })
     }
 
+    #[pyo3(signature = (password=None))]
+    fn get_hotkeypub(&self, password: Option<String>) -> PyResult<PyKeypair> {
+        let keypair = self
+            .inner
+            .get_hotkeypub(password)
+            .map_err(|e| PyErr::new::<PyKeyFileError, _>(e))?;
+        Ok(PyKeypair { inner: keypair })
+    }
+
     #[pyo3(signature = (keypair, encrypt=true, overwrite=false, save_coldkey_to_env=false, coldkey_password=None))]
     fn set_coldkey(
         &mut self,
@@ -1052,6 +1061,18 @@ except argparse.ArgumentError:
             .map_err(|e| PyErr::new::<PyKeyFileError, _>(e))
     }
 
+    #[pyo3(signature = (keypair, encrypt=false, overwrite=false))]
+    fn set_hotkeypub(
+        &mut self,
+        keypair: PyKeypair,
+        encrypt: bool,
+        overwrite: bool,
+    ) -> PyResult<()> {
+        self.inner
+            .set_hotkeypub(keypair.inner, encrypt, overwrite)
+            .map_err(|e| PyErr::new::<PyKeyFileError, _>(e))
+    }
+
     // Getters
     #[getter(coldkey)]
     fn coldkey_py_property(&self) -> PyResult<PyKeypair> {
@@ -1077,6 +1098,14 @@ except argparse.ArgumentError:
         Ok(PyKeypair { inner: keypair })
     }
 
+    #[getter(hotkeypub)]
+    fn hotkeypub_py_property(&self) -> PyResult<PyKeypair> {
+        let keypair = self.inner.hotkeypub_property().map_err(|e| {
+            PyErr::new::<PyKeyFileError, _>(format!("Failed to get hotkeypub: {:?}", e))
+        })?;
+        Ok(PyKeypair { inner: keypair })
+    }
+
     #[getter]
     fn coldkey_file(&self) -> PyResult<PyKeyfile> {
         self.inner
@@ -1101,6 +1130,14 @@ except argparse.ArgumentError:
             .map_err(|e| PyErr::new::<PyKeyFileError, _>(e))
     }
 
+    #[getter]
+    fn hotkeypub_file(&self) -> PyResult<PyKeyfile> {
+        self.inner
+            .hotkeypub_file()
+            .map(|inner| PyKeyfile { inner })
+            .map_err(|e| PyErr::new::<PyKeyFileError, _>(e))
+    }
+
     #[getter]
     fn name(&self) -> String {
         self.inner.get_name()
@@ -1218,6 +1255,20 @@ except argparse.ArgumentError:
             })
     }
 
+    #[pyo3(text_signature = "($self)")]
+    fn unlock_hotkeypub(&mut self) -> PyResult<PyKeypair> {
+        self.inner
+            .unlock_hotkeypub()
+            .map(|inner| PyKeypair { inner })
+            .map_err(|e| match e {
+                KeyFileError::DecryptionError(_) => PyErr::new::<PyPasswordError, _>(format!(
+                    "Decryption failed: {}",
+                    e.to_string()
+                )),
+                _ => PyErr::new::<PyKeyFileError, _>(format!("Failed to unlock hotkey: {:?}", e)),
+            })
+    }
+
     #[pyo3(
         name = "create_new_coldkey",
         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:
             inner: self.inner.clone(),
         })
     }
+
+    #[pyo3(signature = (ss58_address=None, public_key=None, overwrite=None))]
+    fn regenerate_hotkeypub(
+        &mut self,
+        ss58_address: Option<String>,
+        public_key: Option<String>,
+        overwrite: Option<bool>,
+    ) -> PyResult<Self> {
+        let new_inner_wallet = self
+            .inner
+            .regenerate_hotkeypub(ss58_address, public_key, overwrite.unwrap_or(false))
+            .map_err(|e| PyErr::new::<PyKeyFileError, _>(e))?;
+        self.inner = new_inner_wallet;
+        Ok(Wallet {
+            inner: self.inner.clone(),
+        })
+    }
 }