Repository: Project-MONAI/MONAI
File: tests/test_zarr_avg_merger.py (Python)

Code before:
            TEST_CASE_1_DEFAULT_DTYPE,
            TEST_CASE_2_DEFAULT_DTYPE,
            TEST_CASE_3_DEFAULT_DTYPE,
            TEST_CASE_4_DEFAULT_DTYPE,
            TEST_CASE_5_VALUE_DTYPE,
            TEST_CASE_6_COUNT_DTYPE,
            TEST_CASE_7_COUNT_VALUE_DTYPE,
            TEST_CASE_8_DTYPE,
            TEST_CASE_9_LARGER_SHAPE,
            TEST_CASE_10_DIRECTORY_STORE,
            TEST_CASE_11_MEMORY_STORE,
            TEST_CASE_12_CHUNKS,
            TEST_CASE_13_COMPRESSOR_LZ4,
            TEST_CASE_14_COMPRESSOR_PICKLE,
            TEST_CASE_15_COMPRESSOR_LZMA,
            TEST_CASE_16_WITH_LOCK,
            TEST_CASE_17_WITHOUT_LOCK,
        ]
    )
    def test_zarr_avg_merger_patches(self, arguments, patch_locations, expected):
        if "compressor" in arguments:
            if arguments["compressor"] != "default":
                arguments["compressor"] = numcodecs.registry.codec_registry[arguments["compressor"].lower()]()
        if "value_compressor" in arguments:
            if arguments["value_compressor"] != "default":
                arguments["value_compressor"] = numcodecs.registry.codec_registry[
                    arguments["value_compressor"].lower()
                ]()
        if "count_compressor" in arguments:
            if arguments["count_compressor"] != "default":
                arguments["count_compressor"] = numcodecs.registry.codec_registry[
                    arguments["count_compressor"].lower()
                ]()
        merger = ZarrAvgMerger(**arguments)
        for pl in patch_locations:
            merger.aggregate(pl[0], pl[1])
        output = merger.finalize()
        if "value_dtype" in arguments:
            self.assertTrue(merger.get_values().dtype, arguments["value_dtype"])
        if "count_dtype" in arguments:
            self.assertTrue(merger.get_counts().dtype, arguments["count_dtype"])
        # check for multiple call of finalize
        self.assertIs(output, merger.finalize())
        # check if the result is matching the expectation
        assert_allclose(output[:], expected.numpy())

    def test_zarr_avg_merger_finalized_error(self):
        with self.assertRaises(ValueError):
            merger = ZarrAvgMerger(merged_shape=(1, 3, 2, 3))
            merger.finalize()
            merger.aggregate(torch.zeros(1, 3, 2, 2), (3, 3))

Review comment (suggestion):
I'd suggest putting `codec_reg = numcodecs.registry.codec_registry` at the top of this method then using that variable in the three places the registry is needed just to shorten the lines a bit.

Code after:
            TEST_CASE_1_DEFAULT_DTYPE,
            TEST_CASE_2_DEFAULT_DTYPE,
            TEST_CASE_3_DEFAULT_DTYPE,
            TEST_CASE_4_DEFAULT_DTYPE,
            TEST_CASE_5_VALUE_DTYPE,
            TEST_CASE_6_COUNT_DTYPE,
            TEST_CASE_7_COUNT_VALUE_DTYPE,
            TEST_CASE_8_DTYPE,
            TEST_CASE_9_LARGER_SHAPE,
            TEST_CASE_10_DIRECTORY_STORE,
            TEST_CASE_11_MEMORY_STORE,
            TEST_CASE_12_CHUNKS,
            TEST_CASE_13_COMPRESSOR_LZ4,
            TEST_CASE_14_COMPRESSOR_PICKLE,
            TEST_CASE_15_COMPRESSOR_LZMA,
            TEST_CASE_16_WITH_LOCK,
            TEST_CASE_17_WITHOUT_LOCK,
        ]
    )
    def test_zarr_avg_merger_patches(self, arguments, patch_locations, expected):
        codec_reg = numcodecs.registry.codec_registry
        if "compressor" in arguments:
            if arguments["compressor"] != "default":
                arguments["compressor"] = codec_reg[arguments["compressor"].lower()]()
        if "value_compressor" in arguments:
            if arguments["value_compressor"] != "default":
                arguments["value_compressor"] = codec_reg[arguments["value_compressor"].lower()]()
        if "count_compressor" in arguments:
            if arguments["count_compressor"] != "default":
                arguments["count_compressor"] = codec_reg[arguments["count_compressor"].lower()]()
        merger = ZarrAvgMerger(**arguments)
        for pl in patch_locations:
            merger.aggregate(pl[0], pl[1])
        output = merger.finalize()
        if "value_dtype" in arguments:
            self.assertTrue(merger.get_values().dtype, arguments["value_dtype"])
        if "count_dtype" in arguments:
            self.assertTrue(merger.get_counts().dtype, arguments["count_dtype"])
        # check for multiple call of finalize
        self.assertIs(output, merger.finalize())
        # check if the result is matching the expectation
        assert_allclose(output[:], expected.numpy())

    def test_zarr_avg_merger_finalized_error(self):
        with self.assertRaises(ValueError):
            merger = ZarrAvgMerger(merged_shape=(1, 3, 2, 3))
            merger.finalize()
            merger.aggregate(torch.zeros(1, 3, 2, 2), (3, 3))

    def test_zarr_avg_merge_none_merged_shape_error(self):
        with self.assertRaises(ValueError):