import collections
import os
from copy import copy
from visidata import vd
from visidata import Sheet, BaseSheet, VisiData, IndexSheet, Path, Progress, TypedExceptionWrapper, TypedWrapper, UNLOADED
vd.option('safe_error', '#ERR', 'error string to use while saving', replay=True)
vd.option('save_encoding', 'utf-8', 'encoding passed to codecs.open when saving a file', replay=True, help=vd.help_encoding)
_compression_formats = {'gz', 'bz2', 'xz', 'lzma', 'zst', 'zstd'}
def parse_filetype(ft):
'Parse filetype string like "json.gz" into (format, compression). Returns (ft, None) for plain types.'
for sep in ('.', '+'):
if sep in ft:
fmt, comp = ft.rsplit(sep, 1)
if comp in _compression_formats:
return fmt, comp
return ft, None
@Sheet.api
def safe_trdict(vs, delimiter=None):
'returns string.translate dictionary for replacing tabs and newlines'
if vs.options.safety_first:
delim = delimiter or vs.options.delimiter
trdict = {
0: '', # strip NUL completely
10: vs.options.tsv_safe_newline, # \n
13: vs.options.tsv_safe_newline, # \r
}
if not delim or ord(delim) in trdict:
vd.fail(f'cannot use delimiter {repr(delim)} with safety_first')
trdict[ord(delim)] = vs.options.tsv_safe_tab # \t
return trdict
return {}
@Sheet.api
def iterdispvals(sheet, *cols, format=False, delimiter=None):
'For each row in sheet, yield OrderedDict of values for given cols. Values are typed if format=False, or a formatted display string if format=True.'
if not cols:
cols = sheet.visibleCols
transformers = collections.OrderedDict() # list of transformers for each column in order
trdict = sheet.safe_trdict(delimiter=delimiter)
for col in cols:
transformers[col] = [ col.type ]
if format:
formatMaker = getattr(col, 'formatter_'+(col.formatter or sheet.options.disp_formatter))
transformers[col].append(formatMaker(col._formatdict))
if trdict:
transformers[col].append(lambda v,trdict=trdict: v.translate(trdict) if isinstance(v, str) else v) #3256
options_safe_error = sheet.options.safe_error
for r in sheet.iterrows('saving'):
dispvals = collections.OrderedDict() # [col] -> value
for col, transforms in transformers.items():
try:
dispval = col.getValue(r)
except Exception as e:
dispval = options_safe_error or str(e)
try:
for t in transforms:
if dispval is None or isinstance(dispval, float) and dispval != dispval:
break
elif isinstance(dispval, TypedExceptionWrapper):
dispval = options_safe_error or str(dispval)
break
elif isinstance(dispval, TypedWrapper):
dispval = ''
break
else:
dispval = t(dispval)
if (dispval is None or isinstance(dispval, float) and dispval != dispval) and format:
dispval = ''
except Exception:
dispval = str(dispval)
dispvals[col] = dispval
yield dispvals
@Sheet.api
def itervals(sheet, *cols, format=False):
for row in sheet.iterdispvals(*cols, format=format):
yield [row[c] for c in cols]
@BaseSheet.api
def getDefaultSaveName(sheet):
src = getattr(sheet, 'source', None)
if hasattr(src, 'scheme') and src.scheme:
return src.name
if isinstance(src, Path):
if src.given == '-':
return f'stdin.{sheet.options.save_filetype}'
if sheet.options.is_set('save_filetype', sheet):
# if save_filetype is over-ridden from default, use it as the extension
return str(src.with_suffix('')) + '.' + sheet.options.save_filetype
return str(src)
else:
return sheet.name+'.'+getattr(sheet, 'filetype', sheet.options.save_filetype)
@VisiData.api
def saveCols(vd, cols):
sheet = cols[0].sheet
vs = copy(sheet)
vs.columns = list(cols)
vs.rows = list(sheet.rows) # copy list to avoid conflict with rows modifications
if len(cols) == 1:
savedcoltxt = cols[0].name + ' column'
else:
savedcoltxt = '%s columns' % len(cols)
path = vd.inputPath('save %s to: ' % savedcoltxt, value=vs.getDefaultSaveName())
vd.saveSheets(path, vs)
[docs]
@VisiData.api
def saveSheets(vd, givenpath, *vsheets, confirm_overwrite=True):
'''Save all *vsheets* to *givenpath*. Async.
Callers should be careful not to call reload() while saveSheets is still running.
Use vd.sync(saveSheets) to wait for the save to finish.'''
if not vsheets: # blank tuple
vd.warning('no sheets to save')
return
if not givenpath.name:
vd.warning('no save path given')
return
unloaded = [ vs for vs in vsheets if vs.rows is UNLOADED ]
vd.sync(*vd.ensureLoaded(unloaded))
# resolve filetype: path option (from inputPath or -f) > extension > save_filetype
path_filetype = givenpath.options.is_set('filetype', givenpath)
if path_filetype:
fmt, compression = parse_filetype(path_filetype.value)
if compression and givenpath.compression is None:
givenpath.compression = compression
else:
fmt = None
givenext = givenpath.ext.lower()
default_ft = vd.options.save_filetype.lower()
vd.clearCaches()
def _find_saver(ft):
return getattr(vsheets[0], 'save_' + ft, None) or getattr(vd, 'save_' + ft, None)
savefunc = filetype = None
requested = [ft for ft in (fmt, givenext) if ft]
for ft in requested:
savefunc = _find_saver(ft)
if savefunc:
filetype = ft
break
if not savefunc:
savefunc = _find_saver(default_ft)
if not savefunc:
vd.fail(f'no saver for {" or ".join(requested + [default_ft])}')
if requested:
vd.confirm(f'no `{requested[0]}` saver, save as {default_ft}? ') #2286
filetype = default_ft
if confirm_overwrite:
vd.confirmOverwrite(givenpath)
disptype = f'{filetype}.{givenpath.compression}' if givenpath.compression else filetype
vd.status('saving %s sheets to %s as %s' % (len(vsheets), givenpath.given, disptype))
if not givenpath.given.endswith('/'): # forcibly specify save individual files into directory by ending path with /
for vs in vsheets:
vs.hasBeenModified = False
# savefuncs(vd, p, *vsheets) will have 2 argcount (*vsheets does not get counted as an arg)
# savefuncs(vd, p, vs) will have 3 argcount (vs counts as an arg, along with vd, path)
if savefunc.__code__.co_argcount == 3 and len(vsheets) > 1:
vd.fail(f'cannot save multiple {filetype} sheets to non-dir')
return vd.execAsync(savefunc, givenpath, *vsheets)
# path is a dir
# save as individual files in the givenpath directory
try:
os.makedirs(givenpath, exist_ok=True)
except FileExistsError:
pass
if not givenpath.is_dir():
vd.fail(f'cannot save multiple {filetype} sheets to non-dir')
def _savefiles(vsheets, givenpath, savefunc, filetype):
for vs in vsheets:
p = Path((givenpath / vs.name).with_suffix('.'+filetype))
savefunc(p, vs)
vs.hasBeenModified = False
vd.status(f'{givenpath} save finished') #2157
return vd.execAsync(_savefiles, vsheets, givenpath, savefunc, filetype)
@VisiData.api
def save_zip(vd, p, *vsheets):
vd.clearCaches()
import zipfile
with vd.TempDir() as tmpdir:
with zipfile.ZipFile(str(p), 'w', zipfile.ZIP_DEFLATED, allowZip64=True, compresslevel=9) as zfp:
for vs in Progress(vsheets):
filetype = vs.options.save_filetype
tmpp = Path(tmpdir)/f'{vs.name}.{filetype}'
savefunc = getattr(vs, 'save_' + filetype, None) or getattr(vd, 'save_' + filetype, None)
savefunc(tmpp, vs)
zfp.write(tmpp, f'{vs.name}.{vs.options.save_filetype}')
[docs]
@VisiData.api
def save_txt(vd, p, *vsheets):
if len(vsheets) == 1 and vsheets[0].nVisibleCols > 1: #2173
return vd.save_tsv(p, vsheets[0])
with p.open(mode='w', encoding=vsheets[0].options.save_encoding) as fp:
for vs in vsheets:
unitsep = p.options.delimiter
rowsep = p.options.row_delimiter
for dispvals in vs.iterdispvals(*vs.visibleCols, format=True):
fp.write(unitsep.join(dispvals.values()))
fp.write(rowsep)
@BaseSheet.api
def rootSheet(sheet):
r = sheet
while isinstance(r.source, BaseSheet):
r = r.source
return r
BaseSheet.addCommand('Ctrl+S', 'save-sheet', 'vd.saveSheets(inputPath("save to: ", value=getDefaultSaveName()), sheet)', 'save current sheet to given filename and filetype')
BaseSheet.addCommand('', 'save-sheet-really', 'vd.saveSheets(Path(getDefaultSaveName()), sheet, confirm_overwrite=False)', 'save current sheet without asking for filename or confirmation')
BaseSheet.addCommand('', 'save-source', 'vd.saveSheets(rootSheet().source, rootSheet())', 'save root sheet to its source')
BaseSheet.addCommand('gCtrl+S', 'save-all', 'vd.saveSheets(inputPath("save all sheets to: "), *vd.stackedSheets)', 'save all sheets to given file or directory)')
IndexSheet.addCommand('gCtrl+S', 'save-selected', 'vd.saveSheets(inputPath("save %d sheets to: " % nSelectedRows, value="_".join(getattr(vs, "name", None) or "blank" for vs in selectedRows)), *selectedRows)', 'save all selected sheets to given file or directory')
Sheet.addCommand('', 'save-col', 'saveCols([cursorCol])', 'save current column only to given filename and filetype')
Sheet.addCommand('', 'save-col-keys', 'saveCols(keyCols + [cursorCol])', 'save key columns and current column to given filename and filetype')
vd.addMenuItems('''
File > Save > current sheet > save-sheet
File > Save > all sheets > save-all
File > Save > current column > save-col
File > Save > keys and current column > save-col-keys
''')