1###############################################################################
2#
3# common.py - utility functions used in many places in CheckM
4#
5###############################################################################
6#                                                                             #
7#    This program is free software: you can redistribute it and/or modify     #
8#    it under the terms of the GNU General Public License as published by     #
9#    the Free Software Foundation, either version 3 of the License, or        #
10#    (at your option) any later version.                                      #
11#                                                                             #
12#    This program is distributed in the hope that it will be useful,          #
13#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
14#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
15#    GNU General Public License for more details.                             #
16#                                                                             #
17#    You should have received a copy of the GNU General Public License        #
18#    along with this program. If not, see <http://www.gnu.org/licenses/>.     #
19#                                                                             #
20###############################################################################
21
22import os
23import errno
24import sys
25import ast
26import logging
27
28import numpy as np
29
30from checkm.defaultValues import DefaultValues
31
32
33def getBinIdsFromOutDir(outDir):
34    """Get bin ids."""
35    binIds = []
36
37    binDir = os.path.join(outDir, 'bins')
38
39    for f in os.listdir(binDir):
40        if os.path.isdir(os.path.join(binDir, f)) and f != 'storage':
41            binIds.append(f)
42
43    return binIds
44
45
46def readDistribution(prefix):
47    """Read distribution file."""
48    distFile = os.path.join(DefaultValues.DISTRIBUTION_DIR, prefix + '.txt')
49    checkFileExists(distFile)
50
51    with open(distFile, 'r') as f:
52        s = f.read()
53        d = ast.literal_eval(s)
54
55    return d
56
57
58def findNearest(array, value):
59    '''Find nearest array element to a given value.'''
60    idx = (np.abs(np.array(array) - value)).argmin()
61    return array[idx]
62
63
64def checkEmptyDir(inputDir):
65    """Check the the specified directory is empty and create it if necessary."""
66    if not os.path.exists(inputDir):
67        makeSurePathExists(inputDir)
68    else:
69        # check if directory is empty
70        files = os.listdir(inputDir)
71        if len(files) != 0:
72            logger = logging.getLogger('timestamp')
73            logger.error('Output directory must be empty: ' + inputDir + '\n')
74            sys.exit(1)
75
76
77def checkFileExists(inputFile):
78    """Check if file exists."""
79    if not os.path.exists(inputFile):
80        logger = logging.getLogger('timestamp')
81        logger.error('Input file does not exists: ' + inputFile + '\n')
82        sys.exit(1)
83
84
85def checkDirExists(inputDir):
86    """Check if directory exists."""
87    if not os.path.exists(inputDir):
88        logger = logging.getLogger('timestamp')
89        logger.error('Input directory does not exists: ' + inputDir + '\n')
90        sys.exit(1)
91
92
93def makeSurePathExists(path):
94    """Create directory if it does not exist."""
95    if not path:
96        return
97
98    try:
99        os.makedirs(path)
100    except OSError as exception:
101        if exception.errno != errno.EEXIST:
102            logger = logging.getLogger('timestamp')
103            logger.error('Specified path does not exist: ' + path + '\n')
104            sys.exit(1)
105
106
107def binIdFromFilename(filename):
108    """Extract bin id from bin filename."""
109    binId = os.path.basename(filename)
110    binId = os.path.splitext(binId)[0]
111
112    return binId
113
114
115def reassignStdOut(outFile):
116    """Redirect standard out to a file."""
117    oldStdOut = sys.stdout
118    if(outFile != ''):
119        try:
120            # redirect stdout to a file
121            sys.stdout = open(outFile, 'w')
122        except:
123            logger = logging.getLogger('timestamp')
124            logger.error("Error diverting stdout to file: " + outFile)
125            sys.exit(1)
126
127    return oldStdOut
128
129
130def restoreStdOut(outFile, oldStdOut):
131    """Redirect standard out back to system standard out."""
132    if(outFile != ''):
133        try:
134            # redirect stdout to a file
135            sys.stdout.close()
136            sys.stdout = oldStdOut
137        except:
138            logger = logging.getLogger('timestamp')
139            logger.error("Error restoring stdout ", outFile)
140            sys.exit(1)
141