Add python wrappers for get_max_long, get_max_short

crates/hyperdrive-math-py/python/hyperdrive_math_py/hyperdrive_math_py.pyi

@@ -30,3 +30,37 @@ class HyperdriveState:
         str
             The spot price as a string representation of a solidity uint256 value.
         """
+    def get_max_long(self, budget: str, maybe_max_iterations: int) -> str:
+        """Gets the max amount of bonds that can be purchased for the given budget.
+
+        Arguments
+        ----------
+        budget : str
+            The account budget in base for making a long.
+        maybe_max_iterations : int
+            The number of iterations to use for the Newtonian method.
+
+        Returns
+        -------
+        str
+            The maximum long as a string representation of a solidity uint256 value.
+        """
+    def get_max_short(
+        self, budget: str, open_share_price: str, maybe_max_iterations: int
+    ) -> str:
+        """Gets the max amount of bonds that can be shorted for the given budget.
+
+        Arguments
+        ----------
+        budget : str
+            The account budget in base for making a short.
+        open_share_price : str
+            The share price of underlying vault.
+        maybe_max_iterations : int
+            The number of iterations to use for the Newtonian method.
+
+        Returns
+        -------
+        str
+            The maximum short as a string representation of a solidity uint256 value.
+        """

crates/hyperdrive-math-py/src/lib.rs

@@ -1,4 +1,5 @@
 use ethers::core::types::{Address, U256};
+use fixed_point::FixedPoint;
 use hyperdrive_wrappers::wrappers::i_hyperdrive::{Fees, PoolConfig, PoolInfo};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
@@ -149,6 +150,39 @@ impl HyperdriveState {
         let result = U256::from(result_fp).to_string();
         return Ok(result);
     }
+
+    pub fn get_max_long(
+        &self,
+        budget: &str,
+        maybe_max_iterations: Option<usize>,
+    ) -> PyResult<String> {
+        let budget_fp = FixedPoint::from(U256::from_dec_str(budget).map_err(|_| {
+            PyErr::new::<PyValueError, _>("Failed to convert budget string to U256")
+        })?);
+        let result_fp = self.state.get_max_long(budget_fp, maybe_max_iterations);
+        let result = U256::from(result_fp).to_string();
+        return Ok(result);
+    }
+
+    pub fn get_max_short(
+        &self,
+        budget: &str,
+        open_share_price: &str,
+        maybe_max_iterations: Option<usize>,
+    ) -> PyResult<String> {
+        let budget_fp = FixedPoint::from(U256::from_dec_str(budget).map_err(|_| {
+            PyErr::new::<PyValueError, _>("Failed to convert budget string to U256")
+        })?);
+        let open_share_price_fp =
+            FixedPoint::from(U256::from_dec_str(open_share_price).map_err(|_| {
+                PyErr::new::<PyValueError, _>("Failed to convert open_share_price string to U256")
+            })?);
+        let result_fp =
+            self.state
+                .get_max_short(budget_fp, open_share_price_fp, maybe_max_iterations);
+        let result = U256::from(result_fp).to_string();
+        return Ok(result);
+    }
 }
 
 /// A pyO3 wrapper for the hyperdrie_math crate.

python/test/test_hyperdrive_math_wrappers.py

@@ -1,4 +1,5 @@
 """Tests for hyperdrive_math.rs wrappers"""
+import pytest
 from hyperdrive_math_py import HyperdriveState
 from hyperdrive_math_py.types import Fees, PoolConfig, PoolInfo
 
@@ -18,18 +19,18 @@
 
 
 sample_pool_info = PoolInfo(
-    share_reserves="1000000000000000000",
-    bond_reserves="2000000000000000000",
-    lp_total_supply="3000000000000000000",
-    share_price="4000000000000000000",
-    longs_outstanding="5000000000000000000",
-    long_average_maturity_time="6000000000000000000",
-    shorts_outstanding="7000000000000000000",
-    short_average_maturity_time="8000000000000000000",
-    short_base_volume="9000000000000000000",
-    withdrawal_shares_ready_to_withdraw="2500000000000000000",
-    withdrawal_shares_proceeds="3500000000000000000",
-    lp_share_price="1000000000000000000",
+    share_reserves=str(int(1_000_000 * 1e18)),
+    bond_reserves=str(int(2_000_000 * 1e18)),
+    lp_total_supply=str(int(3_000_000 * 1e18)),
+    share_price=str(int(1e18)),
+    longs_outstanding="0",
+    long_average_maturity_time="0",
+    shorts_outstanding="0",
+    short_average_maturity_time="0",
+    short_base_volume="0",
+    withdrawal_shares_ready_to_withdraw="0",
+    withdrawal_shares_proceeds="0",
+    lp_share_price=str(int(1e18)),
 )
 
 
@@ -45,3 +46,77 @@ def test_get_spot_price():
     spot_price = state.get_spot_price()
     assert spot_price is not None, "Failed to get spot price."
     assert isinstance(spot_price, str), "Expected spot price to be a string."
+
+
+def test_max_long():
+    """test get_max_long."""
+    state = HyperdriveState(sample_pool_config, sample_pool_info)
+    budget = "1000000000000000000"  # 1 base
+    max_iterations = 20
+    max_long = state.get_max_long(budget, max_iterations)
+    expected_max_long = "1000000000000000000"  # 1 base
+    assert max_long == expected_max_long
+
+
+def test_max_long_fail_conversion():
+    """test get_max_long."""
+    state = HyperdriveState(sample_pool_config, sample_pool_info)
+    max_iterations = 20
+
+    # bad string
+    budget = "asdf"
+    with pytest.raises(ValueError, match="Failed to convert budget string to U256"):
+        state.get_max_long(budget, max_iterations)
+
+    # bad string
+    budget = "1.23"
+    with pytest.raises(ValueError, match="Failed to convert budget string to U256"):
+        state.get_max_long(budget, max_iterations)
+
+
+def test_max_short():
+    """test get_max_short."""
+    state = HyperdriveState(sample_pool_config, sample_pool_info)
+    budget = "10000000000000000000000"  # 10k base
+    open_share_price = "1000000000000000000"  # 1 base
+    max_iterations = 20
+    max_short = state.get_max_short(budget, open_share_price, max_iterations)
+    expected_max_short = "2583754033693357393077"  # apprx 2583 base
+    assert max_short == expected_max_short
+
+
+def test_max_short_fail_conversion():
+    """test get_max_short."""
+    state = HyperdriveState(sample_pool_config, sample_pool_info)
+    open_share_price = "1000000000000000000"  # 1 base
+    max_iterations = 20
+
+    # bad string
+    budget = "asdf"
+    with pytest.raises(ValueError, match="Failed to convert budget string to U256"):
+        state.get_max_short(budget, open_share_price, max_iterations)
+
+    # bad string
+    budget = "1.23"
+    with pytest.raises(ValueError, match="Failed to convert budget string to U256"):
+        state.get_max_short(budget, open_share_price, max_iterations)
+
+    budget = "10000000000000000000000"  # 10k base
+    # bad string
+    open_share_price = "asdf"
+    with pytest.raises(
+        ValueError, match="Failed to convert open_share_price string to U256"
+    ):
+        state.get_max_short(budget, open_share_price, max_iterations)
+
+
+def test_max_short_fail_budge():
+    """test get_max_short."""
+    state = HyperdriveState(sample_pool_config, sample_pool_info)
+    open_share_price = "1000000000000000000"  # 1 base
+    max_iterations = 20
+
+    # too small, max short requires too much
+    budget = "100000000000000000000"  # 100 base
+    with pytest.raises(BaseException, match="max short exceeded budget"):
+        state.get_max_short(budget, open_share_price, max_iterations)