Refactor imports and clean up code formatting across multiple test files

Author: rhoadesScholar

src/cellmap_data/__init__.py

@@ -22,8 +22,8 @@
 from .image import CellMapImage
 from .image_writer import ImageWriter
 from .multidataset import CellMapMultiDataset
-from .subdataset import CellMapSubset
 from .mutable_sampler import MutableSubsetRandomSampler
+from .subdataset import CellMapSubset
 
 __all__ = [
     "CellMapBaseDataset",

src/cellmap_data/image.py

@@ -9,10 +9,7 @@
 import xarray
 import xarray_tensorstore as xt
 import zarr
-from pydantic_ome_ngff.v04.multiscale import (
-    MultiscaleGroupAttrs,
-    MultiscaleMetadata,
-)
+from pydantic_ome_ngff.v04.multiscale import MultiscaleGroupAttrs, MultiscaleMetadata
 from pydantic_ome_ngff.v04.transform import Scale, Translation, VectorScale
 from scipy.spatial.transform import Rotation as rot
 from xarray_ome_ngff.v04.multiscale import coords_from_transforms

tests/test_base_classes.py

@@ -1,8 +1,9 @@
 """Tests for base abstract classes."""
 
+from abc import ABC
+
 import pytest
 import torch
-from abc import ABC
 
 from cellmap_data.base_dataset import CellMapBaseDataset
 from cellmap_data.base_image import CellMapImageBase
@@ -18,57 +19,57 @@ def test_cannot_instantiate_abstract_class(self):
 
     def test_incomplete_implementation_raises_error(self):
         """Test that incomplete implementations cannot be instantiated."""
-        
+
         # Missing all abstract methods
         class IncompleteDataset(CellMapBaseDataset):
             pass
-        
+
         with pytest.raises(TypeError, match="Can't instantiate abstract class"):
             IncompleteDataset()
-        
+
         # Missing some abstract methods
         class PartialDataset(CellMapBaseDataset):
             @property
             def class_counts(self):
                 return {}
-            
+
             @property
             def class_weights(self):
                 return {}
-        
+
         with pytest.raises(TypeError, match="Can't instantiate abstract class"):
             PartialDataset()
 
     def test_complete_implementation_can_be_instantiated(self):
         """Test that complete implementations can be instantiated."""
-        
+
         class CompleteDataset(CellMapBaseDataset):
             def __init__(self):
                 self.classes = ["class1", "class2"]
                 self.input_arrays = {"raw": {}}
                 self.target_arrays = {"labels": {}}
-            
+
             @property
             def class_counts(self):
                 return {"class1": 100.0, "class2": 200.0}
-            
+
             @property
             def class_weights(self):
                 return {"class1": 0.67, "class2": 0.33}
-            
+
             @property
             def validation_indices(self):
                 return [0, 1, 2]
-            
+
             def to(self, device, non_blocking=True):
                 return self
-            
+
             def set_raw_value_transforms(self, transforms):
                 pass
-            
+
             def set_target_value_transforms(self, transforms):
                 pass
-        
+
         # Should not raise
         dataset = CompleteDataset()
         assert isinstance(dataset, CellMapBaseDataset)
@@ -83,11 +84,11 @@ def set_target_value_transforms(self, transforms):
     def test_attributes_are_defined(self):
         """Test that expected attributes are defined in the base class."""
         # Check type annotations exist
-        assert hasattr(CellMapBaseDataset, '__annotations__')
+        assert hasattr(CellMapBaseDataset, "__annotations__")
         annotations = CellMapBaseDataset.__annotations__
-        assert 'classes' in annotations
-        assert 'input_arrays' in annotations
-        assert 'target_arrays' in annotations
+        assert "classes" in annotations
+        assert "input_arrays" in annotations
+        assert "target_arrays" in annotations
 
 
 class TestCellMapImageBase:
@@ -100,52 +101,52 @@ def test_cannot_instantiate_abstract_class(self):
 
     def test_incomplete_implementation_raises_error(self):
         """Test that incomplete implementations cannot be instantiated."""
-        
+
         # Missing all abstract methods
         class IncompleteImage(CellMapImageBase):
             pass
-        
+
         with pytest.raises(TypeError, match="Can't instantiate abstract class"):
             IncompleteImage()
-        
+
         # Missing some abstract methods
         class PartialImage(CellMapImageBase):
             @property
             def bounding_box(self):
                 return {"x": (0, 100), "y": (0, 100)}
-            
+
             @property
             def sampling_box(self):
                 return {"x": (10, 90), "y": (10, 90)}
-        
+
         with pytest.raises(TypeError, match="Can't instantiate abstract class"):
             PartialImage()
 
     def test_complete_implementation_can_be_instantiated(self):
         """Test that complete implementations can be instantiated."""
-        
+
         class CompleteImage(CellMapImageBase):
             def __getitem__(self, center):
                 return torch.zeros((1, 64, 64))
-            
+
             @property
             def bounding_box(self):
                 return {"x": (0.0, 100.0), "y": (0.0, 100.0)}
-            
+
             @property
             def sampling_box(self):
                 return {"x": (10.0, 90.0), "y": (10.0, 90.0)}
-            
+
             @property
             def class_counts(self):
                 return 1000.0
-            
+
             def to(self, device, non_blocking=True):
                 pass
-            
+
             def set_spatial_transforms(self, transforms):
                 pass
-        
+
         # Should not raise
         image = CompleteImage()
         assert isinstance(image, CellMapImageBase)
@@ -161,59 +162,59 @@ def set_spatial_transforms(self, transforms):
 
     def test_class_counts_supports_dict_return_type(self):
         """Test that class_counts can return a dictionary."""
-        
+
         class MultiClassImage(CellMapImageBase):
             def __getitem__(self, center):
                 return torch.zeros((1, 64, 64))
-            
+
             @property
             def bounding_box(self):
                 return {"x": (0.0, 100.0)}
-            
+
             @property
             def sampling_box(self):
                 return {"x": (10.0, 90.0)}
-            
+
             @property
             def class_counts(self):
                 return {"class1": 500.0, "class2": 300.0, "class3": 200.0}
-            
+
             def to(self, device, non_blocking=True):
                 pass
-            
+
             def set_spatial_transforms(self, transforms):
                 pass
-        
+
         image = MultiClassImage()
         counts = image.class_counts
         assert isinstance(counts, dict)
         assert counts == {"class1": 500.0, "class2": 300.0, "class3": 200.0}
 
     def test_bounding_box_can_be_none(self):
         """Test that bounding_box property can return None."""
-        
+
         class UnboundedImage(CellMapImageBase):
             def __getitem__(self, center):
                 return torch.zeros((1, 64, 64))
-            
+
             @property
             def bounding_box(self):
                 return None
-            
+
             @property
             def sampling_box(self):
                 return None
-            
+
             @property
             def class_counts(self):
                 return 1000.0
-            
+
             def to(self, device, non_blocking=True):
                 pass
-            
+
             def set_spatial_transforms(self, transforms):
                 pass
-        
+
         image = UnboundedImage()
         assert image.bounding_box is None
         assert image.sampling_box is None

tests/test_dataset_edge_cases.py

@@ -1,9 +1,10 @@
 """Tests for CellMapDataset edge cases and special methods."""
 
 import pickle
+
+import numpy as np
 import pytest
 import torch
-import numpy as np
 
 from cellmap_data import CellMapDataset, CellMapMultiDataset
 
@@ -104,6 +105,7 @@ def test_executor_handles_fork(self, minimal_dataset):
 
         # Simulate a fork by changing the PID tracking
         import os
+
         dataset._executor_pid = os.getpid() + 1
 
         # Access executor again - should create new one
@@ -219,10 +221,10 @@ def test_class_counts_property(self, minimal_dataset):
         assert isinstance(counts, dict)
         # class_counts structure has changed - it's now nested with 'totals'
         # Check that the totals key exists and has class entries
-        if 'totals' in counts:
+        if "totals" in counts:
             for cls in dataset.classes:
                 # Class names might have _bg suffix
-                assert any(cls in key for key in counts['totals'].keys())
+                assert any(cls in key for key in counts["totals"].keys())
         else:
             # Old structure - direct class keys
             for cls in dataset.classes:
@@ -249,7 +251,7 @@ def test_validation_indices_property(self, minimal_dataset):
         indices = dataset.validation_indices
 
         # Should be a sequence
-        assert hasattr(indices, '__iter__')
+        assert hasattr(indices, "__iter__")
 
     def test_2d_array_creates_multidataset(self, tmp_path):
         """Test that 2D array without slicing axis triggers special handling."""
@@ -285,7 +287,7 @@ def test_2d_array_creates_multidataset(self, tmp_path):
         # Should create some kind of dataset (either regular or multi)
         # The key is that it doesn't raise an error
         assert dataset is not None
-        assert hasattr(dataset, '__getitem__')
+        assert hasattr(dataset, "__getitem__")
 
     def test_set_raw_value_transforms(self, minimal_dataset):
         """Test setting raw value transforms."""

tests/test_empty_image_writer.py

@@ -4,10 +4,11 @@
 Tests empty image handling and image writing functionality.
 """
 
+import os
+from pathlib import Path
+
 import pytest
 from upath import UPath
-from pathlib import Path
-import os
 
 from cellmap_data import EmptyImage, ImageWriter
 

tests/test_helpers.py

@@ -11,12 +11,8 @@
 import numpy as np
 import zarr
 from pydantic_ome_ngff.v04.axis import Axis
-from pydantic_ome_ngff.v04.multiscale import (
-    MultiscaleMetadata,
-)
-from pydantic_ome_ngff.v04.multiscale import (
-    Dataset as MultiscaleDataset,
-)
+from pydantic_ome_ngff.v04.multiscale import Dataset as MultiscaleDataset
+from pydantic_ome_ngff.v04.multiscale import MultiscaleMetadata
 from pydantic_ome_ngff.v04.transform import VectorScale
 
 

tests/test_image_edge_cases.py

@@ -1,8 +1,8 @@
 """Tests for CellMapImage edge cases and special methods."""
 
+import numpy as np
 import pytest
 import torch
-import numpy as np
 
 from cellmap_data import CellMapImage
 

tests/test_multidataset_datasplit.py

@@ -6,11 +6,7 @@
 
 import pytest
 
-from cellmap_data import (
-    CellMapDataset,
-    CellMapDataSplit,
-    CellMapMultiDataset,
-)
+from cellmap_data import CellMapDataset, CellMapDataSplit, CellMapMultiDataset
 
 from .test_helpers import create_test_dataset
 

tests/test_utils.py

@@ -7,10 +7,7 @@
 import numpy as np
 import torch
 
-from cellmap_data.utils.misc import (
-    get_sliced_shape,
-    torch_max_value,
-)
+from cellmap_data.utils.misc import get_sliced_shape, torch_max_value
 
 
 class TestUtilsMisc: