diff --git a/gridfinder/gridfinder.py b/gridfinder/gridfinder.py
index 6968e48..b67ce3f 100644
--- a/gridfinder/gridfinder.py
+++ b/gridfinder/gridfinder.py
@@ -14,7 +14,7 @@
 from affine import Affine
 from IPython.display import Markdown, display
 
-from gridfinder._util import save_raster
+from gridfinder.util import save_raster
 
 sys.setrecursionlimit(100000)
 
@@ -84,7 +84,7 @@ def optimise(
     jupyter: bool = False,
     animate: bool = False,
     affine: Optional[Affine] = None,
-    animate_path: Optional[str] = None,
+    animate_path: str = "",
     silent: bool = False,
 ) -> np.ndarray:
     """Run the Dijkstra algorithm for the supplied arrays.
@@ -103,6 +103,7 @@ def optimise(
         on-grid point. Values of 0 imply that cell is part of an MV grid line.
     """
 
+
     max_i = costs.shape[0]
     max_j = costs.shape[1]
 
@@ -141,6 +142,7 @@ def zero_and_heap_path(loc: Tuple[int, int]) -> None:
     counter = 0
     progress = 0
     max_cells = targets.shape[0] * targets.shape[1]
+    handle = None
     if jupyter:
         handle = display(Markdown(""), display_id=True)
 
@@ -201,13 +203,15 @@ def zero_and_heap_path(loc: Tuple[int, int]) -> None:
                         if int(progress_new) > int(progress):
                             progress = progress_new
                             message = f"{progress:.2f} %"
-                            if jupyter:
+                            if jupyter and handle:
                                 handle.update(message)
                             elif not silent:
                                 print(message)
                             if animate:
                                 i = int(progress)
                                 path = os.path.join(animate_path, f"arr{i:03d}.tif")
+                                if not affine:
+                                    raise ValueError("Must provide an affine when animate=True")  # NoQA
                                 save_raster(path, dist, affine)
 
     return dist
diff --git a/gridfinder/prepare.py b/gridfinder/prepare.py
index 120a770..9e14cfb 100644
--- a/gridfinder/prepare.py
+++ b/gridfinder/prepare.py
@@ -19,7 +19,7 @@
 from rasterio.warp import Resampling, reproject
 from scipy import signal
 
-from gridfinder._util import clip_raster, save_raster
+from gridfinder.util import clip_raster, save_raster
 
 
 def clip_rasters(
diff --git a/pyproject.toml b/pyproject.toml
index 46b1826..4717a6f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -87,14 +87,11 @@ select = [
 [tool.ruff.lint.isort]
 known-first-party = ["gridfinder"]
 
-[tool.mypy]
-show_error_codes = true
-show_error_context = true
-show_traceback = true
-disallow_untyped_defs = true
-disallow_any_unimported = true
-no_implicit_optional = true
-check_untyped_defs = true
-warn_return_any = true
-warn_unused_ignores = true
-warn_unreachable = true
+[tool.pyright]
+include = ["gridfinder"]
+reportMissingImports = true
+reportMissingParameterType = true
+reportUnnecessaryTypeIgnoreComment = true
+reportDeprecated = true
+pythonVersion = "3.12"
+pythonPlatform = "Linux"