plotlive 0.1.0__py3-none-any.whl

plotlive/__init__.py

@@ -0,0 +1,18 @@
+"""
+plotlive — Interactive matplotlib-compatible graphs rendered with pygame.
+
+Usage:
+    import plotlive.pyplot as plt
+    import numpy as np
+
+    plt.plot([1, 2, 3], [4, 5, 6], label='data')
+    plt.legend()
+    plt.show()
+"""
+from .figure import Figure
+from .axes import Axes
+from . import pyplot
+from .pyplot import show, figure, subplots
+
+__all__ = ['Figure', 'Axes', 'pyplot', 'show', 'figure', 'subplots']
+__version__ = '0.1.0'

plotlive/_jupyter.py

@@ -0,0 +1,127 @@
+"""
+plotlive._jupyter — inline display helpers for Jupyter notebooks.
+
+This module is imported lazily inside pyplot.show() only when a Jupyter
+kernel is detected, so it never affects the normal pygame path.
+"""
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .figure import Figure
+
+
+def is_jupyter() -> bool:
+    """Return True when running inside a Jupyter notebook (not terminal IPython)."""
+    try:
+        from IPython import get_ipython  # type: ignore[import]
+        ip = get_ipython()
+        return ip is not None and type(ip).__name__ == 'ZMQInteractiveShell'
+    except Exception:
+        return False
+
+
+def show_figure(fig: Figure) -> None:
+    """Render figure as PNG and display it inline via IPython."""
+    import pygame
+    import tempfile
+    import os
+    from IPython.display import display, Image  # type: ignore[import]
+    from .renderer import FigureRenderer
+
+    if not pygame.get_init():
+        pygame.init()
+
+    surf = FigureRenderer(fig).render()
+    tmp = None
+    try:
+        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
+            tmp = f.name
+        pygame.image.save(surf, tmp)
+        with open(tmp, 'rb') as f:
+            png_data = f.read()
+    finally:
+        if tmp:
+            try:
+                os.unlink(tmp)
+            except OSError:
+                pass
+    display(Image(data=png_data, format='png'))
+
+
+def show_animation(fig: Figure, set_current_fig) -> None:
+    """Export animation as GIF (or MP4 fallback) and display it inline via IPython.
+
+    set_current_fig is a callable(fig) that updates pyplot's global state so
+    that gcf()/gca() resolve correctly inside state-machine update functions.
+    """
+    import pygame
+    import tempfile
+    import os
+    from IPython.display import display  # type: ignore[import]
+
+    if not pygame.get_init():
+        pygame.init()
+
+    # Keep gcf()/gca() pointing at this figure while frames are rendered
+    set_current_fig(fig)
+
+    anim = fig._animation
+    if anim is None:
+        return
+
+    # --- GIF via Pillow (preferred: self-contained, works everywhere) ---
+    try:
+        import PIL  # noqa: F401 — availability check only
+        from IPython.display import Image  # type: ignore[import]
+        tmp = None
+        try:
+            with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as f:
+                tmp = f.name
+            anim.save(tmp, writer='pillow')
+            with open(tmp, 'rb') as f:
+                gif_data = f.read()
+        finally:
+            if tmp:
+                try:
+                    os.unlink(tmp)
+                except OSError:
+                    pass
+        display(Image(data=gif_data, format='gif'))
+        return
+    except ImportError:
+        pass
+
+    # --- MP4 via imageio + ffmpeg (fallback) ---
+    try:
+        import imageio  # noqa: F401
+        tmp = None
+        try:
+            with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as f:
+                tmp = f.name
+            anim.save(tmp, writer='ffmpeg')
+            with open(tmp, 'rb') as f:
+                video_data = f.read()
+        finally:
+            if tmp:
+                try:
+                    os.unlink(tmp)
+                except OSError:
+                    pass
+        import base64
+        from IPython.display import HTML  # type: ignore[import]
+        b64 = base64.b64encode(video_data).decode('ascii')
+        display(HTML(
+            f'<video controls style="max-width:100%">'
+            f'<source type="video/mp4" src="data:video/mp4;base64,{b64}">'
+            f'</video>'
+        ))
+        return
+    except (ImportError, Exception):
+        pass
+
+    # --- Last resort: first frame as static image ---
+    anim._call_func(0)
+    for ax in fig.axes:
+        ax._dirty = True
+    show_figure(fig)

