diff --git a/mesures/f5d.py b/mesures/f5d.py index d432839..4c8f56e 100644 --- a/mesures/f5d.py +++ b/mesures/f5d.py @@ -3,6 +3,7 @@ from mesures.headers import F5D_HEADER as COLUMNS from mesures.f5 import F5, DTYPES from mesures.utils import check_line_terminator_param +from zipfile import ZipFile import os import pandas as pd @@ -80,10 +81,11 @@ def writer(self): :return: file path of generated F5D File """ existing_files = os.listdir('/tmp') - if existing_files: - versions = [int(f.split('.')[1]) for f in existing_files if self.filename.split('.')[0] in f] - if versions: - self.version = max(versions) + 1 + if self.default_compression != 'zip': + if existing_files: + versions = [int(f.split('.')[1]) for f in existing_files if self.filename.split('.')[0] in f and '.zip' not in f] + if versions: + self.version = max(versions) + 1 file_path = os.path.join('/tmp', self.filename) kwargs = {'sep': ';', @@ -95,5 +97,14 @@ def writer(self): if self.default_compression: kwargs.update({'compression': self.default_compression}) - self.file.to_csv(file_path, **kwargs) + if kwargs.get('compression', False) == 'zip': + self.default_compression = False + zipped_file = ZipFile(os.path.join('/tmp', self.zip_filename), 'w') + file_path = os.path.join('/tmp', self.filename) + kwargs.update({'compression': False}) + self.file.to_csv(file_path) + zipped_file.write(file_path, arcname=os.path.basename(file_path)) + file_path = zipped_file.filename + else: + self.file.to_csv(file_path, **kwargs) return file_path diff --git a/spec/generation_files_spec.py b/spec/generation_files_spec.py index c3463bc..5eb1458 100644 --- a/spec/generation_files_spec.py +++ b/spec/generation_files_spec.py @@ -675,6 +675,13 @@ def get_sample_cups45_data(): expected = 'ES0012345678912345670F;2020/01/01 01:00;0;0;0;0;0;0;0;1;0;FE20214444' assert f.file[f.columns].to_csv(sep=';', header=None, index=False).split('\n')[0] == expected + with it('gets expected content when uses ZIP compression'): + data = SampleData().get_sample_f5d_data() + f = F5D(data, compression='zip') + res = f.writer() + expected = 'ES0012345678912345670F;2020/01/01 01:00;0;0;0;0;0;0;0;1;0;FE20214444' + assert f.file[f.columns].to_csv(sep=';', header=None, index=False).split('\n')[0] == expected + assert '.zip' in res with description('An F1'): with it('is instance of F1 Class'):