Source code for schrodinger.application.watermap.utils

"""
General WaterMap utilities

Copyright Schrodinger, LLC. All rights reserved.
"""

import os
import zipfile

from schrodinger import structure


def archive_wm(wm_fname, out_path=None, trj=False):

    if not os.path.exists(wm_fname):
        raise OSError('"%s" does not exist.' % wm_fname)

    # If there is no watermap property in the file then this is not
    # a valid watermap file
    st_reader = structure.StructureReader(wm_fname)
    for st in st_reader:
        try:
            if st.property['s_watermap_type'] == 'watermap':
                break
        except:
            pass
    else:
        raise OSError('Watermap structure does not exist in "%s".' % wm_fname)

    base_dir = os.path.dirname(wm_fname)
    base_name = os.path.basename(wm_fname)

    cwd = os.getcwd()
    if not base_dir:
        base_dir = os.getcwd()

    os.chdir(base_dir)

    trj_list = []
    for k in list(st.property):
        if k == 's_watermap_idx_file' or k == 's_watermap_trj_file':
            fname = st.property[k]
            trj_list.append(fname)
            del st.property[k]

    file_list = []
    file_list.append(base_name)
    for k in list(st.property):
        if k.find('_file') > 0:
            fname = st.property[k]
            if os.path.exists(fname):
                file_list.append(fname)
            else:
                print('"%s" is not found. Skipping to archive this file.' %
                      fname)

    # If the out_path is just a directory we will add the
    # jobname to the output .zip, otherwise we would end up
    # with a file /path/to/_wm.zip
    if out_path:
        if os.path.isdir(out_path):
            out_path = os.path.join(out_path,
                                    base_name[:base_name.find('_wm.mae')])
        if not out_path.endswith('_wm.zip'):
            archive_fname = out_path + '_wm.zip'
        else:
            archive_fname = out_path
    else:
        archive_fname = base_name[:base_name.find('_wm.mae')] + '_wm.zip'

    # Compress the file
    with zipfile.ZipFile(archive_fname, 'w') as ar_file:
        for f in file_list:
            ar_file.write(f)
        if trj:
            for f in trj_list:
                ar_file.write(f)

    os.chdir(cwd)

    return archive_fname