pybind: Add S2CellId iteration (children, cells, ranges)

src/python/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@pybind11_bazel//:build_defs.bzl", "pybind_extension", "pybind_library")
-load("@rules_cc//cc:defs.bzl", "cc_library")
+load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
 load("@rules_python//python:defs.bzl", "py_library", "py_test")
 
 package(default_visibility = ["//visibility:private"])
@@ -115,15 +115,36 @@ pybind_library(
     ],
 )
 
+cc_library(
+    name = "s2cell_id_range",
+    srcs = ["s2cell_id_range.cc"],
+    hdrs = ["s2cell_id_range.h"],
+    deps = [
+        "//:s2",
+        "@abseil-cpp//absl/log:absl_check",
+    ],
+)
+
 pybind_library(
     name = "s2cell_id_bindings",
     srcs = ["s2cell_id_bindings.cc"],
     deps = [
+        ":s2cell_id_range",
         "//:s2",
         "@abseil-cpp//absl/strings",
     ],
 )
 
+cc_test(
+    name = "s2cell_id_range_test",
+    srcs = ["s2cell_id_range_test.cc"],
+    deps = [
+        ":s2cell_id_range",
+        "//:s2",
+        "//s2/testing:test_main",
+    ],
+)
+
 pybind_library(
     name = "s2cell_bindings",
     srcs = ["s2cell_bindings.cc"],

src/python/README.md

@@ -79,6 +79,10 @@ The Python bindings follow the C++ API closely but with Pythonic conventions:
 - C++ functions that accept or return a vector object use a Python tuple (of length matching the vector dimension)
 - Array indexing operators (e.g., `point[0]`) are not currently supported
 
+**Iterators:**
+- Some classes expose methods that return iterables (e.g., `S2CellId.children()`, `S2CellId.cells(level)`).
+- Iterables support forward iteration, `len()`, indexing, slicing, `in`, and `reversed()` unless otherwise noted.
+
 **Serialization:**
 - The C++ Encoder/Decoder serialization functions are not currently supported
 

src/python/s2cell_id_bindings.cc

@@ -2,11 +2,15 @@
 #include <pybind11/operators.h>
 #include <pybind11/stl.h>
 
+#include <cstdint>
+#include <limits>
 #include <sstream>
+#include <stdexcept>
 #include <vector>
 
 #include "absl/strings/str_cat.h"
 #include "s2/s2cell_id.h"
+#include "python/s2cell_id_range.h"
 #include "s2/s2latlng.h"
 
 namespace py = pybind11;
@@ -68,6 +72,71 @@ void MaybeThrowPositionOutOfRange(uint64_t pos) {
   }
 }
 
+// Raise OverflowError if n exceeds Py_ssize_t; otherwise return it cast.
+// On 32-bit Python, cells() ranges above level ~15 exceed ssize_t max.
+Py_ssize_t LenOrThrow(int64_t n) {
+  if (n > std::numeric_limits<Py_ssize_t>::max()) {
+    throw std::overflow_error(absl::StrCat(
+        "S2CellIdRange size ", n, " exceeds Py_ssize_t max (",
+        std::numeric_limits<Py_ssize_t>::max(),
+        ") on this platform. On 32-bit Python, cells() ranges above "
+        "level ~15 exceed ssize_t. Use .size() for the full int64 count."));
+  }
+  return static_cast<Py_ssize_t>(n);
+}
+
+// Normalize a Python-style index (negative counts from end) against a range
+// of length n and return the 0-based index.  Raises IndexError if the index
+// is still out of [0, n) after normalization.
+// https://docs.python.org/3/reference/datamodel.html#object.__getitem__
+int64_t IndexOrThrow(int64_t i, int64_t n) {
+  if (i < 0) i += n;
+  if (i < 0 || i >= n) throw py::index_error("index out of range");
+  return i;
+}
+
+// Raise ValueError if the slice step is anything other than 1 (or None).
+// S2CellIdRange slices must be contiguous; reversed() covers step=-1.
+void ValidateStepIsOne(py::slice s) {
+  py::object step_obj = s.attr("step");
+  if (step_obj.is_none()) return;
+  int64_t step = step_obj.cast<int64_t>();
+  if (step != 1) {
+    throw py::value_error(absl::StrCat(
+        "S2CellIdRange only supports slices with step 1 (got step ",
+        step, ")"));
+  }
+}
+
+// Resolve a Python slice against a range of length n and return the clamped
+// [start, stop) bounds.  Slice indices are silently clamped per:
+// https://docs.python.org/3/reference/datamodel.html#sequences
+//
+// py::slice::compute() is not used because it requires Py_ssize_t, which is
+// 32-bit on 32-bit platforms; S2CellIdRange can have up to 2^60 cells so
+// indices are handled as int64_t directly.
+std::pair<int64_t, int64_t> ClampedSlice(int64_t n, py::slice s) {
+  auto clamp = [n](int64_t i) -> int64_t {
+    if (i < 0) i += n;
+    if (i < 0) return int64_t{0};
+    if (i > n) return n;
+    return i;
+  };
+  py::object start_obj = s.attr("start");
+  py::object stop_obj = s.attr("stop");
+  int64_t start = start_obj.is_none() ? int64_t{0} : clamp(start_obj.cast<int64_t>());
+  int64_t stop = stop_obj.is_none() ? n : clamp(stop_obj.cast<int64_t>());
+  if (stop < start) stop = start;
+  return {start, stop};
+}
+
+// Dereference an optional S2CellId from an iterator's next(), or raise
+// StopIteration if the iterator is exhausted.
+S2CellId NextOrStop(std::optional<S2CellId> result) {
+  if (!result) throw py::stop_iteration();
+  return *result;
+}
+
 }  // namespace
 
 void bind_s2cell_id(py::module& m) {
@@ -226,6 +295,32 @@ void bind_s2cell_id(py::module& m) {
            }, py::arg("position"),
            "Return the immediate child at the given position (0..3).\n\n"
            "Raises ValueError if this is a leaf cell or position is out of range.")
