lint and format volara and tests

Author: will

tests/test_dataset.py

@@ -1,13 +1,14 @@
+from pathlib import Path
+
+import numpy as np
 import pytest
 import yaml
-import numpy as np
 import zarr
-from pathlib import Path
 from funlib.geometry import Coordinate
-
-from volara.datasets import Raw, Affs, LSD, Labels
 from pydantic import ValidationError
 
+from volara.datasets import LSD, Affs, Labels, Raw
+
 
 @pytest.fixture
 def zarr_store(tmp_path):
@@ -77,7 +78,7 @@ def test_affs_serialization(tmp_path):
     yaml_data = f"""
     dataset_type: affs
     store: {tmp_path / "affs.zarr"}
-    neighborhood: 
+    neighborhood:
       - [0, 1, 0]
       - [0, 0, 1]
     """
@@ -109,15 +110,15 @@ def test_lazy_channel_slicing(zarr_store):
     assert data.shape == (10, 10)
     assert np.all(data == 1.0)  # Original data was 1s
 
-    ds = Raw(store=zarr_store, channels=[0, [0,2,4,6,8]])
+    ds = Raw(store=zarr_store, channels=[0, [0, 2, 4, 6, 8]])
 
     arr = ds.array()
     data = arr[:]
 
     assert data.shape == (5, 10)
     assert np.all(data == 1.0)  # Original data was 1s
 
-    ds = Raw(store=zarr_store, channels=[0, 0, [0,2,4,6,8]])
+    ds = Raw(store=zarr_store, channels=[0, 0, [0, 2, 4, 6, 8]])
 
     arr = ds.array()
     data = arr[:]

volara/blockwise/blockwise.py

@@ -399,7 +399,7 @@ def process_func(block):
             yield task
 
     def get_benchmark_logger(self) -> BenchmarkLogger:
