Repository: Project-MONAI/MONAI
File: tests/transforms/croppad/test_pad_nd_dtypes.py (Python)

Code before:
            patch.object(F, "_pt_pad", wraps=F._pt_pad) as mock_pt,
            patch.object(F, "_np_pad", wraps=F._np_pad) as mock_np,
        ):
            out = pad_nd(img, to_pad, mode="constant", value=0)

        self.assertTrue(mock_pt.called)
        self.assertFalse(mock_np.called)
        self.assertEqual(out.dtype, img.dtype)
        self.assertEqual(out.shape, (1, 6, 8))

    def test_pad_falls_back_to_np_if_pt_raises(self):
        """Test that pad_nd falls back to NumPy when PyTorch raises NotImplementedError."""
        img = torch.ones((1, 4, 4), dtype=torch.bool)
        to_pad = [(0, 0), (1, 1), (2, 2)]
        with (
            patch.object(F, "_pt_pad", new=Mock(side_effect=NotImplementedError("no"))) as mock_pt,
            patch.object(F, "_np_pad", wraps=F._np_pad) as mock_np,
        ):
            out = pad_nd(img, to_pad, mode="constant", value=0)

        self.assertTrue(mock_pt.called)
        self.assertTrue(mock_np.called)
        self.assertEqual(out.dtype, img.dtype)
        self.assertEqual(out.shape, (1, 6, 8))

    @parameterized.expand([torch.bool, torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8, torch.float32])
    def test_pad_dtype_no_error_and_dtype_preserved(self, dtype):
        """Test that pad_nd handles various dtypes without error and preserves dtype."""
        img = torch.ones((1, 4, 4), dtype=dtype)
        to_pad = [(0, 0), (1, 1), (2, 2)]
        out = pad_nd(img, to_pad, mode="constant", value=0)

        self.assertEqual(out.shape, (1, 6, 8))
        self.assertEqual(out.dtype, img.dtype)

    @parameterized.expand(
        [
            ("constant", torch.bool),
            ("constant", torch.int8),
            ("constant", torch.float32),
            ("reflect", torch.bool),
            ("reflect", torch.int8),
            ("reflect", torch.float32),
            ("replicate", torch.bool),
            ("replicate", torch.int8),
            ("replicate", torch.float32),
        ]
    )
    def test_pad_multiple_modes_dtype_preserved(self, mode, dtype):
        """Test that pad_nd preserves dtype across multiple padding modes."""
        img = torch.ones((1, 4, 4), dtype=dtype)

Review comment (suggestion):
Often it's cleaner to put arguments to `expand` into a variable at the top of the file, here and in the next usage it just is cleaner to read the tests.

Code after:
    ("reflect", torch.float32),
    ("replicate", torch.bool),
    ("replicate", torch.int8),
    ("replicate", torch.float32),
]


class TestPadNdDtypes(unittest.TestCase):
    def test_pad_uses_pt_for_bool(self):
        """Test that pad_nd uses PyTorch backend for bool dtype in constant mode."""
        img = torch.ones((1, 4, 4), dtype=torch.bool)
        to_pad = [(0, 0), (1, 1), (2, 2)]
        with (
            patch.object(F, "_pt_pad", wraps=F._pt_pad) as mock_pt,
            patch.object(F, "_np_pad", wraps=F._np_pad) as mock_np,
        ):
            out = pad_nd(img, to_pad, mode="constant", value=0)

        self.assertTrue(mock_pt.called)
        self.assertFalse(mock_np.called)
        self.assertEqual(out.dtype, img.dtype)
        self.assertEqual(out.shape, (1, 6, 8))

    def test_pad_falls_back_to_np_if_pt_raises(self):
        """Test that pad_nd falls back to NumPy when PyTorch raises NotImplementedError."""
        img = torch.ones((1, 4, 4), dtype=torch.bool)
        to_pad = [(0, 0), (1, 1), (2, 2)]
        with (
            patch.object(F, "_pt_pad", new=Mock(side_effect=NotImplementedError("no"))) as mock_pt,
            patch.object(F, "_np_pad", wraps=F._np_pad) as mock_np,
        ):
            out = pad_nd(img, to_pad, mode="constant", value=0)

        self.assertTrue(mock_pt.called)
        self.assertTrue(mock_np.called)
        self.assertEqual(out.dtype, img.dtype)
        self.assertEqual(out.shape, (1, 6, 8))

    @parameterized.expand(DTYPES)
    def test_pad_dtype_no_error_and_dtype_preserved(self, dtype):
        """Test that pad_nd handles various dtypes without error and preserves dtype.
        Args:
            dtype: Input dtype under test.
        """
        img = torch.ones((1, 4, 4), dtype=dtype)
        to_pad = [(0, 0), (1, 1), (2, 2)]
        out = pad_nd(img, to_pad, mode="constant", value=0)

        self.assertEqual(out.shape, (1, 6, 8))
        self.assertEqual(out.dtype, img.dtype)