plotlive/_parsers.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+import numpy as np
+
+_MARKERS = set('os^v<>Dp+x.*phH1234')
+# Order matters: parse multi-char linestyles before single-char
+_LINESTYLES = ['-.', '--', '-', ':']
+_LS_MARKERS = set('-')  # chars that could be either linestyle or marker start
+
+
+def _parse_fmt(fmt: str) -> dict:
+    """
+    Parse matplotlib format string like 'b--o' into property dict.
+    Returns subset of: color, linestyle, marker.
+    """
+    if not fmt:
+        return {}
+
+    result: dict = {}
+    s = fmt
+
+    # Extract linestyle first (multi-char before single-char)
+    for ls in _LINESTYLES:
+        if ls in s:
+            result['linestyle'] = ls
+            s = s.replace(ls, '', 1)
+            break
+
+    # Remaining characters: color and/or marker
+    color_chars = set('brgykwmc')
+    for ch in list(s):
+        if ch in color_chars and 'color' not in result:
+            result['color'] = ch
+        elif ch in _MARKERS and 'marker' not in result:
+            result['marker'] = ch
+
+    return result
+
+
+def _parse_plot_args(args: tuple) -> list[tuple[np.ndarray, np.ndarray, str]]:
+    """
+    Parse positional args to plot() into list of (x, y, fmt) tuples.
+
+    Patterns:
+      (y,)           -> (arange(len(y)), y, '')
+      (y, fmt)       -> (arange(len(y)), y, fmt)  if fmt is str
+      (x, y)         -> (x, y, '')
+      (x, y, fmt)    -> (x, y, fmt)
+      (x, y, fmt, x2, y2, fmt2, ...) -> multiple
+    """
+    segments: list[tuple[np.ndarray, np.ndarray, str]] = []
+    i = 0
+    args_list = list(args)
+
+    while i < len(args_list):
+        a0 = args_list[i]
+
+        # Detect string (format) at position i
+        if isinstance(a0, str):
+            raise ValueError(f"Unexpected string argument at position {i}: {a0!r}")
+
+        x0 = np.asarray(a0, dtype=float)
+
+        if i + 1 >= len(args_list):
+            # Only one array: treat as y
+            segments.append((np.arange(len(x0), dtype=float), x0, ''))
+            i += 1
+            continue
+
+        a1 = args_list[i + 1]
+
+        if isinstance(a1, str):
+            # (y, fmt)
+            segments.append((np.arange(len(x0), dtype=float), x0, a1))
+            i += 2
+            continue
+
+        # a1 is array-like: (x, y, ...)
+        y0 = np.asarray(a1, dtype=float)
+
+        if i + 2 < len(args_list) and isinstance(args_list[i + 2], str):
+            fmt = args_list[i + 2]
+            segments.append((x0, y0, fmt))
+            i += 3
+        else:
+            segments.append((x0, y0, ''))
+            i += 2
+
+    return segments

plotlive/animation.py

