Use C implementation for seek_index

Author: jeromekelleher

python/_tskitmodule.c

@@ -10658,6 +10658,29 @@ Tree_seek(Tree *self, PyObject *args)
     return ret;
 }
 
+static PyObject *
+Tree_seek_index(Tree *self, PyObject *args)
+{
+    PyObject *ret = NULL;
+    tsk_id_t index = 0;
+    int err;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    if (!PyArg_ParseTuple(args, "O&", tsk_id_converter, &index)) {
+        goto out;
+    }
+    err = tsk_tree_seek_index(self->tree, index, 0);
+    if (err != 0) {
+        handle_library_error(err);
+        goto out;
+    }
+    ret = Py_BuildValue("");
+out:
+    return ret;
+}
+
 static PyObject *
 Tree_clear(Tree *self)
 {
@@ -11796,6 +11819,10 @@ static PyMethodDef Tree_methods[] = {
         .ml_meth = (PyCFunction) Tree_seek,
         .ml_flags = METH_VARARGS,
         .ml_doc = "Seeks to the tree at the specified position" },
+    { .ml_name = "seek_index",
+        .ml_meth = (PyCFunction) Tree_seek_index,
+        .ml_flags = METH_VARARGS,
+        .ml_doc = "Seeks to the tree at the specified index" },
     { .ml_name = "clear",
         .ml_meth = (PyCFunction) Tree_clear,
         .ml_flags = METH_NOARGS,

python/tests/test_lowlevel.py

@@ -2968,6 +2968,16 @@ def test_seek_errors(self):
             with pytest.raises(_tskit.LibraryError):
                 tree.seek(bad_pos)
 
+    def test_seek_index_errors(self):
+        ts = self.get_example_tree_sequence()
+        tree = _tskit.Tree(ts)
+        for bad_type in ["", "x", {}]:
+            with pytest.raises(TypeError):
+                tree.seek_index(bad_type)
+        for bad_index in [-1, 10**6]:
+            with pytest.raises(_tskit.LibraryError):
+                tree.seek_index(bad_index)
+
     def test_root_threshold(self):
         for ts in self.get_example_tree_sequences():
             tree = _tskit.Tree(ts)

python/tskit/trees.py

@@ -820,7 +820,6 @@ def seek_index(self, index):
 
         .. include:: substitutions/linear_traversal_warning.rst
 
-
         :param int index: The tree index to seek to.
         :raises IndexError: If an index outside the acceptable range is provided.
         """
@@ -829,12 +828,7 @@ def seek_index(self, index):
             index += num_trees
         if index < 0 or index >= num_trees:
             raise IndexError("Index out of bounds")
-        # This should be implemented in C efficiently using the indexes.
-        # No point in complicating the current implementation by trying
-        # to seek from the correct direction.
-        self.first()
-        while self.index != index:
-            self.next()
+        self._ll_tree.seek_index(index)
 
     def seek(self, position):
         """