[CheckpointWrapper] Use correct extension during atomic writing

bob/pipelines/wrappers.py

@@ -4,6 +4,7 @@ import os
 import tempfile
 
 from functools import partial
+from pathlib import Path
 
 import cloudpickle
 import dask.bag
@@ -394,11 +395,15 @@ class CheckpointWrapper(BaseWrapper, TransformerMixin):
         to_save = getattr(sample, self.sample_attribute)
         for _ in range(self.attempts):
             try:
+
                 dirname = os.path.dirname(path)
                 os.makedirs(dirname, exist_ok=True)
 
                 # Atomic writing
-                with tempfile.NamedTemporaryFile(dir=dirname, delete=False) as f:
+                extension = "".join(Path(path).suffixes)
+                with tempfile.NamedTemporaryFile(
+                    dir=dirname, delete=False, suffix=extension
+                ) as f:
                     self.save_func(to_save, f.name)
                     os.replace(f.name, path)