Code documentation

python/pyngraph/function.cpp

@@ -54,22 +54,39 @@ void regclass_pyngraph_Function(py::module m)
         return "<" + class_name + ": '" + self.get_friendly_name() + "' (" + shape + ")>";
     });
     function.def_static("from_capsule", [](py::object* capsule) {
+        // get the underlying PyObject* which is a PyCapsule pointer
         auto* pybind_capsule_ptr = capsule->ptr();
+        // extract the pointer stored in the PyCapsule under the name CAPSULE_NAME
         auto* capsule_ptr = PyCapsule_GetPointer(pybind_capsule_ptr, CAPSULE_NAME);
-        auto* ngraph_function = static_cast<std::shared_ptr<ngraph::Function>*>(capsule_ptr);
 
-        return *ngraph_function;
+        auto* ngraph_function = static_cast<std::shared_ptr<ngraph::Function>*>(capsule_ptr);
+        if (ngraph_function)
+        {
+            return *ngraph_function;
+        }
+        else
+        {
+            return nullptr;
+        }
     });
     function.def_static("to_capsule", [](std::shared_ptr<ngraph::Function>& ngraph_function) {
+        // create a shared pointer on the heap before putting it in the capsule
+        // this secures the lifetime of the object transferred by the capsule
         auto* sp_copy = new std::shared_ptr<ngraph::Function>(ngraph_function);
-        auto pybind_capsule = py::capsule(sp_copy, CAPSULE_NAME, [](PyObject* capsule) {
+
+        // a destructor callback that will delete the heap allocated shared_ptr
+        // when the capsule is destructed
+        auto sp_deleter = [](PyObject* capsule) {
             auto* capsule_ptr = PyCapsule_GetPointer(capsule, CAPSULE_NAME);
             auto* function_sp = static_cast<std::shared_ptr<ngraph::Function>*>(capsule_ptr);
             if (function_sp)
             {
                 delete function_sp;
             }
-        });
+        };
+
+        // put the shared_ptr in a new capsule under the same name as in "from_capsule"
+        auto pybind_capsule = py::capsule(sp_copy, CAPSULE_NAME, sp_deleter);
 
         return pybind_capsule;
     });