+      .def("children", [](S2CellId self, py::object level_obj) {
+               MaybeThrowIfLeaf(self);
+               S2CellId begin, end;
+               if (level_obj.is_none()) {
+                 begin = self.child_begin();
+                 end = self.child_end();
+               } else {
+                 int level = level_obj.cast<int>();
+                 MaybeThrowLevelOutOfRange(level, self.level() + 1,
+                                           S2CellId::kMaxLevel);
+                 begin = self.child_begin(level);
+                 end = self.child_end(level);
+               }
+               return S2CellIdRange{begin, end};
+           }, py::arg("level") = py::none(),
+           "Return a range over the children of this cell.\n\n"
+           "With no argument, returns the 4 immediate children.\n"
+           "With a level argument, returns all descendants at that level.\n"
+           "Raises ValueError if this is a leaf cell or level is out of range.")
+      .def_static("cells", [](int level) {
+               MaybeThrowLevelOutOfRange(level, 0, S2CellId::kMaxLevel);
+               return S2CellIdRange{S2CellId::Begin(level),
+                                    S2CellId::End(level)};
+           }, py::arg("level"),
+           "Return a range over all cells at the given level across all 6 faces.\n\n"
+           "Warning: the number of cells grows as 6 * 4^level.")
       .def("edge_neighbors", [](S2CellId self) {
                S2CellId neighbors[4];
                self.GetEdgeNeighbors(neighbors);
@@ -279,4 +374,59 @@ void bind_s2cell_id(py::module& m) {
   cls.attr("MAX_LEVEL")    = S2CellId::kMaxLevel;
   cls.attr("MAX_POSITION") = S2CellId::kMaxPosition;
   cls.attr("NUM_FACES")    = S2CellId::kNumFaces;
+
+  py::class_<S2CellIdRange>(m, "S2CellIdRange",
+      "A range of S2CellIds at the same level along the Hilbert curve.\n\n"
+      "Supports len(), indexing, slicing, iteration, reversed(), and 'in'.\n"
+      "Returned from S2CellId.children() and S2CellId.cells().\n\n"
+      "Slicing: only step=1 is supported. Use reversed() for step=-1.\n"
+      "Other step values raise ValueError.")
+      .def("__len__",      [](const S2CellIdRange& self) {
+        return LenOrThrow(self.size());
+      })
+      .def("size", &S2CellIdRange::size,
+           "Return the number of cells in the range as a 64-bit integer.\n\n"
+           "Equivalent to len() on 64-bit platforms; use this method when\n"
+           "the range may exceed Py_ssize_t (e.g. very large cells() ranges\n"
+           "on 32-bit Python).")
+      .def("__getitem__",  [](const S2CellIdRange& self, int64_t i) {
+        return self.at(IndexOrThrow(i, self.size()));
+      }, py::arg("index"))
+      .def("__getitem__",  [](const S2CellIdRange& self, py::slice s) {
+        ValidateStepIsOne(s);
+        auto [start, stop] = ClampedSlice(self.size(), s);
+        return self.slice(start, stop);
+      }, py::arg("slice"))
+      .def("__contains__", &S2CellIdRange::contains, py::arg("cell"))
+      .def("__eq__", [](const S2CellIdRange& self, const S2CellIdRange& other) {
+        return self.begin == other.begin && self.end == other.end;
+      }, py::arg("other"))
+      .def("__repr__", [](const S2CellIdRange& self) {
+        std::ostringstream oss;
+        oss << "S2CellIdRange(begin=" << self.begin.id()
+            << ", end=" << self.end.id() << ")";
+        return oss.str();
+      })
+      .def("__iter__", [](const S2CellIdRange& self) {
+        return S2CellIdForwardIterator{self.begin, self.end};
+      })
+      .def("__reversed__", [](const S2CellIdRange& self) {
+        return S2CellIdReverseIterator{self.end, self.begin};
+      });
+
+  py::class_<S2CellIdForwardIterator>(m, "_S2CellIdForwardIterator", py::module_local())
+      .def("__iter__", [](S2CellIdForwardIterator& self) -> S2CellIdForwardIterator& {
+        return self;
+      })
+      .def("__next__", [](S2CellIdForwardIterator& self) {
+        return NextOrStop(self.next());
+      });
+
+  py::class_<S2CellIdReverseIterator>(m, "_S2CellIdReverseIterator", py::module_local())
+      .def("__iter__", [](S2CellIdReverseIterator& self) -> S2CellIdReverseIterator& {
+        return self;
+      })
+      .def("__next__", [](S2CellIdReverseIterator& self) {
+        return NextOrStop(self.next());
+      });
 }

src/python/s2cell_id_range.cc

@@ -0,0 +1,53 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS-IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "python/s2cell_id_range.h"
+
+#include "absl/log/absl_check.h"
+
+int64_t S2CellIdRange::size() const {
+  ABSL_DCHECK(begin == end || begin.level() == end.level())
+      << "S2CellIdRange: begin and end must be at the same level";
+  return end.distance_from_begin() - begin.distance_from_begin();
+}
+
+S2CellId S2CellIdRange::at(int64_t i) const {
+  ABSL_DCHECK_GE(i, 0);
+  ABSL_DCHECK_LT(i, size());
+  return begin.advance(i);
+}
+
+S2CellIdRange S2CellIdRange::slice(int64_t start, int64_t stop) const {
+  ABSL_DCHECK_GE(start, 0);
+  ABSL_DCHECK_LE(start, stop);
+  ABSL_DCHECK_LE(stop, size());
+  return S2CellIdRange{begin.advance(start), begin.advance(stop)};
+}
+
+bool S2CellIdRange::contains(S2CellId cell) const {
+  return cell >= begin && cell < end && cell.level() == begin.level();
+}
+
+std::optional<S2CellId> S2CellIdForwardIterator::next() {
+  if (cur == end) return std::nullopt;
+  S2CellId result = cur;
+  cur = cur.next();
+  return result;
+}
+
+std::optional<S2CellId> S2CellIdReverseIterator::next() {
+  if (cur == begin) return std::nullopt;
+  cur = cur.prev();
+  return cur;
+}

src/python/s2cell_id_range.h

@@ -0,0 +1,74 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS-IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Helpers for S2CellId sequence iteration used by Python bindings.
+//
+// S2CellIdRange represents a contiguous range of S2CellIds at a fixed level,
+// with forward and reverse iterators.  These types are exposed to pybind11 to
+// support Python's sequence protocol (len, indexing, slicing, iteration,
+// reversed) on S2CellId ranges returned by methods such as children() and
+// cells_at_level().
+
+#ifndef PYTHON_S2CELL_ID_RANGE_H_
+#define PYTHON_S2CELL_ID_RANGE_H_
+
+#include <cstdint>
+#include <optional>
+
+#include "s2/s2cell_id.h"
+
+// A contiguous half-open range [begin, end) of S2CellIds at the same level
+// along the Hilbert curve.  All operations are O(1).  The range may be empty
+// (begin == end).
+struct S2CellIdRange {
+  S2CellId begin;
+  S2CellId end;
+
+  // Return the number of cells in the range.
+  int64_t size() const;
+
+  // Return the cell at 0-based offset i.  i must be in [0, size()); behavior
+  // is undefined otherwise.
+  S2CellId at(int64_t i) const;
+
+  // Return the sub-range [start, stop).  start and stop must satisfy
+  // 0 <= start <= stop <= size(); behavior is undefined otherwise.
+  S2CellIdRange slice(int64_t start, int64_t stop) const;
+
+  // Return true if cell is in the range and at the same level as begin.
+  bool contains(S2CellId cell) const;
+};
+
+// Forward iterator over an S2CellIdRange.
+struct S2CellIdForwardIterator {
+  S2CellId cur;
+  S2CellId end;
+
+  // Return the next cell and advance, or std::nullopt if exhausted.
+  std::optional<S2CellId> next();
+};
+
+// Reverse iterator over an S2CellIdRange.  Construct with cur = range.end;
+// next() decrements cur before yielding, so the empty-range case (begin ==
+// end) is handled without a separate flag.
+struct S2CellIdReverseIterator {
+  S2CellId cur;    // one-past-current; starts at range.end
+  S2CellId begin;  // stop when cur reaches here
+
+  // Return the next cell (in reverse order) and advance, or std::nullopt if
+  // exhausted.
+  std::optional<S2CellId> next();
+};
+
+#endif  // PYTHON_S2CELL_ID_RANGE_H_