From cc43f9339f91011fdb795ea01c07d5122f6179c6 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Tue, 27 Jul 2021 15:58:18 -0700
Subject: [PATCH] feat(locators): implement last,nth (#7870)

---
 docs/src/api/class-locator.md         | 13 +++++++++++++
 src/client/locator.ts                 | 10 +++++++++-
 src/server/injected/injectedScript.ts | 28 +++++++++++++++++++++++----
 src/server/selectors.ts               |  2 +-
 tests/page/locator-query.spec.ts      | 21 +++++++++++++++++++-
 types/types.d.ts                      | 11 +++++++++++
 6 files changed, 78 insertions(+), 7 deletions(-)

diff --git a/docs/src/api/class-locator.md b/docs/src/api/class-locator.md
index b12bb78b52..40e0b4b1b7 100644
--- a/docs/src/api/class-locator.md
+++ b/docs/src/api/class-locator.md
@@ -534,6 +534,11 @@ Returns whether the element is [visible](./actionability.md#visible).
 
 ### option: Locator.isVisible.timeout = %%-input-timeout-%%
 
+## method: Locator.last
+- returns: <[Locator]>
+
+Returns locator to the last matching element.
+
 ## method: Locator.locator
 - returns: <[Locator]>
 
@@ -542,6 +547,14 @@ The method finds an element matching the specified selector in the `Locator`'s s
 
 ### param: Locator.locator.selector = %%-find-selector-%%
 
+## method: Locator.nth
+- returns: <[Locator]>
+
+Returns locator to the n-th matching element.
+
+### param: Locator.nth.index
+- `index` <[int]>
+
 ## async method: Locator.press
 
 Focuses the element, and then uses [`method: Keyboard.down`] and [`method: Keyboard.up`].
diff --git a/src/client/locator.ts b/src/client/locator.ts
index 7fb8fb4f8f..7830cdf04f 100644
--- a/src/client/locator.ts
+++ b/src/client/locator.ts
@@ -99,7 +99,15 @@ export class Locator implements api.Locator {
   }
 
   first(): Locator {
-    return new Locator(this._frame, this._selector + ' >> _first=true');
+    return new Locator(this._frame, this._selector + ' >> _nth=first');
+  }
+
+  last(): Locator {
+    return new Locator(this._frame, this._selector + ` >> _nth=last`);
+  }
+
+  nth(index: number): Locator {
+    return new Locator(this._frame, this._selector + ` >> _nth=${index}`);
   }
 
   async focus(options?: TimeoutOptions): Promise<void> {
diff --git a/src/server/injected/injectedScript.ts b/src/server/injected/injectedScript.ts
index 9162746952..90f9e6a02a 100644
--- a/src/server/injected/injectedScript.ts
+++ b/src/server/injected/injectedScript.ts
@@ -75,6 +75,7 @@ export class InjectedScript {
     this._engines.set('css', this._createCSSEngine());
     this._engines.set('_first', { queryAll: () => [] });
     this._engines.set('_visible', { queryAll: () => [] });
+    this._engines.set('_nth', { queryAll: () => [] });
 
     for (const { name, engine } of customEngines)
       this._engines.set(name, engine);
@@ -110,11 +111,30 @@ export class InjectedScript {
     if (index === selector.parts.length)
       return roots;
 
-    if (selector.parts[index].name === '_first')
-      return roots.slice(0, 1);
+    const part = selector.parts[index];
+    if (part.name === '_nth') {
+      let filtered: ElementMatch[] = [];
+      if (part.body === 'first') {
+        filtered = roots.slice(0, 1);
+      } else if (part.body === 'last') {
+        if (roots.length)
+          filtered = roots.slice(roots.length - 1);
+      } else {
+        if (typeof selector.capture === 'number')
+          throw new Error(`Can't query n-th element in a request with the capture.`);
+        const nth = +part.body;
+        const set = new Set<Element>();
+        for (const root of roots) {
+          set.add(root.element);
+          if (nth + 1 === set.size)
+            filtered = [root];
+        }
+      }
+      return this._querySelectorRecursively(filtered, selector, index + 1, queryCache);
+    }
 
-    if (selector.parts[index].name === '_visible') {
-      const visible = Boolean(selector.parts[index].body);
+    if (part.name === '_visible') {
+      const visible = Boolean(part.body);
       return roots.filter(match => visible === isVisible(match.element));
     }
 
diff --git a/src/server/selectors.ts b/src/server/selectors.ts
index 40e5ca84b9..01bc029450 100644
--- a/src/server/selectors.ts
+++ b/src/server/selectors.ts
@@ -43,7 +43,7 @@ export class Selectors {
       'data-testid', 'data-testid:light',
       'data-test-id', 'data-test-id:light',
       'data-test', 'data-test:light',
-      '_visible', '_first'
+      '_visible', '_nth'
     ]);
     this._engines = new Map();
   }
diff --git a/tests/page/locator-query.spec.ts b/tests/page/locator-query.spec.ts
index 986fbf2794..cbfb620456 100644
--- a/tests/page/locator-query.spec.ts
+++ b/tests/page/locator-query.spec.ts
@@ -17,7 +17,7 @@
 
 import { test as it, expect } from './pageTest';
 
-it('should respect first()', async ({page}) => {
+it('should respect first() and last()', async ({page}) => {
   await page.setContent(`
   <section>
     <div><p>A</p></div>
@@ -27,4 +27,23 @@ it('should respect first()', async ({page}) => {
   expect(await page.locator('div >> p').count()).toBe(6);
   expect(await page.locator('div').locator('p').count()).toBe(6);
   expect(await page.locator('div').first().locator('p').count()).toBe(1);
+  expect(await page.locator('div').last().locator('p').count()).toBe(3);
+});
+
+it('should respect nth()', async ({page}) => {
+  await page.setContent(`
+  <section>
+    <div><p>A</p></div>
+    <div><p>A</p><p>A</p></div>
+    <div><p>A</p><p>A</p><p>A</p></div>
+  </section>`);
+  expect(await page.locator('div >> p').nth(0).count()).toBe(1);
+  expect(await page.locator('div').nth(1).locator('p').count()).toBe(2);
+  expect(await page.locator('div').nth(2).locator('p').count()).toBe(3);
+});
+
+it('should throw on capture w/ nth()', async ({page}) => {
+  await page.setContent(`<section><div><p>A</p></div></section>`);
+  const e = await page.locator('*css=div >> p').nth(0).click().catch(e => e);
+  expect(e.message).toContain(`Can't query n-th element`);
 });
diff --git a/types/types.d.ts b/types/types.d.ts
index b819851800..9a81c8e95b 100644
--- a/types/types.d.ts
+++ b/types/types.d.ts
@@ -7592,6 +7592,11 @@ export interface Locator {
     timeout?: number;
   }): Promise<boolean>;
 
+  /**
+   * Returns locator to the last matching element.
+   */
+  last(): Locator;
+
   /**
    * The method finds an element matching the specified selector in the `Locator`'s subtree. See
    * [Working with selectors](https://playwright.dev/docs/selectors) for more details.
@@ -7599,6 +7604,12 @@ export interface Locator {
    */
   locator(selector: string): Locator;
 
+  /**
+   * Returns locator to the n-th matching element.
+   * @param index
+   */
+  nth(index: number): Locator;
+
   /**
    * Focuses the element, and then uses [keyboard.down(key)](https://playwright.dev/docs/api/class-keyboard#keyboard-down)
    * and [keyboard.up(key)](https://playwright.dev/docs/api/class-keyboard#keyboard-up).