Merge pull request #55 from rocklabs-io/async_request

add async function

ic/agent.py

@@ -37,14 +37,26 @@ def query_endpoint(self, canister_id, data):
         ret = self.client.query(canister_id, data)
         return cbor2.loads(ret)
 
+    async def query_endpoint_async(self, canister_id, data):
+        ret = await self.client.query_async(canister_id, data)
+        return cbor2.loads(ret)
+
     def call_endpoint(self, canister_id, request_id, data):
         self.client.call(canister_id, request_id, data)
         return request_id
 
+    async def call_endpoint_async(self, canister_id, request_id, data):
+        await self.client.call_async(canister_id, request_id, data)
+        return request_id
+
     def read_state_endpoint(self, canister_id, data):
         result = self.client.read_state(canister_id, data)
         return result
 
+    async def read_state_endpoint_async(self, canister_id, data):
+        result = await self.client.read_state_async(canister_id, data)
+        return result
+
     def query_raw(self, canister_id, method_name, arg, return_type = None, effective_canister_id = None):
         req = {
             'request_type': "query",
@@ -61,6 +73,22 @@ def query_raw(self, canister_id, method_name, arg, return_type = None, effective
         elif result['status'] == 'rejected':
             return result['reject_message']
 
+    async def query_raw_async(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,
+            'ingress_expiry': self.get_expiry_date()
+        }
+        _, data = sign_request(req, self.identity)
+        result = await self.query_endpoint_async(canister_id if effective_canister_id is None else effective_canister_id, data)
+        if result['status'] == 'replied':
+            return decode(result['reply']['arg'], return_type)
+        elif result['status'] == 'rejected':
+            return result['reject_message']
+
     def update_raw(self, canister_id, method_name, arg, return_type = None, effective_canister_id = None, **kwargs):
         req = {
             'request_type': "call",
@@ -82,7 +110,26 @@ def update_raw(self, canister_id, method_name, arg, return_type = None, effectiv
         else:
             raise Exception('Timeout to poll result, current status: ' + str(status))
 
-
+    async def update_raw_async(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,
+            'ingress_expiry': self.get_expiry_date()
+        }
+        req_id, data = sign_request(req, self.identity)
+        eid = canister_id if effective_canister_id is None else effective_canister_id
+        _ = await self.call_endpoint_async(eid, req_id, data)
+        # print('update.req_id:', req_id.hex())
+        status, result = await self.poll_async(eid, req_id, **kwargs)
+        if status == 'rejected':
+            raise Exception('Rejected: ' + result.decode())
+        elif status == 'replied': 
+            return decode(result, return_type)
+        else:
+            raise Exception('Timeout to poll result, current status: ' + str(status))        
 
     def read_state_raw(self, canister_id, paths):
         req = {
@@ -101,6 +148,23 @@ def read_state_raw(self, canister_id, paths):
         cert = cbor2.loads(d['certificate'])
         return cert
 
+    async def read_state_raw_async(self, canister_id, paths):
+        req = {
+            'request_type': 'read_state',
+            'sender': self.identity.sender().bytes,
+            'paths': paths, 
+            'ingress_expiry': self.get_expiry_date(),
+        }
+        _, data = sign_request(req, self.identity)
+        ret = await self.read_state_endpoint_async(canister_id, data)
+        if ret == b'Invalid path requested.':
+            raise ValueError('Invalid path requested!')
+        elif ret == b'Could not parse body as read request: invalid type: byte array, expected a sequence':
+            raise ValueError('Could not parse body as read request: invalid type: byte array, expected a sequence')
+        d = cbor2.loads(ret)
+        cert = cbor2.loads(d['certificate'])
+        return cert
+
     def request_status_raw(self, canister_id, req_id):
         paths = [
             ['request_status'.encode(), req_id],
@@ -112,6 +176,17 @@ def request_status_raw(self, canister_id, req_id):
         else:
             return status.decode(), cert
 
+    async def request_status_raw_async(self, canister_id, req_id):
+        paths = [
+            ['request_status'.encode(), req_id],
+        ]
+        cert = await self.read_state_raw_async(canister_id, paths)
+        status = lookup(['request_status'.encode(), req_id, 'status'.encode()], cert)
+        if (status == None):
+            return status, cert
+        else:
+            return status.decode(), cert
+
     def poll(self, canister_id, req_id, delay=1, timeout=float('inf')):
         status = None
         for _ in wait(delay, timeout):
@@ -129,3 +204,21 @@ def poll(self, canister_id, req_id, delay=1, timeout=float('inf')):
             return status, msg
         else:
             return status, _
+
+    async def poll_async(self, canister_id, req_id, delay=1, timeout=float('inf')):
+        status = None
+        for _ in wait(delay, timeout):
+            status, cert = await self.request_status_raw_async(canister_id, req_id)
+            if status == 'replied' or status == 'done' or status  == 'rejected':
+                break
+
+        if status == 'replied':
+            path = ['request_status'.encode(), req_id, 'reply'.encode()]
+            res = lookup(path, cert)
+            return status, res
+        elif status == 'rejected':
+            path = ['request_status'.encode(), req_id, 'reject_message'.encode()]
+            msg = lookup(path, cert)
+            return status, msg
+        else:
+            return status, _

ic/canister.py

@@ -33,6 +33,7 @@ def __init__(self, agent, canister_id, candid=None):
             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))
+            setattr(self, name + '_async', CaniterMethodAsync(agent, canister_id, name, method.argTypes, method.retTypes, anno))
 
 class CaniterMethod:
     def __init__(self, agent, canister_id, name, args, rets, anno = None):
@@ -73,3 +74,43 @@ def __call__(self, *args, **kwargs):
             return res
 
         return list(map(lambda item: item["value"], res))
+
+class CaniterMethodAsync:
+    def __init__(self, agent, canister_id, name, args, rets, anno = None):
+        self.agent = agent
+        self.canister_id = canister_id
+        self.name = name
+        self.args = args
+        self.rets = rets
+
+        self.anno = anno
+
+    async def __call__(self, *args, **kwargs):
+        if len(args) != len(self.args):
+            raise ValueError("Arguments length not match")
+        arguments = []
+        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 = await self.agent.query_raw_async(
+                self.canister_id,
+                self.name, 
+                encode(arguments),
+                self.rets,
+                effective_cansiter_id
+                )
+        else:
+            res = await self.agent.update_raw_async(
+                self.canister_id,
+                self.name, 
+                encode(arguments),
+                self.rets,
+                effective_cansiter_id
+            )
+
+        if type(res) is not list:
+            return res
+
+        return list(map(lambda item: item["value"], res))

ic/client.py

@@ -1,6 +1,6 @@
 # http client
 
-import requests
+import httpx
 
 class Client:
     def __init__(self, url = "https://ic0.app"):
@@ -9,23 +9,51 @@ def __init__(self, url = "https://ic0.app"):
     def query(self, canister_id, data):
         endpoint = self.url + '/api/v2/canister/' + canister_id + '/query'
         headers = {'Content-Type': 'application/cbor'}
-        ret = requests.post(endpoint, data, headers=headers)
+        ret = httpx.post(endpoint, data = data, headers=headers)
         return ret.content
 
     def call(self, canister_id, req_id, data):
         endpoint = self.url + '/api/v2/canister/' + canister_id + '/call'
         headers = {'Content-Type': 'application/cbor'}
-        ret = requests.post(endpoint, data, headers=headers)
+        ret = httpx.post(endpoint, data = data, headers=headers)
         return req_id
 
     def read_state(self, canister_id, data):
         endpoint = self.url + '/api/v2/canister/' + canister_id + '/read_state'
         headers = {'Content-Type': 'application/cbor'}
-        ret = requests.post(endpoint, data, headers=headers)
+        ret = httpx.post(endpoint, data = data, headers=headers)
         return ret.content
 
     def status(self):
         endpoint = self.url + '/api/v2/status'
-        ret = requests.get(endpoint)
+        ret = httpx.get(endpoint)
         print('client.status:', ret.text)
         return ret.content
+
+    async def query_async(self, canister_id, data):
+        async with httpx.AsyncClient() as client:
+            endpoint = self.url + '/api/v2/canister/' + canister_id + '/query'
+            headers = {'Content-Type': 'application/cbor'}
+            ret = await client.post(endpoint, data = data, headers=headers)
+            return ret.content
+
+    async def call_async(self, canister_id, req_id, data):
+        async with httpx.AsyncClient() as client:
+            endpoint = self.url + '/api/v2/canister/' + canister_id + '/call'
+            headers = {'Content-Type': 'application/cbor'}
+            await client.post(endpoint, data = data, headers=headers)
+            return req_id
+
+    async def read_state_async(self, canister_id, data):
+        async with httpx.AsyncClient() as client:
+            endpoint = self.url + '/api/v2/canister/' + canister_id + '/read_state'
+            headers = {'Content-Type': 'application/cbor'}
+            ret = await client.post(endpoint, data = data, headers=headers)
+            return ret.content
+
+    async def status_async(self):
+        async with httpx.AsyncClient() as client:
+            endpoint = self.url + '/api/v2/status'
+            ret = await client.get(endpoint)
+            print('client.status:', ret.text)
+            return ret.content

setup.py

@@ -17,7 +17,7 @@
     author_email = '[email protected]',
     keywords = 'dfinity ic agent',
     install_requires = [
-        'requests>=2.22.0', 
+        'httpx>=0.22.0',
         'ecdsa>=0.18.0b2', 
         'cbor2>=5.4.2', 
         'leb128>=1.0.4', 

test_agent.py

@@ -1,3 +1,4 @@
+import asyncio
 from ic.agent import *
 from ic.identity import *
 from ic.client import *
@@ -8,6 +9,7 @@
 print('principal:', Principal.self_authenticating(iden.der_pubkey))
 ag = Agent(iden, client)
 
+start = time.time()
 # query token totalSupply
 ret = ag.query_raw("gvbup-jyaaa-aaaah-qcdwa-cai", "totalSupply", encode([]))
 print('totalSupply:', ret)
@@ -36,3 +38,38 @@
             ])
         )
 print('result: ', ret)
+
+t = time.time()
+print("sync call elapsed: ", t - start)
+
+async def test_async():
+    ret = await ag.query_raw_async("gvbup-jyaaa-aaaah-qcdwa-cai", "totalSupply", encode([]))
+    print('totalSupply:', ret)
+
+    # query token name
+    ret = await ag.query_raw_async("gvbup-jyaaa-aaaah-qcdwa-cai", "name", encode([]))
+    print('name:', ret)
+
+    # query token balance of user
+    ret = await ag.query_raw_async(
+            "gvbup-jyaaa-aaaah-qcdwa-cai",
+            "balanceOf",
+            encode([
+                {'type': Types.Principal, 'value': iden.sender().bytes}
+            ])
+        )
+    print('balanceOf:', ret)
+
+    # transfer 100 tokens to blackhole
+    ret = await ag.update_raw_async(
+            "gvbup-jyaaa-aaaah-qcdwa-cai",
+            "transfer",
+            encode([
+                {'type': Types.Principal, 'value': 'aaaaa-aa'},
+                {'type': Types.Nat, 'value': 10000000000}
+                ])
+            )
+    print('result: ', ret)
+
+asyncio.run(test_async())
+print("sync call elapsed: ", time.time() - t)

test_canister.py

@@ -1,3 +1,4 @@
+import asyncio
 from ic.canister import Canister
 from ic.client import Client
 from ic.identity import Identity
@@ -435,4 +436,18 @@
         'include_status': [1]
     }
 )
-print(res)
+print(res)
+
+async def async_test():
+  res = await governance.list_proposals_async(
+    {
+        'include_reward_status': [], 
+        'before_proposal': [],
+        'limit': 100, 
+        'exclude_topic': [], 
+        'include_status': [1]
+    }
+  )
+  print(res)
+
+asyncio.run(async_test())