Add a clear function

dirac-institute · Feb 12, 2024 · 9d52bce · 9d52bce
1 parent 2709a29
commit 9d52bce
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 3 deletions.
diff --git a/src/kbmod/search/pydocs/stack_search_docs.h b/src/kbmod/search/pydocs/stack_search_docs.h
@@ -157,11 +157,32 @@ static const auto DOC_StackSearch_prepare_psi_phi = R"doc(
   )doc";
 
 static const auto DOC_StackSearch_get_results = R"doc(
-  todo
+  Get a batch of cached results [start, finish).
+
+  Parameters
+  ----------
+  start : `int`
+      The (inclusive) starting index of the results to retrieve.
+  end : `int`
+      The (exclusive) ending index of the results to retrieve.
+
+  Returns
+  -------
+  results : `List`
+      A list of ``Trajectory`` objects for the cached results.
   )doc";
 
 static const auto DOC_StackSearch_set_results = R"doc(
-  todo
+  Set the cached results. Used for testing.
+
+  Parameters
+  ----------
+  new_results : `List`
+      The list of results to store.
+  )doc";
+
+static const auto DOC_StackSearch_clear_results = R"doc(
+  Clear the cached results.
   )doc";
 
 static const auto DOC_StackSearch_evaluate_single_trajectory = R"doc(

diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp
@@ -265,6 +265,7 @@ std::vector<Trajectory> StackSearch::get_results(int start, int count) {
 
 // This function is used only for testing by injecting known result trajectories.
 void StackSearch::set_results(const std::vector<Trajectory>& new_results) { results = new_results; }
+void StackSearch::clear_results() { results.clear(); }
 
 #ifdef Py_PYTHON_H
 static void stack_search_bindings(py::module& m) {
@@ -303,7 +304,8 @@ static void stack_search_bindings(py::module& m) {
             .def("prepare_psi_phi", &ks::prepare_psi_phi, pydocs::DOC_StackSearch_prepare_psi_phi)
             .def("clear_psi_phi", &ks::clear_psi_phi, pydocs::DOC_StackSearch_clear_psi_phi)
             .def("get_results", &ks::get_results, pydocs::DOC_StackSearch_get_results)
-            .def("set_results", &ks::set_results, pydocs::DOC_StackSearch_set_results);
+            .def("set_results", &ks::set_results, pydocs::DOC_StackSearch_set_results)
+            .def("clear_results", &ks::clear_results, pydocs::DOC_StackSearch_clear_results);
 }
 #endif /* Py_PYTHON_H */
 

diff --git a/src/kbmod/search/stack_search.h b/src/kbmod/search/stack_search.h
@@ -63,6 +63,7 @@ class StackSearch {
 
     // Helper functions for testing
     void set_results(const std::vector<Trajectory>& new_results);
+    void clear_results();
 
     virtual ~StackSearch(){};