Skip to content

Commit

Permalink
Fix segfault when printing closed video_stream
Browse files Browse the repository at this point in the history
  • Loading branch information
WyattBlue committed Oct 20, 2024
1 parent 5bbd3b1 commit 49b7c14
Show file tree
Hide file tree
Showing 7 changed files with 59 additions and 10 deletions.
7 changes: 5 additions & 2 deletions av/audio/stream.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ class _Layout:

class AudioStream(Stream):
codec_context: AudioCodecContext
def encode(self, frame: AudioFrame | None = None) -> list[Packet]: ...
def decode(self, packet: Packet | None = None) -> list[AudioFrame]: ...

# From codec context
frame_size: int
sample_rate: int
Expand All @@ -27,5 +30,5 @@ class AudioStream(Stream):
type: Literal["audio"]
format: _Format
layout: _Layout
def encode(self, frame: AudioFrame | None = None) -> list[Packet]: ...
def decode(self, packet: Packet | None = None) -> list[AudioFrame]: ...

def close(self, strict: bool = True) -> None: ...
4 changes: 2 additions & 2 deletions av/codec/context.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,8 @@ cdef class CodecContext:
cdef _setup_decoded_frame(self, Frame, Packet)

# Implemented by base for the generic send/recv API.
# Note that the user cannot send without recieving. This is because
# _prepare_frames_for_encode may expand a frame into multiple (e.g. when
# Note that the user cannot send without receiving. This is because
# `_prepare_frames_for_encode` may expand a frame into multiple (e.g. when
# resampling audio to a higher rate but with fixed size frames), and the
# send/recv buffer may be limited to a single frame. Ergo, we need to flush
# the buffer as often as possible.
Expand Down
9 changes: 7 additions & 2 deletions av/codec/context.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -193,10 +193,11 @@ cdef class CodecContext:

@property
def extradata(self):
if self.ptr is NULL:
return None
if self.ptr.extradata_size > 0:
return <bytes>(<uint8_t*>self.ptr.extradata)[:self.ptr.extradata_size]
else:
return None
return None

@extradata.setter
def extradata(self, data):
Expand All @@ -222,10 +223,14 @@ cdef class CodecContext:

@property
def is_encoder(self):
if self.ptr is NULL:
return False
return lib.av_codec_is_encoder(self.ptr.codec)

@property
def is_decoder(self):
if self.ptr is NULL:
return False
return lib.av_codec_is_decoder(self.ptr.codec)

cpdef open(self, bint strict=True):
Expand Down
1 change: 0 additions & 1 deletion av/container/streams.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ class StreamContainer:
other: tuple[Stream, ...]

def __init__(self) -> None: ...
def add_stream(self, stream: Stream) -> None: ...
def __len__(self) -> int: ...
def __iter__(self) -> Iterator[Stream]: ...
@overload
Expand Down
4 changes: 4 additions & 0 deletions av/video/codeccontext.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ cdef class VideoCodecContext(CodecContext):

@property
def width(self):
if self.ptr is NULL:
return 0
return self.ptr.width

@width.setter
Expand All @@ -87,6 +89,8 @@ cdef class VideoCodecContext(CodecContext):

@property
def height(self):
if self.ptr is NULL:
return 0
return self.ptr.height

@height.setter
Expand Down
9 changes: 6 additions & 3 deletions av/video/stream.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,11 @@ class VideoStream(Stream):
sample_aspect_ratio: Fraction | None
display_aspect_ratio: Fraction | None
codec_context: VideoCodecContext

def encode(self, frame: VideoFrame | None = None) -> list[Packet]: ...
def encode_lazy(self, frame: VideoFrame | None = None) -> Iterator[Packet]: ...
def decode(self, packet: Packet | None = None) -> list[VideoFrame]: ...

# from codec context
format: VideoFormat
width: int
Expand All @@ -36,6 +41,4 @@ class VideoStream(Stream):
colorspace: int
type: Literal["video"]

def encode(self, frame: VideoFrame | None = None) -> list[Packet]: ...
def encode_lazy(self, frame: VideoFrame | None = None) -> Iterator[Packet]: ...
def decode(self, packet: Packet | None = None) -> list[VideoFrame]: ...
def close(self, strict: bool = True) -> None: ...
35 changes: 35 additions & 0 deletions tests/test_streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,41 @@ def test_selection(self) -> None:
data = container.streams.data[0]
assert data == container.streams.best("data")

def test_printing_closed_video_stream(self) -> None:
input_ = av.open(
fate_suite("amv/MTV_high_res_320x240_sample_Penguin_Joke_MTV_from_WMV.amv")
)
container = av.open("out.mkv", "w")

video_stream = container.add_stream("h264", rate=30)
# encoder = video_stream.codec.name + ""

video_stream.width = input_.streams.video[0].width
video_stream.height = input_.streams.video[0].height
video_stream.pix_fmt = "yuv420p"

for frame in input_.decode(video=0):
container.mux(video_stream.encode(frame))
break

encoder = "libx264"
repr = f"{video_stream}"
assert repr.startswith(f"<av.VideoStream #0 {encoder}, yuv420p 160x120 at ")
assert repr.endswith(">")

# repr = f"{video_stream}"
# assert repr.startswith(f"<av.VideoStream #0 {encoder}, yuv420p 160x120 at ")
# assert repr.endswith(">")

video_stream.close()

repr = f"{video_stream}"
assert repr.startswith(f"<av.VideoStream #0 {encoder}, yuv420p 0x0 at ")
assert repr.endswith(">")

container.close()
input_.close()

# def test_side_data(self) -> None:
# container = av.open(fate_suite("mov/displaymatrix.mov"))
# video = container.streams.video[0]
Expand Down

0 comments on commit 49b7c14

Please sign in to comment.