68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import logging
|
|
import time
|
|
from pathlib import Path
|
|
from threading import Thread
|
|
from typing import Callable
|
|
|
|
from matplotlib import animation, pyplot as plt
|
|
|
|
|
|
def animate(
|
|
fig, frames: list[int] | None, **kwargs
|
|
) -> Callable[..., animation.FuncAnimation]:
|
|
"""Decorator-style wrapper for FuncAnimation."""
|
|
# Why isn't this how the function is exposed by default?
|
|
return lambda func: animation.FuncAnimation(fig, func, frames, **kwargs)
|
|
|
|
# writer = animation.writers["ffmpeg"](fps=15, metadata={"artist": "Me"})
|
|
|
|
def bindfig(fig):
|
|
"""
|
|
Create a function which is compatible with the signature of `pyplot.figure`,
|
|
but alters the already-existing figure `fig` instead.
|
|
"""
|
|
def ret(**kwargs):
|
|
for i, j in kwargs.items():
|
|
if i == "figsize":
|
|
if j is not None:
|
|
fig.set_figwidth(j[0])
|
|
fig.set_figheight(j[1])
|
|
continue
|
|
fig.__dict__["set_" + i](j)
|
|
return fig
|
|
return ret
|
|
|
|
class SympyAnimationWrapper:
|
|
"""
|
|
Context manager for binding a figure for animation.
|
|
"""
|
|
def __init__(self, filename: Path | str):
|
|
self._filename = filename if isinstance(filename, Path) else Path(filename)
|
|
self._current_figure = plt.gcf()
|
|
|
|
self._figure_temp = plt.figure
|
|
|
|
def __enter__(self):
|
|
plt.figure = bindfig(self._current_figure)
|
|
return self.animate
|
|
|
|
def __exit__(self, *_):
|
|
plt.figure = self._figure_temp
|
|
|
|
def animate(self, *args, **kwargs):
|
|
def wrapper(func):
|
|
ret = animate(self._current_figure, *args, **kwargs)(func)
|
|
current_save = ret.save
|
|
def save(*args, **kwargs):
|
|
if self._filename.exists():
|
|
logging.error("Ignoring saving animation '%s': file already exists", self._filename)
|
|
return
|
|
thread = Thread(target=lambda: time.sleep(1) or plt.close())
|
|
thread.run()
|
|
current_save(self._filename, *args, **kwargs)
|
|
ret.save = save
|
|
|
|
return ret
|
|
|
|
return wrapper
|