Move file I/O out of the event loop (#1)

* Use aiofiles for file I/O

* Move all I/O out of the event loop

* Delete commented out code

Author: JeffLIrion

ppadb/device_async.py

@@ -1,10 +1,25 @@
+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:.*)"
@@ -30,16 +45,15 @@ async def _push(self, src, dest, mode, progress):
             await sync.push(src, dest, mode, progress)
 
     async def push(self, src, dest, mode=0o644, progress=None):
-        if not os.path.exists(src):
+        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 os.path.isfile(src):
+        if isfile:
             await self._push(src, dest, mode, progress)
 
-        elif os.path.isdir(src):
-            basename = os.path.basename(src)
-
-            for root, dirs, files in os.walk(src):
+        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))

ppadb/sync_async/__init__.py

@@ -1,13 +1,31 @@
+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
 
@@ -18,14 +36,11 @@ 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)
         """
-        if not os.path.exists(src):
-            raise FileNotFoundError("Can't find the source file {}".format(src))
+        exists, timestamp, total_size = await get_running_loop().run_in_executor(None, _get_src_info, src)
 
-        stat = os.stat(src)
-
-        timestamp = int(stat.st_mtime)
+        if not exists:
+            raise FileNotFoundError("Can't find the source file {}".format(src))
 
-        total_size = os.path.getsize(src)
         sent_size = 0
 
         # SEND
@@ -34,9 +49,9 @@ async def push(self, src, dest, mode, progress=None):
         await self._send_str(Protocol.SEND, args)
 
         # DATA
-        with open(src, 'rb') as stream:
+        async with aiofiles.open(src, 'rb') as stream:
             while True:
-                chunk = stream.read(self.DATA_MAX_LENGTH)
+                chunk = await stream.read(self.DATA_MAX_LENGTH)
                 if not chunk:
                     break
 
@@ -56,13 +71,13 @@ async def pull(self, src, dest):
         await self._send_str(Protocol.RECV, src)
 
         # DATA
-        with open(dest, 'wb') as stream:
+        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()
-                    stream.write(data)
+                    await stream.write(data)
                     continue
 
                 if flag == Protocol.DONE:

setup.py

@@ -31,6 +31,7 @@
     license='MIT license',
     packages=find_packages(exclude=["*.test", "*.test.*", "test.*", "test"]),
     install_requires=[],
+    extras_require={"async": ["aiofiles>=0.4.0"]},
     keywords="adb",
     classifiers=classifiers,
 )

test_async/patchers.py

@@ -1,6 +1,6 @@
 """Patches for async socket functionality."""
 
-from contextlib import contextmanager
+from contextlib import asynccontextmanager
 from unittest.mock import patch
 
 try:
@@ -13,6 +13,39 @@ async def __call__(self, *args, **kwargs):
             return super(AsyncMock, self).__call__(*args, **kwargs)
 
 
+def async_mock_open(read_data=""):
+    class AsyncMockFile:
+        def __init__(self, read_data):
+            self.read_data = read_data
+            _async_mock_open.written = read_data[:0]
+
+        async def read(self, size=-1):
+            if size == -1:
+                ret = self.read_data
+                self.read_data = self.read_data[:0]
+                return ret
+
+            n = min(size, len(self.read_data))
+            ret = self.read_data[:n]
+            self.read_data = self.read_data[n:]
+            return ret
+
+        async def write(self, b):
+            if _async_mock_open.written:
+                _async_mock_open.written += b
+            else:
+                _async_mock_open.written = b
+
+    @asynccontextmanager
+    async def _async_mock_open(*args, **kwargs):
+        try:
+            yield AsyncMockFile(read_data)
+        finally:
+            pass
+
+    return _async_mock_open
+
+
 class FakeStreamWriter:
     def close(self):
         pass

test_async/test_device_async.py

@@ -4,6 +4,7 @@
 
 
 import asyncio
+from contextlib import asynccontextmanager
 import os
 import sys
 import unittest
@@ -16,13 +17,15 @@
 from ppadb.sync_async import SyncAsync
 
 from .async_wrapper import awaiter
-from .patchers import FakeStreamReader, FakeStreamWriter, async_patch
+from .patchers import FakeStreamReader, FakeStreamWriter, async_mock_open, async_patch
 
 
 PNG_IMAGE = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\n\x00\x00\x00\n\x08\x06\x00\x00\x00\x8d2\xcf\xbd\x00\x00\x00\x04sBIT\x08\x08\x08\x08|\x08d\x88\x00\x00\x00\tpHYs\x00\x00\x0fa\x00\x00\x0fa\x01\xa8?\xa7i\x00\x00\x00\x0eIDAT\x18\x95c`\x18\x05\x83\x13\x00\x00\x01\x9a\x00\x01\x16\xca\xd3i\x00\x00\x00\x00IEND\xaeB`\x82'
 
 PNG_IMAGE_NEEDS_REPLACING = PNG_IMAGE[:5] + b'\r' + PNG_IMAGE[5:]
 
+FILEDATA = b'Ohayou sekai.\nGood morning world!'
+
 
 class TestDevice(unittest.TestCase):
     @awaiter
@@ -72,7 +75,7 @@ def progress(*args, **kwargs):
             pass
 
         filedata = b'Ohayou sekai.\nGood morning world!'
-        with patch('os.path.exists', return_value=True), patch('os.path.isfile', return_value=True), patch('ppadb.device_async.open', mock_open(read_data=filedata)), patch('os.stat', return_value=os.stat_result((123,) * 10)), patch('ppadb.sync_async.open', mock_open(read_data=filedata)):
+        with patch('os.path.exists', return_value=True), patch('os.path.isfile', return_value=True), patch('os.stat', return_value=os.stat_result((123,) * 10)), patch('ppadb.sync_async.aiofiles.open', async_mock_open(FILEDATA)):
             with async_patch('asyncio.open_connection', return_value=(FakeStreamReader(), FakeStreamWriter())):
                 with async_patch('{}.FakeStreamReader.read'.format(__name__), side_effect=[b'OKAY', b'OKAY', b'OKAY', PNG_IMAGE_NEEDS_REPLACING, b'', b'OKAY']):
                     await self.device.push('src', 'dest', progress=progress)
@@ -87,14 +90,14 @@ async def test_push_dir(self):
     async def test_pull(self):
         with async_patch('asyncio.open_connection', return_value=(FakeStreamReader(), FakeStreamWriter())):
             with async_patch('{}.FakeStreamReader.read'.format(__name__), side_effect=[b'OKAY', b'OKAY', b'DATA', SyncAsync._little_endian(4), b'TEST', b'DONE', b'OKAY']):
-                with patch('ppadb.sync_async.open', mock_open()):
+                with patch('ppadb.sync_async.aiofiles.open', async_mock_open(FILEDATA)):
                     await self.device.pull('src', 'dest')
 
     @awaiter
     async def test_pull_fail(self):
         with async_patch('asyncio.open_connection', return_value=(FakeStreamReader(), FakeStreamWriter())):
             with async_patch('{}.FakeStreamReader.read'.format(__name__), side_effect=[b'OKAY', b'OKAY', b'FAIL', SyncAsync._little_endian(4), b'TEST', b'DONE', b'OKAY']):
-                with patch('ppadb.sync_async.open', mock_open()):
+                with patch('ppadb.sync_async.aiofiles.open', async_mock_open(FILEDATA)):
                     await self.device.pull('src', 'dest')