From bea39c45e08303fc8d1049ba9db582c1f45c9239 Mon Sep 17 00:00:00 2001
From: Georges Berenger <gberenger@meta.com>
Date: Fri, 17 Jan 2025 13:32:12 -0800
Subject: [PATCH] Context manager fixes

Summary: This diff addresses some context manager methods issues.

Differential Revision: D68341672

fbshipit-source-id: a3e920ca9c4b850e7100810e3f65a99a22a1541c
---
 csrc/reader/VRSReader.cpp | 7 ++-----
 pyvrs/reader.py           | 2 +-
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/csrc/reader/VRSReader.cpp b/csrc/reader/VRSReader.cpp
index 8ece1ac..40d4b1c 100644
--- 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
index 7a9bd57..2149d12 100755
--- 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()