From 333d98d359c1c62628e727bb6cd43a487babec59 Mon Sep 17 00:00:00 2001
From: Anton Oboleninov <anton.oboleninov@saritasa.com>
Date: Tue, 11 Feb 2025 15:39:14 +0700
Subject: [PATCH] Add `__getitem__` to `XPathLocator`

Add ability to get related xpath
locators by index
---
 cspell.config.yaml                |  2 ++
 docs/CHANGELOG.rst                |  4 ++++
 pomcorn/locators/base_locators.py | 22 ++++++++++++++++++++++
 pyproject.toml                    |  2 +-
 tests/locators/test_get_item.py   | 20 ++++++++++++++++++++
 5 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 tests/locators/test_get_item.py

diff --git a/cspell.config.yaml b/cspell.config.yaml
index eab864e..0e3d389 100644
--- a/cspell.config.yaml
+++ b/cspell.config.yaml
@@ -35,6 +35,8 @@ words:
   - yaspeller
 
     # Code words
+  - argnames
+  - argvalues
   - docstrings
   - conftest
   - kwargs
diff --git a/docs/CHANGELOG.rst b/docs/CHANGELOG.rst
index d2455cb..1ce457d 100644
--- a/docs/CHANGELOG.rst
+++ b/docs/CHANGELOG.rst
@@ -3,6 +3,10 @@ Version history
 
 We follow `Semantic Versions <https://semver.org/>`_.
 
+0.8.5 (10.02.25)
+*******************************************************************************
+- Add ability to get related xpath locators by index for ``XPathLocator``
+
 0.8.4 (30.01.25)
 *******************************************************************************
 - Add escaping single and double quotes in the: ``ElementWithTextLocator``,
diff --git a/pomcorn/locators/base_locators.py b/pomcorn/locators/base_locators.py
index 955d713..c5578fc 100644
--- a/pomcorn/locators/base_locators.py
+++ b/pomcorn/locators/base_locators.py
@@ -152,6 +152,28 @@ def __or__(self, other: XPathLocator) -> XPathLocator:
         """
         return XPathLocator(query=f"({self.query} | {other.query})")
 
+    def __getitem__(self, index: int) -> XPathLocator:
+        """Allow get related xpath locator by index.
+
+        Example:
+            div_locator = XPathLocator("//div") --> `//div`
+            div_locator[0]  --> `(//div)[1]`   # xpath numeration starts with 1
+            div_locator[-1] --> `(//div)[last()]`
+
+        """
+        query = f"({self.query})"
+
+        if index >= 0:
+            # `+1` is used here because numeration in xpath starts with 1
+            query += f"[{index + 1}]"
+        elif index == -1:
+            # To avoid ugly locators with `...)[last() - 0]`
+            query += "[last()]"
+        else:
+            query += f"[last() - {abs(index + 1)}]"
+
+        return XPathLocator(query)
+
     def __bool__(self) -> bool:
         """Return whether query of current locator is empty or not."""
         return bool(self.related_query)
diff --git a/pyproject.toml b/pyproject.toml
index 839c254..07fc5f5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pomcorn"
-version = "0.8.4"
+version = "0.8.5"
 description = "Base implementation of Page Object Model"
 authors = [
   "Saritasa <pypi@saritasa.com>",
diff --git a/tests/locators/test_get_item.py b/tests/locators/test_get_item.py
new file mode 100644
index 0000000..a15dab8
--- /dev/null
+++ b/tests/locators/test_get_item.py
@@ -0,0 +1,20 @@
+import pytest
+
+from pomcorn.locators.base_locators import XPathLocator
+
+TEST_QUERY = "//span[text()='Users'][last()]"
+
+
+@pytest.mark.parametrize(
+    argnames=["index", "result"],
+    argvalues=[
+        [-2, f"({TEST_QUERY})[last() - 1]"],
+        [-1, f"({TEST_QUERY})[last()]"],
+        [0, f"({TEST_QUERY})[1]"],
+        [1, f"({TEST_QUERY})[2]"],
+    ],
+)
+def test_getting_related_xpath_locators_by_index(index: int, result: str):
+    """Check that getting related xpath locators by index works correct."""
+    locator = XPathLocator(TEST_QUERY)
+    assert locator[index].query == result