From 08a6cb6bfa289ecd9313702b39a1dd4f3037b557 Mon Sep 17 00:00:00 2001
From: Yuval Ben-Arie <yuval1010@gmail.com>
Date: Thu, 16 Feb 2023 00:47:26 +0200
Subject: [PATCH] Add Support for QuerySet slicing operator

---
 CHANGELOG.rst          |  1 +
 CONTRIBUTORS.rst       |  1 +
 tests/test_queryset.py | 52 ++++++++++++++++++++++++++++++++++++++++++
 tortoise/queryset.py   | 33 +++++++++++++++++++++++++++
 4 files changed, 87 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 82ac4175b..8f4cdd892 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -17,6 +17,7 @@ Added
 - Add __eq__ method to Q to more easily test dynamically-built queries (#1506)
 - Added PlainToTsQuery function for postgres (#1347)
 - Allow field's default keyword to be async function (#1498)
+- Add support for queryset slicing. (#1341)
 
 Fixed
 ^^^^^
diff --git a/CONTRIBUTORS.rst b/CONTRIBUTORS.rst
index d944b8665..5e3811e83 100644
--- a/CONTRIBUTORS.rst
+++ b/CONTRIBUTORS.rst
@@ -53,6 +53,7 @@ Contributors
 * Paul Serov ``@thakryptex``
 * Stanislav Zmiev ``@Ovsyanka83``
 * Waket Zheng ``@waketzheng``
+* Yuval Ben-Arie ``@yuvalbenarie``
 
 Special Thanks
 ==============
diff --git a/tests/test_queryset.py b/tests/test_queryset.py
index 720447de2..a71e5c7ff 100644
--- a/tests/test_queryset.py
+++ b/tests/test_queryset.py
@@ -69,6 +69,58 @@ async def test_offset_negative(self):
         with self.assertRaisesRegex(ParamsError, "Offset should be non-negative number"):
             await IntFields.all().offset(-10)
 
+    async def test_slicing_start_and_stop(self) -> None:
+        sliced_queryset = IntFields.all().order_by("intnum")[1:5]
+        manually_sliced_queryset = IntFields.all().order_by("intnum").offset(1).limit(4)
+        self.assertSequenceEqual(await sliced_queryset, await manually_sliced_queryset)
+
+    async def test_slicing_only_limit(self) -> None:
+        sliced_queryset = IntFields.all().order_by("intnum")[:5]
+        manually_sliced_queryset = IntFields.all().order_by("intnum").limit(5)
+        self.assertSequenceEqual(await sliced_queryset, await manually_sliced_queryset)
+
+    async def test_slicing_only_offset(self) -> None:
+        sliced_queryset = IntFields.all().order_by("intnum")[5:]
+        manually_sliced_queryset = IntFields.all().order_by("intnum").offset(5)
+        self.assertSequenceEqual(await sliced_queryset, await manually_sliced_queryset)
+
+    async def test_slicing_count(self) -> None:
+        queryset = IntFields.all().order_by("intnum")[1:5]
+        self.assertEqual(await queryset.count(), 4)
+
+    def test_slicing_negative_values(self) -> None:
+        with self.assertRaisesRegex(
+            expected_exception=ParamsError,
+            expected_regex="Slice start should be non-negative number or None.",
+        ):
+            _ = IntFields.all()[-1:]
+
+        with self.assertRaisesRegex(
+            expected_exception=ParamsError,
+            expected_regex="Slice stop should be non-negative number greater that slice start, "
+            "or None.",
+        ):
+            _ = IntFields.all()[:-1]
+
+    def test_slicing_stop_before_start(self) -> None:
+        with self.assertRaisesRegex(
+            expected_exception=ParamsError,
+            expected_regex="Slice stop should be non-negative number greater that slice start, "
+            "or None.",
+        ):
+            _ = IntFields.all()[2:1]
+
+    async def test_slicing_steps(self) -> None:
+        sliced_queryset = IntFields.all().order_by("intnum")[::1]
+        manually_sliced_queryset = IntFields.all().order_by("intnum")
+        self.assertSequenceEqual(await sliced_queryset, await manually_sliced_queryset)
+
+        with self.assertRaisesRegex(
+            expected_exception=ParamsError,
+            expected_regex="Slice steps should be 1 or None.",
+        ):
+            _ = IntFields.all()[::2]
+
     async def test_join_count(self):
         tour = await Tournament.create(name="moo")
         await MinRelation.create(tournament=tour)
diff --git a/tortoise/queryset.py b/tortoise/queryset.py
index 074e7868d..7dcae6f6a 100644
--- a/tortoise/queryset.py
+++ b/tortoise/queryset.py
@@ -456,6 +456,39 @@ def offset(self, offset: int) -> "QuerySet[MODEL]":
             queryset._limit = 1000000
         return queryset
 
+    def __getitem__(self, key: slice) -> "QuerySet[MODEL]":
+        """
+        Query offset and limit for Queryset.
+
+        :raises ParamsError: QuerySet indices must be slices.
+
+        :raises ParamsError: Slice steps should be 1 or None.
+
+        :raises ParamsError: Slice start should be non-negative number or None.
+
+        :raises ParamsError: Slice stop should be non-negative number greater that slice start,
+        or None.
+        """
+        if not isinstance(key, slice):
+            raise ParamsError("QuerySet indices must be slices.")
+
+        if not (key.step is None or (isinstance(key.step, int) and key.step == 1)):
+            raise ParamsError("Slice steps should be 1 or None.")
+
+        start = key.start if key.start is not None else 0
+
+        if not isinstance(start, int) or start < 0:
+            raise ParamsError("Slice start should be non-negative number or None.")
+        if key.stop is not None and (not isinstance(key.stop, int) or key.stop <= start):
+            raise ParamsError(
+                "Slice stop should be non-negative number greater that slice start, or None.",
+            )
+
+        queryset = self.offset(start)
+        if key.stop:
+            queryset = queryset.limit(key.stop - start)
+        return queryset
+
     def distinct(self) -> "QuerySet[MODEL]":
         """
         Make QuerySet distinct.