diff --git a/vk_api/exceptions.py b/vk_api/exceptions.py index 648dc5ea..ce0406c4 100644 --- a/vk_api/exceptions.py +++ b/vk_api/exceptions.py @@ -155,3 +155,7 @@ class VkAudioException(Exception): class VkAudioUrlDecodeError(VkAudioException): pass + + +class VkToolsException(VkApiError): + pass diff --git a/vk_api/tools.py b/vk_api/tools.py index 8e69cc9f..d50d8efd 100644 --- a/vk_api/tools.py +++ b/vk_api/tools.py @@ -9,6 +9,7 @@ import sys +from .exceptions import ApiError, VkToolsException from .execute import VkFunction @@ -30,7 +31,7 @@ def __init__(self, vk): self.vk = vk def get_all_iter(self, method, max_count, values=None, key='items', - limit=None, stop_fn=None): + limit=None, stop_fn=None, negative_offset=False): """ Получить все элементы. Работает в методах, где в ответе есть count и items или users. За один запрос получает max_count * 25 элементов @@ -54,30 +55,36 @@ def get_all_iter(self, method, max_count, values=None, key='items', :param stop_fn: функция, отвечающая за выход из цикла :type stop_fn: func + + :param negative_offset: True если offset должен быть отрицательный + :type negative_offset: bool """ values = values.copy() if values else {} + values['count'] = max_count - items_count = 0 offset = 0 + items_count = 0 + count = None while True: - response = vk_get_all_items( - self.vk, method, values, key, max_count, offset - ) - - items = response.get('items') - offset = response.get('offset') - - if items is None or response.get('count') is None: - break # Error - + try: + response = vk_get_all_items( + self.vk, method, key, values, count, offset, + offset_mul=-1 if negative_offset else 1 + ) + except ApiError: + raise VkToolsException( + 'Can\'t load items. Check access to requested items' + ) + + items = response["items"] items_count += len(items) for item in items: yield item - if offset >= response['count']: + if not response['more']: break if limit and items_count >= limit: @@ -86,8 +93,11 @@ def get_all_iter(self, method, max_count, values=None, key='items', if stop_fn and stop_fn(items): break + count = response['count'] + offset = response['offset'] + def get_all(self, method, max_count, values=None, key='items', limit=None, - stop_fn=None): + stop_fn=None, negative_offset=False): """ Использовать только если нужно загрузить все объекты в память. Eсли вы можете обрабатывать объекты по частям, то лучше @@ -97,13 +107,16 @@ def get_all(self, method, max_count, values=None, key='items', limit=None, все данные в память """ - items = list(self.get_all_iter(method, max_count, values, key, limit, - stop_fn)) + items = list( + self.get_all_iter( + method, max_count, values, key, limit, stop_fn, negative_offset + ) + ) return {'count': len(items), key: items} def get_all_slow_iter(self, method, max_count, values=None, key='items', - limit=None, stop_fn=None): + limit=None, stop_fn=None, negative_offset=False): """ Получить все элементы (без использования execute) Работает в методах, где в ответе есть count и items или users @@ -126,34 +139,55 @@ def get_all_slow_iter(self, method, max_count, values=None, key='items', :param stop_fn: функция, отвечающая за выход из цикла :type stop_fn: func + + :param negative_offset: True если offset должен быть отрицательный + :type negative_offset: bool """ values = values.copy() if values else {} - values['count'] = max_count - response = self.vk.method(method, values) - count = response['count'] - items_count = 0 + offset_mul = -1 if negative_offset else 1 - for offset in range(max_count, count + 1, max_count): - values['offset'] = offset + offset = max_count if negative_offset else 0 + count = None + items_count = 0 + + while count is None or offset < count: + values['offset'] = offset * offset_mul response = self.vk.method(method, values) - items = response[key] + + new_count = response['count'] + + count_diff = (new_count - count) if count is not None else 0 + + if count_diff < 0: + offset += count_diff + count = new_count + continue + + response_items = response[key] + items = response_items[count_diff:] items_count += len(items) for item in items: yield item + if len(response_items) < max_count - count_diff: + break + if limit and items_count >= limit: break if stop_fn and stop_fn(items): break + offset += max_count + count = new_count + def get_all_slow(self, method, max_count, values=None, key='items', - limit=None, stop_fn=None): + limit=None, stop_fn=None, negative_offset=False): """ Использовать только если нужно загрузить все объекты в память. Eсли вы можете обрабатывать объекты по частям, то лучше @@ -164,34 +198,54 @@ def get_all_slow(self, method, max_count, values=None, key='items', """ items = list( - self.get_all_slow_iter(method, max_count, values, key, limit, - stop_fn) + self.get_all_slow_iter( + method, max_count, values, key, limit, stop_fn, negative_offset + ) ) return {'count': len(items), key: items} vk_get_all_items = VkFunction( - args=('method', 'values', 'key', 'max_count', 'start_offset'), - clean_args=('method', 'max_count', 'start_offset'), + args=('method', 'key', 'values', 'count', 'offset', 'offset_mul'), + clean_args=('method', 'key', 'offset', 'offset_mul'), code=''' - var max_count = %(max_count)s, - offset = %(start_offset)s, - key = %(key)s; - - var params = {count: max_count, offset: offset} + %(values)s; - - var r = API.%(method)s(params), - items = r[key], - i = 1; - - while(i < 25 && offset + max_count <= r.count) { - offset = offset + max_count; - params.offset = offset; - - items = items + API.%(method)s(params)[key]; - - i = i + 1; + var params = %(values)s, + calls = 0, + items = [], + count = %(count)s, + offset = %(offset)s, + ri; + + while(calls < 25) { + calls = calls + 1; + + params.offset = offset * %(offset_mul)s; + var response = API.%(method)s(params), + new_count = response.count, + count_diff = (count == null ? 0 : new_count - count); + + if (count_diff < 0) { + offset = offset + count_diff; + } else { + ri = response.%(key)s; + items = items + ri.slice(count_diff); + offset = offset + params.count + count_diff; + if (ri.length < params.count) { + calls = 99; + } + } + + count = new_count; + + if (count != null && offset >= count) { + calls = 99; + } }; - return {count: r.count, items: items, offset: offset + max_count}; + return { + count: count, + items: items, + offset: offset, + more: calls != 99 + }; ''')