vortex-nwp 2.0.0b1__py3-none-any.whl

vortex/tools/addons.py

@@ -0,0 +1,297 @@
+"""
+Abstract classes for System addons.
+"""
+
+from collections import defaultdict
+
+from bronx.fancies import loggers
+from bronx.syntax.decorators import nicedeco
+import footprints
+
+from vortex.config import get_from_config_w_default
+from vortex.layout import contexts
+from vortex.tools.env import Environment
+from vortex.tools.systems import OSExtended
+
+logger = loggers.getLogger(__name__)
+
+#: No automatic export
+__all__ = []
+
+
+class Addon(footprints.FootprintBase):
+    """Root class for any :class:`Addon` system subclasses."""
+
+    _abstract = True
+    _collector = ('addon',)
+    _footprint = dict(
+        info = 'Default add-on',
+        attr = dict(
+            kind = dict(),
+            sh = dict(
+                type     = OSExtended,
+                alias    = ('shell',),
+                access   = 'rwx-weak',
+            ),
+            env = dict(
+                type     = Environment,
+                optional = True,
+                default  = None,
+                access   = 'rwx',
+                doc_visibility = footprints.doc.visibility.ADVANCED
+            ),
+            cfginfo = dict(
+                optional = True,
+                default  = '[kind]',
+                doc_visibility = footprints.doc.visibility.ADVANCED
+            ),
+            cmd = dict(
+                optional = True,
+                default  = None,
+                access   = 'rwx',
+            ),
+            path = dict(
+                optional = True,
+                default  = None,
+                access   = 'rwx',
+            ),
+            cycle = dict(
+                optional = True,
+                default  = None,
+                access   = 'rwx',
+            ),
+            toolkind = dict(
+                optional = True,
+                default  = None
+            ),
+        )
+    )
+
+    def __init__(self, *args, **kw):
+        """Abstract Addon initialisation."""
+        logger.debug('Abstract Addon init %s', self.__class__)
+        super().__init__(*args, **kw)
+        self.sh.extend(self)
+        self._context_cache = defaultdict(dict)
+        self._cmd_xperms_cache = set()
+        if self.env is None:
+            self.env = Environment(active=False, clear=True)
+        clsenv = self.__class__.__dict__
+        for k in [x for x in clsenv.keys() if x.isupper()]:
+            self.env[k] = clsenv[k]
+        if self.path is None:
+            self.path = get_from_config_w_default(
+                section="nwp-tools", key=self.kind, default=None,
+            )
+
+    @classmethod
+    def in_shell(cls, shell):
+        """Grep any active instance of that class in the specified shell."""
+        lx = [x for x in shell.search if isinstance(x, cls)]
+        return lx[0] if lx else None
+
+    def _query_context(self):
+        """Return the path and cmd for the current context.
+
+        Results are cached so that the context's localtracker is explored only once.
+
+        .. note:: We use the localtracker instead of the sequence because, in
+            multistep jobs, the localtracker is preserved between steps. It's
+            less elegant but it plays nice with MTOOL.
+        """
+        ctxtag = contexts.Context.tag_focus()
+        if ctxtag not in self._context_cache and self.toolkind is not None:
+            ltrack = contexts.current().localtracker
+            # NB: 'str' is important because local might be in unicode...
+            candidates = [str(self.sh.path.realpath(local))
+                          for local, entry in ltrack.items()
+                          if (entry.latest_rhdict('get').get('resource', dict()).get('kind', '') ==
+                              self.toolkind)]
+            if candidates:
+                realpath = candidates.pop()
+                self._context_cache[ctxtag] = dict(path=self.sh.path.dirname(realpath),
+                                                   cmd=self.sh.path.basename(realpath))
+        return self._context_cache[ctxtag]
+
+    @property
+    def actual_path(self):
+        """The path that should be used in the current context."""
+        infos = self._query_context()
+        ctxpath = infos.get('path', None)
+        return self.path if ctxpath is None else ctxpath
+
+    @property
+    def actual_cmd(self):
+        """The cmd that should be used in the current context."""
+        infos = self._query_context()
+        ctxcmd = infos.get('cmd', None)
+        return self.cmd if ctxcmd is None else ctxcmd
+
+    def _spawn_commons(self, cmd, **kw):
+        """Internal method setting local environment and calling standard shell spawn."""
+
+        # Is there a need for an interpreter ?
+        if 'interpreter' in kw:
+            cmd.insert(0, kw.pop('interpreter'))
+        else:
+            # The first element of the command line needs to be executable
+            if cmd[0] not in self._cmd_xperms_cache:
+                self._cmd_xperms_cache.add(cmd[0])
+                self.sh.xperm(cmd[0], force=True)
+
+        # Overwrite global module env values with specific ones
+        with self.sh.env.clone() as localenv:
+
+            localenv.verbose(True, self.sh)
+            localenv.update(self.env)
+
+            # Check if a pipe is requested
+            inpipe = kw.pop('inpipe', False)
+
+            # Ask the attached shell to run the addon command
+            if inpipe:
+                kw.setdefault('stdout', True)
+                rc = self.sh.popen(cmd, **kw)
+            else:
+                rc = self.sh.spawn(cmd, **kw)
+
+        return rc
+
+    def _spawn(self, cmd, **kw):
+        """Internal method setting local environment and calling standard shell spawn."""
+
+        # Insert the actual tool command as first argument
+        cmd.insert(0, self.actual_cmd)
+        if self.actual_path is not None:
+            cmd[0] = self.actual_path + '/' + cmd[0]
+
+        return self._spawn_commons(cmd, **kw)
+
+    def _spawn_wrap(self, cmd, **kw):
+        """Internal method setting local environment and calling standard shell spawn."""
+
+        # Insert the tool path before the first argument
+        if self.actual_path is not None:
+            cmd[0] = self.actual_path + '/' + cmd[0]
+
+        return self._spawn_commons(cmd, **kw)
+
+
+class FtrawEnableAddon(Addon):
+    """Root class for any :class:`Addon` system subclasses that needs to override rawftput."""
+
+    _abstract = True
+    _footprint = dict(
+        info = 'Default add-on with rawftput support.',
+        attr = dict(
+            rawftshell = dict(
+                info     = "Path to ftserv's concatenation shell",
+                optional = True,
+                default  = None,
+                access   = 'rwx',
+                doc_visibility = footprints.doc.visibility.GURU,
+            ),
+        )
+    )
+
+    def __init__(self, *args, **kw):
+        """Abstract Addon initialisation."""
+        logger.debug('Abstract Addon init %s', self.__class__)
+        super().__init__(*args, **kw)
+        # If needed, look in the config file for the rawftshell
+        if self.rawftshell is None:
+            self.rawftshell = get_from_config_w_default(
+                section="rawftshell", key=self.kind, default=None,
+            )
+
+
+class AddonGroup(footprints.FootprintBase):
+    """Root class for any :class:`AddonGroup` system subclasses.
+
+    An AddonGroup is not really an Addon... it just loads a bunch of other
+    Addons or AddonGroups into the current shell.
+    """
+
+    _abstract = True
+    _collector = ('addon',)
+    _footprint = dict(
+        info = 'Default add-on group',
+        attr = dict(
+            kind = dict(),
+            sh = dict(
+                type     = OSExtended,
+                alias    = ('shell',),
+            ),
+            env = dict(
+                type     = Environment,
+                optional = True,
+                default  = None,
+                doc_visibility = footprints.doc.visibility.ADVANCED,
+            ),
+            cycle = dict(
+                optional = True,
+                default  = None,
+            ),
+            verboseload = dict(
+                optional = True,
+                default  = True,
+                type     = bool,
+            ),
+        )
+    )
+
+    _addonslist = None
+
+    def __init__(self, *args, **kw):
+        """Abstract Addon initialisation."""
+        logger.debug('Abstract Addon init %s', self.__class__)
+        super().__init__(*args, **kw)
+        self._addons_load()
+
+    def _addons_load(self):
+        if self._addonslist is None:
+            raise RuntimeError("the _addonslist classe variable must be overriden.")
+        self._load_addons_from_list(self._addonslist)
+
+    def _load_addons_from_list(self, addons):
+        if self.verboseload:
+            logger.info("Loading the %s Addons group.", self.kind)
+        for addon in addons:
+            _shadd = footprints.proxy.addon(kind=addon, sh=self.sh, env=self.env,
+                                            cycle=self.cycle, verboseload=self.verboseload)
+            if self.verboseload:
+                logger.info("%s Addon is: %s", addon, repr(_shadd))
+
+
+def require_external_addon(*addons):
+    """
+    A method decorator usable in addons, that will check if addons listed in
+    **addons** are properly loaded in the parent System object.
+
+    If not, a :class:`RuntimeError` exception will be raised.
+    """
+    @nicedeco
+    def r_addon_decorator(method):
+
+        def decorated(self, *kargs, **kwargs):
+            # Create a cache in self... ugly but efficient !
+            if not hasattr(self, '_require_external_addon_check_cache'):
+                setattr(self, '_require_external_addon_check_cache', set())
+            ko_addons = set()
+            loaded_addons = None
+            for addon in addons:
+                if addon in self._require_external_addon_check_cache:
+                    continue
+                if loaded_addons is None:
+                    loaded_addons = self.sh.loaded_addons()
+                if addon in loaded_addons:
+                    self._require_external_addon_check_cache.add(addon)
+                else:
+                    ko_addons.add(addon)
+            if ko_addons:
+                raise RuntimeError('The following addons are needed to use the {:s} method: {:s}'
+                                   .format(method.__name__, ', '.join(ko_addons)))
+            return method(self, *kargs, **kwargs)
+
+        return decorated
+    return r_addon_decorator

