Working on a better timeout

adb_shell/adb_device.py

@@ -257,7 +257,7 @@ def connect(self, rsa_keys=None, transport_timeout_s=None, auth_timeout_s=consta
     #                                 Services                                #
     #                                                                         #
     # ======================================================================= #
-    def _service(self, service, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True):
+    def _service(self, service, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None, decode=True):
         """Send an ADB command to the device.
 
         Parameters
@@ -271,6 +271,8 @@ def _service(self, service, command, transport_timeout_s=None, read_timeout_s=co
             and :meth:`BaseTransport.bulk_write() <adb_shell.transport.base_transport.BaseTransport.bulk_write>`
         read_timeout_s : float
             The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`AdbDevice._read`
+        timeout_s : float, None
+            The total time in seconds to wait for the ADB command to finish
         decode : bool
             Whether to decode the output to utf8 before returning
 
@@ -280,7 +282,7 @@ def _service(self, service, command, transport_timeout_s=None, read_timeout_s=co
             The output of the ADB command as a string if ``decode`` is True, otherwise as bytes.
 
         """
-        adb_info = _AdbTransactionInfo(None, None, transport_timeout_s, read_timeout_s)
+        adb_info = _AdbTransactionInfo(None, None, transport_timeout_s, read_timeout_s, timeout_s)
         if decode:
             return b''.join(self._streaming_command(service, command, adb_info)).decode('utf8')
         return b''.join(self._streaming_command(service, command, adb_info))
@@ -317,7 +319,7 @@ def _streaming_service(self, service, command, transport_timeout_s=None, read_ti
             for line in stream:
                 yield line
 
-    def root(self, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S):
+    def root(self, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None):
         """Gain root access.
 
         The device must be rooted in order for this to work.
@@ -329,14 +331,16 @@ def root(self, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_T
             and :meth:`BaseTransport.bulk_write() <adb_shell.transport.base_transport.BaseTransport.bulk_write>`
         read_timeout_s : float
             The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`AdbDevice._read`
+        timeout_s : float, None
+            The total time in seconds to wait for the ADB command to finish
 
         """
         if not self.available:
             raise exceptions.AdbConnectionError("ADB command not sent because a connection to the device has not been established.  (Did you call `AdbDevice.connect()`?)")
 
-        self._service(b'root', b'', transport_timeout_s, read_timeout_s, False)
+        self._service(b'root', b'', transport_timeout_s, read_timeout_s, timeout_s, False)
 
-    def shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True):
+    def shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None, decode=True):
         """Send an ADB shell command to the device.
 
         Parameters
@@ -348,6 +352,8 @@ def shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFA
             and :meth:`BaseTransport.bulk_write() <adb_shell.transport.base_transport.BaseTransport.bulk_write>`
         read_timeout_s : float
             The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`AdbDevice._read`
+        timeout_s : float, None
+            The total time in seconds to wait for the ADB command to finish
         decode : bool
             Whether to decode the output to utf8 before returning
 
@@ -360,7 +366,7 @@ def shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFA
         if not self.available:
             raise exceptions.AdbConnectionError("ADB command not sent because a connection to the device has not been established.  (Did you call `AdbDevice.connect()`?)")
 
-        return self._service(b'shell', command.encode('utf8'), transport_timeout_s, read_timeout_s, decode)
+        return self._service(b'shell', command.encode('utf8'), transport_timeout_s, read_timeout_s, timeout_s, decode)
 
     def streaming_shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True):
         """Send an ADB shell command to the device, yielding each line of output.
@@ -869,6 +875,8 @@ def _read_until_close(self, adb_info):
             The data that was read by :meth:`AdbDevice._read_until`
 
         """
+        start = time.time()
+
         while True:
             cmd, data = self._read_until([constants.CLSE, constants.WRTE], adb_info)
 
@@ -879,6 +887,10 @@ def _read_until_close(self, adb_info):
 
             yield data
 
+            # Make sure the ADB command has not timed out
+            if adb_info.timeout_s is not None and time.time() - start > adb_info.timeout_s:
+                raise exceptions.AdbTimeoutError("The command did not complete within {} seconds".format(adb_info.timeout_s))
+
     def _send(self, msg, adb_info):
         """Send a message to the device.
 

adb_shell/adb_device_async.py

