Improving the Python bindings, particularly the map/pair wrappers

Signed-off-by: David Shah <davey1576@gmail.com>

3 files changed, 176 insertions, 21 deletions

diff --git a/common/pybindings.cc b/common/pybindings.cc
index 454a571c..24269b55 100644
--- a/common/pybindings.cc
+++ b/common/pybindings.cc
@@ -21,10 +21,13 @@
 #include "chip.h"
 #include "design.h"
 #include "emb.h"
+#include "jsonparse.h"
 
 // include after design.h/chip.h
 #include "pybindings.h"
 
+#include <fstream>
+
 // Required to determine concatenated module name (which differs for different
 // archs)
 #define PASTER(x, y) x##_##y
@@ -43,6 +46,24 @@ bool operator==(const PortRef &a, const PortRef &b)
     return (a.cell == b.cell) && (a.port == b.port);
 }
 
+// Load a JSON file into a design
+void parse_json_shim(std::string filename, Design &d)
+{
+    std::ifstream inf(filename);
+    if (!inf)
+        throw std::runtime_error("failed to open file " + filename);
+    std::istream *ifp = &inf;
+    parse_json_file(ifp, filename, &d);
+}
+
+// Create a new Chip and load design from json file
+Design load_design_shim(std::string filename, ChipArgs args)
+{
+    Design d(args);
+    parse_json_shim(filename, d);
+    return d;
+}
+
 BOOST_PYTHON_MODULE(MODULE_NAME)
 {
     class_<GraphicElement>("GraphicElement")
@@ -101,6 +122,9 @@ BOOST_PYTHON_MODULE(MODULE_NAME)
     WRAP_MAP(decltype(Design::nets), "IdNetMap");
     WRAP_MAP(decltype(Design::cells), "IdCellMap");
 
+    def("parse_json", parse_json_shim);
+    def("load_design", load_design_shim);
+
     arch_wrap_python();
 }
 
diff --git a/common/pycontainers.h b/common/pycontainers.h
index baf60229..423e36f5 100644
--- a/common/pycontainers.h
+++ b/common/pycontainers.h
@@ -30,6 +30,8 @@
 
 using namespace boost::python;
 
+inline void KeyError() { PyErr_SetString(PyExc_KeyError, "Key not found"); }
+
 /*
 A wrapper for a Pythonised nextpnr Iterator. The actual class wrapped is a
 pair<Iterator, Iterator> containing (current, end)
@@ -82,15 +84,145 @@ struct range_wrapper
         class_<T>(range_name, no_init).def("__iter__", iter);
         iterator_wrapper<iterator_t, P>().wrap(iter_name);
     }
+
     typedef iterator_wrapper<iterator_t, P> iter_wrap;
 };
 
 #define WRAP_RANGE(t) range_wrapper<t##Range>().wrap(#t "Range", #t "Iterator")
 
 /*
+Wrapper for a pair, allows accessing either using C++-style members (.first and
+.second) or as a Python iterable and indexable object
+ */
+template <typename T1, typename T2> struct pair_wrapper
+{
+    typedef std::pair<T1, T2> T;
+
+    struct pair_iterator_wrapper
+    {
+        static object next(std::pair<T &, int> &iter)
+        {
+            if (iter.second == 0) {
+                iter.second++;
+                return object(iter.first.first);
+            } else if (iter.second == 1) {
+                iter.second++;
+                return object(iter.first.second);
+            } else {
+                PyErr_SetString(PyExc_StopIteration, "End of range reached");
+                boost::python::throw_error_already_set();
+                // Should be unreachable, but prevent control may reach end of
+                // non-void
+                throw std::runtime_error("unreachable");
+            }
+        }
+
+        static void wrap(const char *python_name)
+        {
+            class_<std::pair<T &, int>>(python_name, no_init)
+                    .def("__next__", next);
+        }
+    };
+
+    static object get(T &x, int i)
+    {
+        if ((i >= 2) || (i < 0))
+            KeyError();
+        return (i == 1) ? object(x.second) : object(x.first);
+    }
+
+    static void set(T &x, int i, object val)
+    {
+        if ((i >= 2) || (i < 0))
+            KeyError();
+        if (i == 0)
+            x.first = extract<T1>(val);
+        if (i == 1)
+            x.second = extract<T2>(val);
+    }
+
+    static int len(T &x) { return 2; }
+
+    static std::pair<T &, int> iter(T &x)
+    {
+        return std::make_pair(boost::ref(x), 0);
+    };
+
+    static void wrap(const char *pair_name, const char *iter_name)
+    {
+        pair_iterator_wrapper::wrap(iter_name);
+        class_<T>(pair_name, no_init)
+                .def("__iter__", iter)
+                .def("__len__", len)
+                .def("__getitem__", get)
+                .def("__setitem__", set, with_custodian_and_ward<1, 2>())
+                .def_readwrite("first", &T::first)
+                .def_readwrite("second", &T::second);
+    }
+};
+
+/*
+Special case of above for map key/values
+ */
+template <typename T1, typename T2> struct map_pair_wrapper
+{
+    typedef std::pair<T1, T2> T;
+
+    struct pair_iterator_wrapper
+    {
+        static object next(std::pair<T &, int> &iter)
+        {
+            if (iter.second == 0) {
+                iter.second++;
+                return object(iter.first.first);
+            } else if (iter.second == 1) {
+                iter.second++;
+                return object(iter.first.second);
+            } else {
+                PyErr_SetString(PyExc_StopIteration, "End of range reached");
+                boost::python::throw_error_already_set();
+                // Should be unreachable, but prevent control may reach end of
+                // non-void
+                throw std::runtime_error("unreachable");
+            }
+        }
+
+        static void wrap(const char *python_name)
+        {
+            class_<std::pair<T &, int>>(python_name, no_init)
+                    .def("__next__", next);
+        }
+    };
+
+    static object get(T &x, int i)
+    {
+        if ((i >= 2) || (i < 0))
+            KeyError();
+        return (i == 1) ? object(x.second) : object(x.first);
+    }
+
+    static int len(T &x) { return 2; }
+
+    static std::pair<T &, int> iter(T &x)
+    {
+        return std::make_pair(boost::ref(x), 0);
+    };
+
+    static void wrap(const char *pair_name, const char *iter_name)
+    {
+        pair_iterator_wrapper::wrap(iter_name);
+        class_<T>(pair_name, no_init)
+                .def("__iter__", iter)
+                .def("__len__", len)
+                .def("__getitem__", get)
+                .def_readonly("first", &T::first)
+                .def_readwrite("second", &T::second);
+    }
+};
+
+/*
 Wrapper for a map, either an unordered_map, regular map or dict
  */
