(feat) Improved the logic in the gas limit estimator component to consider the cost of each included element for batch messages

Author: abel

pyinjective/core/broadcaster.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from decimal import Decimal
-from typing import Callable, List, Optional, Tuple
+from typing import List, Optional
 
 import math
 from google.protobuf import any_pb2
@@ -9,7 +9,7 @@
 from pyinjective.async_client import AsyncClient
 from pyinjective.composer import Composer
 from pyinjective.constant import Network
-from pyinjective.proto.cosmos.authz.v1beta1 import tx_pb2 as cosmos_authz_tx_pb
+from pyinjective.core.gas_limit_estimator import GasLimitEstimator
 
 
 class BroadcasterAccountConfig(ABC):
@@ -287,52 +287,25 @@ async def configure_gas_fee_for_transaction(
 
 
 class MessageBasedTransactionFeeCalculator(TransactionFeeCalculator):
-    DEFAULT_GAS_LIMIT = 150_000
-    DEFAULT_EXCHANGE_GAS_LIMIT = 100_000
+    TRANSACTION_GAS_LIMIT = 60_000
 
     def __init__(
             self,
             client: AsyncClient,
             composer: Composer,
-            gas_price: Optional[int] = None,
-            base_gas_limit: Optional[int] = None,
-            base_exchange_gas_limit: Optional[int] = None):
+            gas_price: Optional[int] = None):
         self._client = client
         self._composer = composer
         self._gas_price = gas_price or self.DEFAULT_GAS_PRICE
-        self._base_gas_limit = base_gas_limit or self.DEFAULT_GAS_LIMIT
-        self._base_exchange_gas_limit = base_exchange_gas_limit or self.DEFAULT_EXCHANGE_GAS_LIMIT
-
-        self._gas_limit_per_message_type_rules: List[Tuple[Callable, Decimal]] = [
-            (
-                lambda message: "MsgPrivilegedExecuteContract" in self._message_type(message=message),
-                Decimal("6") * self._base_gas_limit
-            ),
-            (
-                lambda message: "MsgExecuteContract" in self._message_type(message=message),
-                Decimal("2.5") * self._base_gas_limit
-            ),
-            (
-                lambda message: "wasm." in self._message_type(message=message),
-                Decimal("1.5") * self._base_gas_limit
-            ),
-            (
-                lambda message: "exchange." in self._message_type(message=message),
-                Decimal("1") * self._base_exchange_gas_limit
-            ),
-            (
-                lambda message: self._is_governance_message(message=message),
-                Decimal("15") * self._base_gas_limit
-            ),
-        ]
 
     async def configure_gas_fee_for_transaction(
             self,
             transaction: Transaction,
             private_key: PrivateKey,
             public_key: PublicKey,
     ):
-        transaction_gas_limit = math.ceil(self._calculate_gas_limit(messages=transaction.msgs))
+        messages_gas_limit = math.ceil(self._calculate_gas_limit(messages=transaction.msgs))
+        transaction_gas_limit = messages_gas_limit + self.TRANSACTION_GAS_LIMIT
 
         fee = [
             self._composer.Coin(
@@ -351,28 +324,11 @@ def _message_type(self, message: any_pb2.Any) -> str:
             message_type = message.DESCRIPTOR.full_name
         return message_type
 
-    def _is_governance_message(self, message: any_pb2.Any) -> bool:
-        message_type = self._message_type(message=message)
-        return "gov." in message_type and ("MsgDeposit" in message_type or "MsgSubmitProposal" in message_type)
-
     def _calculate_gas_limit(self, messages: List[any_pb2.Any]) -> int:
         total_gas_limit = Decimal("0")
 
         for message in messages:
-            applying_rule = next(
-                (rule_tuple for rule_tuple in self._gas_limit_per_message_type_rules
-                 if rule_tuple[0](message)),
-                None
-            )
-
-            if applying_rule is None:
-                total_gas_limit += self._base_gas_limit
-            else:
-                total_gas_limit += applying_rule[1]
-
-            if self._message_type(message=message).endswith("MsgExec"):
-                exec_message = cosmos_authz_tx_pb.MsgExec.FromString(message.value)
-                sub_messages_limit = self._calculate_gas_limit(messages=exec_message.msgs)
-                total_gas_limit += sub_messages_limit
+            estimator = GasLimitEstimator.for_message(message=message)
+            total_gas_limit += estimator.gas_limit()
 
         return math.ceil(total_gas_limit)