Merge pull request #46 from cloudblue/LITE-24467_fix_slicing_for_resourceset

LITE-24467 Fix slicing for ResourceSet

Author: Francesco Faraone

connect/client/models/iterators.py

@@ -40,7 +40,7 @@ def _load(self):
             self._results_iterator = iter(self._rs._results)
             self._loaded = True
 
-    def __next__(self):
+    def __next__(self):  # noqa: CCR001
         self._load()
 
         if not self._rs._results:
@@ -51,9 +51,14 @@ def __next__(self):
             if (
                 self._rs._content_range is None
                 or self._rs._content_range.last >= self._rs._content_range.count - 1
+                or (self._rs._slice and self._rs._content_range.last >= self._rs._slice.stop - 1)
             ):
                 raise
             self._config['params']['offset'] += self._config['params']['limit']
+            if self._rs._slice:
+                items_to_fetch = self._rs._slice.stop - self._rs._content_range.last - 1
+                if items_to_fetch < self._config['params']['limit']:
+                    self._config['params']['limit'] = items_to_fetch
             results, cr = self._execute_request()
             if not results:
                 raise
@@ -86,7 +91,7 @@ async def _load(self):
             self._results_iterator = iter(self._rs._results)
             self._loaded = True
 
-    async def __anext__(self):
+    async def __anext__(self):  # noqa: CCR001
         await self._load()
 
         if not self._rs._results:
@@ -97,9 +102,14 @@ async def __anext__(self):
             if (
                 self._rs._content_range is None
                 or self._rs._content_range.last >= self._rs._content_range.count - 1
+                or (self._rs._slice and self._rs._content_range.last >= self._rs._slice.stop - 1)
             ):
                 raise StopAsyncIteration
             self._config['params']['offset'] += self._config['params']['limit']
+            if self._rs._slice:
+                items_to_fetch = self._rs._slice.stop - self._rs._content_range.last - 1
+                if items_to_fetch < self._config['params']['limit']:
+                    self._config['params']['limit'] = items_to_fetch
             results, cr = await self._execute_request()
             if not results:
                 raise StopAsyncIteration

connect/client/models/resourceset.py

