add GetBattorPathFromPhoneSerial to find_usb_devices

devil/devil/utils/find_usb_devices.py

Index: devil/devil/utils/find_usb_devices.py
diff --git a/devil/devil/utils/find_usb_devices.py b/devil/devil/utils/find_usb_devices.py
index 4982e46a66e0734c5a8b7db89ae950ccbcfc7a5c..522af4d0033d1c4ba5e22d56c2692310625bd48d 100755
--- a/devil/devil/utils/find_usb_devices.py
+++ b/devil/devil/utils/find_usb_devices.py
@@ -6,8 +6,12 @@
 import re
 import sys
 import argparse
+import json
+import collections
 
+from devil import base_error
 from devil.utils import cmd_helper
+from devil.utils import usb_hubs
 from devil.utils import lsusb
 
 # Note: In the documentation below, "virtual port" refers to the port number
@@ -53,6 +57,15 @@ def IsBattor(tty_string, device_tree_map):
   node = device_tree_map[bus].FindDeviceNumber(device)
   return 'Future Technology Devices International' in node.desc
 
+def GetBattorSerialNumbers(device_tree_map):
+  for x in GetTTYList():
+    if IsBattor(x, device_tree_map):
+      (bus, device) = GetBusDeviceFromTTY(x)
+      devnode = device_tree_map[bus].FindDeviceNumber(device)
+      yield devnode.serial
+
+class BattorError(base_error.BaseError):
+  pass
 
 # Class to identify nodes in the USB topology. USB topology is organized as
 # a tree.
@@ -311,83 +324,12 @@ def GetBusNumberToDeviceTreeMap(fast=False):
   return tree
 
 
-class HubType(object):
-  def __init__(self, id_func, port_mapping):
-    """Defines a type of hub.
-
-    Args:
-      id_func: [USBNode -> bool] is a function that can be run on a node
-        to determine if the node represents this type of hub.
-      port_mapping: [dict(int:(int|dict))] maps virtual to physical port
-        numbers. For instance, {3:1, 1:2, 2:3} means that virtual port 3
-        corresponds to physical port 1, virtual port 1 corresponds to physical
-        port 2, and virtual port 2 corresponds to physical port 3. In the
-        case of hubs with "internal" topology, this is represented by nested
-        maps. For instance, {1:{1:1,2:2},2:{1:3,2:4}} means, e.g. that the
-        device plugged into physical port 3 will show up as being connected
-        to port 1, on a device which is connected to port 2 on the hub.
-    """
-    self._id_func = id_func
-    # v2p = "virtual to physical" ports
-    self._v2p_port = port_mapping
-
-  def IsType(self, node):
-    """Determines if the given Node is a hub of this type.
-
-    Args:
-      node: [USBNode] Node to check.
-    """
-    return self._id_func(node)
-
-  def GetPhysicalPortToNodeTuples(self, node):
-    """Gets devices connected to the physical ports on a hub of this type.
-
-    Args:
-      node: [USBNode] Node representing a hub of this type.
-
-    Yields:
-      A series of (int, USBNode) tuples giving a physical port
-      and the USBNode connected to it.
-
-    Raises:
-      ValueError: If the given node isn't a hub of this type.
-    """
-    if self.IsType(node):
-      for res in self._GppHelper(node, self._v2p_port):
-        yield res
-    else:
-      raise ValueError('Node must be a hub of this type')
-
-  def _GppHelper(self, node, mapping):
-    """Helper function for GetPhysicalPortToNodeMap.
-
-    Gets devices connected to physical ports, based on device tree
-    rooted at the given node and the mapping between virtual and physical
-    ports.
-
-    Args:
-      node: [USBNode] Root of tree to search for devices.
-      mapping: [dict] Mapping between virtual and physical ports.
-
-    Yields:
-      A series of (int, USBNode) tuples giving a physical port
-      and the Node connected to it.
-    """
-    for (virtual, physical) in mapping.iteritems():
-      if node.HasPort(virtual):
-        if isinstance(physical, dict):
-          for res in self._GppHelper(node.PortToDevice(virtual), physical):
-            yield res
-        else:
-          yield (physical, node.PortToDevice(virtual))
-
-
 def GetHubsOnBus(bus, hub_types):
   """Scans for all hubs on a bus of given hub types.
 
   Args:
     bus: [USBNode] Bus object.
-    hub_types: [iterable(HubType)] Possible types of hubs.
+    hub_types: [iterable(usb_hubs.HubType)] Possible types of hubs.
 
   Yields:
     Sequence of tuples representing (hub, type of hub)
@@ -402,7 +344,7 @@ def GetPhysicalPortToNodeMap(hub, hub_type):
   """Gets physical-port:node mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: node}
