[python/matplotlib] Add Text, simplify Dots

python/port/genhdr/qstrdefs.in.h

@@ -78,6 +78,7 @@ Q(axis)
 Q(plot)
 Q(pyplot)
 Q(show)
+Q(text)
 
 // Turtle QSTRs
 Q(turtle)

python/port/mod/matplotlib/modpyplot.cpp

@@ -2,11 +2,13 @@ extern "C" {
 #include "modpyplot.h"
 }
 #include <assert.h>
+#include <escher/palette.h>
 #include "port.h"
 #include "plot_controller.h"
 
 Matplotlib::PlotStore * sPlotStore = nullptr;
 Matplotlib::PlotController * sPlotController = nullptr;
+static int paletteIndex = 0; // FIXME: Needs to be reset at some point
 
 // Internal functions
 
@@ -67,10 +69,24 @@ mp_obj_t modpyplot_plot(mp_obj_t x, mp_obj_t y) {
   mp_obj_get_array(x, &xLength, &xItems);
   mp_obj_get_array(y, &yLength, &yItems);
   if (xLength != yLength) {
-    mp_raise_msg(&mp_type_ValueError, "x and y must have same dimension");
+    mp_raise_ValueError("x and y must have same dimension");
   }
 
-  sPlotStore->addDots(x, y);
+  KDColor color = Palette::DataColor[paletteIndex++]; // FIXME: Share overflow routine
+  for (size_t i=0; i<xLength; i++) {
+    sPlotStore->addDot(xItems[i], yItems[i], color);
+  }
+
+  return mp_const_none;
+}
+
+mp_obj_t modpyplot_text(mp_obj_t x, mp_obj_t y, mp_obj_t s) {
+  // Input parameter validation
+  mp_obj_get_float(x);
+  mp_obj_get_float(y);
+  mp_obj_str_get_str(s);
+
+  sPlotStore->addText(x, y, s);
 
   return mp_const_none;
 }

python/port/mod/matplotlib/modpyplot.h

@@ -5,6 +5,7 @@ void modpyplot_gc_collect();
 
 mp_obj_t modpyplot_axis(mp_obj_t arg);
 mp_obj_t modpyplot_plot(mp_obj_t x, mp_obj_t y);
+mp_obj_t modpyplot_text(mp_obj_t x, mp_obj_t y, mp_obj_t s);
 
 // axis(*args, emit=True, **kwargs)
 //mp_obj_t grid();

python/port/mod/matplotlib/modpyplot_table.c

@@ -4,6 +4,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_0(modpyplot___init___obj, modpyplot___init__);
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(modpyplot_axis_obj, modpyplot_axis);
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(modpyplot_plot_obj, modpyplot_plot);
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(modpyplot_show_obj, modpyplot_show);
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(modpyplot_text_obj, modpyplot_text);
 
 STATIC const mp_rom_map_elem_t modpyplot_module_globals_table[] = {
   { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_pyplot) },
@@ -11,6 +12,7 @@ STATIC const mp_rom_map_elem_t modpyplot_module_globals_table[] = {
   { MP_ROM_QSTR(MP_QSTR_axis), MP_ROM_PTR(&modpyplot_axis_obj) },
   { MP_ROM_QSTR(MP_QSTR_plot), MP_ROM_PTR(&modpyplot_plot_obj) },
   { MP_ROM_QSTR(MP_QSTR_show), MP_ROM_PTR(&modpyplot_show_obj) },
+  { MP_ROM_QSTR(MP_QSTR_text), MP_ROM_PTR(&modpyplot_text_obj) },
 };
 
 STATIC MP_DEFINE_CONST_DICT(modpyplot_module_globals, modpyplot_module_globals_table);

python/port/mod/matplotlib/plot_store.cpp

@@ -11,99 +11,79 @@ PlotStore::PlotStore() : Shared::InteractiveCurveViewRange(),
 
 void PlotStore::flush() {
   m_dots = mp_obj_new_list(0, nullptr);
+  m_texts = mp_obj_new_list(0, nullptr);
 }
 
