This repository was archived by the owner on Jan 19, 2025. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_base.py
More file actions
65 lines (44 loc) · 1.36 KB
/
test_base.py
File metadata and controls
65 lines (44 loc) · 1.36 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
try:
from math import prod
except:
from functools import reduce
from operator import mul
prod = lambda seq: reduce(mul, seq, 1)
import tempfile
from os import path
import numpy as np
from matplotlib import pyplot as plt
from ffmpegio import plugins, open as ffopen
import pytest
@pytest.fixture(scope="module")
def x():
yield np.arange(0, 2 * np.pi, 0.01)
@pytest.fixture(scope="module")
def plot_figure(x):
h = 1080
w = 1920
dpi = 120
fig, ax = plt.subplots(figsize=(w / dpi, h / dpi), dpi=dpi)
(line,) = ax.plot(x, np.sin(x))
yield (h, w, 4), fig, line
plt.close(fig)
def test_hooks(plot_figure):
figsize, fig, _ = plot_figure
hook = plugins.get_hook()
assert hook.video_info(obj=fig) == (figsize, "|u1")
assert len(hook.video_bytes(obj=fig)) == prod(figsize)
def test_video(x, plot_figure):
_, fig, line = plot_figure
interval = 20 # delay in milliseconds
save_count = 50 # number of frames
def animate(i):
line.set_ydata(np.sin(x + i / 50)) # update the data.
return line
with tempfile.TemporaryDirectory() as tmpdir, ffopen(
path.join(tmpdir, "output.mp4"), "wv", 1e3 / interval
) as f:
for n in range(save_count):
animate(n) # update figure
f.write(fig) # write new frame
if __name__ == "__main__":
test_video()