- utils.prod to extend math.prod to earlier python

tikuma-lsuhsc · tikuma-lsuhsc · commit 4c9fcc6d7ea6 · 2022-02-13T12:45:31.000-06:00
- use utils.get_samplesize on all occasions
diff --git a/src/ffmpegio/image.py b/src/ffmpegio/image.py
@@ -44,7 +44,7 @@ def _run_read(*args, shape=None, pix_fmt_in=None, s_in=None, show_log=None, **kw
         if out.returncode:
             raise FFmpegError(out.stderr)
 
-    nbytes = utils.dtype_itemsize(dtype) * prod(shape)
+    nbytes = utils.get_samplesize(shape, dtype)
 
     return plugins.get_hook().bytes_to_video(
         b=out.stdout[-nbytes:], dtype=dtype, shape=shape
diff --git a/src/ffmpegio/plugins/rawdata_bytes.py b/src/ffmpegio/plugins/rawdata_bytes.py
@@ -1,4 +1,4 @@
-from math import prod
+from ..utils import get_samplesize
 from pluggy import HookimplMarker
 
 hookimpl = HookimplMarker("ffmpegio")
@@ -71,7 +71,7 @@ def bytes_to_video(b: bytes, dtype: str, shape: tuple[int, int, int]) -> object:
     return {
         "buffer": b,
         "dtype": dtype,
-        "shape": (len(b) // (prod(shape) * int(dtype[-1])), *shape),
+        "shape": (len(b) // get_samplesize(shape, dtype), *shape),
     }
 
 
@@ -92,5 +92,5 @@ def bytes_to_audio(b: bytes, dtype: str, shape: tuple[int]) -> object:
     return {
         "buffer": b,
         "dtype": dtype,
-        "shape": (len(b) // (prod(shape) * int(dtype[-1])), *shape),
+        "shape": (len(b) // get_samplesize(shape, dtype), *shape),
     }
diff --git a/src/ffmpegio/streams/SimpleStreams.py b/src/ffmpegio/streams/SimpleStreams.py
@@ -1,5 +1,4 @@
 import logging
-from math import prod
 from .. import utils, configure, ffmpegprocess, probe, plugins
 
 # from ..utils import bytes_to_ndarray as _as_array
@@ -719,14 +718,15 @@ def _get_output_info(self, timeout):
 
     def _start_reader(self):
 
+        self._bps_out = utils.get_samplesize(self.shape, self.dtype)
+        self._bps_in = utils.get_samplesize(self.shape_in, self.dtype_in)
+        self._out2in = self.rate / self.rate_in
+
         # start the FFmpeg output reader
+        self._reader.itemsize = self._bps_out
         self._reader.stdout = self._proc.stdout
-        self._reader.itemsize = utils.get_samplesize(self.shape, self.dtype)
         self._reader.start()
 
-        self._bps_out = prod(self.shape) * int(self.dtype[-1])
-        self._bps_in = prod(self.shape_in) * int(self.dtype_in[-1])
-        self._out2in = self.rate / self.rate_in
         self._reader_needs_info = False
 
     def close(self):
diff --git a/src/ffmpegio/utils/__init__.py b/src/ffmpegio/utils/__init__.py
@@ -1,4 +1,12 @@
-from math import cos, prod, radians, sin
+try:
+    from math import prod
+except:
+    from functools import reduce
+    from operator import mul
+
+    prod = lambda seq: reduce(mul, seq, 1)
+
+from math import cos, radians, sin
 import re, fractions
 from .. import caps
 
diff --git a/tests/test_ffmpegprocess.py b/tests/test_ffmpegprocess.py
@@ -1,8 +1,5 @@
-import pytest
-
-import logging, time
+import logging
 from ffmpegio import probe, configure, ffmpegprocess, utils
-import numpy as np
 
 # logging.basicConfig(level=logging.DEBUG)
 
@@ -115,8 +112,7 @@ def progress(*args):
         ffmpeg_args, pix_fmt_in, s_in, r_in
     )
 
-    out = np.empty(shape, dtype)
-    blksize = out.itemsize * np.prod(shape)  # 1 element
+    samplesize = utils.get_samplesize(shape,dtype)
 
     with ffmpegprocess.Popen(
         ffmpeg_args,
@@ -126,10 +122,8 @@ def progress(*args):
         for k in range(15):
             print(proc.stderr.readline())
         for j in range(10):
-            memoryview(out).cast("b")[:] = memoryview(proc.stdout.read(blksize)).cast(
-                "b"
-            )
-            print(j, out.shape, out.dtype)
+            out = proc.stdout.read(samplesize)
+            assert len(out)==samplesize
 
 
 if __name__ == "__main__":
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
@@ -1,4 +1,4 @@
-from math import prod
+from ffmpegio.utils import prod
 from ffmpegio import plugins
 
 
diff --git a/tests/test_simplestreams.py b/tests/test_simplestreams.py
@@ -1,14 +1,8 @@
-from distutils.debug import DEBUG
-import fractions
-import io
-from math import prod
-from time import sleep
 import ffmpegio
 import tempfile, re
 from os import path
-import numpy as np
 import logging
-from ffmpegio import streams
+from ffmpegio import streams, utils
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -30,9 +24,17 @@ def test_read_video():
 
 def test_read_write_video():
     fs, F = ffmpegio.video.read(url, t=1)
-    bps = prod(F["shape"][1:]) * int(F["dtype"][-1])
-    F0 = {"buffer": F["buffer"][:bps], "shape": F["shape"], "dtype": F["dtype"]}
-    F1 = {"buffer": F["buffer"][bps:], "shape": F["shape"], "dtype": F["dtype"]}
+    bps = utils.get_samplesize(F["shape"][-3:], F["dtype"])
+    F0 = {
+        "buffer": F["buffer"][:bps],
+        "shape": (1, *F["shape"][1:]),
+        "dtype": F["dtype"],
+    }
+    F1 = {
+        "buffer": F["buffer"][bps:],
+        "shape": (F["shape"][0] - 1, *F["shape"][1:]),
+        "dtype": F["dtype"],
+    }
 
     with tempfile.TemporaryDirectory() as tmpdirname:
         out_url = path.join(tmpdirname, re.sub(r"\..*?$", outext, path.basename(url)))
@@ -45,7 +47,7 @@ def test_read_audio(caplog):
     # caplog.set_level(logging.DEBUG)
 
     fs, x = ffmpegio.audio.read(url)
-    bps = prod(x["shape"][1:]) * int(x["dtype"][-1])
+    bps = utils.get_samplesize(x["shape"][-1:], x["dtype"])
 
     with ffmpegio.open(url, "ra", show_log=True, blocksize=1024 ** 2) as f:
         # x = f.read(1024)
diff --git a/tests/test_video.py b/tests/test_video.py
@@ -1,8 +1,6 @@
-from math import prod
-from ffmpegio import video, probe
+from ffmpegio import video, probe, utils
 import tempfile, re
 from os import path
-import numpy as np
 
 
 def test_create():
@@ -50,7 +48,6 @@ def test_read_write():
     with tempfile.TemporaryDirectory() as tmpdirname:
         out_url = path.join(tmpdirname, re.sub(r"\..*?$", outext, path.basename(url)))
         print(out_url)
-        print(prod(A["shape"]))
         video.write(out_url, fs, A)
         print(probe.video_streams_basic(out_url))
         fs, A = video.read(out_url, vframes=10)
@@ -73,14 +70,14 @@ def test_read():
     N = 5
     fs, B = video.read(url, vframes=10, ss_in=float(n0 / fs))
     print(B["shape"], A["shape"])
-    nbytes = prod(A["shape"][1:]) * int(A["dtype"][-1])
-    assert A["buffer"][n0 * nbytes :] == B['buffer'][: (10 - n0) * nbytes]
+    nbytes = utils.get_samplesize(A["shape"][-3:], A["dtype"])
+    assert A["buffer"][n0 * nbytes :] == B["buffer"][: (10 - n0) * nbytes]
 
     fs, C = video.read(url, ss_in=float(n0 / fs), t_in=float(N / fs))
 
     print(C["shape"])
-    assert A["buffer"][n0 * nbytes :(n0+N)*nbytes] == C['buffer']
-    
+    assert A["buffer"][n0 * nbytes : (n0 + N) * nbytes] == C["buffer"]
+
     # fs, D = video.read(url, ss_in=n0, t_in=N, units='frames')
     # assert np.array_equal(D, C)
 
@@ -113,7 +110,6 @@ def test_two_pass_write():
     # test_create()
     from ffmpegio import configure, utils, ffmpegprocess
     from ffmpegio.utils import log as log_utils
-    import numpy as np
     from pprint import pprint
     import re
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-from math import prod`
	`1`	`+from ffmpegio.utils import prod`
`2`	`2`	`from ffmpegio import plugins`
`3`	`3`
`4`	`4`