trellis: add pybind bindings for template classes

4 files changed, 62 insertions, 12 deletions

diff --git a/gr-trellis/include/gnuradio/trellis/encoder.h b/gr-trellis/include/gnuradio/trellis/encoder.h
index 1e573d935..96e84e1c9 100644
--- a/gr-trellis/include/gnuradio/trellis/encoder.h
+++ b/gr-trellis/include/gnuradio/trellis/encoder.h
@@ -29,9 +29,7 @@ class TRELLIS_API encoder : virtual public sync_block
 public:
     typedef std::shared_ptr<encoder<IN_T, OUT_T>> sptr;
 
-    static sptr make(const fsm& FSM, int ST);
-
-    static sptr make(const fsm& FSM, int ST, int K);
+    static sptr make(const fsm& FSM, int ST, int K = 0);
 
     virtual fsm FSM() const = 0;
     virtual int ST() const = 0;
diff --git a/gr-trellis/lib/encoder_impl.cc b/gr-trellis/lib/encoder_impl.cc
index e12e5928b..235367ca6 100644
--- a/gr-trellis/lib/encoder_impl.cc
+++ b/gr-trellis/lib/encoder_impl.cc
@@ -20,16 +20,11 @@ namespace gr {
 namespace trellis {
 
 template <class IN_T, class OUT_T>
-typename encoder<IN_T, OUT_T>::sptr encoder<IN_T, OUT_T>::make(const fsm& FSM, int ST)
-{
-    return gnuradio::get_initial_sptr(new encoder_impl<IN_T, OUT_T>(FSM, ST, 0, false));
-}
-
-template <class IN_T, class OUT_T>
 typename encoder<IN_T, OUT_T>::sptr
 encoder<IN_T, OUT_T>::make(const fsm& FSM, int ST, int K)
 {
-    return gnuradio::get_initial_sptr(new encoder_impl<IN_T, OUT_T>(FSM, ST, K, true));
+    return gnuradio::get_initial_sptr(
+        new encoder_impl<IN_T, OUT_T>(FSM, ST, K, K >= 0 ? false : true));
 }
 
 template <class IN_T, class OUT_T>
diff --git a/gr-trellis/python/trellis/bindings/encoder_python.cc b/gr-trellis/python/trellis/bindings/encoder_python.cc
index bfa1c5a63..fef8ca881 100644
--- a/gr-trellis/python/trellis/bindings/encoder_python.cc
+++ b/gr-trellis/python/trellis/bindings/encoder_python.cc
@@ -19,4 +19,34 @@ namespace py = pybind11;
 // pydoc.h is automatically generated in the build directory
 #include <encoder_pydoc.h>
 
-void bind_encoder(py::module& m) {}
+template <class IN_T, class OUT_T>
+void bind_encoder_template(py::module& m, const char* classname)
+{
+    using encoder = gr::trellis::encoder<IN_T, OUT_T>;
+
+    py::class_<encoder,
+               gr::sync_block,
+               gr::block,
+               gr::basic_block,
+               std::shared_ptr<encoder>>(m, classname)
+        .def(py::init(&gr::trellis::encoder<IN_T, OUT_T>::make),
+             py::arg("FSM"),
+             py::arg("ST"),
+             py::arg("K") = 0)
+        .def("FSM", &gr::trellis::encoder<IN_T, OUT_T>::FSM)
+        .def("ST", &gr::trellis::encoder<IN_T, OUT_T>::ST)
+        .def("K", &gr::trellis::encoder<IN_T, OUT_T>::K)
+        .def("set_FSM", &gr::trellis::encoder<IN_T, OUT_T>::set_FSM)
+        .def("set_ST", &gr::trellis::encoder<IN_T, OUT_T>::set_ST)
+        .def("set_K", &gr::trellis::encoder<IN_T, OUT_T>::set_K);
+}
+
+void bind_encoder(py::module& m)
+{
+    bind_encoder_template<std::uint8_t, std::uint8_t>(m, "encoder_bb");
+    bind_encoder_template<std::uint8_t, std::int16_t>(m, "encoder_bs");
+    bind_encoder_template<std::uint8_t, std::int32_t>(m, "encoder_bi");
+    bind_encoder_template<std::int16_t, std::int16_t>(m, "encoder_ss");
+    bind_encoder_template<std::int16_t, std::int32_t>(m, "encoder_si");
+    bind_encoder_template<std::int32_t, std::int32_t>(m, "encoder_ii");
+}
diff --git a/gr-trellis/python/trellis/bindings/viterbi_python.cc b/gr-trellis/python/trellis/bindings/viterbi_python.cc
index 32ba3f3f2..3729ed125 100644
--- a/gr-trellis/python/trellis/bindings/viterbi_python.cc
+++ b/gr-trellis/python/trellis/bindings/viterbi_python.cc
@@ -19,4 +19,31 @@ namespace py = pybind11;
 // pydoc.h is automatically generated in the build directory
 #include <viterbi_pydoc.h>
 
-void bind_viterbi(py::module& m) {}
+template <typename T>
+void bind_viterbi_template(py::module& m, const char* classname)
+{
+    using viterbi = gr::trellis::viterbi<T>;
+
+    py::class_<viterbi, gr::block, gr::basic_block, std::shared_ptr<viterbi>>(m,
+                                                                              classname)
+        .def(py::init(&gr::trellis::viterbi<T>::make),
+             py::arg("FSM"),
+             py::arg("K"),
+             py::arg("S0"),
+             py::arg("SK"))
+        .def("FSM", &viterbi::FSM)
+        .def("K", &viterbi::K)
+        .def("S0", &viterbi::S0)
+        .def("SK", &viterbi::SK)
+        .def("set_FSM", &viterbi::set_FSM)
+        .def("set_K", &viterbi::set_K)
+        .def("set_S0", &viterbi::set_S0)
+        .def("set_SK", &viterbi::set_SK);
+}
+
+void bind_viterbi(py::module& m)
+{
+    bind_viterbi_template<std::uint8_t>(m, "viterbi_b");
+    bind_viterbi_template<std::int16_t>(m, "viterbi_s");
+    bind_viterbi_template<std::int32_t>(m, "viterbi_i");
+}