Merge 943aedb into 1a950d5

apitools/base/py/list_pager.py

@@ -17,12 +17,51 @@
 """A helper function that executes a series of List queries for many APIs."""
 
 from apitools.base.py import encoding
+import six
 
 __all__ = [
     'YieldFromList',
 ]
 
 
+def _GetattrNested(message, attribute):
+    """Gets a possibly nested attribute.
+
+    Same as getattr() if attribute is a string;
+    if attribute is a tuple, returns the nested attribute referred to by
+    the fields in the tuple as if they were a dotted accessor path.
+
+    (ex _GetattrNested(msg, ('foo', 'bar', 'baz')) gets msg.foo.bar.baz
+    """
+    if isinstance(attribute, six.string_types):
+        return getattr(message, attribute)
+    elif len(attribute) == 0:
+        return message
+    else:
+        return _GetattrNested(getattr(message, attribute[0]), attribute[1:])
+
+
+def _SetattrNested(message, attribute, value):
+    """Sets a possibly nested attribute.
+
+    Same as setattr() if attribute is a string;
+    if attribute is a tuple, sets the nested attribute referred to by
+    the fields in the tuple as if they were a dotted accessor path.
+
+    (ex _SetattrNested(msg, ('foo', 'bar', 'baz'), 'v') sets msg.foo.bar.baz
+    to 'v'
+    """
+    if isinstance(attribute, six.string_types):
+        return setattr(message, attribute, value)
+    elif len(attribute) < 1:
+        raise ValueError("Need an attribute to set")
+    elif len(attribute) == 1:
+        return setattr(message, attribute[0], value)
+    else:
+        return setattr(_GetattrNested(message, attribute[:-1]),
+                       attribute[-1], value)
+
+
 def YieldFromList(
         service, request, global_params=None, limit=None, batch_size=100,
         method='List', field='items', predicate=None,
@@ -45,21 +84,23 @@ def YieldFromList(
       method: str, The name of the method used to fetch resources.
       field: str, The field in the response that will be a list of items.
       predicate: lambda, A function that returns true for items to be yielded.
-      current_token_attribute: str, The name of the attribute in a
+      current_token_attribute: str or tuple, The name of the attribute in a
           request message holding the page token for the page being
-          requested.
-      next_token_attribute: str, The name of the attribute in a
-          response message holding the page token for the next page.
-      batch_size_attribute: str, The name of the attribute in a
+          requested. If a tuple, path to attribute.
+      next_token_attribute: str or tuple, The name of the attribute in a
+          response message holding the page token for the next page. If a
+          tuple, path to the attribute.
+      batch_size_attribute: str or tuple, The name of the attribute in a
           response message holding the maximum number of results to be
           returned. None if caller-specified batch size is unsupported.
+          If a tuple, path to the attribute.
 
     Yields:
       protorpc.message.Message, The resources listed by the service.
 
     """
     request = encoding.CopyProtoMessage(request)
-    setattr(request, current_token_attribute, None)
+    _SetattrNested(request, current_token_attribute, None)
     while limit is None or limit:
         if batch_size_attribute:
             # On Py3, None is not comparable so min() below will fail.
@@ -72,10 +113,10 @@ def YieldFromList(
                 request_batch_size = None
             else:
                 request_batch_size = min(batch_size, limit or batch_size)
-            setattr(request, batch_size_attribute, request_batch_size)
+            _SetattrNested(request, batch_size_attribute, request_batch_size)
         response = getattr(service, method)(request,
                                             global_params=global_params)
-        items = getattr(response, field)
+        items = _GetattrNested(response, field)
         if predicate:
             items = list(filter(predicate, items))
         for item in items:
@@ -85,7 +126,7 @@ def YieldFromList(
             limit -= 1
             if not limit:
                 return
-        token = getattr(response, next_token_attribute)
+        token = _GetattrNested(response, next_token_attribute)
         if not token:
             return
-        setattr(request, current_token_attribute, token)
+        _SetattrNested(request, current_token_attribute, token)

apitools/base/py/list_pager_test.py

@@ -27,6 +27,32 @@
 from samples.iam_sample.iam_v1 import iam_v1_messages as iam_messages
 
 
+class Example(object):
+    def __init__(self):
+        self.a = 'aaa'
+        self.b = 'bbb'
+        self.c = 'ccc'
+
+
+class GetterSetterTest(unittest2.TestCase):
+
+    def testGetattrNested(self):
+        o = Example()
+        self.assertEqual(list_pager._GetattrNested(o, 'a'), 'aaa')
+        self.assertEqual(list_pager._GetattrNested(o, ('a',)), 'aaa')
+        o.b = Example()
+        self.assertEqual(list_pager._GetattrNested(o, ('b', 'c')), 'ccc')
+
+    def testSetattrNested(self):
+        o = Example()
+        list_pager._SetattrNested(o, 'b', Example())
+        self.assertEqual(o.b.a, 'aaa')
+        list_pager._SetattrNested(o, ('b', 'a'), 'AAA')
+        self.assertEqual(o.b.a, 'AAA')
+        list_pager._SetattrNested(o, ('c',), 'CCC')
+        self.assertEqual(o.c, 'CCC')
+
+
 class ListPagerTest(unittest2.TestCase):
 
     def _AssertInstanceSequence(self, results, n):