diff --git a/python/src/add_liquidity.py b/python/src/add_liquidity.py index 16fdefb..427d176 100644 --- a/python/src/add_liquidity.py +++ b/python/src/add_liquidity.py @@ -4,6 +4,7 @@ _to_raw_undo_rate_round_up, _compute_and_charge_aggregate_swap_fees, _get_single_input_index, + _require_unbalanced_liquidity_enabled ) from src.base_pool_math import ( compute_add_liquidity_unbalanced, @@ -48,6 +49,7 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_ updated_balances_live_scaled18[i] = a if add_liquidity_input["kind"] == Kind.UNBALANCED.value: + _require_unbalanced_liquidity_enabled(pool_state) amounts_in_scaled18 = max_amounts_in_scaled18 computed = compute_add_liquidity_unbalanced( updated_balances_live_scaled18, @@ -63,6 +65,7 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_ swap_fee_amounts_scaled18 = computed["swap_fee_amounts"] elif add_liquidity_input["kind"] == Kind.SINGLE_TOKEN_EXACT_OUT.value: + _require_unbalanced_liquidity_enabled(pool_state) token_index = _get_single_input_index(max_amounts_in_scaled18) amounts_in_scaled18 = max_amounts_in_scaled18 bpt_amount_out = add_liquidity_input["min_bpt_amount_out_raw"] diff --git a/python/src/remove_liquidity.py b/python/src/remove_liquidity.py index 776dab5..27d7eb2 100644 --- a/python/src/remove_liquidity.py +++ b/python/src/remove_liquidity.py @@ -4,6 +4,7 @@ _get_single_input_index, _to_raw_undo_rate_round_down, _compute_and_charge_aggregate_swap_fees, + _require_unbalanced_liquidity_enabled, ) from src.base_pool_math import ( compute_proportional_amounts_out, @@ -64,6 +65,7 @@ def remove_liquidity( remove_liquidity_input["max_bpt_amount_in_raw"], ) elif remove_liquidity_input["kind"] == RemoveKind.SINGLE_TOKEN_EXACT_IN.value: + _require_unbalanced_liquidity_enabled(pool_state) bpt_amount_in = remove_liquidity_input["max_bpt_amount_in_raw"] amounts_out_scaled18 = min_amounts_out_scaled18 token_out_index = _get_single_input_index( @@ -83,6 +85,7 @@ def remove_liquidity( amounts_out_scaled18[token_out_index] = computed["amount_out_with_fee"] swap_fee_amounts_scaled18 = computed["swap_fee_amounts"] elif remove_liquidity_input["kind"] == RemoveKind.SINGLE_TOKEN_EXACT_OUT.value: + _require_unbalanced_liquidity_enabled(pool_state) amounts_out_scaled18 = min_amounts_out_scaled18 token_out_index = _get_single_input_index( remove_liquidity_input["min_amounts_out_raw"] diff --git a/python/src/utils.py b/python/src/utils.py index 060c660..45fef8f 100644 --- a/python/src/utils.py +++ b/python/src/utils.py @@ -114,3 +114,8 @@ def _get_single_input_index(max_amounts_in): raise ValueError("All zero inputs for single token add") return input_index + + +def _require_unbalanced_liquidity_enabled(pool_state): + if not pool_state.get("supportsUnbalancedLiquidity", True): + raise ValueError("DoesNotSupportUnbalancedLiquidity")