Added PyProgram methods and tests

Differential Revision: D77388495

Author: Conarnar

extension/pybindings/pybindings.cpp

@@ -1012,18 +1012,47 @@ struct PyProgram final {
       program_(load_program(
         loader_.get(),
         program_verification)) {}
+
   static std::unique_ptr<PyProgram> load_from_buffer(
     const py::bytes& buffer,
     Program::Verification program_verification =
         Program::Verification::Minimal) {
     return std::make_unique<PyProgram>(buffer, program_verification);
   }
+
   static std::unique_ptr<PyProgram> load_from_file(
     const std::string& path,
     Program::Verification program_verification =
         Program::Verification::Minimal) {
     return std::make_unique<PyProgram>(path, program_verification);
   }
+
+  PyProgram(const PyProgram&) = delete;
+  PyProgram& operator=(const PyProgram&) = delete;
+  PyProgram(PyProgram&&) = default;
+  PyProgram& operator=(PyProgram&&) = default;
+
+  size_t num_methods() const {
+    return program_->num_methods();
+  }
+
+  std::string get_method_name(size_t method_index) const {
+    Result<const char*> res = program_->get_method_name(method_index);
+    THROW_IF_ERROR(
+      res.error(),
+      "Failed get method name, error: 0x:%" PRIx32,
+      static_cast<uint32_t>(res.error()));
+    return std::string(res.get());
+  }
+
+  std::string get_output_flattening_encoding(std::string method_name) const {
+    Result<const char*> res = program_->get_output_flattening_encoding(method_name.c_str());
+    THROW_IF_ERROR(
+      res.error(),
+      "Failed get output flattening encoding, error: 0x:%" PRIx32,
+      static_cast<uint32_t>(res.error()));
+    return std::string(res.get());
+  }
  private:
   std::unique_ptr<DataLoader> loader_;
   std::unique_ptr<Program> program_;
@@ -1231,7 +1260,13 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
       py::arg("program_verification") =
           Program::Verification::Minimal,
       call_guard);
-  py::class_<PyProgram>(m, "ExecuTorchProgram");
+  py::class_<PyProgram>(m, "ExecuTorchProgram")
+      .def("num_methods", &PyProgram::num_methods, call_guard)
+      .def(
+          "get_method_name",
+          &PyProgram::get_method_name,
+          py::arg("method_index"),
+          call_guard);
 }
 
 namespace {

extension/pybindings/test/make_test.py

@@ -168,6 +168,7 @@ def make_test(  # noqa: C901
     subfunction of wrapper.
     """
     load_fn: Callable = runtime._load_for_executorch_from_buffer
+    load_prog_fn: Callable = runtime._load_program_from_buffer
 
     def wrapper(tester: unittest.TestCase) -> None:
         ######### TEST CASES #########
@@ -474,6 +475,24 @@ def test_unsupported_input_type(tester):
             # This should raise a Python error, not hit a fatal assert in the C++ code.
             tester.assertRaises(RuntimeError, executorch_module, inputs)
 
+        def test_program_methods_one(tester):
+            exported_program, _ = create_program(ModuleAdd())
+            executorch_program = load_prog_fn(exported_program.buffer)
+            tester.assertEqual(executorch_program.num_methods(), 1)
+            tester.assertEqual(executorch_program.get_method_name(0), "forward")
+
+        def test_program_methods_multi(tester):
+            exported_program, _ = create_program(ModuleMulti())
+            executorch_program = load_prog_fn(exported_program.buffer)
+            tester.assertEqual(executorch_program.num_methods(), 2)
+            tester.assertEqual(executorch_program.get_method_name(0), "forward")
+            tester.assertEqual(executorch_program.get_method_name(1), "forward2")
+
+        def test_program_method_index_out_of_bounds(tester):
+            exported_program, _ = create_program(ModuleMulti())
+            executorch_program = load_prog_fn(exported_program.buffer)
+            tester.assertRaises(RuntimeError, executorch_program.get_method_name, 2)
+
         ######### RUN TEST CASES #########
         test_e2e(tester)
         test_multiple_entry(tester)
@@ -490,5 +509,8 @@ def test_unsupported_input_type(tester):
         test_bad_name(tester)
         test_verification_config(tester)
         test_unsupported_input_type(tester)
+        test_program_methods_one(tester)
+        test_program_methods_multi(tester)
+        test_program_method_index_out_of_bounds(tester)
 
     return wrapper