Add Document get_index_from_location / get_location_from_index

CHANGELOG.md

@@ -14,6 +14,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Fixed `DataTable` not scrolling to rows that were just added https://github.com/Textualize/textual/pull/3552
 - Fixed cache bug with `DataTable.update_cell` https://github.com/Textualize/textual/pull/3551
 
+### Added
+
+- Add Document `get_index_from_location` / `get_location_from_index` https://github.com/Textualize/textual/pull/3410
+
 ### Changed
 
 - Buttons will now display multiple lines, and have auto height https://github.com/Textualize/textual/pull/3539

src/textual/document/_document.py

@@ -320,6 +320,42 @@ def line_count(self) -> int:
         """Returns the number of lines in the document."""
         return len(self._lines)
 
+    def get_index_from_location(self, location: Location) -> int:
+        """Given a location, returns the index from the document's text.
+
+        Args:
+            location: The location in the document.
+
+        Returns:
+            The index in the document's text.
+        """
+        row, column = location
+        index = row * len(self.newline) + column
+        for line_index in range(row):
+            index += len(self.get_line(line_index))
+        return index
+
+    def get_location_from_index(self, index: int) -> Location:
+        """Given an index in the document's text, returns the corresponding location.
+
+        Args:
+            index: The index in the document's text.
+
+        Returns:
+            The corresponding location.
+        """
+        column_index = 0
+        newline_length = len(self.newline)
+        for line_index in range(self.line_count):
+            next_column_index = (
+                column_index + len(self.get_line(line_index)) + newline_length
+            )
+            if index < next_column_index:
+                return (line_index, index - column_index)
+            elif index == next_column_index:
+                return (line_index + 1, 0)
+            column_index = next_column_index
+
     def get_line(self, index: int) -> str:
         """Returns the line with the given index from the document.
 

tests/document/test_document.py

@@ -98,3 +98,42 @@ def test_get_selected_text_no_newline_at_end_of_file_windows():
     document = Document(TEXT_WINDOWS)
     selection = document.get_text_range((0, 0), (2, 0))
     assert selection == TEXT_WINDOWS
+
+
+@pytest.mark.parametrize(
+    "text", [TEXT, TEXT_NEWLINE, TEXT_WINDOWS, TEXT_WINDOWS_NEWLINE]
+)
+def test_index_from_location(text):
+    document = Document(text)
+    lines = text.split(document.newline)
+    assert document.get_index_from_location((0, 0)) == 0
+    assert document.get_index_from_location((0, len(lines[0]))) == len(lines[0])
+    assert document.get_index_from_location((1, 0)) == len(lines[0]) + len(
+        document.newline
+    )
+    assert document.get_index_from_location((len(lines) - 1, len(lines[-1]))) == len(
+        text
+    )
+
+
+@pytest.mark.parametrize(
+    "text", [TEXT, TEXT_NEWLINE, TEXT_WINDOWS, TEXT_WINDOWS_NEWLINE]
+)
+def test_location_from_index(text):
+    document = Document(text)
+    lines = text.split(document.newline)
+    assert document.get_location_from_index(0) == (0, 0)
+    assert document.get_location_from_index(len(lines[0])) == (0, len(lines[0]))
+    if len(document.newline) > 1:
+        assert document.get_location_from_index(len(lines[0]) + 1) == (
+            0,
+            len(lines[0]) + 1,
+        )
+    assert document.get_location_from_index(len(lines[0]) + len(document.newline)) == (
+        1,
+        0,
+    )
+    assert document.get_location_from_index(len(text)) == (
+        len(lines) - 1,
+        len(lines[-1]),
+    )