add GetBattorPathFromPhoneSerial to find_usb_devices

devil/devil/utils/find_usb_devices.py

 #!/usr/bin/python
 # Copyright 2016 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.
 
 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
 # as observed by the system (e.g. by usb-devices) and "physical port" refers
 # to the physical numerical label on the physical port e.g. on a USB hub.
 # The mapping between virtual and physical ports is not always the identity
 # (e.g. the port labeled "1" on a USB hub does not always show up as "port 1"
 # when you plug something into it) but, as far as we are aware, the mapping
 # between virtual and physical ports is always the same for a given
 # model of USB hub. When "port number" is referenced without specifying, it
@@ skipping 25 matching lines @@
 
 def GetBattorList(device_tree_map):
   return [x for x in GetTTYList() if IsBattor(x, device_tree_map)]
 
 
 def IsBattor(tty_string, device_tree_map):
   (bus, device) = GetBusDeviceFromTTY(tty_string)
   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.
 class USBNode(object):
   def __init__(self):
     self._port_to_node = {}
 
   @property
   def desc(self):
     raise NotImplementedError
@@ skipping 178 matching lines @@
 
 
 _T_LINE_REGEX = re.compile(r'T:  Bus=(?P<bus>\d{2}) Lev=(?P<lev>\d{2}) '
                            r'Prnt=(?P<prnt>\d{2,3}) Port=(?P<port>\d{2}) '
                            r'Cnt=(?P<cnt>\d{2}) Dev#=(?P<dev>.{3}) .*')
 
 _S_LINE_REGEX = re.compile(r'S:  SerialNumber=(?P<serial>.*)')
 _LSUSB_BUS_DEVICE_RE = re.compile(r'^Bus (\d{3}) Device (\d{3}): (.*)')
 
 
-def GetBusNumberToDeviceTreeMap(fast=False):
+def GetBusNumberToDeviceTreeMap(fast=True):
   """Gets devices currently attached.
 
   Args:
     fast [bool]: whether to do it fast (only get description, not
     the whole dictionary, from lsusb)
 
   Returns:
     map of {bus number: bus object}
     where the bus object has all the devices attached to it in a tree.
   """
@@ skipping 39 matching lines @@
     match = _S_LINE_REGEX.match(line)
     if match:
       if bus_num == -1:
         raise ValueError('S line appears before T line in input file')
       # put the serial number in the device
       tree[bus_num].FindDeviceNumber(device_num).serial = match.group('serial')
 
   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)
   """
   for device in bus.AllNodes():
     for hub_type in hub_types:
       if hub_type.IsType(device):
         yield (device, hub_type)
 
 
 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}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(hub)
   return {port: device for (port, device) in port_device}
 
 
 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)}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(hub)
   return {port: (device.bus_num, device.device_num)
           for (port, device) in port_device}
 
 
 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)}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(hub)
   return {port: device.serial
           for (port, device) in port_device
           if device.serial}
 
 
 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)}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(device)
   bus_device_to_tty = GetBusDeviceToTTYMap()
   return {port: bus_device_to_tty[(device.bus_num, device.device_num)]
           for (port, device) in port_device
           if (device.bus_num, device.device_num) in bus_device_to_tty}
 
 
 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
 
   Yields:
     Sequence of dicts of {physical port: device} where the type of
     device depends on the ident keyword. Each dict is a separate hub.
   """
   if device_tree_map is None:
     device_tree_map = GetBusNumberToDeviceTreeMap(fast=fast)
@@ skipping 71 matching lines @@
   return result
 
 
 # This dictionary described the mapping between physical and
 # virtual ports on a Plugable 7-Port Hub (model USB2-HUB7BC).
 # Keys are the virtual ports, values are the physical port.
 # The entry 4:{1:4, 2:3, 3:2, 4:1} indicates that virtual port
 # 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
   print '==== DEVICE DISPLAY ===='
-  device_trees = GetBusNumberToDeviceTreeMap(fast=True)
+  device_trees = GetBusNumberToDeviceTreeMap()
   for device_tree in device_trees.values():
     device_tree.Display()
   print
 
   # 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):

[jbudorick, 2016/03/30 19:28:58]

I should've brought this up on whatever review introduced the BattOr utility
functions above, but better late than never: this module contains a mix of
generic USB logic and highly specific BattOr logic. Can we move the BattOr logic
into its own module?

