From be4bbc8d9e293f549ff630ff27005ac5f3067762 Mon Sep 17 00:00:00 2001
From: Almar Klein <almar@almarklein.org>
Date: Wed, 13 Mar 2024 15:21:45 +0100
Subject: [PATCH] Lazy import av

---
 fastplotlib/layouts/_frame/_toolbar.py |  2 +-
 fastplotlib/layouts/_record_mixin.py   | 25 +++++++++++++------------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/fastplotlib/layouts/_frame/_toolbar.py b/fastplotlib/layouts/_frame/_toolbar.py
index 94410b8ea..6a0485655 100644
--- a/fastplotlib/layouts/_frame/_toolbar.py
+++ b/fastplotlib/layouts/_frame/_toolbar.py
@@ -1,4 +1,4 @@
-from fastplotlib.layouts._subplot import Subplot
+from .._subplot import Subplot
 
 
 class ToolBar:
diff --git a/fastplotlib/layouts/_record_mixin.py b/fastplotlib/layouts/_record_mixin.py
index e3a491915..e3bfdeba5 100644
--- a/fastplotlib/layouts/_record_mixin.py
+++ b/fastplotlib/layouts/_record_mixin.py
@@ -3,12 +3,17 @@
 from multiprocessing import Queue, Process
 from time import time
 
-try:
-    import av
-except ImportError:
-    HAS_AV = False
-else:
-    HAS_AV = True
+
+def _get_av():
+    try:
+        import av
+    except ImportError:
+        raise ModuleNotFoundError(
+            "Recording to video file requires `av`:\n"
+            "https://github.com/PyAV-Org/PyAV"
+        ) from None
+    else:
+        return av
 
 
 class VideoWriterAV(Process):
@@ -28,6 +33,7 @@ def __init__(
         super().__init__()
         self.queue = queue
 
+        av = _get_av()
         self.container = av.open(path, mode="w")
 
         self.stream = self.container.add_stream(codec, rate=fps, options=options)
@@ -45,6 +51,7 @@ def __init__(
         self.stream.pix_fmt = pixel_format
 
     def run(self):
+        av = _get_av()
         while True:
             if self.queue.empty():  # no frame to write
                 continue
@@ -177,12 +184,6 @@ def record_start(
 
         """
 
-        if not HAS_AV:
-            raise ModuleNotFoundError(
-                "Recording to video file requires `av`:\n"
-                "https://github.com/PyAV-Org/PyAV"
-            )
-
         if Path(path).exists():
             raise FileExistsError(f"File already exists at given path: {path}")