✨ Implement Sake getAbi without proxy method

Ackee-Blockchain · Nov 12, 2024 · d79b322 · d79b322
1 parent a6261c1
commit d79b322
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 3 deletions.
diff --git a/wake/lsp/methods.py b/wake/lsp/methods.py
@@ -142,3 +142,4 @@ class RequestMethodEnum(StrEnum):
     SAKE_GET_BALANCES = "wake/sake/getBalances"
     SAKE_SET_BALANCES = "wake/sake/setBalances"
     SAKE_GET_ABI = "wake/sake/getAbi"
+    SAKE_GET_ABI_WITH_PROXY = "wake/sake/getAbiWithProxy"
diff --git a/wake/lsp/sake.py b/wake/lsp/sake.py
@@ -160,11 +160,21 @@ class SakeLoadStateParams(SakeParams):
 
 class SakeGetAbiParams(SakeParams):
     address: str
+    chain_id: int
 
 
 class SakeGetAbiResult(SakeResult):
     name: str
     abi: List
+
+
+class SakeGetAbiWithProxyParams(SakeParams):
+    address: str
+
+
+class SakeGetAbiWithProxyResult(SakeResult):
+    name: str
+    abi: List
     proxy_name: Optional[str]
     proxy_abi: Optional[List]
     implementation_address: Optional[str]
@@ -788,8 +798,38 @@ async def set_label(self, params: SakeSetLabelParams) -> SakeResult:
         except Exception as e:
             raise LspError(ErrorCodes.InternalError, str(e)) from None
 
-    @chain_connected
     async def get_abi(self, params: SakeGetAbiParams) -> SakeGetAbiResult:
+        try:
+            try:
+                name, abi = get_name_abi_from_explorer(params.address, params.chain_id)
+                return SakeGetAbiResult(success=True, name=name, abi=abi)
+            except AbiNotFound as e:
+                if e.api_key_name is not None:
+                    raise LspError(
+                        ErrorCodes.AbiNotFound,
+                        f"ABI not found for {params.address}",
+                        {
+                            "apiKeyName": e.api_key_name,
+                            "tomlUsed": self.lsp_context.use_toml
+                            and self.lsp_context.toml_path.is_file(),
+                            "configPath": str(
+                                self.lsp_context.config.global_config_path
+                            ),
+                        },
+                    )
+                else:
+                    raise LspError(
+                        ErrorCodes.AbiNotFound,
+                        f"ABI not found for {params.address}",
+                        {},
+                    )
+        except Exception as e:
+            raise LspError(ErrorCodes.InternalError, str(e)) from None
+
+    @chain_connected
+    async def get_abi_with_proxy(
+        self, params: SakeGetAbiWithProxyParams
+    ) -> SakeGetAbiWithProxyResult:
         def info_from_address(address: str) -> Tuple[str, List]:
             fqn = get_fqn_from_address(logic_contract.address, "latest", chain)
             if fqn is not None:
@@ -837,7 +877,7 @@ def info_from_address(address: str) -> Tuple[str, List]:
             name, abi = info_from_address(str(logic_contract.address))
 
             if contract != logic_contract:
-                return SakeGetAbiResult(
+                return SakeGetAbiWithProxyResult(
                     success=True,
                     name=name,
                     abi=abi,
@@ -848,7 +888,7 @@ def info_from_address(address: str) -> Tuple[str, List]:
             else:
                 proxy_name, proxy_abi = info_from_address(params.address)
 
-                return SakeGetAbiResult(
+                return SakeGetAbiWithProxyResult(
                     success=True,
                     name=name,
                     abi=abi,

diff --git a/wake/lsp/server.py b/wake/lsp/server.py
@@ -427,6 +427,14 @@ def __init__(
                 ),
                 SakeGetAbiParams,
             ),
+            RequestMethodEnum.SAKE_GET_ABI_WITH_PROXY: (
+                lambda params: (
+                    self.__sake_context.get_abi_with_proxy(  # pyright: ignore reportAttributeAccessIssue
+                        params,
+                    )
+                ),
+                SakeGetAbiParams,
+            ),
         }
 
         self.__notification_mapping = {