+  """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>},
+  {'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.
+
+  Writes a JSON file consisting of a list of items of the following form:
+  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
+  {'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:
+      filename: Name of file to write.
+  """
+  result = []
+  for (phone, battor) in GenerateSerialMap(hub_types).iteritems():
+    result.append({'phone': phone, 'battor': battor})
+  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 dict of:
+  {<phone serial 1>: <battor serial 1>,
+   <phone serial 2>: <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

[jbudorick, 2016/03/30 19:28:58]

nit: add that this defaults to ['plugable_7port']

[alexandermont, 2016/03/30 23:35:53]

Done

+  """
+  hub_types = [usb_hubs.GetHubType(x)
+               for x in hub_types or ['plugable_7port']]
+  devtree = GetBusNumberToDeviceTreeMap()
+  battor_serials = list(GetBattorSerialNumbers(devtree))
+  p2serial = GetAllPhysicalPortToSerialMaps(hub_types,

[jbudorick, 2016/03/30 19:28:58]

nit: rename p2serial

[alexandermont, 2016/03/30 23:35:52]

Done

+                                            device_tree_map=devtree)
+  port_to_devices = collections.defaultdict(list)

[jbudorick, 2016/03/30 19:28:57]

Rather than a list of serials, wdyt about having the values in this dict be a
collections.namedtuple with two string fields, phone_serial and battor_serial?

[alexandermont, 2016/03/30 23:35:52]

See below.

+  result = {}
+  for hub in p2serial:
+    for (port, serial) in hub.iteritems():
+      port_to_devices[port].append(serial)

[jbudorick, 2016/03/30 19:28:58]

Following from the above, you'd then check whether the serial is a battor serial
here. If so, you'd set the serial as the battor_serial if one hasn't already be
sent (and raise a BattorError for duplicate battor serial numbers if it has). If
not, you'd set the serial as the phone_serial if one hasn't been set (and raise
a BattorError for multiple phones if it has).

You'd then go through port_to_devices and, for each entry with a valid
battor_serial, you'd:
 - ensure that phone_serial is not None and raise a BattorError if not
 - set result[phone_serial] = battor_serial

I think that'd be more clear than what you have here.

[alexandermont, 2016/03/30 23:35:53]

This won't work if the value in the dictionary is a namedtuple, since a
namedtuple is immutable (so you couldn't set the phone_serial once, and the
battor_serial a different time). But you could do this with a list of length 2.
Note that it is possible to have a mutable tuple if you use a specialized class
(e.g. recordclass) but it is probably not worth introducing an additional
dependency just for this.

I have also put more comments to try to explain this function better.

+  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 GetBattorPathFromPhoneSerial(serial, serial_map=None):
+  """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.
+
+  (1) If serial_map is a dict, serial_map is treated as a dictionary mapping

[jbudorick, 2016/03/30 19:28:57]

Why is this handling two arguments in one...? Handle them separately, i.e., have
serial_map and serial_map_file (or revise the call sites s.t. this only receives
one or the other).

[alexandermont, 2016/03/30 23:35:52]

Done

+  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 a string, serial_map is treated as the name of a
+  phone-to-BattOr mapping file (generated with GenerateSerialMapFile)
+  and this will be loaded and used as the dict to map port numbers to
+  BattOr serial numbers.
+
+  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()
+  all_battors = GetBattorList(devtree)
+  if len(all_battors) == 1:
+    return '/dev/' + all_battors[0]
+
+  if isinstance(serial_map, str):
+    serial_map = ReadSerialMapFile(serial_map)
+
+  if not isinstance(serial_map, dict):
+    raise ValueError('Incorrect data type for serial_map')
+
+  if not serial:
+    raise BattorError('Two or more BattOrs connected, no serial provided')
+
+  tty_string = _PhoneToPathMap(serial, serial_map, devtree)
+
+  if not tty_string:
+    raise BattorError('No device with given serial number detected.')
+
+  if IsBattor(tty_string, devtree):
+    return '/dev/' + tty_string
+  else:
+    raise BattorError('Device with given serial number is  not a BattOr.')
+
 def parse_options(argv):
   """Parses and checks the command-line options.
 
   Returns:
     A tuple containing the options structure and a list of categories to
     be traced.
   """
   USAGE = '''./find_usb_devices [--help]
     This script shows the mapping between USB devices and port numbers.
     Clients are not intended to call this script from the command line.
     Clients are intended to call the functions in this script directly.
     For instance, GetAllPhysicalPortToSerialMaps(...)
     Running this script with --help will display this message.
     Running this script without --help will display information about
     devices attached, TTY mapping, and serial number mapping,
     for testing purposes. See design document for API documentation.
   '''
   parser = argparse.ArgumentParser(usage=USAGE)
   return parser.parse_args(argv[1:])
 
 def main():
   parse_options(sys.argv)
   TestUSBTopologyScript()
 
 if __name__ == "__main__":
   sys.exit(main())