Lazy import av (#432)

almarklein · web-flow · commit fe7417980a83 · 2024-03-14T04:36:22.000-04:00
diff --git 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
@@ -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}")
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-from fastplotlib.layouts._subplot import Subplot`
	`1`	`+from .._subplot import Subplot`
`2`	`2`
`3`	`3`
`4`	`4`	`class ToolBar:`