-inline void KeyError() { PyErr_SetString(PyExc_KeyError, "Key not found"); }
 
 template <typename T> struct map_wrapper
 {
@@ -119,11 +251,10 @@ template <typename T> struct map_wrapper
     }
 
     static void wrap(const char *map_name, const char *kv_name,
-                     const char *iter_name)
+                     const char *kv_iter_name, const char *iter_name)
     {
-        class_<KV>(kv_name)
-                .def_readonly("first", &KV::first)
-                .def_readwrite("second", &KV::second);
+        map_pair_wrapper<typename KV::first_type,
+                         typename KV::second_type>::wrap(kv_name, kv_iter_name);
         typedef range_wrapper<T, return_value_policy<copy_non_const_reference>>
                 rw;
         typename rw::iter_wrap().wrap(iter_name);
@@ -136,6 +267,7 @@ template <typename T> struct map_wrapper
 };
 
 #define WRAP_MAP(t, name)                                                      \
-    map_wrapper<t>().wrap(#name, #name "KeyValue", #name "Iterator")
+    map_wrapper<t>().wrap(#name, #name "KeyValue", #name "KeyValueIter",       \
+                          #name "Iterator")
 
 #endif
diff --git a/python/dump_design.py b/python/dump_design.py
index 1850d509..d95f041b 100644
--- a/python/dump_design.py
+++ b/python/dump_design.py
@@ -1,26 +1,25 @@
 # Run ./nextpnr-ice40 --json ice40/blinky.json --run python/dump_design.py
-for cell in sorted(design.cells, key=lambda x: x.first):
-    print("Cell {} : {}".format(cell.first, cell.second.type))
+for cell, cinfo in sorted(design.cells, key=lambda x: x.first):
+    print("Cell {} : {}".format(cell, cinfo.type))
     print("\tPorts:")
-    for port in sorted(cell.second.ports, key=lambda x: x.first):
-        dir = (" <-- ", " --> ", " <-> ")[int(port.second.type)]
-        if port.second.net is not None:
-            print("\t\t{} {} {}".format(port.first, dir, port.second.net.name))
+    for port, pinfo in sorted(cinfo.ports, key=lambda x: x.first):
+        dir = (" <-- ", " --> ", " <-> ")[int(pinfo.type)]
+        if pinfo.net is not None:
+            print("\t\t{} {} {}".format(port, dir, pinfo.net.name))
 
-    if len(cell.second.attrs) > 0:
+    if len(cinfo.attrs) > 0:
         print("\tAttrs:")
-        for attr in cell.second.attrs:
-            print("\t\t{}: {}".format(attr.first, attr.second))
+        for attr, val in cinfo.attrs:
+            print("\t\t{}: {}".format(attr, val))
 
-    if len(cell.second.params) > 0:
+    if len(cinfo.params) > 0:
         print("\tParams:")
-        for param in cell.second.params:
-            val = param.second
+        for param, val in cinfo.params:
             if val.isdigit():
                 val = bin(int(val))[2:]
                 val = "{}'b{}".format(len(val), val)
-            print("\t\t{}: {}".format(param.first, val))
+            print("\t\t{}: {}".format(param, val))
 
-    if not cell.second.bel.nil():
-        print("\tBel: {}".format(chip.getBelName(cell.second.bel)))
+    if not cinfo.bel.nil():
+        print("\tBel: {}".format(chip.getBelName(cinfo.bel)))
     print()