Unverified Commit 402939ed authored by moto's avatar moto Committed by GitHub
Browse files

Add path-like object support to StreamReader/Writer (#3608)

parent ac63c454
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Dict, Iterator, Optional, Tuple, TypeVar, Union
import torch
......@@ -446,7 +447,7 @@ class StreamReader:
For the detailed usage of this class, please refer to the tutorial.
Args:
src (str, file-like object): The media source.
src (str, path-like or file-like object): The media source.
If string-type, it must be a resource indicator that FFmpeg can
handle. This includes a file path, URL, device identifier or
filter expression. The supported value depends on the FFmpeg found
......@@ -512,17 +513,15 @@ class StreamReader:
def __init__(
self,
src: Union[str, BinaryIO],
src: Union[str, Path, BinaryIO],
format: Optional[str] = None,
option: Optional[Dict[str, str]] = None,
buffer_size: int = 4096,
):
if isinstance(src, str):
self._be = _StreamReader(src, format, option)
elif hasattr(src, "read"):
if hasattr(src, "read"):
self._be = _StreamReaderFileObj(src, format, option, buffer_size)
else:
raise ValueError("`src` must be either a string or file-like object.")
self._be = _StreamReader(str(src), format, option)
i = self._be.find_best_audio_stream()
self._default_audio_stream = None if i < 0 else i
......
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Dict, Optional, Union
import torch
......@@ -132,7 +133,7 @@ class StreamWriter:
"""Encode and write audio/video streams chunk by chunk
Args:
dst (str or file-like object): The destination where the encoded data are written.
dst (str, path-like or file-like object): The destination where the encoded data are written.
If string-type, it must be a resource indicator that FFmpeg can
handle. The supported value depends on the FFmpeg found in the system.
......@@ -184,16 +185,14 @@ class StreamWriter:
def __init__(
self,
dst: Union[str, BinaryIO],
dst: Union[str, Path, BinaryIO],
format: Optional[str] = None,
buffer_size: int = 4096,
):
if isinstance(dst, str):
self._s = _StreamWriter(dst, format)
elif hasattr(dst, "write"):
if hasattr(dst, "write"):
self._s = _StreamWriterFileObj(dst, format, buffer_size)
else:
raise ValueError("`dst` must be either a string or a file-like object.")
self._s = _StreamWriter(str(dst), format)
self._is_open = False
@_format_common_args
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment