utils.py

import subprocess

import numpy as np
from autosubmit.config.log import Log
from cdo import Cdo
from nco import Nco
import tempfile
import os
import pwd
import shutil


class Utils(object):

    nco = Nco()
    cdo = Cdo()

    @staticmethod
    def setminmax(filename, variable_list):

        if isinstance(variable_list, basestring):
            variable_list = variable_list.split()

        Log.info('Getting max and min values for {0}', ' '.join(variable_list))

        dataset = Utils.nco.readCdf(filename)
        for variable in variable_list:
            var = dataset.variables[variable]
            values = [np.max(var), np.min(var)]
            Utils.nco.ncatted(input=filename, output=filename,
                              options='-h -a valid_max,{0},m,f,{1}'.format(variable, values[0]))
            Utils.nco.ncatted(input=filename, output=filename,
                              options='-h -a valid_min,{0},m,f,{1}'.format(variable, values[1]))

    @staticmethod
    def rename_variable(filename, old_name, new_name, must_exist=True, rename_dimension=False):
        Utils.rename_variables(filename, {old_name: new_name}, must_exist, rename_dimension)

    @staticmethod
    def rename_variables(filepath, dic_names, must_exist=True, rename_dimension=False):
        handler = Utils.nco.openCdf(filepath)

        for old_name, new_name in dic_names.items():
            if old_name in handler.variables:
                handler.renameVariable(old_name, new_name)
            elif must_exist:
                raise Exception("Variable {0} does not exist in file {1}".format(old_name, filepath))

            if rename_dimension:
                if old_name in handler.dimensions:
                    handler.renameDimension(old_name, new_name)
                elif must_exist:
                    raise Exception("Variable {0} does not exist in file {1}".format(old_name, filepath))

        handler.close()

    @staticmethod
    def move_file(source, destiny):
        if not os.path.exists(os.path.dirname(destiny)):
            os.makedirs(os.path.dirname(destiny))
        shutil.move(source, destiny)

    @staticmethod
    def execute_shell_command(line, log_level=Log.DEBUG):
        if isinstance(line, basestring):
            line = line.split()
        process = subprocess.Popen(line, stdout=subprocess.PIPE)
        for line in process.communicate():
            if not line:
                continue
            Log.log.log(log_level, line)
        if process.returncode != 0:
            raise Exception('Error executing {0}\n Return code: {1}', ' '.join(line), process.returncode)
        return line


class TempFile(object):

    autoclean = True
    files = list()
    scratch_folder = os.path.join('/scratch', pwd.getpwuid(os.getuid())[0])
    prefix = 'temp'

    @staticmethod
    def get(filename=None, clean=True):
        if filename:
            path = os.path.join(TempFile.scratch_folder, filename)
        else:
            path = tempfile.mkstemp(dir=TempFile.scratch_folder, prefix=TempFile.prefix, suffix='.nc')[1]

        if clean:
            TempFile.files.append(path)

        return path

    @staticmethod
    def clean():
        for temp_file in TempFile.files:
            if os.path.exists(temp_file):
                os.remove(temp_file)
        TempFile.files = list()