<commit_before>"""Various miscellaneous functions"""
__author__ = "<NAME>, and <NAME>"
__date__ = "01 Jan 2016"
__version__ = "1.0"

import pickle
import glob
import logging
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import warnings
import zipfile

import ccp4
import exit_util
import pdb_edit

from ample.constants import SHARE_DIR, AMPLEDIR, I2DIR

CCP4 = ccp4.CCP4()
SCRIPT_EXT = '.bat' if sys.platform.startswith('win') else '.sh'
EXE_EXT = '.exe' if sys.platform.startswith('win') else ''
SCRIPT_HEADER = '' if sys.platform.startswith('win') else '#!/bin/bash'

class FileNotFoundError(Exception): pass

# ample_util is used before anything else so there is no logger available
# and we need to a Null handler
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

def amoptd_fix_path(optd, newroot):
    """Update all the paths in an AMPLE results dictionary to be rooted at newroot

    Parameters
    ----------
    optd: dict
      AMPLE results dictionary
    newroot: str
       Path to the AMPLE root directory (topdir containing MRBUMP dir etc.)
    """
    #oldroot = os.sep.join(optd['work_dir'].split(os.sep)[:-1])
    oldroot = optd['work_dir']
    for k in ['benchmark_dir',
              'native_pdb',
              'native_pdb_std',
              'fasta']:
        if k in optd and isinstance(optd[k], str):
            optd[k] = optd[k].replace(oldroot, newroot)

    MRBUMP_FILE_KEYS = ['PHASER_logfile', 'PHASER_pdbout', 'PHASER_mtzout',
                        'REFMAC_logfile', 'REFMAC_pdbout', 'REFMAC_mtzout',
                        'BUCC_logfile', 'BUCC_pdbout', 'BUCC_mtzout',
                        'ARP_logfile', 'ARP_pdbout', 'ARP_mtzout',
                        'SHELXE_logfile', 'SHELXE_pdbout', 'SHELXE_mtzout',
                        'SXRBUCC_logfile', 'SXRBUCC_pdbout','SXRBUCC_mtzout',
                        'SXRARP_logfile', 'SXRARP_pdbout', 'SXRARP_mtzout']
    if 'mrbump_results' in optd:
        for r in optd['mrbump_results']:
            for k in MRBUMP_FILE_KEYS:
                if k in r and isinstance(r[k], str):
                    old = r[k]
                    warnings.warn("FIX MRBUMP BUG buccaneer refine.pdb vs refined.pdb")
                    new = r[k].replace(oldroot, newroot)
                    #logger.info('Changing amopt entry %s from: %s to: %s', k, old, new)
                    r[k] = new
    return optd



def extract_and_validate_models(amoptd, sequence=None, single=True, allsame=True):
    """Extract models given to AMPLE from arguments in the amoptd and validate
    that they are suitable

    Parameters
    ----------
    amoptd : dict
       AMPLE options dictionary
    sequence : str
       single-letter protein sequence - if given a check will be made that all
       models are of this sequence
    single : bool
       if True check each pdb only contains a single model
    allsame : bool
       only extract a file if the suffix is in the list

    """
    
    def is_tar_archive(filename):
        tar_suffixes = ['.tar.gz', '.tgz', '.tar.bz', '.tar.bz2', '.tbz']
        return any(map(lambda x: filename.endswith(x), tar_suffixes))

    def is_zip_archive(filename):
        zip_suffixes = ['.zip']
        return any(map(lambda x: filename.endswith(x), zip_suffixes))
    
    def is_pdb_file(filename):
        return any(map(lambda x: filename.endswith(x), pdb_suffixes))

    def path_to_quark_alldecoy(pdb_files):
        QUARK_DECOY_NAME = 'alldecoy.pdb'
        for pdb in pdb_files:
            if os.path.basename(pdb) == QUARK_DECOY_NAME:
                return pdb
        return None

    filepath = amoptd['models']
    models_dir = amoptd['models_dir']
    pdb_suffixes = ['.pdb', '.PDB']

    if os.path.isfile(filepath):
        basename = os.path.basename(filepath).lower()
        models_dir_tmp = os.path.join(amoptd['work_dir'], 'models.tmp')
        os.mkdir(models_dir_tmp)
        # Extract /copy any pdb files into models_dir_tmp
        if is_tar_archive(basename):
            pdb_files = extract_tar(filepath, models_dir_tmp, suffixes=pdb_suffixes)
        elif is_zip_archive(basename):
            pdb_files = extract_zip(filepath, models_dir_tmp, suffixes=pdb_suffixes)
        elif is_pdb_file(basename):
            shutil.copy2(filepath, models_dir_tmp)
            pdb_files = [os.path.join(models_dir_tmp, filepath)]
        else:
            raise RuntimeError("Do not know how to handle input models file: {}".format(filepath))
        # See if we have an alldecoy.pdb file
        quark_decoy = path_to_quark_alldecoy(pdb_files)
        if quark_decoy:
            split_quark_alldecoy(quark_decoy, models_dir)
            amoptd['quark_models'] = True
            shutil.rmtree(models_dir_tmp) # delete as contains uneeded files extracted from archive
        else:
            # Rename models directory as it now contains our pdb files
            os.rename(models_dir_tmp, models_dir)
    elif os.path.isdir(filepath):
        models_dir = filepath

    if not pdb_edit.check_pdb_directory(models_dir, sequence=sequence, single=single, allsame=allsame):
        msg = "Problem importing pdb files - please check the log for more information"
        exit_util.exit_error(msg)

    amoptd['models_dir'] = models_dir
    return glob.glob(os.path.join(models_dir, "*.pdb"))

