Add context manager methods to FileReader and VRSReader objects

Summary: Context manager style is particular nice with file objects, let's make that work with VRSReader and FileReader, updating existing tests.

Reviewed By: hanghu

Differential Revision: D68307514

fbshipit-source-id: 989ba765505e7cc750697d7f30ab5f19fa2cd400

csrc/reader/VRSReader.cpp

@@ -1105,7 +1105,12 @@ bool OssVRSReader::purgeFileCache() {
 #if IS_VRS_OSS_CODE()
 void pybind_vrsreader(py::module& m) {
   py::class_<PyVRSReader>(m, "Reader")
+      .def(py::init<>())
       .def(py::init<bool>())
+      .def(py::init<const string&>())
+      .def(py::init<const PyFileSpec&>())
+      .def(py::init<const string&, bool>())
+      .def(py::init<const PyFileSpec&, bool>())
       .def("open", py::overload_cast<const string&>(&PyVRSReader::open))
       .def("open", py::overload_cast<const PyFileSpec&>(&PyVRSReader::open))
       .def("close", &PyVRSReader::close)
@@ -1196,7 +1201,18 @@ void pybind_vrsreader(py::module& m) {
           &PyVRSReader::prefetchRecordSequence,
           py::arg("sequence"),
           py::arg("clearSequence") = true)
-      .def("purge_file_cache", &PyVRSReader::purgeFileCache);
+      .def("purge_file_cache", &PyVRSReader::purgeFileCache)
+      .def(
+          "__enter__",
+          [&](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(); },
+          "Exit the runtime context related to this VRSReader");
 }
 #endif
 } // namespace pyvrs

csrc/reader/VRSReader.h

@@ -127,6 +127,21 @@ class OssVRSReader : public VRSReaderBase {
       : streamPlayer_{*this}, autoReadConfigurationRecord_{autoReadConfigurationRecord} {
     init();
   }
+  OssVRSReader() : OssVRSReader(false) {}
+  explicit OssVRSReader(const string& path) : OssVRSReader(false) {
+    open(path);
+  }
+  OssVRSReader(const string& path, bool autoReadConfigurationRecord)
+      : OssVRSReader(autoReadConfigurationRecord) {
+    open(path);
+  }
+  explicit OssVRSReader(const PyFileSpec& spec) : OssVRSReader(false) {
+    open(spec);
+  }
+  OssVRSReader(const PyFileSpec& spec, bool autoReadConfigurationRecord)
+      : OssVRSReader(autoReadConfigurationRecord) {
+    open(spec);
+  }
 
   ~OssVRSReader() override {
     close();

pyvrs/reader.py

@@ -204,6 +204,14 @@ def _init_reader(
         self._reader = reader_cls(auto_read_configuration_records)
         self._open_files(multi_path, file_spec)
 
+    def __enter__(self) -> "VRSReader":
+        """Context manager entry point."""
+        return self
+
+    def __exit__(self) -> None:
+        """Context manager exit point."""
+        self.close()
+
     def _path_to_file_spec(
         self, path: PathType, multi_path: bool
     ) -> Union[FileSpec, List[FileSpec]]: