Repository: just-work/fffw

Issue:
Support computing frames count for video metadata
It's useful to track encoding progress and to restore partial video metadata

Comments:

Patch:
diff --git a/fffw/encoding/filters.py b/fffw/encoding/filters.py
index f42e5b2..194424c 100644
--- a/fffw/encoding/filters.py
+++ b/fffw/encoding/filters.py
@@ -278,8 +278,18 @@ class Trim(AutoFilter):
                 scenes.append(Scene(stream=scene.stream, start=start,
                                     duration=end - start))
 
-        return replace(meta, start=self.start, duration=self.end,
-                       scenes=scenes, streams=streams)
+        kwargs = {
+            'start': self.start,
+            'duration': self.end,
+            'scenes': scenes,
+            'streams': streams
+        }
+        interval = cast(TS, self.end) - cast(TS, self.start)
+        if isinstance(meta, AudioMeta):
+            kwargs['samples'] = round(meta.sampling_rate * interval)
+        if isinstance(meta, VideoMeta):
+            kwargs['frames'] = round(meta.frame_rate * interval)
+        return replace(meta, **kwargs)
 
 
 @dataclass
@@ -353,6 +363,7 @@ class Concat(Filter):
         duration = TS(0)
         scenes = []
         streams: List[str] = []
+        frames: int = 0
         for meta in metadata:
             duration += meta.duration
             scenes.extend(meta.scenes)
@@ -361,12 +372,17 @@ class Concat(Filter):
                     # Add all streams for each concatenated metadata and remove
                     # contiguous duplicates.
                     streams.append(stream)
+            if isinstance(meta, VideoMeta):
+                frames += meta.frames
         kwargs = dict(duration=duration, scenes=scenes, streams=streams)
         meta = metadata[0]
         if isinstance(meta, AudioMeta):
             # Recompute samples and sampling rate: sampling rate from first
             # input, samples count corresponds duration.
             kwargs['samples'] = round(meta.sampling_rate * duration)
+        if isinstance(meta, VideoMeta):
+            # Sum frames count from all input streams
+            kwargs['frames'] = frames
         return replace(metadata[0], **kwargs)
 
 
diff --git a/fffw/graph/meta.py b/fffw/graph/meta.py
index c767955..29beac6 100644
--- a/fffw/graph/meta.py
+++ b/fffw/graph/meta.py
@@ -330,6 +330,8 @@ class VideoMeta(Meta):
     """ Display aspect ratio."""
     frame_rate: float
     """ Frames per second."""
+    frames: int
+    """ Number of frames."""
     device: Optional[Device]
     """ Hardware device asociated with current stream."""
 
@@ -346,6 +348,9 @@ class VideoMeta(Meta):
         else:
             assert str(self.dar) == 'nan'
 
+        interval = float(self.duration - self.start)
+        assert abs(self.frames - interval * self.frame_rate) <= 1
+
 
 @dataclass
 class AudioMeta(Meta):
@@ -369,11 +374,8 @@ class AudioMeta(Meta):
         return AUDIO
 
     def validate(self) -> None:
-        duration = self.duration.total_seconds()
-        if duration != 0:
-            assert abs(self.sampling_rate - self.samples / duration) < 0.001
-        else:
-            assert self.sampling_rate == 0
+        interval = float(self.duration - self.start)
+        assert abs(self.samples - interval * self.sampling_rate) <= 1
 
 
 def audio_meta_data(**kwargs: Any) -> AudioMeta:
@@ -437,6 +439,7 @@ def video_meta_data(**kwargs: Any) -> VideoMeta:
         par=par,
         dar=dar,
         frame_rate=frame_rate,
+        frames=frames,
         device=None,
     )
 
diff --git a/fffw/wrapper/base.py b/fffw/wrapper/base.py
index b595a09..d8326b8 100644
--- a/fffw/wrapper/base.py
+++ b/fffw/wrapper/base.py
@@ -5,13 +5,78 @@ from asyncio.subprocess import Process
 from dataclasses import dataclass
 from logging import getLogger
 from types import TracebackType
-from typing import Tuple, List, Any, Optional, cast, Callable, Union, TextIO, \
-    Type
+from typing import Tuple, List, Any, Optional, cast, Callable, Union, TextIO
+from typing import Type, AsyncIterator
 
 from fffw.wrapper.helpers import quote, ensure_binary, ensure_text
 from fffw.wrapper.params import Params
 
 
+class UniversalLineReader:
+    """
+    Reads bytes from asyncio.StreamReader and splits it to lines with either
+    CR or LF, or even CRLF.
+
+    https://docs.python.org/3/glossary.html#term-universal-newlines
+
+    >>> # noinspection PyUnresolvedReferences
+    ... line_iter = UniversalLineReader(process.stderr)
+    >>> async for line in line_iter:
+    ...    print(line)
+
+    """
+
+    def __init__(self,
+                 reader: asyncio.StreamReader,
+                 bufsize: int = 10 * io.DEFAULT_BUFFER_SIZE,
+                 blocksize: int = io.DEFAULT_BUFFER_SIZE,
+                 encoding: str = 'utf8') -> None:
+        """
+        :param reader: asynchronous stream reader, i.e. stdout/stderr of
+            asyncio Process instance.
+        :param bufsize: max buffer size
+        :param blocksize: read block size
+        :param encoding: text encoding
+        """
+        self.blocksize = blocksize
+        self.reader = reader
+        self.bufsize = bufsize
+        self.buffer = b''
+        self.encoding = encoding
+        self.at_eof = False
+
+    def __aiter__(self) -> AsyncIterator[str]:
+        return self.readlines()
+
+    async def readlines(self) -> AsyncIterator[str]:
+        while not self.at_eof:
+            # StreamReader supports only LF line separator. This leads to buffer
+            # overrun when it contains only CR-terminated lines. Thus, we read
+            # blocks manually and then split it to lines with universal line
+            # separator.
+            block = await self.reader.read(self.blocksize)
+            # empty read means that stream is closed
+            self.at_eof = len(block) == 0
+
+            # checking max buffer size
+            buffered = len(self.buffer) + len(block)
+            if buffered > self.bufsize:
+                raise asyncio.LimitOverrunError("buffer overrun", buffered)
+
+            self.buffer += block
+            if self.buffer:
+                # Split buffer to line with any of CR, LF of CRLF separators
+                # Last line is buffered to handle the case when CRLF sequence is
+                # being split to subsequent reads.
+                [*lines, self.buffer] = self.buffer.splitlines(keepends=True)
+                for line in lines:
+                    yield line.decode(self.encoding)
+        # We always leave one non-empty line above, but stream may be empty. In
+        # this case we don't want to yield empty line.
+        if self.buffer:
+            yield self.buffer.decode(self.encoding)
+
+
 class Runner:
     """ Wrapper for Popen process for non-blocking streams handling."""
 
@@ -69,7 +134,7 @@ class Runner:
         """
         Handles read from stdout/stderr.
 
-        Reads lines from reader and feeds it to callback. Values, filtered by
+        Reads lines from stream and feeds it to callback. Values, filtered by
         callback, are written to output buffer.
 
         :param reader: Process.stdout or Process.stderr instance
@@ -78,14 +143,10 @@ class Runner:
         """
         if callback is None or reader is None:
             return
-        while True:
-            line = await reader.readline()
-            if line:
-                data = callback(line.decode())
-                if data:
-                    output.write(data)
-            else:
-                break
+        async for line in UniversalLineReader(reader):
+            data = callback(line)
+            if data:
+                output.write(data)
 
     @staticmethod
     async def write(writer: Optional[asyncio.StreamWriter],


Test patch:
diff --git a/tests/test_graph.py b/tests/test_graph.py
index e8c25d9..b85e73b 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from typing import cast
 from unittest import TestCase
 
-from fffw.encoding import inputs, outputs
+from fffw.encoding import inputs, outputs, codecs
 from fffw.encoding.complex import FilterComplex
 from fffw.encoding.filters import *
 from fffw.graph import *
@@ -34,17 +34,22 @@ class FilterGraphTestCase(TestCase):
             dar=1.777777778,
             par=1.0,
             duration=300.0,
+            frame_rate=10.0,
+            frame_count=3000
         )
         self.audio_metadata = audio_meta_data(
             duration=200.0,
             sampling_rate=48000,
-            samples_count=200*48000)
+            samples_count=200 * 48000)
 
         self.source = inputs.Input(
             input_file='input.mp4',
             streams=(inputs.Stream(VIDEO, meta=self.video_metadata),
                      inputs.Stream(AUDIO, meta=self.audio_metadata)))
-        self.output = outputs.output_file('output.mp4')
+        self.output = outputs.output_file(
+            'output.mp4',
+            codecs.VideoCodec('libx264'),
+            codecs.AudioCodec('libfdk_aac'))
         self.input_list = inputs.InputList((self.source,))
         self.output_list = outputs.OutputList((self.output,))
         self.fc = FilterComplex(self.input_list, self.output_list)
@@ -237,13 +242,16 @@ class FilterGraphTestCase(TestCase):
         self.assertEqual(vm.width, self.video_metadata.width)
         self.assertEqual(vm.height, self.video_metadata.height)
 
-    def test_concat_metadata(self):
+    def test_concat_video_metadata(self):
         """
         Concat filter sums stream duration
 
         $ ffmpeg -y -i first.mp4 -i second.mp4 -filter_complex concat test.mp4
         """
-        vs = inputs.Stream(VIDEO, meta=video_meta_data(duration=1000.0))
+        video_meta = video_meta_data(duration=1000.0,
+                                     frame_count=10000,
+                                     frame_rate=10.0)
+        vs = inputs.Stream(VIDEO, meta=video_meta)
         self.input_list.append(inputs.input_file('second.mp4', vs))
         concat = vs | Concat(VIDEO)
         self.source | concat
@@ -253,12 +261,15 @@ class FilterGraphTestCase(TestCase):
         vm = cast(VideoMeta, self.output.codecs[0].get_meta_data())
         self.assertEqual(self.video_metadata.duration + vs.meta.duration,
                          vm.duration)
+        self.assertEqual(self.video_metadata.frames + video_meta.frames,
+                         vm.frames)
 
     def test_concat_audio_metadata(self):
         """
         Concat filter sums samples count for audio streams.
         """
-        audio_meta = audio_meta_data(duration=1000.0, sampling_rate=24000,
+        audio_meta = audio_meta_data(duration=1000.0,
+                                     sampling_rate=24000,
                                      samples_count=24000 * 1000)
         a = inputs.Stream(AUDIO, meta=audio_meta)
         self.input_list.append(inputs.input_file('second.mp4', a))
@@ -273,7 +284,7 @@ class FilterGraphTestCase(TestCase):
         self.assertEqual(round(am.duration * audio_meta.sampling_rate),
                          am.samples)
 
-    def test_trim_metadata(self):
+    def test_video_trim_metadata(self):
         """
         Trim filter sets start and changes stream duration.
         $ ffmpeg -y -i source.mp4 -vf trim=start=3:end=4 -an test.mp4
@@ -285,6 +296,21 @@ class FilterGraphTestCase(TestCase):
         vm = cast(VideoMeta, self.output.codecs[0].get_meta_data())
         self.assertEqual(vm.start, TS(3.0))
         self.assertEqual(vm.duration, TS(4.0))
+        self.assertEqual(vm.frames, 1.0 * vm.frame_rate)
+
+    def test_audio_trim_metadata(self):
+        """
+        Trim filter sets start and changes stream duration.
+        $ ffmpeg -y -i source.mp4 -af atrim=start=3:end=4 -vn test.mp4
+
+        Note that resulting video has 3 seconds of frozen frame at 00:00:03.000,
+        total duration is 4.
+        """
+        self.source | Trim(AUDIO, start=3.0, end=4.0) > self.output
+        am = cast(AudioMeta, self.output.codecs[1].get_meta_data())
+        self.assertEqual(am.start, TS(3.0))
+        self.assertEqual(am.duration, TS(4.0))
+        self.assertEqual(am.samples, 1.0 * am.sampling_rate)
 
     def test_setpts_metadata(self):
         """
diff --git a/tests/test_meta_data.py b/tests/test_meta_data.py
index 8f63ce8..801af85 100644
--- a/tests/test_meta_data.py
+++ b/tests/test_meta_data.py
@@ -229,6 +229,7 @@ class MetaDataTestCase(TestCase):
             par=1.0,
             dar=1.778,
             frame_rate=50.0,
+            frames=337,
             device=None,
         )
         self.assertEqual(expected, video)
