import logging
import os
import zipfile
from xmm.util.file import smart_open
from ..base import PipelineStep
logger = logging.getLogger(__name__)
[docs]class ZipStep(PipelineStep):
"""
Zip the exported file.
TODO: The original file will not be deleted in the export directory.
"""
def process_step(self, state, context):
return state, context
def cleanup(self, context, exception=None):
filename = context.get('write', {}).get('filename')
if not isinstance(filename, str) or not os.path.isfile(filename):
return context
filename_zip = os.path.splitext(filename)[0] + '.zip'
with smart_open(filename_zip, 'wb') as fp:
with zipfile.ZipFile(fp, 'w', zipfile.ZIP_DEFLATED) as zip_file:
logging.info(repr(os.path.basename(filename)))
zip_file.write(filename, os.path.basename(filename))
context['write']['filename'] = filename_zip
return context