Merge pull request Swind#47 from JeffLIrion/async-pull-request

Implement async functionality

Author: Swind

ppadb/client_async.py

@@ -0,0 +1,21 @@
+from ppadb.command.host_async import HostAsync
+from ppadb.connection_async import ConnectionAsync
+
+
+class ClientAsync(HostAsync):
+    def __init__(self, host='127.0.0.1', port=5037):
+        self.host = host
+        self.port = port
+
+    async def create_connection(self, timeout=None):
+        conn = ConnectionAsync(self.host, self.port, timeout)
+        return await conn.connect()
+
+    async def device(self, serial):
+        devices = await self.devices()
+
+        for device in devices:
+            if device.serial == serial:
+                return device
+
+        return None

ppadb/command/host_async/__init__.py

@@ -0,0 +1,28 @@
+from ppadb.device_async import DeviceAsync
+
+
+class HostAsync:
+    CONNECT_RESULT_PATTERN = "(connected to|already connected)"
+
+    OFFLINE = "offline"
+    DEVICE = "device"
+    BOOTLOADER = "bootloader"
+
+    async def _execute_cmd(self, cmd):
+        async with await self.create_connection() as conn:
+            await conn.send(cmd)
+            return await conn.receive()
+
+    async def devices(self):
+        cmd = "host:devices"
+        result = await self._execute_cmd(cmd)
+
+        devices = []
+
+        for line in result.split('\n'):
+            if not line:
+                break
+
+            devices.append(DeviceAsync(self, line.split()[0]))
+
+        return devices

ppadb/command/transport_async/__init__.py

@@ -0,0 +1,40 @@
+import logging
+import re
+import time
+
+
+class TransportAsync:
+    async def transport(self, connection):
+        cmd = "host:transport:{}".format(self.serial)
+        await connection.send(cmd)
+
+        return connection
+
+    async def shell(self, cmd, timeout=None):
+        conn = await self.create_connection(timeout=timeout)
+
+        cmd = "shell:{}".format(cmd)
+        await conn.send(cmd)
+
+        result = await conn.read_all()
+        await conn.close()
+        return result.decode('utf-8')
+
+    async def sync(self):
+        conn = await self.create_connection()
+
+        cmd = "sync:"
+        await conn.send(cmd)
+
+        return conn
+
+    async def screencap(self):
+        async with await self.create_connection() as conn:
+            cmd = "shell:/system/bin/screencap -p"
+            await conn.send(cmd)
+            result = await conn.read_all()
+
+        if result and len(result) > 5 and result[5] == 0x0d:
+            return result.replace(b'\r\n', b'\n')
+        else:
+            return result

ppadb/connection_async.py

