[python/matplotlib] Add a Dot iterator

python/port/mod/matplotlib/plot_store.cpp

@@ -1,4 +1,5 @@
 #include "plot_store.h"
+#include <escher/palette.h>
 
 namespace Matplotlib {
 
@@ -18,49 +19,91 @@ void PlotStore::addDots(mp_obj_t x, mp_obj_t y) {
   mp_obj_list_append(m_dots, tuple);
 }
 
-PlotStore::Dot PlotStore::dotAtIndex(int i) {
+PlotStore::DotIterator PlotStore::DotIterator::Begin(mp_obj_t dots) {
+  DotIterator it;
+  mp_obj_list_get(dots, &(it.m_numberOfTuples), &(it.m_tuples));
+  if (it.m_numberOfTuples > 0) {
+    it.m_tupleIndex = 0;
+    it.loadValues();
+  }
+  return it;
+}
+
+PlotStore::DotIterator PlotStore::DotIterator::End(mp_obj_t dots) {
+  DotIterator it;
+  mp_obj_list_get(dots, &(it.m_numberOfTuples), &(it.m_tuples));
+  if (it.m_numberOfTuples > 0) {
+    it.m_tupleIndex = it.m_numberOfTuples;
+    it.m_valueIndex = 0;
+  }
+  return it;
+}
+
+PlotStore::Dot PlotStore::DotIterator::operator*() {
+  return PlotStore::Dot(
+    mp_obj_get_float(m_xValues[m_valueIndex]),
+    mp_obj_get_float(m_yValues[m_valueIndex]),
+    Palette::DataColor[m_tupleIndex] // FIXME: Share the "looping" routing
+  );
+};
+
+bool PlotStore::DotIterator::operator!=(const DotIterator & it) const {
+  return (m_tupleIndex != it.m_tupleIndex || m_valueIndex != it.m_valueIndex);
+};
+
+PlotStore::DotIterator & PlotStore::DotIterator::operator++() {
+  if (m_valueIndex < m_numberOfValues - 1) {
+    m_valueIndex++;
+  } else if (m_tupleIndex < m_numberOfTuples - 1) {
+    m_tupleIndex++;
+    loadValues();
+  } else {
+    m_tupleIndex = m_numberOfTuples;
+    m_valueIndex = 0;
+  }
+
+  return *this;
+}
+
+
+void PlotStore::DotIterator::loadValues() {
+  mp_obj_t tuple = m_tuples[m_tupleIndex];
+
+  mp_obj_t * coordinates;
+  mp_obj_get_array_fixed_n(tuple, 2, &coordinates);
+
+  mp_obj_get_array(coordinates[0], &m_numberOfValues, &m_xValues);
+  mp_obj_get_array(coordinates[1], &m_numberOfValues, &m_yValues);
+
+  m_valueIndex = 0;
+}
+
+void PlotStore::forEachDot(PlotStore::DotCallback callback) {
   size_t numberOfTuples;
   mp_obj_t * tuples;
   mp_obj_list_get(m_dots, &numberOfTuples, &tuples);
 
-  mp_obj_t firstTuple = tuples[0];
-  size_t numberOfObjects;
-  mp_obj_t * objects;
-  mp_obj_tuple_get(firstTuple, &numberOfObjects, &objects);
+  for (size_t t=0; t<numberOfTuples; t++) {
+    mp_obj_t tuple = tuples[t];
+    mp_obj_t * coordinates;
+    mp_obj_get_array_fixed_n(tuple, 2, &coordinates);
 
-  mp_obj_t x = objects[0];
-  mp_obj_t y = objects[1];
+    size_t xLength, yLength;
+    mp_obj_t * xItems, * yItems;
+    mp_obj_get_array(coordinates[0], &xLength, &xItems);
+    mp_obj_get_array(coordinates[1], &yLength, &yItems);
+    assert(xLength == yLength); // Checked in modpyplot_plot
 
-  size_t numberOfX;
-  mp_obj_t * xValues;
-  mp_obj_list_get(x, &numberOfX, &xValues);
-  size_t numberOfY;
-  mp_obj_t * yValues;
-  mp_obj_list_get(y, &numberOfY, &yValues);
+    for (size_t i=0; i<xLength; i++) {
+      PlotStore::Dot dot(
+        mp_obj_get_float(xItems[i]),
+        mp_obj_get_float(yItems[i]),
+        Palette::DataColor[t]
+      );
 
-  float goodX = mp_obj_get_float(xValues[i]);
-  float goodY = mp_obj_get_float(yValues[i]);
-
-  return PlotStore::Dot(goodX, goodY, KDColorRed);
-}
-
-int PlotStore::numberOfDots() {
-  size_t numberOfTuples;
-  mp_obj_t * tuples;
-  mp_obj_list_get(m_dots, &numberOfTuples, &tuples);
-
-  mp_obj_t firstTuple = tuples[0];
-  size_t numberOfObjects;
-  mp_obj_t * objects;
-  mp_obj_tuple_get(firstTuple, &numberOfObjects, &objects);
-
-  mp_obj_t x = objects[0];
-
-  size_t numberOfX;
-  mp_obj_t * xValues;
-  mp_obj_list_get(x, &numberOfX, &xValues);
-
-  return numberOfX;
+      callback(&dot);
+    }
+  }
 }
 
 }