|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | from decimal import Decimal |
4 | | -from typing import TYPE_CHECKING |
| 4 | +from typing import TYPE_CHECKING, Optional |
5 | 5 |
|
6 | 6 | from eth_utils import to_normalized_address |
7 | 7 | from superfluid import CFA_V1, Operation, Web3FlowInfo |
8 | 8 |
|
9 | 9 | from aleph.sdk.exceptions import InsufficientFundsError |
10 | 10 |
|
11 | | -from ..evm_utils import get_super_token_address, to_human_readable_token, to_wei_token |
| 11 | +from ..evm_utils import ( |
| 12 | + FlowUpdate, |
| 13 | + from_wei_token, |
| 14 | + get_super_token_address, |
| 15 | + to_wei_token, |
| 16 | +) |
12 | 17 |
|
13 | 18 | if TYPE_CHECKING: |
14 | 19 | from aleph.sdk.chains.ethereum import ETHAccount |
@@ -52,7 +57,7 @@ def can_start_flow(self, flow: Decimal, block=True) -> bool: |
52 | 57 | if not valid and block: |
53 | 58 | raise InsufficientFundsError( |
54 | 59 | required_funds=float(MIN_FLOW_4H), |
55 | | - available_funds=to_human_readable_token(balance), |
| 60 | + available_funds=float(from_wei_token(balance)), |
56 | 61 | ) |
57 | 62 | return valid |
58 | 63 |
|
@@ -96,3 +101,51 @@ async def update_flow(self, receiver: str, flow: Decimal) -> str: |
96 | 101 | flow_rate=int(to_wei_token(flow)), |
97 | 102 | ), |
98 | 103 | ) |
| 104 | + |
| 105 | + async def manage_flow( |
| 106 | + self, |
| 107 | + receiver: str, |
| 108 | + flow: Decimal, |
| 109 | + update_type: FlowUpdate, |
| 110 | + ) -> Optional[str]: |
| 111 | + """ |
| 112 | + Update the flow of a Superfluid stream between a sender and receiver. |
| 113 | + This function either increases or decreases the flow rate between the sender and receiver, |
| 114 | + based on the update_type. If no flow exists and the update type is augmentation, it creates a new flow |
| 115 | + with the specified rate. If the update type is reduction and the reduction amount brings the flow to zero |
| 116 | + or below, the flow is deleted. |
| 117 | +
|
| 118 | + :param receiver: Address of the receiver in hexadecimal format. |
| 119 | + :param flow: The flow rate to be added or removed (in ether). |
| 120 | + :param update_type: The type of update to perform (augmentation or reduction). |
| 121 | + :return: The transaction hash of the executed operation (create, update, or delete flow). |
| 122 | + """ |
| 123 | + |
| 124 | + # Retrieve current flow info |
| 125 | + flow_info: Web3FlowInfo = await self.account.get_flow(receiver) |
| 126 | + |
| 127 | + current_flow_rate_wei: Decimal = Decimal(flow_info["flowRate"] or "0") |
| 128 | + flow_rate_wei: int = int(to_wei_token(flow)) |
| 129 | + |
| 130 | + if update_type == FlowUpdate.INCREASE: |
| 131 | + if current_flow_rate_wei > 0: |
| 132 | + # Update existing flow by increasing the rate |
| 133 | + new_flow_rate_wei = current_flow_rate_wei + flow_rate_wei |
| 134 | + new_flow_rate_ether = from_wei_token(new_flow_rate_wei) |
| 135 | + return await self.account.update_flow(receiver, new_flow_rate_ether) |
| 136 | + else: |
| 137 | + # Create a new flow if none exists |
| 138 | + return await self.account.create_flow(receiver, flow) |
| 139 | + else: |
| 140 | + if current_flow_rate_wei > 0: |
| 141 | + # Reduce the existing flow |
| 142 | + new_flow_rate_wei = current_flow_rate_wei - flow_rate_wei |
| 143 | + # Ensure to not leave infinitesimal flows |
| 144 | + # Often, there were 1-10 wei remaining in the flow rate, which prevented the flow from being deleted |
| 145 | + if new_flow_rate_wei > 99: |
| 146 | + new_flow_rate_ether = from_wei_token(new_flow_rate_wei) |
| 147 | + return await self.account.update_flow(receiver, new_flow_rate_ether) |
| 148 | + else: |
| 149 | + # Delete the flow if the new flow rate is zero or negative |
| 150 | + return await self.account.delete_flow(receiver) |
| 151 | + return None |
0 commit comments