ostools.py

#******************************************************************************
# * Copyright (c) 2019, XtremeDV. All rights reserved.
# *
# * Licensed under the Apache License, Version 2.0 (the "License");
# * you may not use this file except in compliance with the License.
# * You may obtain a copy of the License at
# *
# * http://www.apache.org/licenses/LICENSE-2.0
# *
# * Unless required by applicable law or agreed to in writing, software
# * distributed under the License is distributed on an "AS IS" BASIS,
# * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# * See the License for the specific language governing permissions and
# * limitations under the License.
# *
# * Author: Jude Zhang, Email: zhajio.1988@gmail.com
# *******************************************************************************
# Copyright (c) 2014-2018, Lars Asplund lars.anders.asplund@gmail.com
"""
Provides operating systems dependent functionality that can be easily
stubbed for testing
"""

from __future__ import print_function

import time
import subprocess
import threading
import psutil
import shutil
import sys
import signal
try:
    # Python 3.x
    from queue import Queue, Empty
except ImportError:
    # Python 2.7
    from Queue import Queue, Empty  # pylint: disable=import-error

from os.path import exists, getmtime, dirname, relpath, splitdrive
import os
import io

import logging
LOGGER = logging.getLogger(__name__)

class ProgramStatus(object):
    """
    Maintain global program status to support graceful shutdown
    """
    def __init__(self):
        self._lock = threading.Lock()
        self._shutting_down = False

    @property
    def is_shutting_down(self):
        with self._lock:  # pylint: disable=not-context-manager
            return self._shutting_down

    def check_for_shutdown(self):
        if self.is_shutting_down:
            raise KeyboardInterrupt

    def shutdown(self):
        with self._lock:  # pylint: disable=not-context-manager
            LOGGER.debug("ProgramStatus.shutdown")
            self._shutting_down = True

    def reset(self):
        with self._lock:  # pylint: disable=not-context-manager
            self._shutting_down = False


PROGRAM_STATUS = ProgramStatus()


class InterruptableQueue(object):
    """
    A Queue which can be interrupted
    """

    def __init__(self):
        self._queue = Queue()

    def get(self):
        """
        Get a value from the queue
        """
        while True:
            PROGRAM_STATUS.check_for_shutdown()
            try:
                return self._queue.get(timeout=0.1)
            except Empty:
                pass

    def put(self, value):
        self._queue.put(value)

    def empty(self):
        return self._queue.empty()


class Process(object):
    """
    A simple process interface which supports asynchronously consuming the stdout and stderr
    of the process while it is running.
    """

    class NonZeroExitCode(Exception):
        pass

    def __init__(self, cmd, cwd=None, env=None):
        self._cmd = cmd
        self._cwd = cwd

        # Create process with new process group
        # Sending a signal to a process group will send it to all children
        # Hopefully this way no orphaned processes will be left behind
        self._process = subprocess.Popen(
            self._cmd,
            cwd=self._cwd,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            shell=True,
            bufsize=0,
            # Create new process group on POSIX, setpgrp does not exist on Windows
            #preexec_fn=os.setsid)
            preexec_fn=os.setpgrp)  # pylint: disable=no-member
        LOGGER.debug("Started process with pid=%i: '%s'", self._process.pid, (" ".join(self._cwd)))

        self._queue = InterruptableQueue()
        self._reader = AsynchronousFileReader(self._process.stdout, self._queue)
        self._reader.start()

    def write(self, *args, **kwargs):
        """ Write to stdin """
        if not self._process.stdin.closed:
            self._process.stdin.write(*args, **kwargs)

    def writeline(self, line):
        """ Write a line to stdin """
        if not self._process.stdin.closed:
            self._process.stdin.write(line + "\n")
            self._process.stdin.flush()

    def next_line(self):
        """
        Return either the next line or the exit code
        """

        if not self._reader.eof():
            # Show what we received from standard output.
            msg = self._queue.get()

            if msg is not None:
                return msg

        retcode = self.wait()
        return retcode

    def wait(self):
        """
        Wait while without completely blocking to avoid
        deadlock when shutting down
        """
        while self._process.poll() is None:
            PROGRAM_STATUS.check_for_shutdown()
            time.sleep(0.05)
            LOGGER.debug("Waiting for process with pid=%i to stop", self._process.pid)
        return self._process.returncode

    def is_alive(self):
        """
        Returns true if alive
        """
        return self._process.poll() is None

    def consume_output(self, callback=print):
        """
        Consume the output of the process.
        The output is interpreted as UTF-8 text.

        @param callback Called for each line of output
        @raises Process.NonZeroExitCode when the process does not exit with code zero
        """

        def default_callback(*args, **kwargs):
            pass

        if not callback:
            callback = default_callback

        while not self._reader.eof():
            line = self._queue.get()
            if line is None:
                break
            if callback(line) is not None:
                return

        retcode = None
        while retcode is None:
            retcode = self.wait()
            if retcode != 0:
                raise Process.NonZeroExitCode

    def terminate(self):
        """
        Terminate the process
        """
        if self._process.poll() is None:
            process = psutil.Process(self._process.pid)
            proc_list = process.children(recursive=True)
            #proc_list.reverse()
            for proc in proc_list:
                proc.kill()
            process.kill()

        # Let's be tidy and join the threads we've started.
        if self._process.poll() is None:
            LOGGER.debug("Terminating process with pid=%i", self._process.pid)
            self._process.terminate()

        if self._process.poll() is None:
            time.sleep(0.05)

        if self._process.poll() is None:
            LOGGER.debug("Killing process with pid=%i", self._process.pid)
            self._process.kill()

        if self._process.poll() is None:
            LOGGER.debug("Waiting for process with pid=%i", self._process.pid)
            self.wait()

        LOGGER.debug("Process with pid=%i terminated with code=%i",
                     self._process.pid,
                     self._process.returncode)

        self._reader.join()
        self._process.stdout.close()
        self._process.stdin.close()

    def __del__(self):
        try:
            self.terminate()
        except KeyboardInterrupt:
            LOGGER.debug("Process.__del__: Ignoring KeyboardInterrupt")