@@ -415,7 +357,7 @@ def GetPhysicalPortToBusDeviceMap(hub, hub_type):
   """Gets physical-port:(bus#, device#) mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: (bus number, device number)}
@@ -429,7 +371,7 @@ def GetPhysicalPortToSerialMap(hub, hub_type):
   """Gets physical-port:serial# mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: serial number)}
@@ -444,7 +386,7 @@ def GetPhysicalPortToTTYMap(device, hub_type):
   """Gets physical-port:tty-string mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: tty-string)}
@@ -460,7 +402,7 @@ def CollectHubMaps(hub_types, map_func, device_tree_map=None, fast=False):
   """Runs a function on all hubs in the system and collects their output.
 
   Args:
-    hub_types: [HubType] List of possible hub types.
+    hub_types: [usb_hubs.HubType] List of possible hub types.
     map_func: [string] Function to run on each hub.
     device_tree: Previously constructed device tree map, if any.
     fast: Whether to construct device tree fast, if not already provided
@@ -552,30 +494,10 @@ def GetBusDeviceToTTYMap():
 # 4 connects to another 'virtual' hub that itself has the
 # virtual-to-physical port mapping {1:4, 2:3, 3:2, 4:1}.
 
-PLUGABLE_7PORT_LAYOUT = {1:7,
-                         2:6,
-                         3:5,
-                         4:{1:4, 2:3, 3:2, 4:1}}
 
 def TestUSBTopologyScript():
   """Test display and hub identification."""
   # Identification criteria for Plugable 7-Port Hub
-  def _is_plugable_7port_hub(node):
-    """Check if a node is a Plugable 7-Port Hub
-    (Model USB2-HUB7BC)
-    The topology of this device is a 4-port hub,
-    with another 4-port hub connected on port 4.
-    """
-    if not isinstance(node, USBDeviceNode):
-      return False
-    if '4-Port HUB' not in node.desc:
-      return False
-    if not node.HasPort(4):
-      return False
-    return '4-Port HUB' in node.PortToDevice(4).desc
-
-  plugable_7port = HubType(_is_plugable_7port_hub,
-                           PLUGABLE_7PORT_LAYOUT)
   print '==== USB TOPOLOGY SCRIPT TEST ===='
 
   # Display devices
@@ -587,19 +509,199 @@ def TestUSBTopologyScript():
 
   # Display TTY information about devices plugged into hubs.
   print '==== TTY INFORMATION ===='