@@ -0,0 +1,104 @@
+import asyncio
+import struct
+import socket
+
+from ppadb.protocol import Protocol
+from ppadb.utils.logger import AdbLogging
+
+logger = AdbLogging.get_logger(__name__)
+
+
+class ConnectionAsync:
+    def __init__(self, host='localhost', port=5037, timeout=None):
+        self.host = host
+        self.port = int(port)
+        self.timeout = timeout
+
+        self.reader = None
+        self.writer = None
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, type, value, traceback):
+        await self.close()
+
+    async def connect(self):
+        logger.debug("Connect to ADB server - %s:%d", self.host, self.port)
+
+        try:
+            if self.timeout:
+                self.reader, self.writer = await asyncio.wait_for(asyncio.open_connection(self.host, self.port), self.timeout)
+            else:
+                self.reader, self.writer = await asyncio.open_connection(self.host, self.port)
+
+        except (OSError, asyncio.TimeoutError) as e:
+            raise RuntimeError("ERROR: connecting to {}:{} {}.\nIs adb running on your computer?".format(self.host, self.port, e))
+
+        return self
+
+    async def close(self):
+        logger.debug("Connection closed...")
+        
+        if self.writer:
+            try:
+                self.writer.close()
+                await self.writer.wait_closed()
+            except OSError:
+                pass
+
+        self.reader = None
+        self.writer = None
+
+    ##############################################################################################################
+    #
+    # Send command & Receive command result
+    #
+    ##############################################################################################################
+    async def _recv(self, length):
+        return await asyncio.wait_for(self.reader.read(length), self.timeout)
+
+    async def _send(self, data):
+        self.writer.write(data)
+        await asyncio.wait_for(self.writer.drain(), self.timeout)
+
+    async def receive(self):
+        nob = int((await self._recv(4)).decode('utf-8'), 16)
+        return (await self._recv(nob)).decode('utf-8')
+
+    async def send(self, msg):
+        msg = Protocol.encode_data(msg)
+        logger.debug(msg)
+        await self._send(msg)
+        return await self._check_status()
+
+    async def _check_status(self):
+        recv = (await self._recv(4)).decode('utf-8')
+        if recv != Protocol.OKAY:
+            error = (await self._recv(1024)).decode('utf-8')
+            raise RuntimeError("ERROR: {} {}".format(repr(recv), error))
+
+        return True
+
+    ##############################################################################################################
+    #
+    # Socket read/write
+    #
+    ##############################################################################################################
+    async def read_all(self):
+        data = bytearray()
+
+        while True:
+            recv = await self._recv(4096)
+            if not recv:
+                break
+            data += recv
+
+        return data
+
+    async def read(self, length=0):
+        data = await self._recv(length)
+        return data
+
+    async def write(self, data):
+        await self._send(data)

ppadb/device_async.py

@@ -0,0 +1,69 @@
+try:
+    from asyncio import get_running_loop
+except ImportError:  # pragma: no cover
+    from asyncio import get_event_loop as get_running_loop  # Python 3.6 compatibility
+
+import re
+import os
+
+from ppadb.command.transport_async import TransportAsync
+from ppadb.sync_async import SyncAsync
+
+
+def _get_src_info(src):
+    exists = os.path.exists(src)
+    isfile = os.path.isfile(src)
+    isdir = os.path.isdir(src)
+    basename = os.path.basename(src)
+    walk = None if not isdir else list(os.walk(src))
+
+    return exists, isfile, isdir, basename, walk
+
+
+class DeviceAsync(TransportAsync):
+    INSTALL_RESULT_PATTERN = "(Success|Failure|Error)\s?(.*)"
+    UNINSTALL_RESULT_PATTERN = "(Success|Failure.*|.*Unknown package:.*)"
+
+    def __init__(self, client, serial):
+        self.client = client
+        self.serial = serial
+
+    async def create_connection(self, set_transport=True, timeout=None):
+        conn = await self.client.create_connection(timeout=timeout)
+
+        if set_transport:
+            await self.transport(conn)
+
+        return conn
+
+    async def _push(self, src, dest, mode, progress):
+        # Create a new connection for file transfer
+        sync_conn = await self.sync()
+        sync = SyncAsync(sync_conn)
+
+        async with sync_conn:
+            await sync.push(src, dest, mode, progress)
+
+    async def push(self, src, dest, mode=0o644, progress=None):
+        exists, isfile, isdir, basename, walk = await get_running_loop().run_in_executor(None, _get_src_info, src)
+        if not exists:
+            raise FileNotFoundError("Cannot find {}".format(src))
+
+        if isfile:
+            await self._push(src, dest, mode, progress)
+
+        elif isdir:
+            for root, dirs, files in walk:
+                root_dir_path = os.path.join(basename, root.replace(src, ""))
+
+                await self.shell("mkdir -p {}/{}".format(dest, root_dir_path))
+
+                for item in files:
+                    await self._push(os.path.join(root, item), os.path.join(dest, root_dir_path, item), mode, progress)
+
+    async def pull(self, src, dest):
+        sync_conn = await self.sync()
+        sync = SyncAsync(sync_conn)
+
+        async with sync_conn:
+            return await sync.pull(src, dest)