<commit_msg>Use native functions to check for zip/tar files<commit_after>"""Various miscellaneous functions"""
__author__ = "<NAME>, and <NAME>"
__date__ = "01 Jan 2016"
__version__ = "1.0"

import pickle
import glob
import logging
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import warnings
import zipfile

import ccp4
import exit_util
import pdb_edit

from ample.constants import SHARE_DIR, AMPLEDIR, I2DIR

CCP4 = ccp4.CCP4()
SCRIPT_EXT = '.bat' if sys.platform.startswith('win') else '.sh'
EXE_EXT = '.exe' if sys.platform.startswith('win') else ''
SCRIPT_HEADER = '' if sys.platform.startswith('win') else '#!/bin/bash'

class FileNotFoundError(Exception): pass

# ample_util is used before anything else so there is no logger available
# and we need to a Null handler
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

def amoptd_fix_path(optd, newroot):
    """Update all the paths in an AMPLE results dictionary to be rooted at newroot

    Parameters
    ----------
    optd: dict
      AMPLE results dictionary
    newroot: str
       Path to the AMPLE root directory (topdir containing MRBUMP dir etc.)
    """
    #oldroot = os.sep.join(optd['work_dir'].split(os.sep)[:-1])
    oldroot = optd['work_dir']
    for k in ['benchmark_dir',
              'native_pdb',
              'native_pdb_std',
              'fasta']:
        if k in optd and isinstance(optd[k], str):
            optd[k] = optd[k].replace(oldroot, newroot)

    MRBUMP_FILE_KEYS = ['PHASER_logfile', 'PHASER_pdbout', 'PHASER_mtzout',
                        'REFMAC_logfile', 'REFMAC_pdbout', 'REFMAC_mtzout',
                        'BUCC_logfile', 'BUCC_pdbout', 'BUCC_mtzout',
                        'ARP_logfile', 'ARP_pdbout', 'ARP_mtzout',
                        'SHELXE_logfile', 'SHELXE_pdbout', 'SHELXE_mtzout',
                        'SXRBUCC_logfile', 'SXRBUCC_pdbout','SXRBUCC_mtzout',
                        'SXRARP_logfile', 'SXRARP_pdbout', 'SXRARP_mtzout']
    if 'mrbump_results' in optd:
        for r in optd['mrbump_results']:
            for k in MRBUMP_FILE_KEYS:
                if k in r and isinstance(r[k], str):
                    old = r[k]
                    warnings.warn("FIX MRBUMP BUG buccaneer refine.pdb vs refined.pdb")
                    new = r[k].replace(oldroot, newroot)
                    #logger.info('Changing amopt entry %s from: %s to: %s', k, old, new)
                    r[k] = new
    return optd



def extract_and_validate_models(amoptd, sequence=None, single=True, allsame=True):
    """Extract models given to AMPLE from arguments in the amoptd and validate
    that they are suitable

    Parameters
    ----------
    amoptd : dict
       AMPLE options dictionary
    sequence : str
       single-letter protein sequence - if given a check will be made that all
       models are of this sequence
    single : bool
       if True check each pdb only contains a single model
    allsame : bool
       only extract a file if the suffix is in the list

    """
    
    def is_pdb_file(filename):
        return any(map(lambda x: filename.endswith(x), pdb_suffixes))

    def path_to_quark_alldecoy(pdb_files):
        QUARK_DECOY_NAME = 'alldecoy.pdb'
        for pdb in pdb_files:
            if os.path.basename(pdb) == QUARK_DECOY_NAME:
                return pdb
        return None

    filepath = amoptd['models']
    models_dir = amoptd['models_dir']
    pdb_suffixes = ['.pdb', '.PDB']

    if os.path.isfile(filepath):
        models_dir_tmp = os.path.join(amoptd['work_dir'], 'models.tmp')
        if not os.path.isdir(models_dir_tmp):
            os.mkdir(models_dir_tmp)
        # Extract /copy any pdb files into models_dir_tmp
        if tarfile.is_tarfile(filepath):
            pdb_files = extract_tar(filepath, models_dir_tmp, suffixes=pdb_suffixes)
        elif zipfile.is_zipfile(filepath):
            pdb_files = extract_zip(filepath, models_dir_tmp, suffixes=pdb_suffixes)
        elif is_pdb_file(filepath):
            shutil.copy2(filepath, models_dir_tmp)
            pdb_files = [os.path.join(models_dir_tmp, filepath)]
        else:
            raise RuntimeError("Do not know how to handle input models file: {}".format(filepath))
        # See if we have an alldecoy.pdb file
        quark_decoy = path_to_quark_alldecoy(pdb_files)
        if quark_decoy:
            split_quark_alldecoy(quark_decoy, models_dir)
            amoptd['quark_models'] = True
            shutil.rmtree(models_dir_tmp) # delete as contains uneeded files extracted from archive
        else:
            # Rename models directory as it now contains our pdb files
            os.rename(models_dir_tmp, models_dir)
    elif os.path.isdir(filepath):
        models_dir = filepath

    if not pdb_edit.check_pdb_directory(models_dir, sequence=sequence, single=single, allsame=allsame):
        msg = "Problem importing pdb files - please check the log for more information"
        exit_util.exit_error(msg)

    amoptd['models_dir'] = models_dir
    return glob.glob(os.path.join(models_dir, "*.pdb"))


{'fraction_non_alphanumeric': 0.09935225757287103, 'fraction_numerical': 0.0029529434177938655, 'mean_word_length': 2.6391681109185443, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1170892', 'n_tokens_mistral': 3522, 'n_tokens_neox': 3364, 'n_words': 880}