From 98f94ee0823208ee934a24488000a065dae9562f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20B=C3=B6hm?= <188768+fb55@users.noreply.github.com>
Date: Mon, 16 May 2022 10:12:45 +0100
Subject: [PATCH] feat: Add a limit argument to `select`

---
 src/index.spec.ts  | 13 +++++++++++++
 src/index.ts       | 35 ++++++++++++++++++++++++++---------
 src/positionals.ts | 19 ++++++++++++++++---
 3 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/src/index.spec.ts b/src/index.spec.ts
index 632020de..6837794e 100644
--- a/src/index.spec.ts
+++ b/src/index.spec.ts
@@ -117,4 +117,17 @@ describe("index", () => {
             select("ul:first li:lt(3)", dom, { context: [dom] })
         ).toHaveLength(2);
     });
+
+    it("should support limit argument", () => {
+        const dom = parseDocument("<p>Paragraph".repeat(10));
+        expect(select("p:even", dom, {}, 1)).toHaveLength(1);
+        expect(select("p:even", dom, {}, 5)).toHaveLength(5);
+        expect(select("p:odd", dom, {}, 1)).toHaveLength(1);
+        expect(select("p:odd", dom, {}, 5)).toHaveLength(5);
+        expect(select("p:lt(5)", dom, {}, 2)).toHaveLength(2);
+        expect(select("p:lt(5)", dom, {}, 6)).toHaveLength(5);
+
+        // Should not use the limit for positionals before the end of the selector
+        expect(select("p:odd + p:eq(5)", dom, {}, 1)).toHaveLength(1);
+    });
 });
diff --git a/src/index.ts b/src/index.ts
index 6dc0d9a9..8360722d 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -188,16 +188,23 @@ function filterBySelector(
         const root = options.root ?? getDocumentRoot(elements[0]);
         const opts = { ...options, context: elements, relativeSelector: false };
         selector.push(SCOPE_PSEUDO);
-        return findFilterElements(root, selector, opts, true);
+        return findFilterElements(root, selector, opts, true, elements.length);
     }
     // Performance optimization: If we don't have to traverse, just filter set.
-    return findFilterElements(elements, selector, options, false);
+    return findFilterElements(
+        elements,
+        selector,
+        options,
+        false,
+        elements.length
+    );
 }
 
 export function select(
     selector: string | ((el: Element) => boolean),
     root: AnyNode | AnyNode[],
-    options: Options = {}
+    options: Options = {},
+    limit = Infinity
 ): Element[] {
     if (typeof selector === "function") {
         return find(root, selector);
@@ -206,12 +213,12 @@ export function select(
     const [plain, filtered] = groupSelectors(parse(selector));
 
     const results: Element[][] = filtered.map((sel) =>
-        findFilterElements(root, sel, options, true)
+        findFilterElements(root, sel, options, true, limit)
     );
 
     // Plain selectors can be queried in a single go
     if (plain.length) {
-        results.push(findElements(root, plain, options, Infinity));
+        results.push(findElements(root, plain, options, limit));
     }
 
     if (results.length === 0) {
@@ -244,17 +251,21 @@ function findFilterElements(
     root: AnyNode | AnyNode[],
     selector: Selector[],
     options: Options,
-    queryForSelector: boolean
+    queryForSelector: boolean,
+    totalLimit: number
 ): Element[] {
     const filterIndex = selector.findIndex(isFilter);
     const sub = selector.slice(0, filterIndex);
     const filter = selector[filterIndex] as CheerioSelector;
+    // If we are at the end of the selector, we can limit the number of elements to retrieve.
+    const partLimit =
+        selector.length - 1 === filterIndex ? totalLimit : Infinity;
 
     /*
      * Set the number of elements to retrieve.
      * Eg. for :first, we only have to get a single element.
      */
-    const limit = getLimit(filter.name, filter.data);
+    const limit = getLimit(filter.name, filter.data, partLimit);
 
     if (limit === 0) return [];
 
@@ -315,10 +326,16 @@ function findFilterElements(
      * Otherwise,
      */
     return remainingSelector.some(isFilter)
-        ? findFilterElements(result, remainingSelector, options, false)
+        ? findFilterElements(
+              result,
+              remainingSelector,
+              options,
+              false,
+              totalLimit
+          )
         : remainingHasTraversal
         ? // Query existing elements to resolve traversal.
-          findElements(result, [remainingSelector], options, Infinity)
+          findElements(result, [remainingSelector], options, totalLimit)
         : // If we don't have any more traversals, simply filter elements.
           filterElements(result, [remainingSelector], options);
 }
diff --git a/src/positionals.ts b/src/positionals.ts
index 426cc228..3c4ee77a 100644
--- a/src/positionals.ts
+++ b/src/positionals.ts
@@ -37,7 +37,11 @@ export function isFilter(s: Selector): s is CheerioSelector {
     return false;
 }
 
-export function getLimit(filter: Filter, data: string | null): number {
+export function getLimit(
+    filter: Filter,
+    data: string | null,
+    partLimit: number
+): number {
     const num = data != null ? parseInt(data, 10) : NaN;
 
     switch (filter) {
@@ -47,10 +51,19 @@ export function getLimit(filter: Filter, data: string | null): number {
         case "eq":
             return isFinite(num) ? (num >= 0 ? num + 1 : Infinity) : 0;
         case "lt":
-            return isFinite(num) ? (num >= 0 ? num : Infinity) : 0;
+            return isFinite(num)
+                ? num >= 0
+                    ? Math.min(num, partLimit)
+                    : Infinity
+                : 0;
         case "gt":
             return isFinite(num) ? Infinity : 0;
-        default:
+        case "odd":
+            return 2 * partLimit;
+        case "even":
+            return 2 * partLimit - 1;
+        case "last":
+        case "not":
             return Infinity;
     }
 }