From 1c5ed8f022847ca53dfe4cfeb7ce6bfdf552c6e5 Mon Sep 17 00:00:00 2001
From: Gigon Bae <gbae@nvidia.com>
Date: Tue, 28 Sep 2021 11:27:53 -0700
Subject: [PATCH] Add __enter__ and __exit__ for CuImage

---
 cpp/include/cucim/cuimage.h        |  2 ++
 cpp/src/cuimage.cpp                | 25 +++++++++++++++++++++----
 cpp/src/filesystem/file_handle.cpp |  2 +-
 python/pybind11/cucim_py.cpp       | 14 ++++++++++++++
 python/pybind11/cucim_pydoc.h      |  9 +++++++++
 5 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/cpp/include/cucim/cuimage.h b/cpp/include/cucim/cuimage.h
index 9e4f0dcc1..18efc028b 100644
--- a/cpp/include/cucim/cuimage.h
+++ b/cpp/include/cucim/cuimage.h
@@ -164,6 +164,8 @@ class EXPORT_VISIBLE CuImage : public std::enable_shared_from_this<CuImage>
 
     void save(std::string file_path) const;
 
+    void close();
+
 private:
     using Mutex = std::mutex;
     using ScopedLock = std::scoped_lock<Mutex>;
diff --git a/cpp/src/cuimage.cpp b/cpp/src/cuimage.cpp
index 6f126deb5..8942c557e 100644
--- a/cpp/src/cuimage.cpp
+++ b/cpp/src/cuimage.cpp
@@ -229,10 +229,7 @@ CuImage::CuImage() : std::enable_shared_from_this<CuImage>()
 CuImage::~CuImage()
 {
     //    printf("[cuCIM] CuImage::~CuImage()\n");
-    if (file_handle_.client_data)
-    {
-        image_formats_->formats[0].image_parser.close(&file_handle_);
-    }
+    close();
     image_formats_ = nullptr; // memory release is handled by the framework
     if (image_metadata_)
     {
@@ -615,6 +612,10 @@ CuImage CuImage::read_region(std::vector<int64_t>&& location,
         // Read region from internal file if image_data_ is nullptr
         if (image_data_ == nullptr)
         {
+            if (file_handle_.fd < 0) // file_handle_ is not opened
+            {
+                throw std::runtime_error("[Error] The image file is closed!");
+            }
             if (!image_formats_->formats[0].image_reader.read(
                     &file_handle_, image_metadata_, &request, image_data, nullptr /*out_metadata*/))
             {
@@ -778,6 +779,10 @@ std::set<std::string> CuImage::associated_images() const
 
 CuImage CuImage::associated_image(const std::string& name, const io::Device& device) const
 {
+    if (file_handle_.fd < 0) // file_handle_ is not opened
+    {
+        throw std::runtime_error("[Error] The image file is closed!");
+    }
     auto it = associated_images_.find(name);
     if (it != associated_images_.end())
     {
@@ -855,6 +860,18 @@ void CuImage::save(std::string file_path) const
         fs.close();
     }
 }
+
+void CuImage::close()
+{
+    if (file_handle_.client_data)
+    {
+        image_formats_->formats[0].image_parser.close(&file_handle_);
+    }
+    file_handle_.cufile = nullptr;
+    file_handle_.path = nullptr;
+    file_handle_.fd = -1;
+}
+
 void CuImage::ensure_init()
 {
     ScopedLock g(mutex_);
diff --git a/cpp/src/filesystem/file_handle.cpp b/cpp/src/filesystem/file_handle.cpp
index 1b3036d94..dfb55c8de 100644
--- a/cpp/src/filesystem/file_handle.cpp
+++ b/cpp/src/filesystem/file_handle.cpp
@@ -21,7 +21,7 @@
 #include "cucim/codec/hash_function.h"
 
 CuCIMFileHandle::CuCIMFileHandle()
-    : fd(0),
+    : fd(-1),
       cufile(nullptr),
       type(FileHandleType::kUnknown),
       path(nullptr),
diff --git a/python/pybind11/cucim_py.cpp b/python/pybind11/cucim_py.cpp
index 65c140178..753cac9bc 100644
--- a/python/pybind11/cucim_py.cpp
+++ b/python/pybind11/cucim_py.cpp
@@ -153,7 +153,21 @@ PYBIND11_MODULE(_cucim, m)
              py::arg("name") = "", //
              py::arg("device") = io::Device()) //
         .def("save", &CuImage::save, doc::CuImage::doc_save, py::call_guard<py::gil_scoped_release>()) //
+        .def("close", &CuImage::close, doc::CuImage::doc_close, py::call_guard<py::gil_scoped_release>()) //
         .def("__bool__", &CuImage::operator bool, py::call_guard<py::gil_scoped_release>()) //
+        .def(
+            "__enter__",
+            [](const std::shared_ptr<CuImage>& cuimg) { //
+                return cuimg; //
+            }, //
+            py::call_guard<py::gil_scoped_release>())
+        .def(
+            "__exit__",
+            [](const std::shared_ptr<CuImage>& cuimg, const py::object& type, const py::object& value,
+               const py::object& traceback) { //
+                cuimg->close(); //
+            }, //
+            py::call_guard<py::gil_scoped_release>())
         .def(
             "__repr__", //
             [](const CuImage& cuimg) { //
diff --git a/python/pybind11/cucim_pydoc.h b/python/pybind11/cucim_pydoc.h
index 1fa4b7af6..8d985535f 100644
--- a/python/pybind11/cucim_pydoc.h
+++ b/python/pybind11/cucim_pydoc.h
@@ -227,6 +227,15 @@ Saves image data to the file path.
 Currently it supports only .ppm file format that can be viewed by `eog` command in Ubuntu.
 )doc")
 
+// void close();
+PYDOC(close, R"doc(
+Closes the file handle.
+
+Once the file handle is closed, the image object (if loaded before) still exists but cannot read additional images
+from the file.
+)doc")
+
+
 // void _set_array_interface(const CuImage& cuimg);
 PYDOC(_set_array_interface, R"doc(
 Add `__array_interface__` or `__cuda_array_interface__` depending on the memory type.