Merge pull request #52 from rocklabs-io/common-cansiter

Common cansiter

Author: ccyanxyz

ic/agent.py

@@ -45,49 +45,41 @@ def read_state_endpoint(self, canister_id, data):
         result = self.client.read_state(canister_id, data)
         return result
 
-    def query_raw(self, canister_id, method_name, *arg):
-        assert len(arg) == 1 or len(arg) == 2
+    def query_raw(self, canister_id, method_name, arg, return_type = None, effective_canister_id = None):
         req = {
             'request_type': "query",
             'sender': self.identity.sender().bytes,
             'canister_id': Principal.from_str(canister_id).bytes if isinstance(canister_id, str) else canister_id.bytes,
             'method_name': method_name,
-            'arg': arg[0],
+            'arg': arg,
             'ingress_expiry': self.get_expiry_date()
         }
         _, data = sign_request(req, self.identity)
-        result = self.query_endpoint(canister_id, data)
+        result = self.query_endpoint(canister_id if effective_canister_id is None else effective_canister_id, data)
         if result['status'] == 'replied':
-            if len(arg) == 1:
-                res = decode(result['reply']['arg'])
-            else:
-                res = decode(result['reply']['arg'], arg[1])
-            return res
+            return decode(result['reply']['arg'], return_type)
         elif result['status'] == 'rejected':
             return result['reject_message']
 
-    def update_raw(self, canister_id, method_name, *arg, **kwargs):
-        assert len(arg) == 1 or len(arg) == 2
+    def update_raw(self, canister_id, method_name, arg, return_type = None, effective_canister_id = None, **kwargs):
         req = {
             'request_type': "call",
             'sender': self.identity.sender().bytes,
             'canister_id': Principal.from_str(canister_id).bytes if isinstance(canister_id, str) else canister_id.bytes,
             'method_name': method_name,
-            'arg': arg[0],
+            'arg': arg,
             'ingress_expiry': self.get_expiry_date()
         }
         req_id, data = sign_request(req, self.identity)
-        _ = self.call_endpoint(canister_id, req_id, data)
+        eid = canister_id if effective_canister_id is None else effective_canister_id
+        _ = self.call_endpoint(eid, req_id, data)
         # print('update.req_id:', req_id.hex())
-        status, result = self.poll(canister_id, req_id, **kwargs)
+        status, result = self.poll(eid, req_id, **kwargs)
         if status != 'replied':
-            return  status
+            return status
         else:
-            if len(arg) == 1:
-                res = decode(result)
-            else:
-                res = decode(result, arg[1])
-            return res
+            return decode(result, return_type)
+            
 
 
     def read_state_raw(self, canister_id, paths):

ic/candid.py

@@ -1,6 +1,7 @@
 # candid.example.py shows how to use candid's en/decode
 
 import leb128
+from collections.abc import Iterable
 from struct import pack,unpack
 from abc import abstractclassmethod, ABCMeta
 from enum import Enum
@@ -528,7 +529,7 @@ def __init__(self, _type: Type):
         self._type = _type
 
     def covariant(self, x):
-        return type(x) == list and not False in list(map(self._type.covariant, x))
+        return isinstance(x, Iterable) and not False in list(map(self._type.covariant, x))
 
     def encodeValue(self, val):
         length = leb128.u.encode(len(val))

ic/canister.py

@@ -1,7 +1,7 @@
 from .parser.DIDEmitter import *;
 from antlr4 import *
 from antlr4.InputStream import InputStream
-from .candid import encode
+from .candid import encode, FuncClass
 
 class Canister:
     def __init__(self, agent, canister_id, candid=None):
@@ -30,8 +30,9 @@ def __init__(self, agent, canister_id, candid=None):
         self.actor = emitter.getActor()
 
         for name, method in self.actor["methods"].items():
-            anno = None if len(method[2]) == 0 else method[2][0]
-            setattr(self, name, CaniterMethod(agent, canister_id, name, method[0], method[1], anno))
+            assert type(method) == FuncClass
+            anno = None if len(method.annotations) == 0 else method.annotations[0]
+            setattr(self, name, CaniterMethod(agent, canister_id, name, method.argTypes, method.retTypes, anno))
 
 class CaniterMethod:
     def __init__(self, agent, canister_id, name, args, rets, anno = None):
@@ -50,19 +51,22 @@ def __call__(self, *args, **kwargs):
         for i, arg in enumerate(args):
             arguments.append({"type": self.args[i], "value": arg})
 
+        effective_cansiter_id = args[0]['canister_id'] if self.canister_id == 'aaaaa-aa' and len(args) > 0 and type(args[0]) == dict and 'canister_id' in args[0] else self.canister_id
         if self.anno == 'query':
             res = self.agent.query_raw(
                 self.canister_id,
                 self.name, 
                 encode(arguments),
-                self.rets
+                self.rets,
+                effective_cansiter_id
                 )
         else:
             res = self.agent.update_raw(
                 self.canister_id,
                 self.name, 
                 encode(arguments),
-                self.rets
+                self.rets,
+                effective_cansiter_id
             )
 
         if type(res) is not list: