diff --git a/hls4ml/writer/catapult_writer.py b/hls4ml/writer/catapult_writer.py index ba7e511995..86181bada0 100755 --- a/hls4ml/writer/catapult_writer.py +++ b/hls4ml/writer/catapult_writer.py @@ -893,11 +893,12 @@ def write_tar(self, model): model (ModelGraph): the hls4ml model. """ - if not os.path.exists(model.config.get_output_dir() + '.tar.gz'): - with tarfile.open(model.config.get_output_dir() + '.tar.gz', mode='w:gz') as archive: + if model.config.get_writer_config().get('WriteTar', False): + tar_path = model.config.get_output_dir() + '.tar.gz' + if os.path.exists(tar_path): + os.remove(tar_path) + with tarfile.open(tar_path, mode='w:gz') as archive: archive.add(model.config.get_output_dir(), recursive=True) - else: - print('Project .tar.gz archive already exists') def write_hls(self, model): self.write_output_dir(model) diff --git a/test/pytest/test_catapult_writer.py b/test/pytest/test_catapult_writer.py new file mode 100644 index 0000000000..cfd6e84b93 --- /dev/null +++ b/test/pytest/test_catapult_writer.py @@ -0,0 +1,33 @@ +from pathlib import Path + +import pytest + +from hls4ml.writer.catapult_writer import CatapultWriter + + +class DummyConfig: + def __init__(self, output_dir, write_tar): + self.output_dir = str(output_dir) + self.writer_config = {'WriteTar': write_tar} + + def get_output_dir(self): + return self.output_dir + + def get_writer_config(self): + return self.writer_config + + +class DummyModel: + def __init__(self, output_dir, write_tar): + self.config = DummyConfig(output_dir, write_tar) + + +@pytest.mark.parametrize('write_tar', [True, False]) +def test_catapult_write_tar_respects_writer_config(tmp_path, write_tar): + output_dir = tmp_path / 'myproject' + output_dir.mkdir() + (output_dir / 'marker.txt').write_text('marker') + + CatapultWriter().write_tar(DummyModel(output_dir, write_tar)) + + assert Path(f'{output_dir}.tar.gz').exists() == write_tar