cdist/cdist/core/python_type.py

import logging
import os
import io
import sys
import re
from cdist import message, Error


class PythonType:
    def __init__(self, env, cdist_object, local, remote, message_prefix=None):
        self.env = env
        self.cdist_object = cdist_object
        self.object_id = cdist_object.object_id
        self.object_name = cdist_object.name
        self.cdist_type = cdist_object.cdist_type
        self.local = local
        self.remote = remote
        self.object_path = cdist_object.absolute_path
        self.type_path = cdist_object.cdist_type.absolute_path
        self.explorer_path = os.path.join(self.object_path, 'explorer')
        self.parameters = cdist_object.parameters
        self.stdin_path = os.path.join(self.object_path, 'stdin')
        self.log = logging.getLogger(
            self.local.target_host[0] + ':' + self.object_name)

        self.message_prefix = message_prefix
        self.message = None

    def get_parameter(self, name):
        return self.parameters.get(name)

    def get_explorer_file(self, name):
        path = os.path.join(self.explorer_path, name)
        return path

    def get_explorer(self, name):
        path = self.get_explorer_file(name)
        with open(path, 'r') as f:
            value = f.read()
            if value:
                value = value.strip()
            return value

    def run_local(self, command, env=None):
        rv = self.local.run(command, env=env, return_output=True)
        if rv:
            rv = rv.rstrip('\n')
        return rv

    def run_remote(self, command, env=None):
        rv = self.remote.run(command, env=env, return_output=True)
        if rv:
            rv = rv.rstrip('\n')
        return rv

    def transfer(self, source, destination):
        self.remote.transfer(source, destination)

    def die(self, msg):
        raise Error("{}: {}".format(self.cdist_object, msg))

    def type_manifest(self):
        pass

    def type_gencode(self):
        pass

    def manifest(self, stdout=None, stderr=None):
        try:
            if self.message_prefix:
                self.message = message.Message(self.message_prefix,
                                               self.local.messages_path)
                self.env.update(self.message.env)
            if stdout is not None:
                stdout_save = sys.stdout
                sys.stdout = stdout
            if stderr is not None:
                stderr_save = sys.stderr
                sys.stderr = stderr
            yield from self.type_manifest()
        finally:
            if self.message:
                self.message.merge_messages()
            if stdout is not None:
                sys.stdout = stdout_save
            if stderr is not None:
                sys.stderr = stderr_save

    def run(self, stdout=None, stderr=None):
        try:
            if self.message_prefix:
                self.message = message.Message(self.message_prefix,
                                               self.local.messages_path)
            if stdout is not None:
                stdout_save = sys.stdout
                sys.stdout = stdout
            if stderr is not None:
                stderr_save = sys.stderr
                sys.stderr = stderr
            return self.type_gencode()
        finally:
            if self.message:
                self.message.merge_messages()
            if stdout is not None:
                sys.stdout = stdout_save
            if stderr is not None:
                sys.stderr = stderr_save

    def send_message(self, msg):
        if self.message:
            with open(self.message.messages_out, 'a') as f:
                print(msg, file=f)

    def receive_message(self, pattern):
        if self.message:
            with open(self.message.messages_in, 'r') as f:
                for line in f:
                    match = re.search(pattern, line)
                    if match:
                        return match
        return None


class ManifestEntry:
    def __init__(self, name, stdin=None, parameters=None):
        self.name = name
        if parameters is None:
            self.parameters = {}
        else:
            self.parameters = parameters
        self.set_stdin(stdin)

    def set_stdin(self, value):
        # If file-like object then read its value.
        if value is not None and isinstance(value, io.IOBase):
            value = value.read()

        # Convert to bytes file-like object.
        if value is None:
            self.stdin = None
        elif isinstance(value, str):
            self.stdin = io.BytesIO(value.encode('utf-8'))
        elif isinstance(value, bytes) or isinstance(value, bytearray):
            self.stdin = io.BytesIO(value)
        else:
            raise TypeError("value must be str, bytes, bytearray, file-like "
                            "object or None")

    def cmd_line(self):
        argv = [self.name, ]
        for param in self.parameters:
            argv.append(param)
            val = self.parameters[param]
            if val:
                argv.append(val)
        return argv

    def __repr__(self):
        return '<ManifestEntry name={}, parameters={}, stdin={}>'.format(
            self.name, self.parameters, self.stdin)