@@ -250,7 +250,7 @@ async def connect(self, rsa_keys=None, transport_timeout_s=None, auth_timeout_s=
     #                                 Services                                #
     #                                                                         #
     # ======================================================================= #
-    async def _service(self, service, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True):
+    async def _service(self, service, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None, decode=True):
         """Send an ADB command to the device.
 
         Parameters
@@ -264,6 +264,8 @@ async def _service(self, service, command, transport_timeout_s=None, read_timeou
             and :meth:`BaseTransportAsync.bulk_write() <adb_shell.transport.base_transport_async.BaseTransportAsync.bulk_write>`
         read_timeout_s : float
             The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`AdbDeviceAsync._read`
+        timeout_s : float, None
+            The total time in seconds to wait for the ADB command to finish
         decode : bool
             Whether to decode the output to utf8 before returning
 
@@ -273,7 +275,7 @@ async def _service(self, service, command, transport_timeout_s=None, read_timeou
             The output of the ADB command as a string if ``decode`` is True, otherwise as bytes.
 
         """
-        adb_info = _AdbTransactionInfo(None, None, transport_timeout_s, read_timeout_s)
+        adb_info = _AdbTransactionInfo(None, None, transport_timeout_s, read_timeout_s, timeout_s)
         if decode:
             return b''.join([x async for x in self._streaming_command(service, command, adb_info)]).decode('utf8')
         return b''.join([x async for x in self._streaming_command(service, command, adb_info)])
@@ -310,7 +312,7 @@ async def _streaming_service(self, service, command, transport_timeout_s=None, r
             async for line in stream:
                 yield line
 
-    async def root(self, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S):
+    async def root(self, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None):
         """Gain root access.
 
         The device must be rooted in order for this to work.
@@ -322,14 +324,16 @@ async def root(self, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_
             and :meth:`BaseTransportAsync.bulk_write() <adb_shell.transport.base_transport_async.BaseTransportAsync.bulk_write>`
         read_timeout_s : float
             The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`AdbDeviceAsync._read`
+        timeout_s : float, None
+            The total time in seconds to wait for the ADB command to finish
 
         """
         if not self.available:
             raise exceptions.AdbConnectionError("ADB command not sent because a connection to the device has not been established.  (Did you call `AdbDeviceAsync.connect()`?)")
 
-        await self._service(b'root', b'', transport_timeout_s, read_timeout_s, False)
+        await self._service(b'root', b'', transport_timeout_s, read_timeout_s, timeout_s, False)
 
-    async def shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True):
+    async def shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None, decode=True):
         """Send an ADB shell command to the device.
 
         Parameters
@@ -341,6 +345,8 @@ async def shell(self, command, transport_timeout_s=None, read_timeout_s=constant
             and :meth:`BaseTransportAsync.bulk_write() <adb_shell.transport.base_transport_async.BaseTransportAsync.bulk_write>`
         read_timeout_s : float
             The total time in seconds to wait for a ``b'CLSE'`` or ``b'OKAY'`` command in :meth:`AdbDeviceAsync._read`
+        timeout_s : float, None
+            The total time in seconds to wait for the ADB command to finish
         decode : bool
             Whether to decode the output to utf8 before returning
 
@@ -353,7 +359,7 @@ async def shell(self, command, transport_timeout_s=None, read_timeout_s=constant
         if not self.available:
             raise exceptions.AdbConnectionError("ADB command not sent because a connection to the device has not been established.  (Did you call `AdbDeviceAsync.connect()`?)")
 
-        return await self._service(b'shell', command.encode('utf8'), transport_timeout_s, read_timeout_s, decode)
+        return await self._service(b'shell', command.encode('utf8'), transport_timeout_s, read_timeout_s, timeout_s, decode)
 
     async def streaming_shell(self, command, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, decode=True):
         """Send an ADB shell command to the device, yielding each line of output.
@@ -862,6 +868,8 @@ async def _read_until_close(self, adb_info):
             The data that was read by :meth:`AdbDeviceAsync._read_until`
 
         """
+        start = time.time()
+
         while True:
             cmd, data = await self._read_until([constants.CLSE, constants.WRTE], adb_info)
 
@@ -872,6 +880,10 @@ async def _read_until_close(self, adb_info):
 
             yield data
 
+            # Make sure the ADB command has not timed out
+            if adb_info.timeout_s is not None and time.time() - start > adb_info.timeout_s:
+                raise exceptions.AdbTimeoutError("The command did not complete within {} seconds".format(adb_info.timeout_s))
+
     async def _send(self, msg, adb_info):
         """Send a message to the device.
 

adb_shell/adb_message.py

@@ -88,7 +88,7 @@ def unpack(message):
     arg1 : int
         TODO
     data_length : int