-        benchmark_db_path = Path("volara_benchmark_logs/benchmark.db")
+        _benchmark_db_path = Path("volara_benchmark_logs/benchmark.db")
         return BenchmarkLogger(
             None,
             task=self.task_name,

volara/blockwise/graph_mws.py

@@ -1,16 +1,14 @@
-from contextlib import contextmanager
-from typing import Annotated, Literal
-import tempfile
-import itertools
 import functools
+import itertools
+import tempfile
+from contextlib import contextmanager
 from pathlib import Path
-import time
-from pprint import pprint
+from typing import Annotated, Literal
 
 import daisy
 import mwatershed as mws
-import numpy as np
 import networkx as nx
+import numpy as np
 from funlib.geometry import Coordinate, Roi
 from pydantic import Field
 
@@ -20,7 +18,6 @@
 from ..dbs import PostgreSQL, SQLite
 from ..utils import PydanticCoordinate
 from .blockwise import BlockwiseTask
-from ..datasets import Labels
 
 DB = Annotated[
     PostgreSQL | SQLite,
@@ -313,7 +310,6 @@ def process_block(block: daisy.Block):
                     both_sides=True,
                 )
 
-                t1 = time.time()
                 edges = []
                 inputs = set(
                     node
@@ -391,15 +387,13 @@ def process_block(block: daisy.Block):
                     in_group = new_seg_to_frag_mapping.setdefault(int(out_frag), set())
                     in_group.add(int(in_frag))
 
-                t1 = time.time()
                 # save the lut to a temporary file for this block
                 block_lut = LUT(
                     path=f"{tmp_path}/{'-'.join([str(o) for o in block.write_roi.offset])}-lut"
                 )
                 lut = np.array([inputs, outputs])
                 block_lut.save(lut, edges=edges)
 
-                t1 = time.time()
                 # read luts and existing super fragments in neighboring blocks
                 existing_luts = [
                     LUT(
@@ -420,7 +414,6 @@ def process_block(block: daisy.Block):
                     + [self.lut],
                 ).load()
 
-                t1 = time.time()
                 frag_seg_mapping: dict[int, int] = {
                     int(k): int(v) for k, v in total_lut.T
                 }
@@ -429,10 +422,8 @@ def process_block(block: daisy.Block):
                     in_group = seg_frag_mapping.setdefault(out_frag, set())
                     in_group.add(in_frag)
 
-                t1 = time.time()
                 out_graph = out_rag_provider.read_graph(block.read_roi)
 
-                t1 = time.time()
                 for out_seg, in_frags in new_seg_to_frag_mapping.items():
                     agglomerated_attrs = {
                         "size": sum(
@@ -454,7 +445,6 @@ def process_block(block: daisy.Block):
 
                     out_graph.add_node(int(out_seg), **agglomerated_attrs)
 
-                t1 = time.time()
                 edges_to_agglomerate = {}
                 for u, v in graph.edges():
                     if (
@@ -601,7 +591,6 @@ def process_block(block: daisy.Block):
                     edge_attrs=list(self.weights.keys()),
                 )
 
-                t1 = time.time()
                 edges = []
                 for u, v, edge_attrs in graph.edges(data=True):
                     if (
@@ -630,15 +619,13 @@ def process_block(block: daisy.Block):
                 else:
                     inputs, outputs = [], []
 
-                t1 = time.time()
                 # save the lut to a temporary file for this block
                 block_lut = LUT(
                     path=f"{tmp_path}/{'-'.join([str(o) for o in block.write_roi.offset])}-lut"
                 )
                 lut = np.array([inputs, outputs])
                 block_lut.save(lut, edges=edges)
 
-                t1 = time.time()
                 # read luts and existing super fragments in neighboring blocks
                 existing_luts = [
                     LUT(
@@ -659,7 +646,6 @@ def process_block(block: daisy.Block):
                     + [self.lut],
                 ).load()
 
-                t1 = time.time()
                 frag_seg_mapping: dict[int, int] = {
                     int(k): int(v) for k, v in total_lut.T
                 }
@@ -668,11 +654,9 @@ def process_block(block: daisy.Block):
                     in_group = seg_frag_mapping.setdefault(out_frag, set())
                     in_group.add(in_frag)
 
-                t1 = time.time()
                 out_graph = out_rag_provider.read_graph(block.read_roi)
                 assert out_graph.number_of_nodes() == 0, out_graph.number_of_nodes
 
-                t1 = time.time()
                 for out_frag in np.unique(outputs):
                     if out_frag is not None and out_frag not in out_graph.nodes:
                         in_group = seg_frag_mapping[out_frag]
@@ -699,7 +683,6 @@ def process_block(block: daisy.Block):
 
                         out_graph.add_node(int(out_frag), **agglomerated_attrs)
 
-                t1 = time.time()
                 edges_to_agglomerate = {}
                 for u, v in graph.edges():
                     if u in frag_seg_mapping and v in frag_seg_mapping:

volara/datasets.py

@@ -242,7 +242,6 @@ def stack(data, other_data):
             arr.lazy_op(lambda data: stack(data, self.stack.array("r").data))  # type: ignore[possibly-missing-attribute]
 
 
-
 class Affs(Dataset):
     """
     Represents a dataset containing affinities.

volara/lut.py

@@ -1,10 +1,10 @@
-from volara.tmp import replace_values
-
-from pathlib import Path
 from collections.abc import Sequence
+from pathlib import Path
 
 import numpy as np
 
+from volara.tmp import replace_values
+
 from .utils import StrictBaseModel
 
 
@@ -38,7 +38,9 @@ def drop(self):
             self.file.unlink()
 
     def save(self, lut: np.ndarray, edges=None):
-        np.savez_compressed(self.file, fragment_segment_lut=lut.astype(int), edges=edges)
+        np.savez_compressed(
+            self.file, fragment_segment_lut=lut.astype(int), edges=edges
+        )
 
     def load(self) -> np.ndarray | None:
         if not self.file.exists():
@@ -76,5 +78,7 @@ def load_iterated(self):
         for lut in self.luts[1:]:
             next_map = lut.load()
             if next_map is not None:
-                starting_map[1] = replace_values(starting_map[1], next_map[0], next_map[1])
+                starting_map[1] = replace_values(
+                    starting_map[1], next_map[0], next_map[1]
+                )
         return starting_map