@@ -266,7 +267,6 @@ class MetaDataTestCase(TestCase):
         self.assertTrue(fields(ExtendedVideoMeta))
 
 
-
 class TimeStampTestCase(TestCase):
     td: timedelta
     ts: meta.TS
diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py
new file mode 100644
index 0000000..ba68852
--- /dev/null
+++ b/tests/test_wrapper.py
@@ -0,0 +1,121 @@
+import asyncio
+import io
+import random
+import sys
+import time
+from dataclasses import dataclass
+from typing import List
+from unittest import TestCase, mock
+
+from fffw.wrapper import param
+from fffw.wrapper.base import UniversalLineReader, BaseWrapper
+
+
+@dataclass
+class Python(BaseWrapper):
+    command = 'python'
+    module: str = param(name='m')
+
+
+def script():
+    line = input()
+    sys.stdout.write(f'stdout: {line}\n')
+    sys.stderr.write(f'stderr: {line}\n')
+    time.sleep(int(line) // 100)
+    return int(line)
+
+
+if __name__ == '__main__':
+    sys.exit(script())
+
+
+class WrapperTestCase(TestCase):
+
+    def test_run_child_process(self):
+        p = Python(module='tests.test_wrapper')
+        ret, out, err = p.run('1')
+        self.assertEqual(ret, 1)
+        self.assertEqual(out, 'stdout: 1\n')
+        self.assertEqual(err, 'stderr: 1\n')
+
+    def test_child_timeout(self):
+        p = Python(module='tests.test_wrapper')
+        ret, out, err = p.run('100', timeout=0.01)
+        self.assertEqual(ret, -9)
+
+
+class UniversalLineReaderTestCase(TestCase):
+    def setUp(self) -> None:
+        self.data = io.BytesIO()
+        self.stream = mock.MagicMock(read=self.read)
+        self.reader = UniversalLineReader(self.stream, bufsize=100, blocksize=5)
+
+    async def read(self, n=-1):
+        return self.data.read(n)
+
+    async def iterate(self) -> List[str]:
+        result = []
+        async for line in self.reader:
+            result.append(line)
+        return result
+
+    def assert_lines(self, lines: List[str]):
+        for line in lines:
+            self.data.write(line.encode('utf-8'))
+        self.data.seek(0)
+        result = asyncio.get_event_loop().run_until_complete(self.iterate())
+        self.assertListEqual(lines, result)
+
+    def test_read_lf(self):
+        lines = []
+        for c in 'abcdefghijklmnopqrstuvwxyz':
+            length = random.randint(0, 15)
+            lines.append(c * length + '\n')
+        for line in lines:
+            self.data.write(line.encode('utf-8'))
+        self.data.seek(0)
+
+        self.assert_lines(lines)
+
+    def test_read_cr(self):
+        lines = []
+        for c in 'abcdefghijklmnopqrstuvwxyz':
+            length = random.randint(0, 15)
+            lines.append(c * length + '\r')
+
+        self.assert_lines(lines)
+
+    def test_read_crlf(self):
+        lines = []
+        for c in 'abcdefghijklmnopqrstuvwxyz':
+            length = random.randint(0, 15)
+            lines.append(c * length + '\r\n')
+
+        self.assert_lines(lines)
+
+    def test_empty_lines(self):
+        lines = [
+            'a\n',
+            '\n',
+            'b\n'
+        ]
+
+        self.assert_lines(lines)
+
+    def test_last_incomplete_line(self):
+        lines = [
+            'aaaaa\n',
+            'b'
+        ]
+
+        self.assert_lines(lines)
+
+    def test_empty_stream(self):
+        self.assert_lines([])
+
+    def test_buffer_overrun(self):
+        max_line = 'a' * self.reader.bufsize
+        lines = [max_line + '\n']
+
+        with self.assertRaises(asyncio.LimitOverrunError):
+            self.assert_lines(lines)