Source code for ocgis.util.zipper
import io
import os
import zipfile
from ocgis.constants import OutputFormatName
class Zipper(object):
def __init__(self, base_path=None):
self.base_path = base_path
def get_zip_stream(self):
buffer = io.BytesIO()
zip = zipfile.ZipFile(buffer, 'w', zipfile.ZIP_DEFLATED)
try:
items = os.listdir(self.base_path)
path = self.base_path
except OSError:
path, items = os.path.split(self.base_path)
items = [items]
for item in items:
filepath = os.path.join(path, item)
zip.write(filepath, arcname=item)
zip.close()
buffer.flush()
zip_stream = buffer.getvalue()
buffer.close()
return (zip_stream)
@classmethod
def compress_path(cls, path, items):
zipf = zipfile.ZipFile(path, 'w', zipfile.ZIP_DEFLATED)
for item in items:
arcname = item.get('arcname')
zipf.write(item['filename'], arcname=arcname)
zipf.close()
return (path)
def get_items(folder):
items = []
for dirpath, dirnames, filenames in os.walk(folder):
for filename in filenames:
if os.path.split(dirpath)[1] == 'shp':
arcname = os.path.join('shp', filename)
else:
arcname = filename
items.append({'filename': os.path.join(dirpath, filename),
'arcname': arcname})
return (items)
def get_zipped_path(path_zip, folder):
items = get_items(folder)
Zipper.compress_path(path_zip, items)
return (path_zip)
[docs]def format_return(ret_path, ops, with_auxiliary_files=False):
"""
Format an OpenClimateGIS path returning an absolute path to a zip file or simply passing `ret_path` through.
>>> import ocgis
>>> ops = ocgis.OcgOperations(...)
>>> ret = ops.execute()
>>> new_path = ocgis.format_return(ret,ops)
:param ret_path: The path returned by :func:`ocgis.OcgOperations.execute`.
:type ret_path: str
:param ops: Instance of :class:`ocgis.OcgOperations`.
:type ops: :class:`ocgis.OcgOperations`
:param with_auxiliary_files: If `True`, zip additional contents of CSV and netCDF returns.
:type with_auxiliary_files: bool
:raises: NotImplementedError
:returns: str
"""
# can do nothing with numpy returns
if ops.output_format == OutputFormatName.OCGIS:
raise NotImplementedError
# the folder containing all output files
folder = os.path.split(ret_path)[0]
# name for output zipfile
path_zip = os.path.join(folder, ops.prefix + '.zip')
if ops.output_format in ['csv', 'nc']:
# add all files
if with_auxiliary_files:
ret = get_zipped_path(path_zip, folder)
# only interested in the file.
else:
ret = ret_path
# otherwise return all files
else:
ret = get_zipped_path(path_zip, folder)
return ret