-        The length of the data sent by the device (used by :meth:`adb_shell.adb_device._read`)
+        The length of the data sent by the device (used by :meth:`adb_shell.adb_device.AdbDevice._read` and :meth:`adb_shell.adb_device_async.AdbDeviceAsync._read`)
     data_checksum : int
         The checksum of the data sent by the device
 
@@ -114,18 +114,18 @@ class AdbMessage(object):
     command : bytes
         A command; examples used in this package include :const:`adb_shell.constants.AUTH`, :const:`adb_shell.constants.CNXN`, :const:`adb_shell.constants.CLSE`, :const:`adb_shell.constants.OPEN`, and :const:`adb_shell.constants.OKAY`
     arg0 : int
-        Usually the local ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` provides :const:`adb_shell.constants.VERSION`, :const:`adb_shell.constants.AUTH_SIGNATURE`, and :const:`adb_shell.constants.AUTH_RSAPUBLICKEY`
+        Usually the local ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` and :meth:`~adb_shell.adb_device_async.AdbDeviceAsync.connect` provide :const:`adb_shell.constants.VERSION`, :const:`adb_shell.constants.AUTH_SIGNATURE`, and :const:`adb_shell.constants.AUTH_RSAPUBLICKEY`
     arg1 : int
-        Usually the remote ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` provides :const:`adb_shell.constants.MAX_ADB_DATA`
+        Usually the remote ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` and :meth:`~adb_shell.adb_device_async.AdbDeviceAsync.connect` provide :const:`adb_shell.constants.MAX_ADB_DATA`
     data : bytes
         The data that will be sent
 
     Attributes
     ----------
     arg0 : int
-        Usually the local ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` provides :const:`adb_shell.constants.VERSION`, :const:`adb_shell.constants.AUTH_SIGNATURE`, and :const:`adb_shell.constants.AUTH_RSAPUBLICKEY`
+        Usually the local ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` and :meth:`~adb_shell.adb_device_async.AdbDeviceAsync.connect` provide :const:`adb_shell.constants.VERSION`, :const:`adb_shell.constants.AUTH_SIGNATURE`, and :const:`adb_shell.constants.AUTH_RSAPUBLICKEY`
     arg1 : int