@@ -0,0 +1,279 @@
+from __future__ import annotations
+
+import inspect
+from typing import Callable
+
+_ANIMATION_EVENT_ID = None
+
+
+def _get_animation_event():
+    global _ANIMATION_EVENT_ID
+    if _ANIMATION_EVENT_ID is None:
+        import pygame
+        _ANIMATION_EVENT_ID = pygame.USEREVENT + 1
+    return _ANIMATION_EVENT_ID
+
+
+# ------------------------------------------------------------------
+# Export helpers
+# ------------------------------------------------------------------
+
+def _surf_to_array(surf):
+    """pygame Surface → (H, W, 3) uint8 numpy array."""
+    import pygame
+    import numpy as np
+    arr = pygame.surfarray.array3d(surf)        # (W, H, 3)
+    return np.ascontiguousarray(arr.transpose(1, 0, 2))  # (H, W, 3)
+
+
+def _save_gif(arrays: list, filename: str, fps: int) -> None:
+    try:
+        from PIL import Image
+    except ImportError:
+        raise ImportError(
+            'GIF export requires Pillow:\n'
+            '  pip install Pillow'
+        )
+    duration_ms = max(20, int(1000 / fps))
+    imgs = [Image.fromarray(arr) for arr in arrays]
+    imgs[0].save(
+        filename,
+        save_all=True,
+        append_images=imgs[1:],
+        loop=0,
+        duration=duration_ms,
+        optimize=False,
+    )
+
+
+def _save_video(arrays: list, filename: str, fps: int) -> None:
+    try:
+        import imageio
+    except ImportError:
+        raise ImportError(
+            'Video export requires imageio:\n'
+            '  pip install imageio[ffmpeg]'
+        )
+    with imageio.get_writer(filename, fps=fps) as writer:
+        for arr in arrays:
+            writer.append_data(arr)
+
+
+def _resolve_backend(filename: str, writer) -> str:
+    """Return 'gif' or 'video', inferring from writer name and file extension."""
+    ext = filename.rsplit('.', 1)[-1].lower() if '.' in filename else ''
+    if isinstance(writer, str):
+        if writer == 'pillow':
+            return 'gif'
+        if writer in ('ffmpeg', 'ffmpeg_file', 'imageio'):
+            return 'video'
+    if ext == 'gif':
+        return 'gif'
+    if ext in ('mp4', 'mov', 'avi', 'webm'):
+        return 'video'
+    raise ValueError(
+        f'Unknown format ".{ext}". '
+        'Supported extensions: .gif  .mp4  .mov  .avi'
+    )
+
+
+def _resolve_frames(frames, save_count) -> list:
+    """Convert matplotlib-style frames arg to a concrete list of frame values."""
+    if frames is None:
+        n = save_count if save_count is not None else 100
+        return list(range(n))
+    if isinstance(frames, int):
+        return list(range(frames))
+    if inspect.isgeneratorfunction(frames):
+        gen = frames()
+        if save_count is not None:
+            return [next(gen) for _ in range(save_count)]
+        return list(gen)
+    return list(frames)
+
+
+class FuncAnimation:
+    """
+    matplotlib-compatible animation driven by a user-supplied function.
+
+    Matches the ``matplotlib.animation.FuncAnimation`` interface exactly::
+
+        from plotlive.animation import FuncAnimation
+        anim = FuncAnimation(fig, update, frames=50, interval=200)
+        anim.save('output.gif')
+
+    Animations start **paused** — press Space to play, ←/→ to step.
+    """
+
+    def __init__(
+        self,
+        fig,
+        func: Callable,
+        frames=None,
+        init_func=None,
+        fargs=None,
+        save_count: int | None = None,
+        *,
+        cache_frame_data: bool = True,
+        interval: int = 200,
+        repeat_delay: int = 0,
+        repeat: bool = True,
+        blit: bool = False,
+    ):
+        self.figure = fig
+        self._func = func
+        self._fargs = tuple(fargs) if fargs is not None else ()
+        self.interval = interval
+        self.repeat = repeat
+
+        self._frame_seq = _resolve_frames(frames, save_count)
+        self.frames = len(self._frame_seq)
+
+        self._displayed_frame: int = -1
+        self._next_frame: int = 0
+        self._playing: bool = False
+        self._finished: bool = False
+
+    def _call_func(self, i: int) -> None:
+        frame = self._frame_seq[i]
+        if self._fargs:
+            self._func(frame, *self._fargs)
+        else:
+            self._func(frame)
+
+    @property
+    def event_id(self) -> int:
+        return _get_animation_event()
+
+    # ------------------------------------------------------------------
+    # Playback control
+    # ------------------------------------------------------------------
+
+    def start(self) -> None:
+        import pygame
+        self._playing = True
+        self._finished = False
+        pygame.time.set_timer(self.event_id, self.interval)
+
+    def pause(self) -> None:
+        import pygame
+        self._playing = False
+        pygame.time.set_timer(self.event_id, 0)
+
+    def resume(self) -> None:
+        if not self._finished:
+            self.start()
+
+    def toggle(self) -> None:
+        if self._playing:
+            self.pause()
+        else:
+            self.resume()
+
+    # ------------------------------------------------------------------
+    # Frame rendering
+    # ------------------------------------------------------------------
+
+    def _show_frame(self, n: int) -> bool:
+        n = max(0, min(n, self.frames - 1))
+        self._call_func(n)
+        self._displayed_frame = n
+        self._next_frame = n + 1
+        for ax in self.figure.axes:
+            ax._dirty = True
+        return True
+
+    def step_forward(self) -> bool:
+        next_n = self._displayed_frame + 1
+        if next_n >= self.frames:
+            return False
+        return self._show_frame(next_n)
+
+    def step_back(self) -> bool:
+        prev_n = max(0, self._displayed_frame - 1) if self._displayed_frame >= 0 else 0
+        return self._show_frame(prev_n)
+
+    # ------------------------------------------------------------------
+    # Export — matches matplotlib's Animation.save() signature
+    # ------------------------------------------------------------------
+
+    def save(
+        self,
+        filename: str,
+        writer=None,
+        fps: int | None = None,
+        dpi: float | None = None,
+        codec: str | None = None,
+        bitrate: int | None = None,
+        extra_args=None,
+        metadata: dict | None = None,
+        extra_anim=None,
+        savefig_kwargs: dict | None = None,
+        *,
+        progress_callback=None,
+    ) -> None:
+        """
+        Render all frames off-screen and write to *filename*.
+
+        Matches ``matplotlib.animation.Animation.save()``::
+
+            anim.save('output.gif', writer='pillow', fps=10)
+            anim.save('output.mp4', writer='ffmpeg', fps=30)
+            anim.save('output.gif')   # writer and fps inferred automatically
+
+        Requirements:
+          GIF   — pip install Pillow
+          Video — pip install imageio[ffmpeg]
+        """
+        import pygame
+        from .renderer import FigureRenderer
+
+        if not pygame.get_init():
+            pygame.init()
+
+        backend = _resolve_backend(filename, writer)  # raises early for bad ext
+        effective_fps = fps if fps is not None else max(1, 1000 // self.interval)
+
+        print(f'Exporting {self.frames} frames → {filename}')
+        arrays = []
+        _step = max(1, self.frames // 10)
+        for i in range(self.frames):
+            self._call_func(i)
+            surf = FigureRenderer(self.figure).render()
+            arrays.append(_surf_to_array(surf))
+            if progress_callback is not None:
+                progress_callback(i + 1, self.frames)
+            if (i + 1) % _step == 0 or i == self.frames - 1:
+                print(f'  {i + 1}/{self.frames}')
+
+        self._call_func(0)  # restore to frame 0 so show() opens at the start
+
+        if backend == 'gif':
+            _save_gif(arrays, filename, effective_fps)
+        else:
+            _save_video(arrays, filename, effective_fps)
+        print(f'Saved → {filename}')
+
+    # ------------------------------------------------------------------
+    # Timer callback (called by the pygame event loop in pyplot.show())
+    # ------------------------------------------------------------------
+
+    def _on_timer(self) -> bool:
+        if not self._playing or self._finished:
+            return False
+
+        self._show_frame(self._next_frame)
+
+        if self._next_frame >= self.frames:
+            if self.repeat:
+                self._next_frame = 0
+                self._displayed_frame = -1
+            else:
+                self._finished = True
+                self.pause()
+
+        return True
+
+
+# Alias so existing code using Animation still works
+Animation = FuncAnimation