From 9c891ca296aebb0f99bf3c8c62e48d545da37c29 Mon Sep 17 00:00:00 2001
From: Steven Johnson <srj@google.com>
Date: Thu, 21 Apr 2022 10:12:29 -0700
Subject: [PATCH] Grab-bag of minor Python fixes

---
 python_bindings/src/PyConciseCasts.cpp | 20 ++++++++++----------
 python_bindings/src/PyExpr.cpp         |  1 +
 python_bindings/src/PyImageParam.cpp   |  2 +-
 python_bindings/src/PyLoopLevel.cpp    |  4 +++-
 python_bindings/stub/PyStubImpl.cpp    | 10 ++++++++++
 5 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/python_bindings/src/PyConciseCasts.cpp b/python_bindings/src/PyConciseCasts.cpp
index 4ee79d703673..ebd646c0934b 100644
--- a/python_bindings/src/PyConciseCasts.cpp
+++ b/python_bindings/src/PyConciseCasts.cpp
@@ -6,34 +6,34 @@ namespace PythonBindings {
 void define_concise_casts(py::module &m) {
     // explicit cast should be tried before
     //  the pybind11::implicitly_convertible<T, Expr> conversion
-    m.def("f64", [](double v) {
+    m.def("f64", [](double v) -> Expr {
         return Expr(v);
     });
-    m.def("f32", [](float v) {
+    m.def("f32", [](float v) -> Expr {
         return Expr(v);
     });
-    m.def("i64", [](int64_t v) {
+    m.def("i64", [](int64_t v) -> Expr {
         return Expr(v);
     });
-    m.def("i32", [](int32_t v) {
+    m.def("i32", [](int32_t v) -> Expr {
         return Expr(v);
     });
-    m.def("i16", [](int16_t v) {
+    m.def("i16", [](int16_t v) -> Expr {
         return Expr(v);
     });
-    m.def("i8", [](int8_t v) {
+    m.def("i8", [](int8_t v) -> Expr {
         return Expr(v);
     });
-    m.def("u64", [](uint64_t v) {
+    m.def("u64", [](uint64_t v) -> Expr {
         return Expr(v);
     });
-    m.def("u32", [](uint32_t v) {
+    m.def("u32", [](uint32_t v) -> Expr {
         return Expr(v);
     });
-    m.def("u16", [](uint16_t v) {
+    m.def("u16", [](uint16_t v) -> Expr {
         return Expr(v);
     });
-    m.def("u8", [](uint8_t v) {
+    m.def("u8", [](uint8_t v) -> Expr {
         return Expr(v);
     });
     // pybind11::implicitly_convertible<T, Expr> conversions
diff --git a/python_bindings/src/PyExpr.cpp b/python_bindings/src/PyExpr.cpp
index 423f2778e5a2..5f061cd8ff77 100644
--- a/python_bindings/src/PyExpr.cpp
+++ b/python_bindings/src/PyExpr.cpp
@@ -48,6 +48,7 @@ void define_expr(py::module &m) {
             .def("__nonzero__", to_bool)
 
             .def("type", &Expr::type)
+            .def("defined", &Expr::defined)
             .def("__repr__", [](const Expr &e) -> std::string {
                 std::ostringstream o;
                 o << "<halide.Expr of type " << halide_type_to_string(e.type()) << ": " << e << ">";
diff --git a/python_bindings/src/PyImageParam.cpp b/python_bindings/src/PyImageParam.cpp
index 1a3e35f50a4d..dfc6a8f66c7c 100644
--- a/python_bindings/src/PyImageParam.cpp
+++ b/python_bindings/src/PyImageParam.cpp
@@ -57,7 +57,7 @@ void define_image_param(py::module &m) {
     auto image_param_class =
         py::class_<ImageParam>(m, "ImageParam", output_image_param_class)
             .def(py::init<>())
-            .def(py::init<Type, int>())
+            .def(py::init<Type, int>(), py::arg("type"), py::arg("dimensions"))
             .def(py::init<Type, int, std::string>(), py::arg("type"), py::arg("dimensions"), py::arg("name"))
             .def("set", &ImageParam::set)
             .def("get", &ImageParam::get)
diff --git a/python_bindings/src/PyLoopLevel.cpp b/python_bindings/src/PyLoopLevel.cpp
index dade7619444d..b2db4b5d035d 100644
--- a/python_bindings/src/PyLoopLevel.cpp
+++ b/python_bindings/src/PyLoopLevel.cpp
@@ -17,7 +17,9 @@ void define_loop_level(py::module &m) {
             .def_static("root", &LoopLevel::root)
             .def("__repr__", [](const LoopLevel &b) -> std::string {
                 std::ostringstream o;
-                o << "<halide.LoopLevel " << (b.defined() ? b.to_string() : "UNDEF") << ">";
+                // b.to_string() fails for locked LoopLevels. Just output something generic.
+                // o << "<halide.LoopLevel " << (b.defined() ? b.to_string() : "UNDEF") << ">";
+                o << "<halide.LoopLevel>";
                 return o.str();
             });
 }
diff --git a/python_bindings/stub/PyStubImpl.cpp b/python_bindings/stub/PyStubImpl.cpp
index b5803cbd7fd1..c0683ca8e42b 100644
--- a/python_bindings/stub/PyStubImpl.cpp
+++ b/python_bindings/stub/PyStubImpl.cpp
@@ -141,6 +141,16 @@ py::object generate_impl(const GeneratorFactory &factory, const GeneratorContext
         } else {
             if (py::isinstance<LoopLevel>(value)) {
                 generator_params[name] = value.cast<LoopLevel>();
+            } else if (py::isinstance<py::list>(value)) {
+                // Convert [hl.UInt(8), hl.Int(16)] -> uint8,int16
+                std::string v;
+                for (auto t : value) {
+                    if (!v.empty()) {
+                        v += ",";
+                    }
+                    v += py::str(t).cast<std::string>();
+                }
+                generator_params[name] = v;
             } else {
                 generator_params[name] = py::str(value).cast<std::string>();
             }