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 __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Dict, Iterator, Optional, Tuple, TypeVar, Union from typing import BinaryIO, Dict, Iterator, Optional, Tuple, TypeVar, Union
import torch import torch
...@@ -446,7 +447,7 @@ class StreamReader: ...@@ -446,7 +447,7 @@ class StreamReader:
For the detailed usage of this class, please refer to the tutorial. For the detailed usage of this class, please refer to the tutorial.
Args: 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 If string-type, it must be a resource indicator that FFmpeg can
handle. This includes a file path, URL, device identifier or handle. This includes a file path, URL, device identifier or
filter expression. The supported value depends on the FFmpeg found filter expression. The supported value depends on the FFmpeg found
...@@ -512,17 +513,15 @@ class StreamReader: ...@@ -512,17 +513,15 @@ class StreamReader:
def __init__( def __init__(
self, self,
src: Union[str, BinaryIO], src: Union[str, Path, BinaryIO],
format: Optional[str] = None, format: Optional[str] = None,
option: Optional[Dict[str, str]] = None, option: Optional[Dict[str, str]] = None,
buffer_size: int = 4096, buffer_size: int = 4096,
): ):
if isinstance(src, str): if hasattr(src, "read"):
self._be = _StreamReader(src, format, option)
elif hasattr(src, "read"):
self._be = _StreamReaderFileObj(src, format, option, buffer_size) self._be = _StreamReaderFileObj(src, format, option, buffer_size)
else: 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() i = self._be.find_best_audio_stream()
self._default_audio_stream = None if i < 0 else i self._default_audio_stream = None if i < 0 else i
......
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Dict, Optional, Union from typing import BinaryIO, Dict, Optional, Union
import torch import torch
...@@ -132,7 +133,7 @@ class StreamWriter: ...@@ -132,7 +133,7 @@ class StreamWriter:
"""Encode and write audio/video streams chunk by chunk """Encode and write audio/video streams chunk by chunk
Args: 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 If string-type, it must be a resource indicator that FFmpeg can
handle. The supported value depends on the FFmpeg found in the system. handle. The supported value depends on the FFmpeg found in the system.
...@@ -184,16 +185,14 @@ class StreamWriter: ...@@ -184,16 +185,14 @@ class StreamWriter:
def __init__( def __init__(
self, self,
dst: Union[str, BinaryIO], dst: Union[str, Path, BinaryIO],
format: Optional[str] = None, format: Optional[str] = None,
buffer_size: int = 4096, buffer_size: int = 4096,
): ):
if isinstance(dst, str): if hasattr(dst, "write"):
self._s = _StreamWriter(dst, format)
elif hasattr(dst, "write"):
self._s = _StreamWriterFileObj(dst, format, buffer_size) self._s = _StreamWriterFileObj(dst, format, buffer_size)
else: else:
raise ValueError("`dst` must be either a string or a file-like object.") self._s = _StreamWriter(str(dst), format)
self._is_open = False self._is_open = False
@_format_common_args @_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