-void PlotStore::addDots(mp_obj_t x, mp_obj_t y) {
-  mp_obj_t items[2] = {x, y};
-  mp_obj_t tuple = mp_obj_new_tuple(2, items);
-  mp_obj_list_append(m_dots, tuple);
-}
+// Iterators
 
-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();
-  }
+template <class T>
+PlotStore::ListIterator<T> PlotStore::ListIterator<T>::Begin(mp_obj_t list) {
+  ListIterator<T> it;
+  mp_obj_list_get(list, &(it.m_numberOfTuples), &(it.m_tuples));
   return it;
 }
 
-PlotStore::DotIterator PlotStore::DotIterator::End(mp_obj_t dots) {
-  DotIterator it;
-  mp_obj_list_get(dots, &(it.m_numberOfTuples), &(it.m_tuples));
+template <class T>
+PlotStore::ListIterator<T> PlotStore::ListIterator<T>::End(mp_obj_t list) {
+  ListIterator<T> it;
+  mp_obj_list_get(list, &(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) {
+template <class T>
+PlotStore::ListIterator<T> & PlotStore::ListIterator<T>::operator++() {
+  if (m_tupleIndex < m_numberOfTuples) {
     m_tupleIndex++;
-    loadValues();
-  } else {
-    m_tupleIndex = m_numberOfTuples;
-    m_valueIndex = 0;
   }
-
   return *this;
 }
 
+template <class T>
+bool PlotStore::ListIterator<T>::operator!=(const PlotStore::ListIterator<T> & it) const {
+  return m_tupleIndex != it.m_tupleIndex;
+};
 
-void PlotStore::DotIterator::loadValues() {
-  mp_obj_t tuple = m_tuples[m_tupleIndex];
+template <class T>
+T PlotStore::ListIterator<T>::operator*() {
+  return T(m_tuples[m_tupleIndex]);
+};
 
-  mp_obj_t * coordinates;
-  mp_obj_get_array_fixed_n(tuple, 2, &coordinates);
+// Dots
 
-  mp_obj_get_array(coordinates[0], &m_numberOfValues, &m_xValues);
-  mp_obj_get_array(coordinates[1], &m_numberOfValues, &m_yValues);
+template class PlotStore::ListIterator<PlotStore::Dot>;
 
-  m_valueIndex = 0;
+PlotStore::Dot::Dot(mp_obj_t tuple) {
+  mp_obj_t * elements;
+  mp_obj_get_array_fixed_n(tuple, 3, &elements);
+  m_x = mp_obj_get_float(elements[0]);
+  m_y = mp_obj_get_float(elements[1]);
+  m_color = KDColor::RGB16(mp_obj_get_int(elements[2]));
 }
 
-void PlotStore::forEachDot(PlotStore::DotCallback callback) {
-  size_t numberOfTuples;
-  mp_obj_t * tuples;
-  mp_obj_list_get(m_dots, &numberOfTuples, &tuples);
+void PlotStore::addDot(mp_obj_t x, mp_obj_t y, KDColor c) {
+  mp_obj_t color = mp_obj_new_int(c);
+  mp_obj_t items[3] = {x, y, color};
+  mp_obj_t tuple = mp_obj_new_tuple(3, items);
+  mp_obj_list_append(m_dots, tuple);
+}
 
-  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);
+// Text
 
-    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
+PlotStore::Text::Text(mp_obj_t tuple) {
+  mp_obj_t * elements;
+  mp_obj_get_array_fixed_n(tuple, 3, &elements);
+  m_x = mp_obj_get_float(elements[0]);
+  m_y = mp_obj_get_float(elements[1]);
+  m_string = mp_obj_str_get_str(elements[2]);
+}
 
-    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]
-      );
-
-      callback(&dot);
-    }
-  }
+void PlotStore::addText(mp_obj_t x, mp_obj_t y, mp_obj_t string) {
+  mp_obj_t items[3] = {x, y, string};
+  mp_obj_t tuple = mp_obj_new_tuple(3, items);
+  mp_obj_list_append(m_texts, tuple);
 }
 
 }