-  for port_map in GetAllPhysicalPortToTTYMaps([plugable_7port],
+  for port_map in GetAllPhysicalPortToTTYMaps([usb_hubs.PLUGABLE_7PORT],
                                               device_tree_map=device_trees):
     print port_map
   print
 
   # Display serial number information about devices plugged into hubs.
   print '==== SERIAL NUMBER INFORMATION ===='
-  for port_map in GetAllPhysicalPortToSerialMaps([plugable_7port],
+  for port_map in GetAllPhysicalPortToSerialMaps([usb_hubs.PLUGABLE_7PORT],
                                                  device_tree_map=device_trees):
     print port_map
+
+  print
+  # Display phone to BattOr map
+  print GenerateSerialMap()
   print ''
+
   return 0
 
+
+def ReadSerialMapFile(filename):
+  """Reads JSON file giving phone-to-battor serial number map.
+
+  Parses a JSON file consisting of a list of items of the following form:
+  [{'phone':<phone serial 1>, 'battor':<battor serial 1>},

[rnephew (Reviews Here), 2016/03/25 00:17:53]

Nit: Spaces here too. See below.

[alexandermont, 2016/03/28 18:47:29]

Done

+  {'phone':<phone serial 2>, 'battor':<battor serial 2>}, ...]
+
+  indicating which phone serial numbers should be matched with
+  which BattOr serial numbers. Returns dictionary of the form:
+
+  {<phone serial 1>: <BattOr serial 1>,
+   <phone serial 2>: <BattOr serial 2>}
+
+  Args:
+      filename: Name of file to read.
+  """
+  result = {}
+  with open(filename, 'r') as infile:
+    in_dict = json.load(infile)
+  for x in in_dict:
+    result[x['phone']] = x['battor']
+  return result
+
+def GenerateSerialMapFile(filename, hub_types=None):
+  """Writes a map of phone serial numbers to BattOr serial numbers to file.
+
+  Args:
+      filename: Name of file to write.
+  """
+  result = []
+  for (phone, battor) in GenerateSerialMap(hub_types).iteritems():
+    result.append({'phone':phone, 'battor':battor})

[rnephew (Reviews Here), 2016/03/25 00:17:53]

nit: 'phone': phone, 'battor': battor

note the spaces.

[alexandermont, 2016/03/28 18:47:29]

Done

+  with open(filename, 'w') as outfile:
+    json.dump(result, outfile)
+
+def GenerateSerialMap(hub_types=None):
+  """Generates a map of phone serial numbers to BattOr serial numbers.
+
+  Generates a JSON file consisting of a list of items of the following form:
+  [{'phone':<phone serial 1>, 'battor':<battor serial 1>},

[rnephew (Reviews Here), 2016/03/25 00:17:52]

Spaces here too.

[alexandermont, 2016/03/28 18:47:29]

Done

+  {'phone':<phone serial 2>, 'battor':<battor serial 2>}, ...]
+
+  indicating which phone serial numbers should be matched with
+  which BattOr serial numbers. Mapping is based on the physical port numbers
+  of the hubs that the BattOrs and phones are connected to.
+
+  Args:
+      hub_types: List of hub types to check for
+  """
+  hub_types = [usb_hubs.GetHubType(x)
+               for x in hub_types or ['plugable_7port']]
+  devtree = GetBusNumberToDeviceTreeMap(fast=True)
+  battor_serials = list(GetBattorSerialNumbers(devtree))
+  p2serial = GetAllPhysicalPortToSerialMaps(hub_types,
+                                            device_tree_map=devtree)
+  port_to_devices = collections.defaultdict(list)
+  result = {}
+  for hub in p2serial:
+    for (port, serial) in hub.iteritems():
+      port_to_devices[port].append(serial)
+  for (port, serial_list) in port_to_devices.iteritems():
+    battor_serial = [x for x in serial_list if x in battor_serials]
+    if len(battor_serial) >= 2:
+      raise BattorError('Duplicate BattOr serial numbers detected')
+    if len(battor_serial) == 1:
+      battor_serial = battor_serial[0]
+      phone_serial = [x for x in serial_list if x != battor_serial]
+      if len(phone_serial) >= 2:
+        raise BattorError('Multiple phones matched with same BattOr')
+      if len(phone_serial) == 0:
+        raise BattorError('BattOr has no matching phone')
+      if battor_serial in result:
+        raise BattorError('Duplicate BattOr serial numbers detected')
+      result[phone_serial[0]] = battor_serial
+  return result
+
+def _PhoneToPathMap(serial, serial_map, devtree):
+  """Maps phone serial number to TTY path, assuming serial map is provided."""
+  battor_serial = serial_map[serial]
+  for tree in devtree.values():
+    for node in tree.AllNodes():
+      if isinstance(node, USBDeviceNode):
+        if node.serial == battor_serial:
+          bus_device_to_tty = GetBusDeviceToTTYMap()
+          bus_device = (node.bus_num, node.device_num)
+          try:
+            return bus_device_to_tty[bus_device]
+          except KeyError:
+            raise BattorError('Device with given serial number not a BattOr '
+                              '(does not have TTY path)')
+
+def _PhoneToPathPorts(serial, hub_types, devtree):
+  """Maps phone serial number to TTY path, assuming hub types are provided."""
+  hub_types = [usb_hubs.GetHubType(x)
+               for x in hub_types or ['plugable_7port']]
+  p2serial = GetAllPhysicalPortToSerialMaps(hub_types,
+                                            device_tree_map=devtree)
+  p2tty = GetAllPhysicalPortToTTYMaps(hub_types,
+                                      device_tree_map=devtree)
+
+  # get the port number of this device
+  port_num = -1
+  for hub in p2serial:
+    for (port, s) in hub.iteritems():
+      if serial == s:
+        port_num = port
+  if port_num == -1:
+    raise BattorError('Device with given serial number not found.')
+
+  # get the tty for this port number
+  tty_string = None
+  for hub in p2tty:
+    x = hub.get(port_num)
+    if x is not None:
+      if tty_string:
+        raise BattorError('Two TTY devices with matching port number.')
+      else:
+        tty_string = x
+  return tty_string
+
+def GetBattorPathFromPhoneSerial(serial, serial_map=None, hub_types=None):

[nednguyen, 2016/03/25 03:05:29]

In which use case a client will not want to use serial_map but use "hub_types"
instead? I usually advocate for simpler API, so unless there is a use case which
people can't use serial_map and have to use hub_types, I would say deprecating
hub_types to simplify your design.

[rnephew (Reviews Here), 2016/03/25 05:03:19]

Thinking more about it, that would most likely happen outside of this function.

[nednguyen, 2016/03/27 16:44:32]

I see. Can we leave the hub_types param out and instead have the callsites
translate hubtypes to serial_map themselves? That's way, it's less ambiguous
about what this function will do when both serial_map & hub_types are present.

[alexandermont, 2016/03/28 18:47:29]

Done

+  """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.
+
+  (1) If serial_map is available, serial_map is treated as a dictionary mapping
+  phone serial numbers to BattOr serial numbers. This function will get the
+  TTY path for the given BattOr serial number.
+
+  (2) If serial_map is None and hub_types is available, hub_types is treated
+  as a list of strings representing the types of hubs that the phones and
+  BattOrs are plugged into (see usb_hubs.py for details). This function will
+  get the BattOr that is connected to the same physical port number,
+  on a different hub, that the phone is connected to.
+
+  (3) If serial_map is None and hub_types is None, then it will assume that the
+  type of hub the BattOrs and phones are plugged into are Plugable USB
+  7-Port Hubs, and then proceed as in (2).
+
+  Args:
+    serial: Serial number of phone connected on the same physical port that
+    the BattOr is connected to.
+    serial_map_file: Map of phone serial numbers to BattOr serial numbers.
+    hub_types: List of hub types to check for. Used only if serial_map_file
+    is None.
+
+  Returns:
+    Device string used to communicate with device.
+
+  Raises:
+    ValueError: If serial number is not given.
+    BattorError: If BattOr not found or unexpected USB topology.
+  """
+  # If there's only one BattOr connected to the system, just use that one.
+  # This allows for use on, e.g., a developer's workstation with no hubs.
+  devtree = GetBusNumberToDeviceTreeMap(fast=True)
+  all_battors = GetBattorList(devtree)
+  if len(all_battors) == 1:
+    return '/dev/' + all_battors[0]
+  if not serial:
+    raise BattorError('Two or more BattOrs connected, no serial provided')
+
+  if serial_map is not None:
+    err_string = 'with given serial number'
+    tty_string = _PhoneToPathMap(serial, serial_map, devtree)
+

[rnephew (Reviews Here), 2016/03/25 05:02:41]

nit: no space here.

[alexandermont, 2016/03/28 18:47:29]

Done

+  else:
+    err_string = 'on matching port'
+    tty_string = _PhoneToPathPorts(serial, hub_types, devtree)
+  if not tty_string:

[rnephew (Reviews Here), 2016/03/25 05:02:41]

space here instead.

[alexandermont, 2016/03/28 18:47:29]

Done

+    raise BattorError('No device %s detected.' % err_string)
+  if IsBattor(tty_string, devtree):
+    return '/dev/' + tty_string
+  else:
+    raise BattorError('Device %s is not a BattOr.' % err_string)
+
 def parse_options(argv):
   """Parses and checks the command-line options.