-        Usually the remote ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` provides :const:`adb_shell.constants.MAX_ADB_DATA`
+        Usually the remote ID, but :meth:`~adb_shell.adb_device.AdbDevice.connect` and :meth:`~adb_shell.adb_device_async.AdbDeviceAsync.connect` provide :const:`adb_shell.constants.MAX_ADB_DATA`
     command : int
         The input parameter ``command`` converted to an integer via :const:`adb_shell.constants.ID_TO_WIRE`
     data : bytes

adb_shell/constants.py

@@ -75,7 +75,7 @@
 SEND = b'SEND'
 STAT = b'STAT'
 
-#: Commands that are recognized by :meth:`adb_shell.adb_device.AdbDevice._read`
+#: Commands that are recognized by :meth:`adb_shell.adb_device.AdbDevice._read` and :meth:`adb_shell.adb_device_async.AdbDeviceAsync._read`
 IDS = (AUTH, CLSE, CNXN, OKAY, OPEN, SYNC, WRTE)
 
 #: A dictionary where the keys are the commands in :const:`IDS` and the values are the keys converted to integers
@@ -84,7 +84,7 @@
 #: A dictionary where the keys are integers and the values are their corresponding commands (type = bytes) from :const:`IDS`
 WIRE_TO_ID = {wire: cmd_id for cmd_id, wire in ID_TO_WIRE.items()}
 
-#: Commands that are recognized by :meth:`adb_shell.adb_device.AdbDevice._filesync_read`
+#: Commands that are recognized by :meth:`adb_shell.adb_device.AdbDevice._filesync_read` and :meth:`adb_shell.adb_device_async.AdbDeviceAsync._filesync_read`
 FILESYNC_IDS = (DATA, DENT, DONE, FAIL, LIST, OKAY, QUIT, RECV, SEND, STAT)
 
 #: A dictionary where the keys are the commands in :const:`FILESYNC_IDS` and the values are the keys converted to integers
@@ -111,8 +111,8 @@
 #: The size of an ADB message
 MESSAGE_SIZE = struct.calcsize(MESSAGE_FORMAT)
 
-#: Default authentication timeout (in s) for :meth:`adb_shell.tcp_transport.TcpTransport.connect`
+#: Default authentication timeout (in s) for :meth:`adb_shell.adb_device.AdbDevice.connect` and :meth:`adb_shell.adb_device_async.AdbDeviceAsync.connect`
 DEFAULT_AUTH_TIMEOUT_S = 10.
 
-#: Default total timeout (in s) for :meth:`adb_shell.adb_device.AdbDevice._read`
+#: Default total timeout (in s) for :meth:`adb_shell.adb_device.AdbDevice._read`, :meth:`adb_shell.adb_device.AdbDevice._read_until`, :meth:`adb_shell.adb_device_async.AdbDeviceAsync._read`, and :meth:`adb_shell.adb_device_async.AdbDeviceAsync._read_until`
 DEFAULT_READ_TIMEOUT_S = 10.

adb_shell/exceptions.py

@@ -38,6 +38,12 @@ class AdbConnectionError(Exception):
     """
 
 
+class AdbTimeoutError(Exception):
+    """ADB command did not complete within the specified time.
+
+    """
+
+
 class DeviceAuthError(Exception):
     """Device authentication failed.
 

adb_shell/hidden_helpers.py

@@ -81,6 +81,15 @@ def _open(name, mode='r'):
 class _AdbTransactionInfo(object):  # pylint: disable=too-few-public-methods
     """A class for storing info and settings used during a single ADB "transaction."
 
+    Note that if ``timeout_s`` is not ``None``, then:
+
+    ::
+
+       self.transport_timeout_s <= self.read_timeout_s <= self.timeout_s
+
+    If ``timeout_s`` is ``None``, the first inequality still applies.
+
+
     Parameters
     ----------
     local_id : int
@@ -94,6 +103,8 @@ class _AdbTransactionInfo(object):  # pylint: disable=too-few-public-methods
         :meth:`BaseTransportAsync.bulk_write() <adb_shell.transport.base_transport_async.BaseTransportAsync.bulk_write>`
     read_timeout_s : float
         The total time in seconds to wait for a command in ``expected_cmds`` in :meth:`AdbDevice._read` and :meth:`AdbDeviceAsync._read`
+    timeout_s : float, None
+        The total time in seconds to wait for the ADB command to finish
 
     Attributes
     ----------
@@ -103,18 +114,21 @@ class _AdbTransactionInfo(object):  # pylint: disable=too-few-public-methods
         The total time in seconds to wait for a command in ``expected_cmds`` in :meth:`AdbDevice._read` and :meth:`AdbDeviceAsync._read`
     remote_id : int
         The ID for the recipient
+    timeout_s : float, None
+        The total time in seconds to wait for the ADB command to finish
     transport_timeout_s : float, None
         Timeout in seconds for sending and receiving packets, or ``None``; see :meth:`BaseTransport.bulk_read() <adb_shell.transport.base_transport.BaseTransport.bulk_read>`,
         :meth:`BaseTransport.bulk_write() <adb_shell.transport.base_transport.BaseTransport.bulk_write>`,
         :meth:`BaseTransportAsync.bulk_read() <adb_shell.transport.base_transport_async.BaseTransportAsync.bulk_read>`, and
         :meth:`BaseTransportAsync.bulk_write() <adb_shell.transport.base_transport_async.BaseTransportAsync.bulk_write>`
 
     """
-    def __init__(self, local_id, remote_id, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S):
+    def __init__(self, local_id, remote_id, transport_timeout_s=None, read_timeout_s=constants.DEFAULT_READ_TIMEOUT_S, timeout_s=None):
         self.local_id = local_id
         self.remote_id = remote_id
-        self.transport_timeout_s = transport_timeout_s
-        self.read_timeout_s = read_timeout_s
+        self.timeout_s = timeout_s
+        self.read_timeout_s = read_timeout_s if self.timeout_s is None else min(read_timeout_s, self.timeout_s)
+        self.transport_timeout_s = self.read_timeout_s if transport_timeout_s is None else min(transport_timeout_s, self.read_timeout_s)
 
 
 class _FileSyncTransactionInfo(object):  # pylint: disable=too-few-public-methods

adb_shell/transport/tcp_transport.py

@@ -95,7 +95,7 @@ def connect(self, transport_timeout_s=None):
         if timeout:
             # Put the socket in non-blocking mode
             # https://docs.python.org/3/library/socket.html#socket.socket.settimeout
-            self._connection.setblocking(0)
+            self._connection.setblocking(False)
 
     def bulk_read(self, numbytes, transport_timeout_s=None):
         """Receive data from the socket.