Context manager fixes

Summary: This diff addresses some context manager methods issues. Differential Revision: D68341672 fbshipit-source-id: a3e920ca9c4b850e7100810e3f65a99a22a1541c
facebookresearch · Jan 17, 2025 · bea39c4 · bea39c4
1 parent 485fe29
commit bea39c4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/csrc/reader/VRSReader.cpp b/csrc/reader/VRSReader.cpp
@@ -1204,14 +1204,11 @@ void pybind_vrsreader(py::module& m) {
       .def("purge_file_cache", &PyVRSReader::purgeFileCache)
       .def(
           "__enter__",
-          [&](PyVRSReader& r) { return &r; },
+          [](PyVRSReader& r) { return &r; },
           "Enter the runtime context related to this VRSReader")
       .def(
           "__exit__",
-          [&](PyVRSReader& r,
-              pybind11::object* exc_type,
-              pybind11::object* exc_value,
-              pybind11::object* traceback) { r.close(); },
+          [](PyVRSReader& r, const py::args&) { r.close(); },
           "Exit the runtime context related to this VRSReader");
 }
 #endif

diff --git a/pyvrs/reader.py b/pyvrs/reader.py
@@ -208,7 +208,7 @@ def __enter__(self) -> "VRSReader":
         """Context manager entry point."""
         return self
 
-    def __exit__(self) -> None:
+    def __exit__(self, type, value, traceback) -> None:
         """Context manager exit point."""
         self.close()