vortex/tools/arm.py

@@ -0,0 +1,76 @@
+"""
+Module needed to work with ARM tools such as Forge.
+"""
+
+from bronx.fancies import loggers
+
+from vortex.util.config import load_template
+
+#: No automatic export
+__all__ = ['ArmForgeTool']
+
+logger = loggers.getLogger(__name__)
+
+
+class ArmForgeTool:
+    """Work with the ARM tools such as DDT & MAP."""
+
+    def __init__(self, ticket):
+        """
+        :param ticket: The current Vortex' session ticket.
+        """
+        self._t = ticket
+        self._sh = self._t.sh
+        self._config = self._sh.default_target.items('armtools')
+        self._ddtpath = self._sh.env.get('VORTEX_ARM_DDT_PATH', None)
+        self._mappath = self._sh.env.get('VORTEX_ARM_MAP_PATH', None)
+        self._forgedir = self._sh.env.get('VORTEX_ARM_FORGE_DIR',
+                                          self.config.get('forgedir', None))
+        self._forgeversion = self._sh.env.get('VORTEX_ARM_FORGE_VERSION',
+                                              self.config.get('forgeversion', 999999))
+        self._forgeversion = int(self._forgeversion)
+        if self._forgedir and self._ddtpath is None:
+            self._ddtpath = self._sh.path.join(self._forgedir, 'bin', 'ddt')
+        if self._forgedir and self._mappath is None:
+            self._mappath = self._sh.path.join(self._forgedir, 'bin', 'map')
+
+    @property
+    def config(self):
+        """The configuration dictionary."""
+        return self._config
+
+    @property
+    def ddtpath(self):
+        """The path to the DDT debuger executable."""
+        if self._ddtpath is None:
+            raise RuntimeError('DDT requested but the DDT path is not configured.')
+        return self._ddtpath
+
+    @property
+    def mappath(self):
+        """The path to the MAP profiler executable."""
+        if self._mappath is None:
+            raise RuntimeError('MAP requested but the MAP path is not configured.')
+        return self._mappath
+
+    def _dump_forge_session(self, sources=(), workdir=None):
+        """Create the ARM Forge's session file to list source directories."""
+        targetfile = 'armforge-vortex-session-file.ddt'
+        if workdir:
+            targetfile = self._sh.path.join(workdir, targetfile)
+        tpl = load_template(self._t, '@armforge-session-conf.tpl',
+                            encoding='utf-8', version=self._forgeversion)
+        sconf = tpl.substitute(sourcedirs='\n'.join(['        <directory>{:s}</directory>'.format(d)
+                                                     for d in sources]))
+        with open(targetfile, 'w') as fhs:
+            fhs.write(sconf)
+        return targetfile
+
+    def ddt_prefix_cmd(self, sources=(), workdir=None):
+        """Generate the prefix command required to start DDT."""
+        if sources:
+            return [self.ddtpath,
+                    '--session={:s}'.format(self._dump_forge_session(sources, workdir=workdir)),
+                    '--connect']
+        else:
+            return [self.ddtpath, '--connect']

vortex/tools/compression.py

@@ -0,0 +1,322 @@
+"""
+Stream/File compression tools.
+
+The user interface for such tools is the :class:`CompressionPipeline`.
+"""
+
+from contextlib import contextmanager
+import io
+import functools
+import operator
+
+import footprints
+from bronx.fancies import loggers
+from vortex.util.iosponge import IoSponge
+
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+class CompressionPipeline:
+    """Main interface to data compression algorithms."""
+
+    def __init__(self, system, compression=''):
+        """
+        :param System system: The system object that will be used to carry out
+            the task.
+        :param str compression: The description of the compression tools that
+            will be used in this compression pipeline.
+            (e.g. 'gzip' will just compress/uncompress using the gzip software,
+            'gzip|bzip2' will use the gzip software piped to the bzip2 software
+            (which is useless), 'gzip&complevel=5' will use the gzip software
+            with a compression factor of 5.)
+
+        :note: See the subclasses of :class:`CompressionUnit` for a description
+            of available compression tools (and their options).
+        """
+        self._units = list()
+        self._sh = system
+        self.description_string = compression
+        for c in [c for c in compression.split('|') if c]:
+            c_raw = c.split('&')
+            ckind = c_raw.pop(0)
+            cargs = dict([arg.split('=', 1) for arg in c_raw])
+            self.add_compression_unit(ckind, **cargs)
+
+    def add_compression_unit(self, unit, **kwargs):
+        """Add a new compression tool to the compression pipeline.
+
+        :param str unit: The kind of the compression tool (see  :class:`CompressionUnit`
+            subclases
+        :param kwargs: Options that will be used during the compression tool
+            initialisation
+        """
+        c_unit = footprints.proxy.compression_unit(kind=unit, **kwargs)
+        if c_unit is None:
+            raise ValueError("The {:s} compression unit could not be found.".
+                             format(unit))
+        self._units.append(c_unit)
+
+    @property
+    def units(self):
+        """The list of compression tools forming the compression pipeline."""
+        return self._units
+
+    @property
+    def _rawftp_shell(self):
+        """The name of the corresponding rawftp specialshell (if relevant)."""
+        if len(self.units) == 1:
+            return self.units[0].rawftp_shell(self._sh)
+        else:
+            return None
+
+    @property
+    def suffix(self):
+        """The suffix usualy associated with this compression pipeline."""
+        s = '.'.join([s.suffix for s in self.units])
+        return '.' + s if s else ''
+
+    @property
+    def compression_factor(self):
+        """The minimal compression factor expected with such a compression pipeline."""
+        return functools.reduce(operator.mul, [s.cfactor for s in self.units], 1.)
+
+    @staticmethod
+    def _inputstream_size(stream):
+        """Find out the size of a seekable input stream."""
+        estimated_size = 0
+        try:
+            stream.seek(0, io.SEEK_END)
+            estimated_size = stream.tell()
+            stream.seek(0)
+        except AttributeError:
+            logger.warning('Could not rewind <source:%s>', str(stream))
+        except OSError:
+            logger.debug('Seek trouble <source:%s>', str(stream))
+        return estimated_size
+
+    @contextmanager
+    def _openstream(self, local, mode='rb'):
+        """If *local* is not an opened file, open it..."""
+        if isinstance(local, str):
+            localfh = open(local, mode)
+            yield localfh
+            localfh.close()
+        elif isinstance(local, io.IOBase):
+            yield local
+        else:
+            raise ValueError("Unknown type for {!s}".format(local))
+
+    def _genericstream_close(self, processes):
+        """Close a list of Popen objects (and look for the returncode)."""
+        for i, p in enumerate(processes):
+            if not self._sh.pclose(p):
+                logger.error("Abnormal return code for one of the processes (#%d)", i)
+
+    @contextmanager
+    def compress2stream(self, local, iosponge=False):
+        """Compress *local* into a pipe or an :class:`IoSponge` object.
+
+        *local* can be an opened file-like object or a filename.
+
+        This method creates a context manager. Example::
+
+            source='myfile'
+            cp = CompressionPipeline(systemobj, 'gzip')
+            ftp = systemobj.ftp('hendrix.meteo.fr')
+            with cp.compress2stream(source) as csource:
+                ftp.put(csource, 'remote_compressedfile')
+
+        When leaving the context, the gzip process that compresses the data will
+        be properly closed
+        """
+        with self._openstream(local) as stream:
+            estimated_size = self._inputstream_size(stream) * self.compression_factor
+            processes = list()
+            lstream = stream
+            for unit in self.units:
+                p = unit.compress(self._sh, lstream)
+                lstream = p.stdout
+                processes.append(p)
+            if iosponge:
+                yield IoSponge(lstream, guessed_size=estimated_size)
+            else:
+                yield lstream
+            self._genericstream_close(processes)
+
+    def _xcopyfileobj(self, in_fh, out_fh):
+        try:
+            self._sh.copyfileobj(in_fh, out_fh)
+        except OSError:
+            return False
+        else:
+            return True
+
+    def compress2file(self, local, destination):
+        """Compress *local* into a file (named *destination*)
+
+        *local* can be an opened file-like object or a filename.
+        *destination* is a filename.
+        """
+        with open(destination, 'wb') as fhout:
+            with self.compress2stream(local) as fhcompressed:
+                return self._xcopyfileobj(fhcompressed, fhout)
+
+    def compress2rawftp(self, local):
+        """
+        Return the name of the rawftp's specialshell that can be used to
+        compress the *local* data. It might return None.
+        """
+        return self._rawftp_shell
+
+    @contextmanager
+    def stream2uncompress(self, destination):
+        """Uncompress piped data to *destination*.
+
+        *destination* can be an opened file-like object or a filename.
+
+        This method creates a context manager. Example::
+
+            destination='mydownloadedfile'
+            cp = CompressionPipeline(systemobj, 'gzip')
+            ftp = systemobj.ftp('hendrix.meteo.fr')
+            with cp.stream2uncompress(destination) as cdestination:
+                ftp.get('remote_compressedfile', cdestination)
+
+        When leaving the context, the gunzip process that uncompresses the data
+        will be properly closed.
+        """
+        with self._openstream(destination, 'wb') as dstream:
+            processes = list()
+            instream = True
+            nunits = len(self.units)
+            for i, unit in enumerate(reversed(self.units)):
+                outstream = dstream if i == nunits - 1 else True
+                p = unit.uncompress(self._sh, instream, outstream)
+                instream = p.stdout
+                processes.append(p)
+            yield processes[0].stdin
+            self._genericstream_close(processes)
+
+    def file2uncompress(self, local, destination):
+        """Uncompress *local* into *destination*.
+
+        *local* is a filename.
+        *destination* can be an opened file-like object or a filename.
+        """
+        with self.stream2uncompress(destination) as fhuncompressed:
+            with open(local, 'rb') as fhcompressed:
+                return self._xcopyfileobj(fhcompressed, fhuncompressed)
+
+
+class CompressionUnit(footprints.FootprintBase):
+    """Defines compress/uncompress methods for a given compression tool."""
+
+    _abstract = True
+    _collector = ('compression_unit',)
+    _footprint = dict(
+        info = 'Abstract Compression Unit',
+        attr = dict(
+            kind = dict(
+                info = "The name of the compression tool.",
+            ),
+            suffix = dict(
+                info = "The usual file extension for this compression tool.",
+                optional = True,
+            ),
+            cfactor = dict(
+                info = "The usual compression factor for this compression tool.",
+                type = float,
+                default = 1.,
+                optional =True,
+            ),
+        ),
+    )
+
+    def rawftp_shell(self, sh):
+        """The rawftp's speciall shell that may carry out a comparable compression."""
+        return None
+
+    def _run_in_pipe(self, sh, cmd, stream, outstream=True):
+        """Run *cmd* with the piped input *stream*."""
+        p = sh.popen(cmd, stdin=stream, stdout=outstream, bufsize=8192)
+        return p
+
+    def compress(self, sh, stream):
+        """Compress the input *stream*. Returns a Popen object."""
+        raise NotImplementedError()
+
+    def uncompress(self, sh, stream, outstream=True):
+        """Uncompress the input *stream*. Returns a Popen object."""
+        raise NotImplementedError()
+
+
+class GzipCompressionUnit(CompressionUnit):
+
+    _footprint = dict(
+        info = 'Compress/Uncompress a stream using gzip',
+        attr = dict(
+            kind = dict(
+                values = ['gzip', 'gz']
+            ),
+            suffix = dict(
+                default = 'gz',
+            ),
+            complevel = dict(
+                info = "The gzip algorithm compression level (see 'man gzip')",
+                type = int,
+                values = range(1, 10),
+                default = 6,
+                optional = True
+            ),
+            cfactor = dict(
+                default = 0.9
+            ),
+        ),
+    )
+
+    def compress(self, sh, stream):
+        """Compress the input *stream*. Returns a Popen object."""
+        return self._run_in_pipe(sh, ['gzip', '--stdout', '-{!s}'.format(self.complevel)],
+                                 stream)
+
+    def uncompress(self, sh, stream, outstream=True):
+        """Uncompress the input *stream*. Returns a Popen object."""
+        return self._run_in_pipe(sh, ['gunzip', '--stdout'], stream, outstream)
+
+
+class Bzip2CompressionUnit(CompressionUnit):
+
+    _footprint = dict(
+        info = 'Compress/Uncompress a stream using bzip2',
+        attr = dict(
+            kind = dict(
+                values = ['bzip2', 'bz2']
+            ),
+            suffix = dict(
+                default = 'bz2',
+            ),
+            complevel = dict(
+                info = "The bzip2 algorithm compression level (see 'man bzip2')",
+                type = int,
+                values = range(1, 10),
+                default = 9,
+                optional = True
+            ),
+            cfactor = dict(
+                default = 0.85,
+            ),
+        ),
+    )
+
+    def compress(self, sh, stream):
+        """Compress the input *stream*. Returns a Popen object."""
+        return self._run_in_pipe(sh, ['bzip2', '--stdout', '-{!s}'.format(self.complevel)],
+                                 stream)
+
+    def uncompress(self, sh, stream, outstream=True):
+        """Uncompress the input *stream*. Returns a Popen object."""
+        return self._run_in_pipe(sh, ['bunzip2', '--stdout'], stream, outstream)