1#!/usr/bin/env python
2#
3# Copyright 2015 Google Inc.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""A helper function that executes a series of List queries for many APIs."""
18
19from apitools.base.py import encoding
20import six
21
22__all__ = [
23    'YieldFromList',
24]
25
26
27def _GetattrNested(message, attribute):
28    """Gets a possibly nested attribute.
29
30    Same as getattr() if attribute is a string;
31    if attribute is a tuple, returns the nested attribute referred to by
32    the fields in the tuple as if they were a dotted accessor path.
33
34    (ex _GetattrNested(msg, ('foo', 'bar', 'baz')) gets msg.foo.bar.baz
35    """
36    if isinstance(attribute, six.string_types):
37        return getattr(message, attribute)
38    elif len(attribute) == 0:
39        return message
40    else:
41        return _GetattrNested(getattr(message, attribute[0]), attribute[1:])
42
43
44def _SetattrNested(message, attribute, value):
45    """Sets a possibly nested attribute.
46
47    Same as setattr() if attribute is a string;
48    if attribute is a tuple, sets the nested attribute referred to by
49    the fields in the tuple as if they were a dotted accessor path.
50
51    (ex _SetattrNested(msg, ('foo', 'bar', 'baz'), 'v') sets msg.foo.bar.baz
52    to 'v'
53    """
54    if isinstance(attribute, six.string_types):
55        return setattr(message, attribute, value)
56    elif len(attribute) < 1:
57        raise ValueError("Need an attribute to set")
58    elif len(attribute) == 1:
59        return setattr(message, attribute[0], value)
60    else:
61        return setattr(_GetattrNested(message, attribute[:-1]),
62                       attribute[-1], value)
63
64
65def YieldFromList(
66        service, request, global_params=None, limit=None, batch_size=100,
67        method='List', field='items', predicate=None,
68        current_token_attribute='pageToken',
69        next_token_attribute='nextPageToken',
70        batch_size_attribute='maxResults',
71        get_field_func=_GetattrNested):
72    """Make a series of List requests, keeping track of page tokens.
73
74    Args:
75      service: apitools_base.BaseApiService, A service with a .List() method.
76      request: protorpc.messages.Message, The request message
77          corresponding to the service's .List() method, with all the
78          attributes populated except the .maxResults and .pageToken
79          attributes.
80      global_params: protorpc.messages.Message, The global query parameters to
81           provide when calling the given method.
82      limit: int, The maximum number of records to yield. None if all available
83          records should be yielded.
84      batch_size: int, The number of items to retrieve per request.
85      method: str, The name of the method used to fetch resources.
86      field: str, The field in the response that will be a list of items.
87      predicate: lambda, A function that returns true for items to be yielded.
88      current_token_attribute: str or tuple, The name of the attribute in a
89          request message holding the page token for the page being
90          requested. If a tuple, path to attribute.
91      next_token_attribute: str or tuple, The name of the attribute in a
92          response message holding the page token for the next page. If a
93          tuple, path to the attribute.
94      batch_size_attribute: str or tuple, The name of the attribute in a
95          response message holding the maximum number of results to be
96          returned. None if caller-specified batch size is unsupported.
97          If a tuple, path to the attribute.
98      get_field_func: Function that returns the items to be yielded. Argument
99          is response message, and field.
100
101    Yields:
102      protorpc.message.Message, The resources listed by the service.
103
104    """
105    request = encoding.CopyProtoMessage(request)
106    _SetattrNested(request, current_token_attribute, None)
107    while limit is None or limit:
108        if batch_size_attribute:
109            # On Py3, None is not comparable so min() below will fail.
110            # On Py2, None is always less than any number so if batch_size
111            # is None, the request_batch_size will always be None regardless
112            # of the value of limit. This doesn't generally strike me as the
113            # correct behavior, but this change preserves the existing Py2
114            # behavior on Py3.
115            if batch_size is None:
116                request_batch_size = None
117            else:
118                request_batch_size = min(batch_size, limit or batch_size)
119            _SetattrNested(request, batch_size_attribute, request_batch_size)
120        response = getattr(service, method)(request,
121                                            global_params=global_params)
122        items = get_field_func(response, field)
123        if predicate:
124            items = list(filter(predicate, items))
125        for item in items:
126            yield item
127            if limit is None:
128                continue
129            limit -= 1
130            if not limit:
131                return
132        token = _GetattrNested(response, next_token_attribute)
133        if not token:
134            return
135        _SetattrNested(request, current_token_attribute, token)
136