From a7559249c9bd362b92e05410cf8187c6161a03c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20B=C3=B6hm?= <me@feedic.com>
Date: Tue, 5 Jan 2021 06:39:47 -0800
Subject: [PATCH] feat(select): Allow a function as the selector

This allows users to specify their own search functions, or pre-compile selectors.
---
 src/index.spec.ts |  5 +++++
 src/index.ts      | 28 +++++++++++++++++++++++++---
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/src/index.spec.ts b/src/index.spec.ts
index 29b0993c..346b37e6 100644
--- a/src/index.spec.ts
+++ b/src/index.spec.ts
@@ -8,6 +8,11 @@ describe("index", () => {
         expect(select("div", dom)).toHaveLength(1);
     });
 
+    it("should find with a function", () => {
+        const dom = parseDOM("<div><p>First<p>Second") as Element[];
+        expect(select((elem) => elem.name === "p", dom)).toHaveLength(2);
+    });
+
     it("should support positionals", () => {
         const dom = parseDOM("<div><p>First<p>Second") as Element[];
         expect(select("p:first", dom)).toMatchInlineSnapshot(`
diff --git a/src/index.ts b/src/index.ts
index 80244668..e9829569 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -125,10 +125,14 @@ function filterParsed(
 }
 
 export function select(
-    selector: string,
+    selector: string | ((el: Element) => boolean),
     root: Element | Element[],
     options: Options = {}
 ): Element[] {
+    if (typeof selector === "function") {
+        return find(root, selector);
+    }
+
     const [plain, filtered] = groupSelectors(parse(selector, options));
 
     const results: Element[][] = filtered.map((sel) =>
@@ -276,6 +280,11 @@ function findFilterElements(
           filterElements(result, [remainingSelector], remainingOpts);
 }
 
+interface CompiledQuery {
+    (el: Element): boolean;
+    shouldTestNextSiblings?: boolean;
+}
+
 function findElements(
     root: Element | Element[],
     sel: Selector[][],
@@ -284,8 +293,21 @@ function findElements(
 ): Element[] {
     if (limit === 0) return [];
 
-    // @ts-expect-error TS seems to mess up the type here ¯\_(ツ)_/¯
-    const query = compileToken<Node, Element>(sel, options, root);
+    const query: CompiledQuery = compileToken<Node, Element>(
+        sel,
+        // @ts-expect-error TS seems to mess up the type here ¯\_(ツ)_/¯
+        options,
+        root
+    );
+
+    return find(root, query, limit);
+}
+
+function find(
+    root: Element | Element[],
+    query: CompiledQuery,
+    limit = Infinity
+): Element[] {
     const elems = prepareContext<Node, Element>(
         root,
         DomUtils,