ppadb/sync_async/__init__.py

@@ -0,0 +1,125 @@
+try:
+    from asyncio import get_running_loop
+except ImportError:  # pragma: no cover
+    from asyncio import get_event_loop as get_running_loop  # Python 3.6 compatibility
+
+import struct
+import os
+
+from ppadb.protocol import Protocol
+from ppadb.sync.stats import S_IFREG
+from ppadb.utils.logger import AdbLogging
+
+import aiofiles
+
+logger = AdbLogging.get_logger(__name__)
+
+
+def _get_src_info(src):
+    exists = os.path.exists(src)
+    if not exists:
+        return exists, None, None
+
+    timestamp = os.stat(src).st_mtime
+    total_size = os.path.getsize(src)
+
+    return exists, timestamp, total_size
+
+
+class SyncAsync:
+    DATA_MAX_LENGTH = 65536
+
+    def __init__(self, connection):
+        self.connection = connection
+
+    async def push(self, src, dest, mode, progress=None):
+        """Push from local path |src| to |dest| on device.
+        :param progress: callback, called with (filename, total_size, sent_size)
+        """
+        exists, timestamp, total_size = await get_running_loop().run_in_executor(None, _get_src_info, src)
+
+        if not exists:
+            raise FileNotFoundError("Can't find the source file {}".format(src))
+
+        sent_size = 0
+
+        # SEND
+        mode = mode | S_IFREG
+        args = "{dest},{mode}".format(dest=dest, mode=mode)
+        await self._send_str(Protocol.SEND, args)
+
+        # DATA
+        async with aiofiles.open(src, 'rb') as stream:
+            while True:
+                chunk = await stream.read(self.DATA_MAX_LENGTH)
+                if not chunk:
+                    break
+
+                sent_size += len(chunk)
+                await self._send_length(Protocol.DATA, len(chunk))
+                await self.connection.write(chunk)
+
+                if progress is not None:
+                    progress(src, total_size, sent_size)
+
+        # DONE
+        await self._send_length(Protocol.DONE, timestamp)
+        await self.connection._check_status()
+
+    async def pull(self, src, dest):
+        # RECV
+        await self._send_str(Protocol.RECV, src)
+
+        # DATA
+        async with aiofiles.open(dest, 'wb') as stream:
+            while True:
+                flag = (await self.connection.read(4)).decode('utf-8')
+
+                if flag == Protocol.DATA:
+                    data = await self._read_data()
+                    await stream.write(data)
+                    continue
+
+                if flag == Protocol.DONE:
+                    await self.connection.read(4)
+                    return
+
+                if flag == Protocol.FAIL:
+                    return (await self._read_data()).decode('utf-8')
+
+    @staticmethod
+    def _integer(little_endian):
+        return struct.unpack("<I", little_endian)
+
+    @staticmethod
+    def _little_endian(n):
+        return struct.pack('<I', n)
+
+    async def _read_data(self):
+        length = self._integer(await self.connection.read(4))[0]
+        data = bytearray()
+        while len(data) < length:
+            data += await self.connection.read(length - len(data))
+        return data
+
+    async def _send_length(self, cmd, length):
+        le_len = self._little_endian(length)
+        data = cmd.encode() + le_len
+
+        logger.debug("Send length: {}".format(data))
+        await self.connection.write(data)
+
+    async def _send_str(self, cmd, args):
+        """
+        Format:
+            {Command}{args length(little endian)}{str}
+        Length:
+            {4}{4}{str length}
+        """
+        logger.debug("{} {}".format(cmd, args))
+        args = args.encode('utf-8')
+
+        le_args_len = self._little_endian(len(args))
+        data = cmd.encode() + le_args_len + args
+        logger.debug("Send string: {}".format(data))
+        await self.connection.write(data)