@@ -34,7 +34,7 @@ def __init__(
         self._results = None
         self._limit = self._client.default_limit or 100
         self._offset = 0
-        self._slice = False
+        self._slice = None
         self._content_range = None
         self._fields = None
         self._search = None
@@ -264,6 +264,21 @@ def _copy(self):
 
         return rs
 
+    def _validate_key(self, key):
+        if not isinstance(key, (int, slice)):
+            raise TypeError('ResourceSet indices must be integers or slices.')
+
+        if isinstance(key, slice) and (key.start is None or key.stop is None):
+            raise ValueError('Both start and stop indexes must be specified.')
+
+        if (not isinstance(key, slice) and (key < 0)) or (
+            isinstance(key, slice) and (key.start < 0 or key.stop < 0)
+        ):
+            raise ValueError('Negative indexing is not supported.')
+
+        if isinstance(key, slice) and not (key.step is None or key.step == 0):
+            raise ValueError('Indexing with step is not supported.')
+
     def help(self):
         """
         Output the ResourceSet documentation to the console.
@@ -301,7 +316,7 @@ def __bool__(self):
 
     def __getitem__(self, key):  # noqa: CCR001
         """
-        If called with and integer index, returns the item
+        If called with slice and integer index, returns the item
         at index ``key``.
 
         If key is a slice, set the pagination limit and offset
@@ -313,19 +328,7 @@ def __getitem__(self, key):  # noqa: CCR001
         :return: The resource at index ``key`` or self if ``key`` is a slice.
         :rtype: dict, ResultSet
         """
-        if not isinstance(key, (int, slice)):
-            raise TypeError('ResourceSet indices must be integers or slices.')
-
-        if isinstance(key, slice) and (key.start is None or key.stop is None):
-            raise ValueError('Both start and stop indexes must be specified.')
-
-        if (not isinstance(key, slice) and (key < 0)) or (
-            isinstance(key, slice) and (key.start < 0 or key.stop < 0)
-        ):
-            raise ValueError('Negative indexing is not supported.')
-
-        if isinstance(key, slice) and not (key.step is None or key.step == 0):
-            raise ValueError('Indexing with step is not supported.')
+        self._validate_key(key)
 
         if self._results is not None:
             return self._results[key]
@@ -339,8 +342,10 @@ def __getitem__(self, key):  # noqa: CCR001
 
         copy = self._copy()
         copy._offset = key.start
-        copy._limit = key.stop - key.start
-        copy._slice = True
+        copy._slice = key
+        if copy._slice.stop - copy._slice.start < copy._limit:
+            copy._limit = copy._slice.stop - copy._slice.start
+
         return copy
 
     def count(self):
@@ -429,7 +434,7 @@ def __bool__(self):
 
     def __getitem__(self, key):  # noqa: CCR001
         """
-        If called with and integer index, returns the item
+        If called with slice and integer index, returns the item
         at index ``key``.
 
         If key is a slice, set the pagination limit and offset
@@ -441,19 +446,7 @@ def __getitem__(self, key):  # noqa: CCR001
         :return: The resource at index ``key`` or self if ``key`` is a slice.
         :rtype: dict, ResultSet
         """
-        if not isinstance(key, (int, slice)):
-            raise TypeError('ResourceSet indices must be integers or slices.')
-
-        if isinstance(key, slice) and (key.start is None or key.stop is None):
-            raise ValueError('Both start and stop indexes must be specified.')
-
-        if (not isinstance(key, slice) and (key < 0)) or (
-            isinstance(key, slice) and (key.start < 0 or key.stop < 0)
-        ):
-            raise ValueError('Negative indexing is not supported.')
-
-        if isinstance(key, slice) and not (key.step is None or key.step == 0):
-            raise ValueError('Indexing with step is not supported.')
+        self._validate_key(key)
 
         if self._results is not None:
             return self._results[key]
@@ -463,8 +456,10 @@ def __getitem__(self, key):  # noqa: CCR001
 
         copy = self._copy()
         copy._offset = key.start
-        copy._limit = key.stop - key.start
-        copy._slice = True
+        copy._slice = key
+        if copy._slice.stop - copy._slice.start < copy._limit:
+            copy._limit = copy._slice.stop - copy._slice.start
+
         return copy
 
     async def count(self):

tests/async_client/test_models.py

@@ -790,21 +790,41 @@ async def test_rs_getitem_not_evaluated(async_client_mock, async_rs_factory):
         rs[0]
 
 
-def test_rs_getitem_slice(mocker, async_client_mock, async_rs_factory):
+@pytest.mark.asyncio
+async def test_rs_getitem_slice(mocker, async_client_mock, async_rs_factory):
     mocker.patch(
-        'connect.client.models.resourceset.parse_content_range',
-        return_value=ContentRange(0, 9, 10),
-    )
-    expected = [{'id': i} for i in range(10)]
-    rs = async_rs_factory(
-        client=async_client_mock(methods=['get']),
+        'connect.client.models.iterators.parse_content_range',
+        return_value=ContentRange(2, 3, 10),
     )
+    rs = async_rs_factory(client=async_client_mock(methods=['get']))
+    expected = [{'id': 2}, {'id': 3}]
     rs._client.get.return_value = expected
     sliced = rs[2:4]
+    results = [item async for item in sliced]
+    assert results == expected
     assert isinstance(sliced, AsyncResourceSet)
-    assert sliced._limit == 2
-    assert sliced._offset == 2
-    rs._client.get.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_rs_getitem_slice_limit(async_mocker, mocker, async_client_mock, async_rs_factory):
+    mocker.patch(
+        'connect.client.models.iterators.parse_content_range',
+        side_effect=[
+            ContentRange(0, 9, 25),
+            ContentRange(10, 19, 25),
+            ContentRange(20, 21, 25),
+        ],
+    )
+    rs = async_rs_factory(client=async_client_mock(methods=['get']))
+    rs._client.get = async_mocker.AsyncMock(side_effect=[
+        [{'id': i} for i in range(10)],
+        [{'id': i + 10} for i in range(10)],
+        [{'id': 20}, {'id': 21}],
+    ])
+    rs._limit = 10
+    sliced = rs[0:22]
+    results = [item async for item in sliced]
+    assert results == [{'id': i} for i in range(22)]
 
 
 def test_rs_getitem_slice_type(async_rs_factory):

tests/client/test_models.py

@@ -698,16 +698,55 @@ def test_rs_getitem(mocker, rs_factory):
 
 def test_rs_getitem_slice(mocker, rs_factory):
     mocker.patch(
-        'connect.client.models.resourceset.parse_content_range',
-        return_value=ContentRange(0, 9, 10),
+        'connect.client.models.iterators.parse_content_range',
+        return_value=ContentRange(2, 3, 10),
     )
-    expected = [{'id': i} for i in range(10)]
+    expected = [{'id': 2}, {'id': 3}]
     rs = rs_factory()
     rs._client.get = mocker.MagicMock(return_value=expected)
     sliced = rs[2:4]
-    assert isinstance(sliced, ResourceSet)
-    assert sliced._limit == 2
-    assert sliced._offset == 2
+    results = [resource for resource in sliced]
+    assert results == expected
+
+
+def test_rs_getitem_slice_more_than_limit(mocker, rs_factory):
+    mocker.patch(
+        'connect.client.models.iterators.parse_content_range',
+        side_effect=[
+            ContentRange(1, 100, 257),
+            ContentRange(101, 101, 257),
+        ],
+    )
+    rs = rs_factory()
+    rs._client.get = mocker.MagicMock(side_effect=[
+        [{'id': i + 1} for i in range(100)],
+        [{'id': 101}],
+    ])
+    sliced = rs[1:102]
+    results = [resource for resource in sliced]
+    assert results == [{'id': i + 1} for i in range(101)]
+
+
+def test_rs_getitem_slice_more_than_limit_no_append(mocker, rs_factory):
+    mocker.patch(
+        'connect.client.models.iterators.parse_content_range',
+        side_effect=[
+            ContentRange(0, 9, 25),
+            ContentRange(10, 19, 25),
+            ContentRange(20, 21, 25),
+        ],
+    )
+    rs = rs_factory()
+    rs._client.get = mocker.MagicMock(side_effect=[
+        [{'id': i} for i in range(10)],
+        [{'id': i + 10} for i in range(10)],
+        [{'id': 20}, {'id': 21}],
+    ])
+    rs._limit = 10
+    rs._client.resourceset_append = False
+    sliced = rs[0:22]
+    results = [resource for resource in sliced]
+    assert results == [{'id': i} for i in range(22)]
 
 
 def test_rs_getitem_slice_type(mocker, rs_factory):