__init__.py 5.37 KB
Newer Older
Clément Haëck's avatar
Clément Haëck committed
1 2
"""General global variables."""

Clément Haëck's avatar
Clément Haëck committed
3
import argparse
Clément Haëck's avatar
Clément Haëck committed
4
import configparser
Clément Haëck's avatar
Clément Haëck committed
5
import itertools
Clément Haëck's avatar
Clément Haëck committed
6
import os
Clément Haëck's avatar
Clément Haëck committed
7
import re
Clément Haëck's avatar
Clément Haëck committed
8
import sys
Clément Haëck's avatar
Clément Haëck committed
9 10 11
from datetime import datetime
from os import path
from typing import List
Clément Haëck's avatar
Clément Haëck committed
12

13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29
# Scripts repository
try:
    root_code = os.environ['SUBMESO_COLOR_CODE_DIR']
except KeyError:
    raise KeyError("SUBMESO_COLOR_CODE_DIR env variable is not defined")

# Data directory
try:
    root_data = os.environ['SUBMESO_COLOR_DATA_DIR']
except KeyError:
    raise KeyError("SUBMESO_COLOR_DATA_DIR env variable is not defined")
if not path.isdir(root_data):
    raise NameError(f"root_data '{root_data}' is not a valid directory.")

# Plot directory
root_plot = path.join(root_data, 'Plots')

Clément Haëck's avatar
Clément Haëck committed
30 31
# Read config file
_config = configparser.ConfigParser()
32
_config_file = path.join(root_code, 'lib', 'conf.ini')
Clément Haëck's avatar
Clément Haëck committed
33 34
if not path.isfile(_config_file):
    raise RuntimeError("Missing configuration file lib/conf.ini")
Clément Haëck's avatar
Clément Haëck committed
35

Clément Haëck's avatar
Clément Haëck committed
36
_config.read(_config_file)
Clément Haëck's avatar
Clément Haëck committed
37

38 39 40 41 42 43 44 45 46 47

def to_bool(x=None):
    if x is None:
        return False
    return _config._convert_to_boolean(x)


_supported_types = {t.__name__: t for t in [int, str, float, type(None)]}
_supported_types['bool'] = to_bool

Clément Haëck's avatar
Clément Haëck committed
48 49 50 51 52 53 54 55
DEFAULTS = {}
DEFAULTS_TYPES = {}
for k in _config.options('defaults'):
    tp_str, val = _config.get('defaults', k).split(',')
    tp = _supported_types[tp_str]
    val = val.strip()
    DEFAULTS[k] = tp(val) if val else tp()
    DEFAULTS_TYPES[k] = tp
56 57


Clément Haëck's avatar
Clément Haëck committed
58 59
def _fixes_as_dict(fixes: List[str]):
    return {k: v for k, v in fixes}
60

Clément Haëck's avatar
Clément Haëck committed
61

62
# Set default fixes
Clément Haëck's avatar
Clément Haëck committed
63 64 65 66 67 68 69
if 'fixes' in DEFAULTS:
    d = DEFAULTS['fixes'].split()
    if len(d) % 2 != 0:
        raise ValueError("Default fixes are not grouped in pairs.")
    d = [[d[i], d[i+1]] for i in range(0, len(d), 2)]
    DEFAULTS['fixes'] = _fixes_as_dict(d)
    DEFAULTS_TYPES['fixes'] = dict
70

71

72 73 74 75 76
# Set default open_mfdataset kwargs
DEFAULTS['open_mf_kw'] = dict()
DEFAULTS_TYPES['open_mf_kw'] = dict


77
def check_output_dir(odir, log=None, file=False):
Clément Haëck's avatar
Clément Haëck committed
78 79 80 81 82 83 84
    """Check if directory exists, if not create it.

    Parameters
    ----------
    odir: str
    log: logging.Logger instance
        If None, no log will be output.
85 86
    file: bool
        If `odir` is actually a file.
Clément Haëck's avatar
Clément Haëck committed
87
    """
88 89
    if file:
        odir = path.dirname(odir)
Clément Haëck's avatar
Clément Haëck committed
90 91 92 93 94 95
    if log is not None:
        log.info('output to %s', odir)
    if not path.isdir(odir):
        if log is not None:
            log.info('creating %s', odir)
        os.makedirs(odir)
Clément Haëck's avatar
Clément Haëck committed
96 97


