Let the CM be on open() (#1410)

Author: NicolasHug

src/torchcodec/encoders/_multi_stream_encoder.py

@@ -89,11 +89,12 @@ def add_audio(
 
     # TODO MultiStreamEncoder: Maybe there should 2 separate methods, one for
     # file, one for file-like.
-    def open(self, dest, *, format: str | None = None) -> None:
+    def open(self, dest, *, format: str | None = None) -> "StreamingEncoder":
         if format is not None:
             _core.streaming_encoder_open_file_like(self._encoder_tensor, format, dest)
         else:
             _core.streaming_encoder_open_file(self._encoder_tensor, str(dest))
+        return self
 
     def close(self) -> None:
         _core.streaming_encoder_close(self._encoder_tensor)

test/test_encoders.py

@@ -1589,8 +1589,7 @@ def test_context_manager(self, tmp_path, method):
         enc, encoder_output, open_kwargs = self._create_encoder(method, tmp_path, "mp4")
         frames = torch.randint(0, 256, (5, 3, 64, 64), dtype=torch.uint8)
         video = enc.add_video(height=64, width=64, frame_rate=30.0)
-        enc.open(**open_kwargs)
-        with enc:
+        with enc.open(**open_kwargs):
             video.write(frames)
 
         # The output is valid and decodable, proving close() was called by __exit__.