class AsynchronousFileReader(threading.Thread):
    """
    Helper class to implement asynchronous reading of a file
    in a separate thread. Pushes read lines on a queue to
    be consumed in another thread.
    """

    def __init__(self, fd, queue, encoding="utf-8"):
        threading.Thread.__init__(self)

        # If Python 3 change encoding of TextIOWrapper to utf-8 ignoring decode errors
        if isinstance(fd, io.TextIOWrapper):
            fd = io.TextIOWrapper(fd.buffer, encoding=encoding, errors="ignore")

        self._fd = fd
        self._queue = queue
        self._encoding = encoding

    def run(self):
        """The body of the thread: read lines and put them on the queue."""
        for line in iter(self._fd.readline, ""):
            if PROGRAM_STATUS.is_shutting_down:
                break

            # Convert string into utf-8 if necessary
            if sys.version_info.major == 2:
                string = line[:-1].decode(encoding=self._encoding, errors="ignore")
            else:
                string = line[:-1]

            self._queue.put(string)
        self._queue.put(None)

    def eof(self):
        """Check whether there is no more content to expect."""
        return not self.is_alive() and self._queue.empty()


def read_file(file_name, encoding="utf-8", newline=None):
    """ To stub during testing """
    try:
        with io.open(file_name, "r", encoding=encoding, newline=newline) as file_to_read:
            data = file_to_read.read()
    except UnicodeDecodeError:
        LOGGER.warning("Could not decode file %s using encoding %s, ignoring encoding errors",
                       file_name, encoding)
        with io.open(file_name, "r", encoding=encoding, errors="ignore", newline=newline) as file_to_read:
            data = file_to_read.read()

    return data

def write_file(file_name, contents, encoding="utf-8"):
    """ To stub during testing """

    path = dirname(file_name)
    if path == "":
        path = "."

    if not file_exists(path):
        os.makedirs(path)

    with io.open(file_name, "wb") as file_to_write:
        file_to_write.write(contents.encode(encoding=encoding))


def file_exists(file_name):
    """ To stub during testing """
    return exists(file_name)


def get_modification_time(file_name):
    """ To stub during testing """
    return getmtime(file_name)


def get_time():
    """ To stub during testing """
    return time.time()


def renew_path(path):
    """
    Ensure path directory exists and is empty
    """
    if exists(path):
        shutil.rmtree(path)
    os.makedirs(path)


def simplify_path(path):
    """
    Return relative path towards current working directory
    unless it is a separate Windows drive
    """
    cwd = os.getcwd()
    drive_cwd = splitdrive(cwd)[0]
    drive_path = splitdrive(path)[0]
    if drive_path == drive_cwd:
        return relpath(path, cwd)

    return path