98
def get_args(args, add_args=None, description=''):
Clément Haëck's avatar
Clément Haëck committed
99 100 101 102 103 104 105 106 107 108 109
    """Get command line arguments.

    With lots of typical arguments.
    Arguments prefixed by a single hyphen:
    `python script.py -region GS -year 2007 -fix d 01`

    Returns a dictionnary of arguments.

    Parameters
    ----------
    args: list of str
Clément Haëck's avatar
Clément Haëck committed
110 111
        List of typical arguments that will automatically added to the parser,
        with the default value taken from the configuration file.
Clément Haëck's avatar
Clément Haëck committed
112
        See the configuration file to see details on available arguments.
Clément Haëck's avatar
Clément Haëck committed
113

Clément Haëck's avatar
Clément Haëck committed
114 115 116 117
        'fixes' is a special case. It adds a '-fix' argument to the command
        line. There can be any number of 'fix' arguments. Each is a couple of
        key (str) and value (str). Each is appended to dictionnary of fixes
        in args['fixes']. See filefinder.Finder.fix_matchers().
Clément Haëck's avatar
Clément Haëck committed
118 119
    add_args: callable
        Function that take an argparse.ArgumentParser and add arguments to it.
Clément Haëck's avatar
Clément Haëck committed
120
        Arguments can be overwritten.
Clément Haëck's avatar
Clément Haëck committed
121

Clément Haëck's avatar
Clément Haëck committed
122
    """
Clément Haëck's avatar
Clément Haëck committed
123 124 125 126 127
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawTextHelpFormatter,
        description=description,
        conflict_handler='resolve')

Clément Haëck's avatar
Clément Haëck committed
128 129 130 131
    for name in args:
        if name not in DEFAULTS:
            raise KeyError(f"'{name}' not in defaults parameters.")
        parser.add_argument('-{}'.format(name),
Clément Haëck's avatar
Clément Haëck committed
132 133
                            type=DEFAULTS_TYPES[name],
                            default=DEFAULTS[name])
Clément Haëck's avatar
Clément Haëck committed
134

Clément Haëck's avatar
Clément Haëck committed
135
    if 'fixes' in args:
Clément Haëck's avatar
Clément Haëck committed
136 137 138
        parser.add_argument(
            '-fix', nargs=2, type=str, action='append',
            default=list(itertools.chain(DEFAULTS['fixes'].items())))
139

Clément Haëck's avatar
Clément Haëck committed
140 141
    if add_args is not None:
        add_args(parser)
142

Clément Haëck's avatar
Clément Haëck committed
143
    dargs = vars(parser.parse_args())
144

Clément Haëck's avatar
Clément Haëck committed
145
    if 'fixes' in args:
Clément Haëck's avatar
Clément Haëck committed
146
        dargs['fixes'] = _fixes_as_dict(dargs['fix'])
147

Clément Haëck's avatar
Clément Haëck committed
148
    if 'date' in dargs:
Clément Haëck's avatar
Clément Haëck committed
149
        date = re.match(r'(\d\d\d\d)[-\s/]?(\d\d)[-\s/]?(\d\d)?',
Clément Haëck's avatar
Clément Haëck committed
150
                        dargs['date'])
Clément Haëck's avatar
Clément Haëck committed
151
        if date is None:
Clément Haëck's avatar
Clément Haëck committed
152 153
            raise KeyError("-date argument not in correct format "
                           "(YYYY[- / ]MM[- / ][DD])")
Clément Haëck's avatar
Clément Haëck committed
154 155 156
        dargs['date_str'] = [g if g is not None else '01'
                             for g in date.groups()]
        dargs['date'] = datetime(*[int(g) for g in dargs['date_str']])
157

Clément Haëck's avatar
Clément Haëck committed
158
    return dargs
Clément Haëck's avatar
Clément Haëck committed
159 160 161 162 163 164 165


def progressbar(it, prefix="", size=60, file=sys.stdout):
    count = len(it)

    def show(j):
        x = int(size*j/count)
Clément Haëck's avatar
Clément Haëck committed
166 167
        file.write("%s[%s%s] %i/%i\r" % (prefix, "#"*x, "."*(size-x),
                                         j, count))
Clément Haëck's avatar
Clément Haëck committed
168 169 170 171 172 173 174
        file.flush()
    show(0)
    for i, item in enumerate(it):
        yield item
        show(i+1)
    file.write("\n")
    file.flush()
175 176


Clément Haëck's avatar
Clément Haëck committed
177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192
def m_start(msg):
    print(msg + '...', end='', flush=True)


def m_end(msg='done'):
    print(msg + '.', flush=True)


def m_next(msg_start, msg_end=None):
    if msg_end is None:
        m_end()
    else:
        m_end(msg_end)
    m_start(msg_start)


193 194 195 196 197 198 199
def fix_time_daily(*datasets):
    out = []
    for ds in datasets:
        out.append(ds.assign_coords(time=ds.time.dt.floor("D")))
    if len(out) == 1:
        return out[0]
    return out