Source code for stetl.filters.fileextractor

# Extracts a file from and archive file like a .zip,
# and saves it as the given file name.
#
# Author: Just van den Broecke (generic and VsiFileExtractor)
# Author: Frank Steggink (ZipFileExtractor)
#
import os.path
from stetl.component import Config
from stetl.filter import Filter
from stetl.util import Util
from stetl.packet import FORMAT

log = Util.get_log('fileextractor')

DEFAULT_BUFFER_SIZE = 1024 * 1024 * 1024




class FileExtractor(Filter):
    """
    Abstract Base Class.
    Extracts a file an archive and saves as the configured file name.

    consumes=FORMAT.any, produces=FORMAT.string
    """

    # Start attribute config meta



    @Config(ptype=str, default=None, required=True)
    def file_path(self):
        """
        File name to write the extracted file to.
        """
        pass




    @Config(ptype=bool, default=True, required=False)
    def delete_file(self):
        """
        Delete the file when the chain has been completed?
        """
        pass




    @Config(ptype=int, default=DEFAULT_BUFFER_SIZE, required=False)
    def buffer_size(self):
        """
        Buffer size for read buffer during extraction.
        """
        pass


    # End attribute config meta

    # Constructor
    def __init__(self, configdict, section, consumes=FORMAT.any, produces=FORMAT.string):
        Filter.__init__(self, configdict, section, consumes=consumes, produces=produces)

    def delete_target_file(self):
        if os.path.isfile(self.file_path):
            os.remove(self.file_path)

    def extract_file(self, packet):
        log.error('Only classes derived from FileExtractor can be used!')



    def invoke(self, packet):

        if packet.data is None:
            log.info("Input data is empty")
            return packet

        # Optionally remove old file
        self.delete_target_file()
        self.extract_file(packet)
        packet.data = self.file_path

        if not os.path.isfile(self.file_path):
            log.warn('Extracted file {} does not exist'.format(self.file_path))
            packet.data = None

        return packet




    def after_chain_invoke(self, packet):
        if not self.delete_file:
            return

        self.delete_target_file()

        return True






class ZipFileExtractor(FileExtractor):
    """
    Extracts a file from a ZIP file, and saves it as the given file name.
    Author: Frank Steggink

    consumes=FORMAT.record, produces=FORMAT.string
    """

    def __init__(self, configdict, section):
        FileExtractor.__init__(self, configdict, section, consumes=FORMAT.record)

    def extract_file(self, packet):

        import zipfile

        with zipfile.ZipFile(packet.data['file_path']) as z:
            with open(self.file_path, 'wb') as f:
                with z.open(packet.data['name']) as zf:
                    while True:
                        buffer = zf.read(self.buffer_size)
                        if not buffer:
                            break
                        f.write(buffer)





class VsiFileExtractor(FileExtractor):
    """
    Extracts a file from a GDAL /vsi path spec, and saves it as the given file name.

    Example paths:
    /vsizip/{/project/nlextract/data/BAG-2.0/BAGNLDL-08112020.zip}/9999STA08112020.zip'
    /vsizip/{/vsizip/{BAGGEM0221L-15022021.zip}/GEM-WPL-RELATIE-15022021.zip}/GEM-WPL-RELATIE-15022021-000001.xml

    See also stetl.inputs.fileinput.VsiZipFileInput that generates these paths.

    Author: Just van den Broecke

    consumes=FORMAT.gdal_vsi_path, produces=FORMAT.string
    """

    def __init__(self, configdict, section):
        FileExtractor.__init__(self, configdict, section, consumes=FORMAT.gdal_vsi_path)

    def extract_file(self, packet):
        from stetl.util import gdal

        # Example input path can be as complex as this:
        #
        vsi_file_path = packet.data
        vsi = None
        vsi_len = 0
        try:
            # gdal.VSIF does not support 'with' so old-school open/close.
            log.info('Extracting {}'.format(vsi_file_path))
            vsi = gdal.VSIFOpenL(vsi_file_path, 'rb')
            with open(self.file_path, 'wb') as f:
                gdal.VSIFSeekL(vsi, 0, 2)
                vsi_len = gdal.VSIFTellL(vsi)
                gdal.VSIFSeekL(vsi, 0, 0)
                read_size = self.buffer_size
                if vsi_len < read_size:
                    read_size = vsi_len

                while True:
                    buffer = gdal.VSIFReadL(1, read_size, vsi)
                    if not buffer or len(buffer) == 0:
                        break
                    f.write(buffer)

        except Exception as e:
            log.error('Cannot extract {} err={}'.format(vsi_file_path, str(e)))
            raise e
        finally:
            if vsi:
                log.info('Extracted {} ok len={} bytes'.format(vsi_file_path, vsi_len))
                gdal.VSIFCloseL(vsi)