python/port/mod/matplotlib/plot_store.h

@@ -14,9 +14,38 @@ public:
   PlotStore();
   void flush();
 
+  // Iterators
+
+  template <class T>
+  class ListIterator {
+  public:
+    static ListIterator Begin(mp_obj_t list);
+    static ListIterator End(mp_obj_t list);
+    T operator*();
+    ListIterator & operator++();
+    bool operator!=(const ListIterator & it) const;
+  private:
+    ListIterator() : m_tupleIndex(0) {}
+    mp_obj_t * m_tuples;
+    size_t m_numberOfTuples;
+    size_t m_tupleIndex;
+  };
+
+  template <class T>
+  class Iterable {
+  public:
+    Iterable(mp_obj_t list) : m_list(list) {}
+    T begin() const { return T::Begin(m_list); }
+    T end() const { return T::End(m_list); }
+  private:
+    mp_obj_t m_list;
+  };
+
+  // Dots
+
   class Dot {
   public:
-    Dot(float x, float y, KDColor color) : m_x(x), m_y(y), m_color(color) {}
+    Dot(mp_obj_t tuple);
     float x() const { return m_x; }
     float y() const { return m_y; }
     KDColor color() const { return m_color; }
@@ -26,52 +55,35 @@ public:
     KDColor m_color;
   };
 
-  class DotIterator {
+  void addDot(mp_obj_t x, mp_obj_t y, KDColor c);
+  Iterable<ListIterator<Dot>> dots() { return Iterable<ListIterator<Dot>>(m_dots); }
+
+  // Texts
+
+  class Text {
   public:
-    static DotIterator Begin(mp_obj_t dots);
-    static DotIterator End(mp_obj_t dots);
-    Dot operator*();
-    bool operator!=(const DotIterator & it) const;
-    DotIterator & operator++();
+    Text(mp_obj_t tuple);
+    float x() const { return m_x; }
+    float y() const { return m_y; }
+    const char * string() const { return m_string; }
   private:
-    void loadValues();
-    mp_obj_t * m_tuples;
-    size_t m_numberOfTuples;
-    size_t m_tupleIndex;
-    mp_obj_t * m_xValues;
-    mp_obj_t * m_yValues;
-    size_t m_numberOfValues;
-    size_t m_valueIndex;
+    float m_x;
+    float m_y;
+    const char * m_string;
   };
 
-  class Dots {
-  public:
-    Dots(mp_obj_t dots) : m_dots(dots) {}
-    DotIterator begin() const { return DotIterator::Begin(m_dots); }
-    DotIterator end() const { return DotIterator::End(m_dots); }
-  private:
-    mp_obj_t m_dots;
-  };
-
-  void addDots(mp_obj_t x, mp_obj_t y);
-  Dots dots() { return Dots(m_dots); }
+  void addText(mp_obj_t x, mp_obj_t y, mp_obj_t string);
+  Iterable<ListIterator<Text>> texts() { return Iterable<ListIterator<Text>>(m_texts); }
 
   void setGrid(bool grid) { m_grid = grid; }
   bool grid() { return m_grid; }
 private:
-  mp_obj_t m_dots; // A list of (x,y), where x and y are lists of numbers
-  bool m_grid;
-
-  /*
-  mp_obj_array_t * m_plots;
-  mp_obj_array_t * m_arrows;
-  mp_obj_array_t * m_scatters;
-  mp_obj_array_t * m_texts;
-  mp_obj_array_t * m_rects;
+  mp_obj_t m_dots; // List of (x, y, color)
+  mp_obj_t m_texts; // List of (x, y, string)
+  mp_obj_t m_rects; // List of (x, y, w, h, color)
+  mp_obj_t m_segments; // List of (x, y, dx, dy, style, color)
 
   bool m_grid;
-  */
-
 };
 
 }