New API for persistent ADB shell

devil/devil/android/sdk/adb_wrapper.py

 # Copyright 2013 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """This module wraps Android's adb tool.
 
 This is a thin wrapper around the adb interface. Any additional complexity
 should be delegated to a higher level (ex. DeviceUtils).
 """
 
 import collections
 import errno
 import logging
 import os
 import re
+import subprocess
 
 from devil import devil_env
 from devil.android import decorators
 from devil.android import device_errors
 from devil.utils import cmd_helper
 from devil.utils import lazy
 from devil.utils import timeout_retry
 
 with devil_env.SysPath(devil_env.DEPENDENCY_MANAGER_PATH):
   import dependency_manager  # pylint: disable=import-error
@@ skipping 39 matching lines @@
 
 
 def _ShouldRetryAdbCmd(exc):
   return not isinstance(exc, device_errors.NoAdbError)
 
 
 DeviceStat = collections.namedtuple('DeviceStat',
                                     ['st_mode', 'st_size', 'st_time'])
 
 
+def _IsExtraneousLine(line, send_cmd):
+  """Determine if a line read from stdout in persistent shell is extraneous.
+
+  The results output to stdout by the persistent shell process
+  (in PersistentShell below) often include "extraneous" lines that are
+  not part of the output of the shell command. These "extraneous" lines
+  do not always appear and are of two forms: shell prompt lines and lines
+  that just duplicate what the input command was. This function
+  detects these extraneous lines. Since all these lines have the
+  original command in them, that is what it detects ror.
+
+  Args:
+      line: Output line to check.
+      send_cmd: Command that was sent to adb persistent shell.
+  """
+  return send_cmd.rstrip() in line
+
+
 class AdbWrapper(object):
   """A wrapper around a local Android Debug Bridge executable."""
 
   _adb_path = lazy.WeakConstant(_FindAdb)
 
   def __init__(self, device_serial):
     """Initializes the AdbWrapper.
 
     Args:
       device_serial: The device serial number as a string.
     """
     if not device_serial:
       raise ValueError('A device serial must be specified')
     self._device_serial = str(device_serial)
 
+  class PersistentShell(object):
+    '''Class to use persistent shell for ADB.
+
+    This class allows a persistent ADB shell to be created, where multiple
+    commands can be passed into it. This avoids the overhead of starting
+    up a new ADB shell for each command.
+
+    Example of use:
+    with pshell as PersistentShell('123456789'):
+        pshell.RunCommand('which ls')
+        pshell.RunCommandAndClose('echo TEST')
+    '''
+    def __init__(self, serial):
+      """Initialization function:
+
+      Args:
+        serial: Serial number of device.
+      """
+      self._cmd = [AdbWrapper.GetAdbPath(), '-s', serial, 'shell']
+      self._process = None
+
+    def __enter__(self):
+      self.Start()
+      self.WaitForReady()
+      return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+      self.Stop()
+
+    def Start(self):
+      """Start the shell."""
+      if self._process is not None:
+        raise RuntimeError('Persistent shell already running.')
+      self._process = subprocess.Popen(self._cmd,
+                                       stdin=subprocess.PIPE,
+                                       stdout=subprocess.PIPE,
+                                       shell=False)
+
+    def WaitForReady(self):
+      """Wait for the shell to be ready after starting.
+
+      Sends an echo command, then waits until it gets a response.
+      """
+      self._process.stdin.write('echo\n')
+      output_line = self._process.stdout.readline()
+      while output_line.rstrip() != '':
+        output_line = self._process.stdout.readline()
+
+    def RunCommand(self, command, close=False):
+      """Runs an ADB command and returns the output.
+
+      Note that there can be approximately 40 ms of additional latency
+      between sending the command and receiving the results if close=False
+      due to the use of Nagle's algorithm in the TCP socket between the
+      adb server and client. To avoid this extra latency, set close=True.
+
+      Args:
+        command: Command to send.
+      Returns:
+        The command output, given as a list of lines, and the exit code
+      """
+
+      if close:
+        def run_cmd(cmd):
+          send_cmd = '( %s ); echo $?; exit;\n' % cmd.rstrip()
+          (output, _) = self._process.communicate(send_cmd)
+          self._process = None
+          for x in output.splitlines():
+            yield x
+
+      else:
+        def run_cmd(cmd):
+          send_cmd = '( %s ); echo DONE:$?;\n' % cmd.rstrip()
+          self._process.stdin.write(send_cmd)
+          while True:
+            output_line = self._process.stdout.readline().rstrip()
+            if output_line[:5] == 'DONE:':
+              yield output_line[5:]
+              break
+            yield output_line
+
+      result = [line for line in run_cmd(command)
+                if not _IsExtraneousLine(line, command)]
+
+      return (result[:-1], int(result[-1]))
+
+    def Stop(self):
+      """Stops the ADB process if it is still running."""
+      if self._process is not None:
+        self._process.stdin.write('exit\n')
+        self._process = None
+
   @classmethod
   def GetAdbPath(cls):
     return cls._adb_path.read()
 
   @classmethod
   def _BuildAdbCmd(cls, args, device_serial, cpu_affinity=None):
     if cpu_affinity is not None:
       cmd = ['taskset', '-c', str(cpu_affinity)]
     else:
       cmd = []
@@ skipping 595 matching lines @@
   @property
   def is_emulator(self):
     return _EMULATOR_RE.match(self._device_serial)
 
   @property
   def is_ready(self):
     try:
       return self.GetState() == _READY_STATE
     except device_errors.CommandFailedError:
       return False