<commit_before># Copyright (c) 2013, Ultrabug
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# includes
################################################################################
import os
import imp
import sys
import argparse
import threading
from threading import Thread

from json import loads
from json import dumps

from time import time
from time import sleep

from datetime import datetime
from datetime import timedelta

from signal import signal
from signal import SIGUSR1

from subprocess import Popen
from subprocess import PIPE
from subprocess import call

from syslog import syslog
from syslog import LOG_ERR
from syslog import LOG_INFO

try:
    # python3
    from queue import Queue
    from queue import Empty
except ImportError:
    # python2
    from Queue import Queue
    from Queue import Empty

# module globals and defaults
################################################################################
CACHE_TIMEOUT = 60
DISABLE_TRANSFORM = False
I3STATUS_CONFIG = '/etc/i3status.conf'
INCLUDE_PATH = '.i3/py3status'
INTERVAL = 1

# functions
################################################################################
def print_line(message):
    """
    Non-buffered printing to stdout
    """
    sys.stdout.write(message + '\n')
    sys.stdout.flush()

def read_line():
    """
    Interrupted respecting reader for stdin
    """
    try:
        line = sys.stdin.readline().strip()
        # i3status sends EOF, or an empty line
        if not line:
            sys.exit(3)
        return line
    except KeyboardInterrupt:
        sys.exit()

def print_output(prefix, output_list, modules_cache):
    """
    Merge current user-classes cache with i3status' output and display on i3bar
    """
    for module in modules_cache:
        index, json = modules_cache[module]
        for j in output_list:
            if j['name'] == json['name']:
                j.update(json)
                break
        else:
            output_list.insert(index, json)

    output = prefix+dumps(output_list)
    print_line(output)
    return output

def i3status_config_reader(config_file):
    """
    i3status.conf reader so we can adapt our code to the i3status config
    """
    in_time = False
    in_general = False
    config = {
        'colors': False,
        'color_good': None,
        'color_bad' : None,
        'color_degraded' : None,
        'color_separator': None,
        'interval': 5,
        'output_format': None,
        'time_format': '%Y-%m-%d %H:%M:%S',
        }
    for line in open(config_file, 'r'):
        line = line.strip(' \t\n\r')
        if line.startswith('general'):
            in_general = True
        elif line.startswith('time'):
            in_time = True
        elif line.startswith('}'):
            in_general = False
            in_time = False
        if in_general and '=' in line:
            key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
            if key in config:
                if value in ['true', 'false']:
                    value = 'True' if value == 'true' else 'False'
                config[key] = eval(value)
        if in_time and '=' in line:
            key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
            if 'time_' + key in config:
                config['time_' + key] = eval(value)
    return config

class I3status(Thread):
    """
    Run i3status in a thread and send its output to a Queue
    """
    def __init__(self, config_file):
        """
        set our useful properties
        """
        Thread.__init__(self)
        self.config_file = config_file
        self.kill = False
        self.queue = Queue()
        self.init = True
        self.started = False

    def stop(self):
        """
        break the i3status loop
        """
        self.kill = True

    def run(self):
        """
        run i3status and queue its output
        """
        i3status_pipe = Popen(
            ['i3status', '-c', self.config_file],
            stdout=PIPE,
            stderr=PIPE,
            )
        self.queue.put(i3status_pipe.stdout.readline())
        self.queue.put(i3status_pipe.stdout.readline())
        while not self.kill:
            line = i3status_pipe.stdout.readline()
            if len(line) > 0:
                self.queue.put(line)
                try:
                    line = line.decode()
                except Exception:
                    pass
                if len(line) > 1 and line.startswith('['):
                    self.init = False
                elif line.startswith(',['):
                    self.started = True
            else:
                break

def process_line(line, **kwargs):
    """
    Main line processor logic
    """
    j = []
    prefix = ''
    if line.startswith('{') and 'version' in line:
        print_line(line.strip('\n'))
    elif line == '[\n':
        print_line(line.strip('\n'))
    else:
        if line.startswith(','):
            line, prefix = line[1:], ','
        elif kwargs['delta'] > 0:
            prefix = ','

        # integrated transformations
        if not DISABLE_TRANSFORM:
            j = transform(loads(line), **kwargs)
        else:
            j = loads(line)

    return (prefix, j)

def transform(j, **kwargs):
    """
    Integrated transformations:
    - update the 'time' object so that it's updated at INTERVAL seconds
    """
    try:
        for item in j:
            # time modification
            if item['name'] in [ 'time', 'tztime' ]:
                time_format = I3STATUS_CONFIG['time_format']
                date = datetime.strptime( item['full_text'], time_format ) \
                    + timedelta(seconds=kwargs['delta'])
                item['full_text'] = date.strftime(time_format)
                if kwargs['delta'] > 0:
                    item['transformed'] = True
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "transformation failed (%s)" % (str(err)))
    return j

class UserModules(Thread):
    """
    This thread executes and caches all of the user's own Py3status classes.
    We run it in a separate thread so that any delay introduced by an external
    class doesn't affect the i3bar updating process.
    """
    def __init__(self, cache_timeout, i3status_conf, include_path, interval):
        """
        set our useful properties
        """
        Thread.__init__(self)
        self.cache = {}
        self.cache_timeout = cache_timeout
        self.classes = {}
        self.i3status_conf = i3status_conf
        self.include_path = include_path
        self.interval = interval
        self.i3status_json = {}
        self.kill = False

    @staticmethod
    def load_from_file(filepath):
        """
        return Py3status user-written classes
        """
        inst = None
        expected_class = 'Py3status'
        mod_name, file_ext = os.path.splitext(os.path.split(filepath)[-1])
        if file_ext.lower() == '.py':
            py_mod = imp.load_source(mod_name, filepath)
            if hasattr(py_mod, expected_class):
                inst = py_mod.Py3status()
        return (mod_name, inst)

    def load_files(self):
        """
        read user-written Py3status class files for dynamic inclusion
        """
        if self.include_path and os.path.isdir(self.include_path):
            for file_name in os.listdir(self.include_path):
                try:
                    module, class_inst = self.load_from_file(
                        self.include_path + file_name
                        )
                    if module and class_inst:
                        self.classes[file_name] = (class_inst, [])
                        for method in dir(class_inst):
                            if not method.startswith('__'):
                                self.classes[file_name][1].append(method)
                except Exception:
                    err = sys.exc_info()[1]
                    syslog(LOG_ERR, "loading %s failed (%s)" \
                        % (file_name, str(err)))

    def execute_classes(self):
        """
        run on every user class included and execute every method on the json,
        then cache it for later usage
        """
        for class_name in sorted( self.classes.keys() ):
            my_class, my_methods = self.classes[class_name]
            for my_method in my_methods:
                try:
                    # handle a cache on user class methods results
                    try:
                        index, result = self.cache[my_method]
                        if time() > result['cached_until']:
                            raise KeyError('cache timeout')
                    except KeyError:
                        # execute the method
                        try:
                            meth = getattr(my_class, my_method)
                            index, result = meth(
                                self.i3status_json,
                                self.i3status_conf,
                                )
                        except Exception:
                            err = sys.exc_info()[1]
                            syslog(LOG_ERR, "user method %s failed (%s)" \
                                % (my_method, str(err)))
                            index, result = (0, {'name': '', 'full_text': ''})

                        # respect user-defined cache timeout for this module
                        if 'cached_until' not in result:
                            result['cached_until'] = time() + self.cache_timeout

                        # validate the response
                        assert isinstance(result, dict), "should return a dict"
                        assert 'full_text' in result, "missing 'full_text' key"
                        assert 'name' in result, "missing 'name' key"
                    finally:
                        self.cache[my_method] = (index, result)
                except Exception:
                    err = sys.exc_info()[1]
                    syslog(LOG_ERR, "injection failed (%s)" % str(err))

    def stop(self):
        """
        break this thread's loop
        """
        self.kill = True

    def clear(self):
        """
        clear user cache
        """
        self.cache = {}

    def run(self):
        """
        periodically execute user classes
        """
        self.load_files()
        while not self.kill:
            self.execute_classes()
            sleep(self.interval)

# main stuff
################################################################################
def main():
    """
    Main logic function
    """
    try:
        # global definition
        global CACHE_TIMEOUT, DISABLE_TRANSFORM
        global I3STATUS_CONFIG, INCLUDE_PATH, INTERVAL

        # command line options
        parser = argparse.ArgumentParser(
            description='The agile, python-powered, i3status wrapper')
        parser = argparse.ArgumentParser(add_help=True)
        parser.add_argument('-c', action="store",
            dest="i3status_conf", type=str,
            default=I3STATUS_CONFIG, help="path to i3status config file")
        parser.add_argument('-d', action="store_true",
            dest="disable_transform", help="disable integrated transformations")
        parser.add_argument('-i', action="store",
            dest="include_path", type=str,
            default=INCLUDE_PATH, help="user-based class include directory")
        parser.add_argument('-n', action="store",
            dest="interval", type=int,
            default=INTERVAL, help="update interval in seconds (default 1 sec)")
        parser.add_argument('-t', action="store",
            dest="cache_timeout", type=int, default=CACHE_TIMEOUT,
            help="default injection cache timeout in seconds (default 60 sec)")
        options = parser.parse_args()

        # configuration and helper variables
        CACHE_TIMEOUT = options.cache_timeout
        DISABLE_TRANSFORM = True if options.disable_transform else False
        I3STATUS_CONFIG = i3status_config_reader(options.i3status_conf)
        INCLUDE_PATH = os.path.abspath( options.include_path ) + '/'
        INTERVAL = options.interval

        # py3status uses only the i3bar protocol
        assert I3STATUS_CONFIG['output_format'] == 'i3bar', \
            'i3status output_format should be set to "i3bar"'
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "py3status init error (%s)" % str(err))
        sys.exit(1)

    try:
        # load user defined modules from a separate thread so that they cannot
        # cause any delay in updating the i3bar
        modules_thread = UserModules(
            CACHE_TIMEOUT,
            I3STATUS_CONFIG,
            INCLUDE_PATH,
            INTERVAL,
            )
        modules_thread.start()

        # spawn a i3status process on a separate thread
        # we will receive its output on a Queue which we will poll for messages
        i3status_thread = I3status(options.i3status_conf)
        i3status_thread.start()

        # SIGUSR1 forces a refresh of the bar both for py3status and i3status,
        # this mimics the USR1 signal handling of i3status (see man i3status)
        def sig_handler(signum, frame):
            """
            Raise a Warning level exception when a user sends a SIGUSR1 signal
            """
            raise UserWarning("received USR1, forcing refresh")
        signal(SIGUSR1, sig_handler)

        # control variables
        forced = False

        # main loop
        while True:
            try:
                # get a timestamp of now
                tst = time()

                # try to read i3status' output for at most INTERVAL seconds else
                # raise the Empty exception
                line = i3status_thread.queue.get(timeout=INTERVAL)
                try:
                    # python3 compatibility code
                    line = line.decode()
                except UnicodeDecodeError:
                    pass

                # i3status first output lines should be processed asap as only
                # the following lines will be processed by py3status
                if i3status_thread.started:
                    if not forced and I3STATUS_CONFIG['interval'] > INTERVAL:
                        # add a calculated sleep honoring py3status refresh
                        # time of the bar every INTERVAL seconds
                        sleep(INTERVAL - float( '{:.2}'.format( time()-tst ) ))
                    else:
                        # reset the SIGUSR1 flag forcing the refresh of the bar
                        forced = False

                # process the output line straight from i3status
                prefix, output_list = process_line(line, delta=0)
                if output_list:
                    modules_thread.i3status_json = output_list
                    print_output(
                        prefix,
                        output_list,
                        modules_thread.cache
                        )
            except Empty:
                # make sure i3status has started before modifying its output
                if i3status_thread.started or not i3status_thread.init:
                    prefix, output_list = process_line(line, delta=INTERVAL)
                    if output_list:
                        modules_thread.i3status_json = output_list
                        line = print_output(
                            prefix,
                            output_list,
                            modules_thread.cache
                            )
                else:
                    syslog(LOG_INFO, "waiting for i3status")
            except UserWarning:
                # SIGUSR1 was received, we also force i3status to refresh by
                # sending it a SIGUSR1 as well then we refresh the bar asap
                msg = sys.exc_info()[1]
                syslog(LOG_INFO, str(msg))
                call(["killall", "-s", "USR1", "i3status"])
                modules_thread.clear()
                forced = True
            except KeyboardInterrupt:
                break
            finally:
                if threading.active_count() < 3:
                    # i3status or modules thread died ? oops
                    break

        i3status_thread.stop()
        modules_thread.stop()
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "py3status error (%s)" % str(err))
        sys.exit(2)

if __name__ == '__main__':
    main()
<commit_msg>Handle killing of user-written modules nicely by calling their kill method if implemented upon exit of py3status, ignore kill methods execution when running normally<commit_after># Copyright (c) 2013, Ultrabug
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# includes
################################################################################
import os
import imp
import sys
import argparse
import threading
from threading import Thread

from json import loads
from json import dumps

from time import time
from time import sleep

from datetime import datetime
from datetime import timedelta

from signal import signal
from signal import SIGUSR1

from subprocess import Popen
from subprocess import PIPE
from subprocess import call

from syslog import syslog
from syslog import LOG_ERR
from syslog import LOG_INFO

try:
    # python3
    from queue import Queue
    from queue import Empty
except ImportError:
    # python2
    from Queue import Queue
    from Queue import Empty

# module globals and defaults
################################################################################
CACHE_TIMEOUT = 60
DISABLE_TRANSFORM = False
I3STATUS_CONFIG = '/etc/i3status.conf'
INCLUDE_PATH = '.i3/py3status'
INTERVAL = 1

# functions
################################################################################
def print_line(message):
    """
    Non-buffered printing to stdout
    """
    sys.stdout.write(message + '\n')
    sys.stdout.flush()

def read_line():
    """
    Interrupted respecting reader for stdin
    """
    try:
        line = sys.stdin.readline().strip()
        # i3status sends EOF, or an empty line
        if not line:
            sys.exit(3)
        return line
    except KeyboardInterrupt:
        sys.exit()

def print_output(prefix, output_list, modules_cache):
    """
    Merge current user-classes cache with i3status' output and display on i3bar
    """
    for module in modules_cache:
        index, json = modules_cache[module]
        for j in output_list:
            if j['name'] == json['name']:
                j.update(json)
                break
        else:
            output_list.insert(index, json)

    output = prefix+dumps(output_list)
    print_line(output)
    return output

def i3status_config_reader(config_file):
    """
    i3status.conf reader so we can adapt our code to the i3status config
    """
    in_time = False
    in_general = False
    config = {
        'colors': False,
        'color_good': None,
        'color_bad' : None,
        'color_degraded' : None,
        'color_separator': None,
        'interval': 5,
        'output_format': None,
        'time_format': '%Y-%m-%d %H:%M:%S',
        }
    for line in open(config_file, 'r'):
        line = line.strip(' \t\n\r')
        if line.startswith('general'):
            in_general = True
        elif line.startswith('time'):
            in_time = True
        elif line.startswith('}'):
            in_general = False
            in_time = False
        if in_general and '=' in line:
            key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
            if key in config:
                if value in ['true', 'false']:
                    value = 'True' if value == 'true' else 'False'
                config[key] = eval(value)
        if in_time and '=' in line:
            key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
            if 'time_' + key in config:
                config['time_' + key] = eval(value)
    return config

class I3status(Thread):
    """
    Run i3status in a thread and send its output to a Queue
    """
    def __init__(self, config_file):
        """
        set our useful properties
        """
        Thread.__init__(self)
        self.config_file = config_file
        self.kill = False
        self.queue = Queue()
        self.init = True
        self.started = False

    def stop(self):
        """
        break the i3status loop
        """
        self.kill = True

    def run(self):
        """
        run i3status and queue its output
        """
        i3status_pipe = Popen(
            ['i3status', '-c', self.config_file],
            stdout=PIPE,
            stderr=PIPE,
            )
        self.queue.put(i3status_pipe.stdout.readline())
        self.queue.put(i3status_pipe.stdout.readline())
        while not self.kill:
            line = i3status_pipe.stdout.readline()
            if len(line) > 0:
                self.queue.put(line)
                try:
                    line = line.decode()
                except Exception:
                    pass
                if len(line) > 1 and line.startswith('['):
                    self.init = False
                elif line.startswith(',['):
                    self.started = True
            else:
                break

def process_line(line, **kwargs):
    """
    Main line processor logic
    """
    j = []
    prefix = ''
    if line.startswith('{') and 'version' in line:
        print_line(line.strip('\n'))
    elif line == '[\n':
        print_line(line.strip('\n'))
    else:
        if line.startswith(','):
            line, prefix = line[1:], ','
        elif kwargs['delta'] > 0:
            prefix = ','

        # integrated transformations
        if not DISABLE_TRANSFORM:
            j = transform(loads(line), **kwargs)
        else:
            j = loads(line)

    return (prefix, j)

def transform(j, **kwargs):
    """
    Integrated transformations:
    - update the 'time' object so that it's updated at INTERVAL seconds
    """
    try:
        for item in j:
            # time modification
            if item['name'] in [ 'time', 'tztime' ]:
                time_format = I3STATUS_CONFIG['time_format']
                date = datetime.strptime( item['full_text'], time_format ) \
                    + timedelta(seconds=kwargs['delta'])
                item['full_text'] = date.strftime(time_format)
                if kwargs['delta'] > 0:
                    item['transformed'] = True
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "transformation failed (%s)" % (str(err)))
    return j

class UserModules(Thread):
    """
    This thread executes and caches all of the user's own Py3status classes.
    We run it in a separate thread so that any delay introduced by an external
    class doesn't affect the i3bar updating process.
    """
    def __init__(self, cache_timeout, i3status_conf, include_path, interval):
        """
        set our useful properties
        """
        Thread.__init__(self)
        self.cache = {}
        self.cache_timeout = cache_timeout
        self.classes = {}
        self.i3status_conf = i3status_conf
        self.include_path = include_path
        self.interval = interval
        self.i3status_json = {}
        self.kill = False

    @staticmethod
    def load_from_file(filepath):
        """
        return Py3status user-written classes
        """
        inst = None
        expected_class = 'Py3status'
        mod_name, file_ext = os.path.splitext(os.path.split(filepath)[-1])
        if file_ext.lower() == '.py':
            py_mod = imp.load_source(mod_name, filepath)
            if hasattr(py_mod, expected_class):
                inst = py_mod.Py3status()
        return (mod_name, inst)

    def load_files(self):
        """
        read user-written Py3status class files for dynamic inclusion
        """
        if self.include_path and os.path.isdir(self.include_path):
            for file_name in os.listdir(self.include_path):
                try:
                    module, class_inst = self.load_from_file(
                        self.include_path + file_name
                        )
                    if module and class_inst:
                        self.classes[file_name] = (class_inst, [])
                        for method in dir(class_inst):
                            if not method.startswith('__'):
                                self.classes[file_name][1].append(method)
                except Exception:
                    err = sys.exc_info()[1]
                    syslog(LOG_ERR, "loading %s failed (%s)" \
                        % (file_name, str(err)))

    def execute_classes(self):
        """
        Run on every user class included and execute every method on the json,
        then cache it for later usage.
        We exclude the 'kill' methods on the classes which are meant to be used
        for convenience only.
        """
        for class_name in sorted( self.classes.keys() ):
            my_class, my_methods = self.classes[class_name]
            for my_method in [ m for m in my_methods if m not in ['kill'] ]:
                try:
                    # handle a cache on user class methods results
                    try:
                        index, result = self.cache[my_method]
                        if time() > result['cached_until']:
                            raise KeyError('cache timeout')
                    except KeyError:
                        # execute the method
                        try:
                            meth = getattr(my_class, my_method)
                            index, result = meth(
                                self.i3status_json,
                                self.i3status_conf,
                                )
                        except Exception:
                            err = sys.exc_info()[1]
                            syslog(LOG_ERR, "user method %s failed (%s)" \
                                % (my_method, str(err)))
                            index, result = (0, {'name': '', 'full_text': ''})

                        # respect user-defined cache timeout for this module
                        if 'cached_until' not in result:
                            result['cached_until'] = time() + self.cache_timeout

                        # validate the response
                        assert isinstance(result, dict), "should return a dict"
                        assert 'full_text' in result, "missing 'full_text' key"
                        assert 'name' in result, "missing 'name' key"
                    finally:
                        self.cache[my_method] = (index, result)
                except Exception:
                    err = sys.exc_info()[1]
                    syslog(LOG_ERR, "injection failed (%s)" % str(err))

    def stop(self):
        """
        call any 'kill' method on Py3status modules and break this thread's loop
        """
        for class_name in sorted( self.classes.keys() ):
            my_class, my_methods = self.classes[class_name]
            if 'kill' in my_methods:
                kill_module = getattr(my_class, 'kill')
                kill_module()
        self.kill = True

    def clear(self):
        """
        clear user cache
        """
        self.cache = {}

    def run(self):
        """
        periodically execute user classes
        """
        self.load_files()
        while not self.kill:
            self.execute_classes()
            sleep(self.interval)

# main stuff
################################################################################
def main():
    """
    Main logic function
    """
    try:
        # global definition
        global CACHE_TIMEOUT, DISABLE_TRANSFORM
        global I3STATUS_CONFIG, INCLUDE_PATH, INTERVAL

        # command line options
        parser = argparse.ArgumentParser(
            description='The agile, python-powered, i3status wrapper')
        parser = argparse.ArgumentParser(add_help=True)
        parser.add_argument('-c', action="store",
            dest="i3status_conf", type=str,
            default=I3STATUS_CONFIG, help="path to i3status config file")
        parser.add_argument('-d', action="store_true",
            dest="disable_transform", help="disable integrated transformations")
        parser.add_argument('-i', action="store",
            dest="include_path", type=str,
            default=INCLUDE_PATH, help="user-based class include directory")
        parser.add_argument('-n', action="store",
            dest="interval", type=int,
            default=INTERVAL, help="update interval in seconds (default 1 sec)")
        parser.add_argument('-t', action="store",
            dest="cache_timeout", type=int, default=CACHE_TIMEOUT,
            help="default injection cache timeout in seconds (default 60 sec)")
        options = parser.parse_args()

        # configuration and helper variables
        CACHE_TIMEOUT = options.cache_timeout
        DISABLE_TRANSFORM = True if options.disable_transform else False
        I3STATUS_CONFIG = i3status_config_reader(options.i3status_conf)
        INCLUDE_PATH = os.path.abspath( options.include_path ) + '/'
        INTERVAL = options.interval

        # py3status uses only the i3bar protocol
        assert I3STATUS_CONFIG['output_format'] == 'i3bar', \
            'i3status output_format should be set to "i3bar"'
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "py3status init error (%s)" % str(err))
        sys.exit(1)

    try:
        # load user defined modules from a separate thread so that they cannot
        # cause any delay in updating the i3bar
        modules_thread = UserModules(
            CACHE_TIMEOUT,
            I3STATUS_CONFIG,
            INCLUDE_PATH,
            INTERVAL,
            )
        modules_thread.start()

        # spawn a i3status process on a separate thread
        # we will receive its output on a Queue which we will poll for messages
        i3status_thread = I3status(options.i3status_conf)
        i3status_thread.start()

        # SIGUSR1 forces a refresh of the bar both for py3status and i3status,
        # this mimics the USR1 signal handling of i3status (see man i3status)
        def sig_handler(signum, frame):
            """
            Raise a Warning level exception when a user sends a SIGUSR1 signal
            """
            raise UserWarning("received USR1, forcing refresh")
        signal(SIGUSR1, sig_handler)

        # control variables
        forced = False

        # main loop
        while True:
            try:
                # get a timestamp of now
                tst = time()

                # try to read i3status' output for at most INTERVAL seconds else
                # raise the Empty exception
                line = i3status_thread.queue.get(timeout=INTERVAL)
                try:
                    # python3 compatibility code
                    line = line.decode()
                except UnicodeDecodeError:
                    pass

                # i3status first output lines should be processed asap as only
                # the following lines will be processed by py3status
                if i3status_thread.started:
                    if not forced and I3STATUS_CONFIG['interval'] > INTERVAL:
                        # add a calculated sleep honoring py3status refresh
                        # time of the bar every INTERVAL seconds
                        sleep(INTERVAL - float( '{:.2}'.format( time()-tst ) ))
                    else:
                        # reset the SIGUSR1 flag forcing the refresh of the bar
                        forced = False

                # process the output line straight from i3status
                prefix, output_list = process_line(line, delta=0)
                if output_list:
                    modules_thread.i3status_json = output_list
                    print_output(
                        prefix,
                        output_list,
                        modules_thread.cache
                        )
            except Empty:
                # make sure i3status has started before modifying its output
                if i3status_thread.started or not i3status_thread.init:
                    prefix, output_list = process_line(line, delta=INTERVAL)
                    if output_list:
                        modules_thread.i3status_json = output_list
                        line = print_output(
                            prefix,
                            output_list,
                            modules_thread.cache
                            )
                else:
                    syslog(LOG_INFO, "waiting for i3status")
            except UserWarning:
                # SIGUSR1 was received, we also force i3status to refresh by
                # sending it a SIGUSR1 as well then we refresh the bar asap
                msg = sys.exc_info()[1]
                syslog(LOG_INFO, str(msg))
                call(["killall", "-s", "USR1", "i3status"])
                modules_thread.clear()
                forced = True
            except KeyboardInterrupt:
                break
            finally:
                if threading.active_count() < 3:
                    # i3status or modules thread died ? oops
                    break

        i3status_thread.stop()
        modules_thread.stop()
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "py3status error (%s)" % str(err))
        sys.exit(2)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># Copyright (c) 2013, Ultrabug
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# includes
################################################################################
import os
import imp
import sys
import argparse
import threading

from json import loads
from json import dumps

from time import time
from time import sleep

from datetime import datetime
from datetime import timedelta

from signal import signal
from signal import SIGUSR1

from subprocess import Popen
from subprocess import PIPE
from subprocess import call

from syslog import syslog
from syslog import LOG_ERR
from syslog import LOG_INFO

try:
	# python3
	from queue import Queue
	from queue import Empty
except ImportError:
	# python2
	from Queue import Queue
	from Queue import Empty

# functions
################################################################################
def print_line(message):
	"""
	Non-buffered printing to stdout
	"""
	sys.stdout.write(message + '\n')
	sys.stdout.flush()

def read_line():
	"""
	Interrupted respecting reader for stdin
	"""
	try:
		line = sys.stdin.readline().strip()
		# i3status sends EOF, or an empty line
		if not line:
			sys.exit(3)
		return line
	except KeyboardInterrupt:
		sys.exit()

def i3status_config_reader():
	"""
	i3status.conf reader so we can adapt our code to the i3status config
	"""
	in_time = False
	in_general = False
	config = {
		'colors': False,
		'color_good': None,
		'color_bad' : None,
		'color_degraded' : None,
		'color_separator': None,
		'interval': 5,
		'output_format': None,
		'time_format': '%Y-%m-%d %H:%M:%S',
		}
	for line in open(I3STATUS_CONFIG_FILE, 'r'):
		line = line.strip(' \t\n\r')
		if line.startswith('general'):
			in_general = True
		elif line.startswith('time'):
			in_time = True
		elif line.startswith('}'):
			in_general = False
			in_time = False
		if in_general and '=' in line:
			key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
			if key in config:
				if value in ['true', 'false']:
					value = 'True' if value == 'true' else 'False'
				config[key] = eval(value)
		if in_time and '=' in line:
			key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
			if 'time_' + key in config:
				config['time_' + key] = eval(value)
	return config

def i3status(message_queue):
	"""
	Execute i3status in a thread and send its output to a Queue to py3status
	"""
	i3status_pipe = Popen(['i3status', '-c', I3STATUS_CONFIG_FILE], stdout=PIPE, stderr=PIPE)
	message_queue.put(i3status_pipe.stdout.readline())
	message_queue.put(i3status_pipe.stdout.readline())
	while True:
		line = i3status_pipe.stdout.readline()
		if len(line) > 0:
			message_queue.put(line)
		else:
			break

def process_line(line, **kwargs):
	"""
	Main line processor logic
	"""
	if line.startswith('{') and 'version' in line:
		print_line(line.strip('\n'))
	elif line == '[\n':
		print_line(line.strip('\n'))
	else:
		prefix = ''
		if line.startswith(','):
			line, prefix = line[1:], ','
		elif kwargs['delta'] > 0:
			prefix = ','

		# integrated transformations
		if not DISABLE_TRANSFORM:
			j = transform(loads(line), **kwargs)
		else:
			j = loads(line)

		# user-based injection and transformation
		j = inject(j)

		print_line(prefix+dumps(j))

def inject(j):
	"""
	Run on every user class included and execute every method on the json,
	then inject the result at the start of the json
	"""
	# inject our own functions' results
	for fn in sorted( USER_CLASSES.keys() ):
		my_class, my_methods = USER_CLASSES[fn]
		for my_method in my_methods:
			try:
				# handle a cache on user class methods results
				try:
					index, result = USER_CACHE[my_method]
					if time() > result['cached_until']:
						raise KeyError('cache timeout')
				except KeyError:
					# execute the method
					try:
						meth = getattr(my_class, my_method)
						index, result = meth(j, I3STATUS_CONFIG)
					except Exception:
						err = sys.exc_info()[1]
						syslog(LOG_ERR, "user method %s failed (%s)" % (my_method, str(err)))
						index, result = (0, {'name': '', 'full_text': ''})

					# respect user-defined cache timeout for this module
					if 'cached_until' not in result:
						result['cached_until'] = time() + CACHE_TIMEOUT

					# validate the response
					assert isinstance(result, dict), "user method didn't return a dict"
					assert 'full_text' in result, "missing 'full_text' key"
					assert 'name' in result, "missing 'name' key"
				finally:
					USER_CACHE[my_method] = (index, result)
					j.insert(index, result)
			except Exception:
				err = sys.exc_info()[1]
				syslog(LOG_ERR, "injection failed (%s)" % str(err))
	return j

def transform(j, **kwargs):
	"""
	Integrated transformations:
	- update the 'time' object so that it's updated at INTERVAL seconds
	- update the 'run_watch' objects so that we rely on the color instead of useless 'yes' or 'no' status
	"""
	try:
		for item in j:
			# time modification
			if item['name'] in [ 'time', 'tztime' ]:
				date = datetime.strptime(item['full_text'], I3STATUS_CONFIG['time_format']) + timedelta(seconds=kwargs['delta'])
				item['full_text'] = date.strftime(I3STATUS_CONFIG['time_format'])
	except Exception:
		err = sys.exc_info()[1]
		syslog(LOG_ERR, "transformation failed (%s)" % (str(err)))
	return j

def load_from_file(filepath):
	"""
	Load Py3status user class for later injection
	"""
	inst = None
	expected_class = 'Py3status'
	mod_name, file_ext = os.path.splitext(os.path.split(filepath)[-1])
	if file_ext.lower() == '.py':
		py_mod = imp.load_source(mod_name, filepath)
		if hasattr(py_mod, expected_class):
			inst = py_mod.Py3status()
	return (mod_name, inst)

# main stuff
################################################################################
def main():
	try:
		# global definition
		global CACHE_TIMEOUT, DISABLE_TRANSFORM, I3STATUS_CONFIG_FILE
		global I3STATUS_CONFIG, INCLUDE_PATH, INTERVAL, USER_CACHE, USER_CLASSES

		# command line options
		PARSER = argparse.ArgumentParser(description='The agile, python-powered, i3status wrapper')
		PARSER = argparse.ArgumentParser(add_help=True)
		PARSER.add_argument('-c', action="store", dest="i3status_conf", type=str, default="/etc/i3status.conf", help="path to i3status config file")
		PARSER.add_argument('-d', action="store_true", dest="disable_transform", help="disable integrated transformations")
		PARSER.add_argument('-i', action="store", dest="include_path", type=str, default='.i3/py3status', help="user-based class include directory")
		PARSER.add_argument('-n', action="store", dest="interval", type=int, default=1, help="polling interval in seconds (default 1 sec)")
		PARSER.add_argument('-t', action="store", dest="cache_timeout", type=int, default=60, help="injection cache timeout in seconds (default 60 sec)")
		OPTS = PARSER.parse_args()

		# configuration and helper variables
		CACHE_TIMEOUT = OPTS.cache_timeout
		DELTA = 0
		DISABLE_TRANSFORM = True if OPTS.disable_transform else False
		FORCED = False
		I3STATUS_CONFIG_FILE = OPTS.i3status_conf
		I3STATUS_CONFIG = i3status_config_reader()
		INCLUDE_PATH = os.path.abspath( OPTS.include_path ) + '/'
		INTERVAL = OPTS.interval
		STARTED = False

		# py3status uses only the i3bar protocol
		assert I3STATUS_CONFIG['output_format'] == 'i3bar', 'unsupported output_format'

		# read user-written Py3status class files for dynamic inclusion
		USER_CACHE = {}
		USER_CLASSES = {}
		if INCLUDE_PATH and os.path.isdir(INCLUDE_PATH):
			for fn in os.listdir(INCLUDE_PATH):
				module, class_inst = load_from_file(INCLUDE_PATH + fn)
				if module and class_inst:
					USER_CLASSES[fn] = (class_inst, [])
					for method in dir(class_inst):
						if not method.startswith('__'):
							USER_CLASSES[fn][1].append(method)

		# spawn a i3status process on a separate thread
		# we will receive its output on a Queue which we will poll for messages
		MESSAGE_QUEUE = Queue()
		I3STATUS_THREAD = threading.Thread(target=i3status, name='i3status', args=(MESSAGE_QUEUE, ))
		I3STATUS_THREAD.start()

		# SIGUSR1 forces a refresh of the bar both for py3status and i3status,
		# this mimics the USR1 signal handling of i3status (see man i3status)
		def sig_handler(signum, frame):
			"""
			Raise a Warning level exception when a user sends a SIGUSR1 signal
			"""
			raise UserWarning("received USR1, forcing refresh")
		signal(SIGUSR1, sig_handler)

		# main loop
		while True:
			try:
				# get a timestamp of now
				TS = time()

				# try to read i3status' output for at most INTERVAL seconds else
				# raise the Empty exception
				LINE = MESSAGE_QUEUE.get(timeout=INTERVAL)
				try:
					# python3 compatibility code
					LINE = LINE.decode()
				except UnicodeDecodeError:
					pass

				# i3status first output lines should be processed asap as only
				# the following lines will be processed by py3status
				if LINE.startswith(',['):
					if not FORCED and I3STATUS_CONFIG['interval'] > INTERVAL:
						# add a calculated sleep honoring py3status refresh
						# time of the bar every INTERVAL seconds
						sleep( INTERVAL - float( '{:.2}'.format( time()-TS ) ) )
					else:
						# reset the SIGUSR1 flag forcing the refresh of the bar
						FORCED = False
					# flag i3status as started
					STARTED = True

				# process the output line and reset the time DELTA
				process_line(LINE, delta=0)
				DELTA = 0
			except Empty:
				# make sure i3status has started before modifying its output
				if STARTED:
					# the DELTA helps us adjust the clock without a syscall
					DELTA += INTERVAL
					process_line(LINE, delta=DELTA)
					if threading.active_count() < 2:
						# i3status died ? oups
						break
				else:
					syslog(LOG_INFO, "waiting for i3status")
			except UserWarning:
				# SIGUSR1 was received, we also force i3status to refresh by
				# sending it a SIGUSR1 as well then we refresh the bar asap
				msg = sys.exc_info()[1]
				syslog(LOG_INFO, str(msg))
				call(["killall", "-s", "USR1", "i3status"])
				USER_CACHE = {}
				FORCED = True
			except KeyboardInterrupt:
				break

		I3STATUS_THREAD.join()
	except Exception:
		err = sys.exc_info()[1]
		syslog(LOG_ERR, "py3status error (%s)" % str(err))
		sys.exit(1)

if __name__ == '__main__':
	main()
<commit_msg>major cleanup, better threading of i3status, enhanced main logic for better responsiveness, drop useless variables, respect PEP8 identation, more PEP8 variable cleanup<commit_after># Copyright (c) 2013, Ultrabug
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# includes
################################################################################
import os
import imp
import sys
import argparse
import threading
from threading import Thread

from json import loads
from json import dumps

from time import time
from time import sleep

from datetime import datetime
from datetime import timedelta

from signal import signal
from signal import SIGUSR1

from subprocess import Popen
from subprocess import PIPE
from subprocess import call

from syslog import syslog
from syslog import LOG_ERR
from syslog import LOG_INFO

try:
    # python3
    from queue import Queue
    from queue import Empty
except ImportError:
    # python2
    from Queue import Queue
    from Queue import Empty

# module globals and defaults
################################################################################
CACHE_TIMEOUT = 60
DISABLE_TRANSFORM = False
I3STATUS_CONFIG = '/etc/i3status.conf'
INCLUDE_PATH = '.i3/py3status'
INTERVAL = 1
USER_CACHE = {}
USER_CLASSES = {}

# functions
################################################################################
def print_line(message):
    """
    Non-buffered printing to stdout
    """
    sys.stdout.write(message + '\n')
    sys.stdout.flush()

def read_line():
    """
    Interrupted respecting reader for stdin
    """
    try:
        line = sys.stdin.readline().strip()
        # i3status sends EOF, or an empty line
        if not line:
            sys.exit(3)
        return line
    except KeyboardInterrupt:
        sys.exit()

def i3status_config_reader(config_file):
    """
    i3status.conf reader so we can adapt our code to the i3status config
    """
    in_time = False
    in_general = False
    config = {
        'colors': False,
        'color_good': None,
        'color_bad' : None,
        'color_degraded' : None,
        'color_separator': None,
        'interval': 5,
        'output_format': None,
        'time_format': '%Y-%m-%d %H:%M:%S',
        }
    for line in open(config_file, 'r'):
        line = line.strip(' \t\n\r')
        if line.startswith('general'):
            in_general = True
        elif line.startswith('time'):
            in_time = True
        elif line.startswith('}'):
            in_general = False
            in_time = False
        if in_general and '=' in line:
            key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
            if key in config:
                if value in ['true', 'false']:
                    value = 'True' if value == 'true' else 'False'
                config[key] = eval(value)
        if in_time and '=' in line:
            key, value = line.split('=')[0].strip(), line.split('=')[1].strip()
            if 'time_' + key in config:
                config['time_' + key] = eval(value)
    return config

class I3status(Thread):
    """
    Run i3status in a thread and send its output to a Queue
    """
    def __init__(self, config_file):
        """
        set our useful properties
        """
        Thread.__init__(self)
        self.config_file = config_file
        self.kill = False
        self.queue = Queue()
        self.init = True
        self.started = False

    def stop(self):
        """
        break the i3status loop
        """
        self.kill = True

    def run(self):
        """
        run i3status and queue its output
        """
        i3status_pipe = Popen(
            ['i3status', '-c', self.config_file],
            stdout=PIPE,
            stderr=PIPE,
            )
        self.queue.put(i3status_pipe.stdout.readline())
        self.queue.put(i3status_pipe.stdout.readline())
        while not self.kill:
            line = i3status_pipe.stdout.readline()
            if len(line) > 0:
                self.queue.put(line)
                try:
                    line = line.decode()
                except Exception:
                    pass
                if len(line) > 1 and line.startswith('['):
                    self.init = False
                elif line.startswith(',['):
                    self.started = True
            else:
                break

def process_line(line, **kwargs):
    """
    Main line processor logic
    """
    if line.startswith('{') and 'version' in line:
        print_line(line.strip('\n'))
    elif line == '[\n':
        print_line(line.strip('\n'))
    else:
        prefix = ''
        if line.startswith(','):
            line, prefix = line[1:], ','
        elif kwargs['delta'] > 0:
            prefix = ','

        # integrated transformations
        if not DISABLE_TRANSFORM:
            j = transform(loads(line), **kwargs)
        else:
            j = loads(line)

        # user-based injection and transformation
        j = inject(j)

        output = prefix+dumps(j)
        print_line(output)
        return output

def inject(j):
    """
    Run on every user class included and execute every method on the json,
    then inject the result at the start of the json
    """
    # inject our own functions' results
    for class_name in sorted( USER_CLASSES.keys() ):
        my_class, my_methods = USER_CLASSES[class_name]
        for my_method in my_methods:
            try:
                # handle a cache on user class methods results
                try:
                    index, result = USER_CACHE[my_method]
                    if time() > result['cached_until']:
                        raise KeyError('cache timeout')
                except KeyError:
                    # execute the method
                    try:
                        meth = getattr(my_class, my_method)
                        index, result = meth(j, I3STATUS_CONFIG)
                    except Exception:
                        err = sys.exc_info()[1]
                        syslog(LOG_ERR, "user method %s failed (%s)" \
                            % (my_method, str(err)))
                        index, result = (0, {'name': '', 'full_text': ''})

                    # respect user-defined cache timeout for this module
                    if 'cached_until' not in result:
                        result['cached_until'] = time() + CACHE_TIMEOUT

                    # validate the response
                    assert isinstance(result, dict), "user should return a dict"
                    assert 'full_text' in result, "missing 'full_text' key"
                    assert 'name' in result, "missing 'name' key"
                finally:
                    USER_CACHE[my_method] = (index, result)
                    j.insert(index, result)
            except Exception:
                err = sys.exc_info()[1]
                syslog(LOG_ERR, "injection failed (%s)" % str(err))
    return j

def transform(j, **kwargs):
    """
    Integrated transformations:
    - update the 'time' object so that it's updated at INTERVAL seconds
    """
    try:
        for item in j:
            # time modification
            if item['name'] in [ 'time', 'tztime' ]:
                time_format = I3STATUS_CONFIG['time_format']
                date = datetime.strptime( item['full_text'], time_format ) \
                    + timedelta(seconds=kwargs['delta'])
                item['full_text'] = date.strftime(time_format)
                if kwargs['delta'] > 0:
                    item['transformed'] = True
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "transformation failed (%s)" % (str(err)))
    return j

def load_from_file(filepath):
    """
    Load Py3status user class for later injection
    """
    inst = None
    expected_class = 'Py3status'
    mod_name, file_ext = os.path.splitext(os.path.split(filepath)[-1])
    if file_ext.lower() == '.py':
        py_mod = imp.load_source(mod_name, filepath)
        if hasattr(py_mod, expected_class):
            inst = py_mod.Py3status()
    return (mod_name, inst)

# main stuff
################################################################################
def main():
    """
    Main logic function
    """
    try:
        # global definition
        global CACHE_TIMEOUT, DISABLE_TRANSFORM
        global I3STATUS_CONFIG, INCLUDE_PATH, INTERVAL, USER_CACHE, USER_CLASSES

        # command line options
        parser = argparse.ArgumentParser(
            description='The agile, python-powered, i3status wrapper')
        parser = argparse.ArgumentParser(add_help=True)
        parser.add_argument('-c', action="store",
            dest="i3status_conf", type=str,
            default=I3STATUS_CONFIG, help="path to i3status config file")
        parser.add_argument('-d', action="store_true",
            dest="disable_transform", help="disable integrated transformations")
        parser.add_argument('-i', action="store",
            dest="include_path", type=str,
            default=INCLUDE_PATH, help="user-based class include directory")
        parser.add_argument('-n', action="store",
            dest="interval", type=int,
            default=INTERVAL, help="update interval in seconds (default 1 sec)")
        parser.add_argument('-t', action="store",
            dest="cache_timeout", type=int, default=CACHE_TIMEOUT,
            help="default injection cache timeout in seconds (default 60 sec)")
        options = parser.parse_args()

        # configuration and helper variables
        CACHE_TIMEOUT = options.cache_timeout
        DISABLE_TRANSFORM = True if options.disable_transform else False
        I3STATUS_CONFIG = i3status_config_reader(options.i3status_conf)
        INCLUDE_PATH = os.path.abspath( options.include_path ) + '/'
        INTERVAL = options.interval

        # py3status uses only the i3bar protocol
        assert I3STATUS_CONFIG['output_format'] == 'i3bar', \
            'i3status output_format should be set to "i3bar"'
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "py3status init error (%s)" % str(err))
        sys.exit(1)

    try:
        USER_CACHE = {}
        USER_CLASSES = {}
        # read user-written Py3status class files for dynamic inclusion
        if INCLUDE_PATH and os.path.isdir(INCLUDE_PATH):
            for file_name in os.listdir(INCLUDE_PATH):
                module, class_inst = load_from_file(INCLUDE_PATH + file_name)
                if module and class_inst:
                    USER_CLASSES[file_name] = (class_inst, [])
                    for method in dir(class_inst):
                        if not method.startswith('__'):
                            USER_CLASSES[file_name][1].append(method)

        # spawn a i3status process on a separate thread
        # we will receive its output on a Queue which we will poll for messages
        i3status_thread = I3status(options.i3status_conf)
        i3status_thread.start()

        # SIGUSR1 forces a refresh of the bar both for py3status and i3status,
        # this mimics the USR1 signal handling of i3status (see man i3status)
        def sig_handler(signum, frame):
            """
            Raise a Warning level exception when a user sends a SIGUSR1 signal
            """
            raise UserWarning("received USR1, forcing refresh")
        signal(SIGUSR1, sig_handler)

        # control variables
        forced = False

        # main loop
        while True:
            try:
                # get a timestamp of now
                tst = time()

                # try to read i3status' output for at most INTERVAL seconds else
                # raise the Empty exception
                line = i3status_thread.queue.get(timeout=INTERVAL)
                try:
                    # python3 compatibility code
                    line = line.decode()
                except UnicodeDecodeError:
                    pass

                # i3status first output lines should be processed asap as only
                # the following lines will be processed by py3status
                if i3status_thread.started:
                    if not forced and I3STATUS_CONFIG['interval'] > INTERVAL:
                        # add a calculated sleep honoring py3status refresh
                        # time of the bar every INTERVAL seconds
                        sleep(INTERVAL - float( '{:.2}'.format( time()-tst ) ))
                    else:
                        # reset the SIGUSR1 flag forcing the refresh of the bar
                        forced = False

                # process the output line straight from i3status
                process_line(line, delta=0)
            except Empty:
                # make sure i3status has started before modifying its output
                if i3status_thread.started or not i3status_thread.init:
                    line = process_line(line, delta=INTERVAL)
                    if threading.active_count() < 2:
                        # i3status died ? oups
                        break
                else:
                    syslog(LOG_INFO, "waiting for i3status")
            except UserWarning:
                # SIGUSR1 was received, we also force i3status to refresh by
                # sending it a SIGUSR1 as well then we refresh the bar asap
                msg = sys.exc_info()[1]
                syslog(LOG_INFO, str(msg))
                call(["killall", "-s", "USR1", "i3status"])
                USER_CACHE = {}
                forced = True
            except KeyboardInterrupt:
                break

        i3status_thread.stop()
    except Exception:
        err = sys.exc_info()[1]
        syslog(LOG_ERR, "py3status error (%s)" % str(err))
        sys.exit(2)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Wed Sep 21 17:08:59 2016

@author: bitzer
"""

import math
import numpy as np
import scipy
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
from warnings import warn
from abc import ABCMeta, abstractmethod

#%% define transformations
class transform(metaclass=ABCMeta):
    """A parameter transformation."""

    def __init__(self, **params):
        self.transformed_range = None

    @abstractmethod
    def transform(self, x):
        return None

    @abstractmethod
    def transformed_pdf(self, y, mu, sigma2):
        return None

    @abstractmethod
    def transformed_mode(self, mu, sigma2):
        return None

    def find_transformed_mode(self, mu, sigma2, **kwargs):
        fun = lambda x: -self.transformed_pdf(x, mu, sigma2)

        res = scipy.optimize.minimize_scalar(fun, **kwargs)

        return res.x

    @abstractmethod
    def transformed_ppf(self, q, mu, sigma2):
        return None

    def approximate_transformed_ppf(self, q, mu, sigma2):
        if np.isscalar(q):
            q = np.array(q)

        x = np.random.normal(mu, math.sqrt(sigma2), 1000)

        return np.percentile(self.transform(x), q * 100)

class identity(transform):
    def __init__(self):
        self.transformed_range = np.r_[-np.inf, np.inf]

    def transform(self, x):
        return x

    def transformed_pdf(self, y, mu, sigma2):
        return scipy.stats.norm.pdf(y, loc=mu, scale=math.sqrt(sigma2))

    def transformed_mode(self, mu, sigma2):
        return mu

    def transformed_ppf(self, q, mu, sigma2):
        return scipy.stats.norm.ppf(q, loc=mu, scale=math.sqrt(sigma2))

class exponential(transform):
    def __init__(self):
        self.transformed_range = np.r_[0, np.inf]

    def transform(self, x):
        return np.exp(x)

    def transformed_pdf(self, y, mu, sigma2):
        return scipy.stats.lognorm.pdf(y, math.sqrt(sigma2), scale=math.exp(mu))

    def transformed_mode(self, mu, sigma2):
        return math.exp(mu - sigma2)

    def transformed_ppf(self, q, mu, sigma2):
        return scipy.stats.lognorm.ppf(q, math.sqrt(sigma2), scale=math.exp(mu))

class gaussprob(transform):

    @property
    def width(self):
        return self._width

    @property
    def shift(self):
        return self._shift

    def __init__(self, width=1.0, shift=0.0):
        self._width = width
        self._shift = shift

        self.transformed_range = np.r_[self.shift, self.width + self.shift]

    def transform(self, x):
        return gaussprob_trans(x, self.width, self.shift)

    def transformed_pdf(self, y, mu, sigma2):
        if sigma2 == 1.0:
            sigma2 -= 1e-15
            warn('subtracted 1e-15 from sigma2=1.0 to avoid division by 0')

        return np.exp((sigma2 - 1) / 2 / sigma2 * (
                      scipy.stats.norm.ppf((y - self.shift) / self.width) -
                      mu / (1 - sigma2)) ** 2 +
                      mu**2 / (1 - sigma2) / 2
                     ) / np.sqrt(sigma2) / self.width

    def transformed_mode(self, mu, sigma2):
        return self.find_transformed_mode(mu, sigma2, method='Bounded',
                                          bounds=self.transformed_range)

    def transformed_ppf(self, q, mu, sigma2):
        if np.isscalar(q):
            q = np.array(q)

        warn('gaussprob transform: No analytic ppf implemented. '
             'Using numeric approximation.')

        return self.approximate_transformed_ppf(q, mu, sigma2)

try:
    from numba import vectorize, float64
    @vectorize([float64(float64, float64, float64)], nopython=True)
    def gaussprob_trans(x, width, shift):
        cdf = (1 + math.erf(x / math.sqrt(2))) / 2

        return cdf * width + shift
except ImportError:
    def gaussprob_trans(x, width, shift):
        return scipy.stats.norm.cdf(x) * width + shift

#%% define parameter container
class parameter_container:

    @property
    def names(self):
        return self.params.name

    def __init__(self):
        self.params = pd.DataFrame(columns=('name', 'transform'))
        self.P = 0
        self.mu = np.array([])
        self.cov = np.array([[]])

    def add_param(self, name, mu, sigma, transform=identity()):
        self.params.loc[self.P] = [name, transform]
        self.P += 1
        self.generate_transformfun()

        self.mu = np.r_[self.mu, mu]
        cov = np.zeros((self.P, self.P))
        cov[:self.P-1, :self.P-1] = self.cov
        self.cov = cov
        self.cov[-1, -1] = sigma ** 2

    def generate_transformfun(self):
        trstr = ""
        for i in range(self.P):
            trstr += ("self.params.loc[%s, 'transform']." % i +
                      "transform(values[:, %s])," % i)
        trstr = trstr[:-1]

        self.transformfun = eval("lambda self, values: np.c_[%s]" % trstr)

    def transform(self, values):
        return self.transformfun(self, values)

    def sample(self, S, mu=None, cov=None):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        return np.random.multivariate_normal(self.mu, self.cov, S)

    def sample_transformed(self, S, mu=None, cov=None):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        return self.transform(self.sample(S))

    def get_transformed_mode(self, mu=None, cov=None):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        mode = np.full(self.P, np.nan)
        for i, par in self.params.iterrows():
            mode[i] = par.transform.transformed_mode(mu[i], cov[i, i])

        return mode

    def plot_param_dist(self, mu=None, cov=None, S=500, q_lower=0.005,
                        q_upper=0.995, only_marginals=False, dist_names=['']):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        if only_marginals:
            fig, axes = plt.subplots(1, self.P)
            for i, par in self.params.iterrows():
                xlim = par.transform.transformed_ppf(np.r_[q_lower, q_upper],
                                                 mu[i], cov[i, i])
                x = np.linspace(xlim[0], xlim[1], 1000)
                axes[i].plot(x, par.transform.transformed_pdf(x, mu[i], cov[i, i]))
                axes[i].set_xlim(xlim)
                axes[i].set_xlabel(par.name)
                if i == 0:
                    axes[i].set_ylabel('pdf')

            return fig, axes
        else:
            samples = np.random.multivariate_normal(mu, cov, S)
            samples = self.transform(samples)
            samples = pd.DataFrame(samples, columns=self.params.name)
            samples['distribution'] = dist_names[0]

            pg = sns.PairGrid(samples, hue='distribution', diag_sharey=False)

            # scatter plot in upper diagonal
            pg = pg.map_upper(plt.scatter, alpha=0.3)

            # correlation of sampels in lower diagonal
            pg = pg.map_lower(plot_corrcoef)

            # fill diagonal with empty axes
            pg = pg.map_diag(lambda x, **kwargs: None)
            # plot analytical pdfs in diagonal
            for i, par in self.params.iterrows():
                xlim = par.transform.transformed_ppf(np.r_[q_lower, q_upper],
                                                 mu[i], cov[i, i])

                x = np.linspace(xlim[0], xlim[1], 1000)
                pg.diag_axes[i].plot(x, par.transform.transformed_pdf(x, mu[i], cov[i, i]))
                pg.diag_axes[i].set_xlim(xlim)

                # also set y-limits of off-diagonal
                if self.P > 1:
                    if i==0:
                        pg.axes[0, 1].set_ylim(xlim)
                    else:
                        pg.axes[i, 0].set_ylim(xlim)

#            pg.add_legend(frameon=True)

            return pg

def plot_corrcoef(x, y, **kwargs):
    corrcoefs = np.corrcoef(x, y)

    ax = plt.gca()
    ax.text(0.5, 0.5, 'R = %4.2f' % corrcoefs[0, 1],
            horizontalalignment='center', verticalalignment='center',
            transform=ax.transAxes, **kwargs)

#    ax.set_axis_off()

#%% some tests
if __name__ == "__main__":
    # test parameter container
    pars = parameter_container()
    pars.add_param('noisestd', 0, 1, transform=exponential())
    pars.add_param('prior', 0.8, 0.5, transform=gaussprob())
    pars.add_param('ndtmean', -5, 2)

    pg = pars.plot_param_dist()

    for i, par in pars.params.iterrows():
        mu = pars.mu[i]
        sigma2 = pars.cov[i, i]
        x = par.transform.transformed_mode(mu, sigma2)
        pg.diag_axes[i].plot(x, par.transform.transformed_pdf(x, mu, sigma2), '*r')

    # function for checking the implemented gaussprobpdf
    def check_gaussprobpdf(mu=0.0, sigma=1.0):
        g_samples = scipy.stats.norm.rvs(loc=mu, scale=sigma, size=10000)
        p_samples = scipy.stats.norm.cdf(g_samples)

        gtr = gaussprob()

        plt.figure()
        ax = sns.distplot(p_samples)
        lower, upper = ax.get_xlim()
        yy = np.linspace(lower, upper, 1000)
        ax.plot(yy, gtr.transformed_pdf(yy, mu, sigma**2))<commit_msg>transform function is unpicklable so account for that<commit_after># -*- coding: utf-8 -*-
"""
Created on Wed Sep 21 17:08:59 2016

@author: bitzer
"""

import math
import numpy as np
import scipy
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
from warnings import warn
from abc import ABCMeta, abstractmethod

#%% define transformations
class transform(metaclass=ABCMeta):
    """A parameter transformation."""

    def __init__(self, **params):
        self.transformed_range = None

    @abstractmethod
    def transform(self, x):
        return None

    @abstractmethod
    def transformed_pdf(self, y, mu, sigma2):
        return None

    @abstractmethod
    def transformed_mode(self, mu, sigma2):
        return None

    def find_transformed_mode(self, mu, sigma2, **kwargs):
        fun = lambda x: -self.transformed_pdf(x, mu, sigma2)

        res = scipy.optimize.minimize_scalar(fun, **kwargs)

        return res.x

    @abstractmethod
    def transformed_ppf(self, q, mu, sigma2):
        return None

    def approximate_transformed_ppf(self, q, mu, sigma2):
        if np.isscalar(q):
            q = np.array(q)

        x = np.random.normal(mu, math.sqrt(sigma2), 1000)

        return np.percentile(self.transform(x), q * 100)

class identity(transform):
    def __init__(self):
        self.transformed_range = np.r_[-np.inf, np.inf]

    def transform(self, x):
        return x

    def transformed_pdf(self, y, mu, sigma2):
        return scipy.stats.norm.pdf(y, loc=mu, scale=math.sqrt(sigma2))

    def transformed_mode(self, mu, sigma2):
        return mu

    def transformed_ppf(self, q, mu, sigma2):
        return scipy.stats.norm.ppf(q, loc=mu, scale=math.sqrt(sigma2))

class exponential(transform):
    def __init__(self):
        self.transformed_range = np.r_[0, np.inf]

    def transform(self, x):
        return np.exp(x)

    def transformed_pdf(self, y, mu, sigma2):
        return scipy.stats.lognorm.pdf(y, math.sqrt(sigma2), scale=math.exp(mu))

    def transformed_mode(self, mu, sigma2):
        return math.exp(mu - sigma2)

    def transformed_ppf(self, q, mu, sigma2):
        return scipy.stats.lognorm.ppf(q, math.sqrt(sigma2), scale=math.exp(mu))

class gaussprob(transform):

    @property
    def width(self):
        return self._width

    @property
    def shift(self):
        return self._shift

    def __init__(self, width=1.0, shift=0.0):
        self._width = width
        self._shift = shift

        self.transformed_range = np.r_[self.shift, self.width + self.shift]

    def transform(self, x):
        return gaussprob_trans(x, self.width, self.shift)

    def transformed_pdf(self, y, mu, sigma2):
        if sigma2 == 1.0:
            sigma2 -= 1e-15
            warn('subtracted 1e-15 from sigma2=1.0 to avoid division by 0')

        return np.exp((sigma2 - 1) / 2 / sigma2 * (
                      scipy.stats.norm.ppf((y - self.shift) / self.width) -
                      mu / (1 - sigma2)) ** 2 +
                      mu**2 / (1 - sigma2) / 2
                     ) / np.sqrt(sigma2) / self.width

    def transformed_mode(self, mu, sigma2):
        return self.find_transformed_mode(mu, sigma2, method='Bounded',
                                          bounds=self.transformed_range)

    def transformed_ppf(self, q, mu, sigma2):
        if np.isscalar(q):
            q = np.array(q)

        warn('gaussprob transform: No analytic ppf implemented. '
             'Using numeric approximation.')

        return self.approximate_transformed_ppf(q, mu, sigma2)

try:
    from numba import vectorize, float64
    @vectorize([float64(float64, float64, float64)], nopython=True)
    def gaussprob_trans(x, width, shift):
        cdf = (1 + math.erf(x / math.sqrt(2))) / 2

        return cdf * width + shift
except ImportError:
    def gaussprob_trans(x, width, shift):
        return scipy.stats.norm.cdf(x) * width + shift

#%% define parameter container
class parameter_container:

    @property
    def names(self):
        return self.params.name

    def __init__(self):
        self.params = pd.DataFrame(columns=('name', 'transform'))
        self.P = 0
        self.mu = np.array([])
        self.cov = np.array([[]])

    def __getstate__(self):
        state = self.__dict__.copy()
        del state['transformfun']
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.generate_transformfun()

    def add_param(self, name, mu, sigma, transform=identity()):
        self.params.loc[self.P] = [name, transform]
        self.P += 1
        self.generate_transformfun()

        self.mu = np.r_[self.mu, mu]
        cov = np.zeros((self.P, self.P))
        cov[:self.P-1, :self.P-1] = self.cov
        self.cov = cov
        self.cov[-1, -1] = sigma ** 2

    def generate_transformfun(self):
        trstr = ""
        for i in range(self.P):
            trstr += ("self.params.loc[%s, 'transform']." % i +
                      "transform(values[:, %s])," % i)
        trstr = trstr[:-1]

        self.transformfun = eval("lambda self, values: np.c_[%s]" % trstr)

    def transform(self, values):
        return self.transformfun(self, values)

    def sample(self, S, mu=None, cov=None):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        return np.random.multivariate_normal(self.mu, self.cov, S)

    def sample_transformed(self, S, mu=None, cov=None):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        return self.transform(self.sample(S))

    def get_transformed_mode(self, mu=None, cov=None):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        mode = np.full(self.P, np.nan)
        for i, par in self.params.iterrows():
            mode[i] = par.transform.transformed_mode(mu[i], cov[i, i])

        return mode

    def plot_param_dist(self, mu=None, cov=None, S=500, q_lower=0.005,
                        q_upper=0.995, only_marginals=False, dist_names=['']):
        if mu is None:
            mu = self.mu
        if cov is None:
            cov = self.cov

        if only_marginals:
            fig, axes = plt.subplots(1, self.P)
            for i, par in self.params.iterrows():
                xlim = par.transform.transformed_ppf(np.r_[q_lower, q_upper],
                                                 mu[i], cov[i, i])
                x = np.linspace(xlim[0], xlim[1], 1000)
                axes[i].plot(x, par.transform.transformed_pdf(x, mu[i], cov[i, i]))
                axes[i].set_xlim(xlim)
                axes[i].set_xlabel(par.name)
                if i == 0:
                    axes[i].set_ylabel('pdf')

            return fig, axes
        else:
            samples = np.random.multivariate_normal(mu, cov, S)
            samples = self.transform(samples)
            samples = pd.DataFrame(samples, columns=self.params.name)
            samples['distribution'] = dist_names[0]

            pg = sns.PairGrid(samples, hue='distribution', diag_sharey=False)

            # scatter plot in upper diagonal
            pg = pg.map_upper(plt.scatter, alpha=0.3)

            # correlation of sampels in lower diagonal
            pg = pg.map_lower(plot_corrcoef)

            # fill diagonal with empty axes
            pg = pg.map_diag(lambda x, **kwargs: None)
            # plot analytical pdfs in diagonal
            for i, par in self.params.iterrows():
                xlim = par.transform.transformed_ppf(np.r_[q_lower, q_upper],
                                                 mu[i], cov[i, i])

                x = np.linspace(xlim[0], xlim[1], 1000)
                pg.diag_axes[i].plot(x, par.transform.transformed_pdf(x, mu[i], cov[i, i]))
                pg.diag_axes[i].set_xlim(xlim)

                # also set y-limits of off-diagonal
                if self.P > 1:
                    if i==0:
                        pg.axes[0, 1].set_ylim(xlim)
                    else:
                        pg.axes[i, 0].set_ylim(xlim)

#            pg.add_legend(frameon=True)

            return pg

def plot_corrcoef(x, y, **kwargs):
    corrcoefs = np.corrcoef(x, y)

    ax = plt.gca()
    ax.text(0.5, 0.5, 'R = %4.2f' % corrcoefs[0, 1],
            horizontalalignment='center', verticalalignment='center',
            transform=ax.transAxes, **kwargs)

#    ax.set_axis_off()

#%% some tests
if __name__ == "__main__":
    # test parameter container
    pars = parameter_container()
    pars.add_param('noisestd', 0, 1, transform=exponential())
    pars.add_param('prior', 0.8, 0.5, transform=gaussprob())
    pars.add_param('ndtmean', -5, 2)

    pg = pars.plot_param_dist()

    for i, par in pars.params.iterrows():
        mu = pars.mu[i]
        sigma2 = pars.cov[i, i]
        x = par.transform.transformed_mode(mu, sigma2)
        pg.diag_axes[i].plot(x, par.transform.transformed_pdf(x, mu, sigma2), '*r')

    # function for checking the implemented gaussprobpdf
    def check_gaussprobpdf(mu=0.0, sigma=1.0):
        g_samples = scipy.stats.norm.rvs(loc=mu, scale=sigma, size=10000)
        p_samples = scipy.stats.norm.cdf(g_samples)

        gtr = gaussprob()

        plt.figure()
        ax = sns.distplot(p_samples)
        lower, upper = ax.get_xlim()
        yy = np.linspace(lower, upper, 1000)
        ax.plot(yy, gtr.transformed_pdf(yy, mu, sigma**2))<|endoftext|>


<commit_before>from pyantidot.exception import NotFoundException
from pyantidot.helpers import high_light
from pyantidot.tools import Bunch

class BunchContainer(object):
    def __init__(self, bunch: Bunch):
        self._bunch = bunch

class Pager(BunchContainer):
    pass

class HighLightText(BunchContainer):
    @property
    def is_match(self) -> bool:
        return 'match' in self._bunch and 'text' not in self._bunch

    def __str__(self) -> str:
        if 'match' in self._bunch:
            return self._bunch.match
        return self._bunch.text

class HighLightTextList(list):
    def __str__(self) -> str:
        return ' '.join(map(str, self))

    def get_highlight_text(self, surround) -> str:
        return high_light(self, surround)

class Content(BunchContainer):
    @property
    def doc_id(self) -> int:
        return int(self._bunch.docId)

    @property
    def uri(self) -> str:
        return self._bunch.uri

    @property
    def title(self) -> HighLightTextList:
        return HighLightTextList([HighLightText(bunch) for bunch in self._bunch.title])

    @property
    def abstract(self) -> HighLightTextList:
        return HighLightTextList([HighLightText(bunch) for bunch in self._bunch.abstract])

    @property
    def relevance(self) -> Bunch:
        return self._bunch.relevance

    @property
    def client_data(self) -> [Bunch]:
        return [bunch for bunch in self._bunch.clientData]

class ReplySetNode(BunchContainer):
    @property
    def key(self) -> object:
        return self._bunch.key

    @property
    def labels(self) -> [Bunch]:
        return [label for label in self._bunch.labels]

    @property
    def items(self) -> int:
        return int(self._bunch['items'])  # We can't use self._bunch.items function here, Bunch.items is a function

class ReplySetFacet(BunchContainer):
    @property
    def layout(self) -> str:
        return self._bunch.layout

    @property
    def type(self) -> str:
        return self._bunch.type

    @property
    def id(self) -> int:
        return self._bunch.id

    @property
    def nodes(self) -> [ReplySetNode]:
        return [ReplySetNode(node) for node in self._bunch.node]

    @property
    def pager(self) -> Pager:
        return Pager(self._bunch.pager)

    @property
    def labels(self) -> [Bunch]:
        return [label for label in self._bunch.labels]

    def node(self, key: object) -> ReplySetNode:
        for node_bunch in self._bunch.node:
            if node_bunch.key == key:
                return ReplySetNode(node_bunch)
        raise NotFoundException()

class ReplySet(BunchContainer):
    @property
    def contents(self) -> [Content]:
        return [Content(bunch) for bunch in self._bunch.content.reply]

    @property
    def meta_bunch(self) -> Bunch:
        return self._bunch.meta

    @property
    def pager(self) -> Pager:
        return Pager(self._bunch.pager)

    @property
    def facets(self) -> [ReplySetFacet]:
        return [ReplySetFacet(bunch) for bunch in self._bunch.facets.facet]

    def content(self, uri: object) -> Content:
        for content_bunch in self._bunch.content.reply:
            if content_bunch.uri == uri:
                return Content(content_bunch)
        raise NotFoundException()

    def facet(self, id: object) -> ReplySetFacet:
        for facet in self._bunch.facets.facet:
            if facet.id == id:
                return ReplySetFacet(facet)
        raise NotFoundException()

class Header(BunchContainer):
    @property
    def query_bunch(self) -> Bunch:
        return self._bunch.query

    @property
    def user_bunch(self) -> Bunch:
        return self._bunch.user

    @property
    def performance_bunch(self) -> Bunch:
        return self._bunch.performance

    @property
    def info_bunch(self) -> Bunch:
        return self._bunch.info

class Response(BunchContainer):
    @property
    def header(self) -> Header:
        return Header(self._bunch.header)

    @property
    def replies(self) -> list:
        return [ReplySet(bunch) for bunch in self._bunch.replySet]

    def reply(self, uri: object) -> ReplySet:
        for reply_set_bunch in self._bunch.replySet:
            if reply_set_bunch.meta.uri == uri:
                return ReplySet(reply_set_bunch)
        raise NotFoundException()
<commit_msg>declare new properties<commit_after>from pyantidot.exception import NotFoundException
from pyantidot.helpers import high_light
from pyantidot.tools import Bunch

class BunchContainer(object):
    def __init__(self, bunch: Bunch):
        self._bunch = bunch

    def get_raw(self):
        return self._bunch

class Pager(BunchContainer):
    @property
    def current_page(self) -> int:
        return int(self._bunch.currentPage)

    @property
    def new_page(self):
        return int(self._bunch.currentPage)

    @property
    def pages(self) -> [int]:
        return self._bunch.page

class HighLightText(BunchContainer):
    @property
    def is_match(self) -> bool:
        return 'match' in self._bunch and 'text' not in self._bunch

    def __str__(self) -> str:
        if 'match' in self._bunch:
            return self._bunch.match
        return self._bunch.text

class HighLightTextList(list):
    def __str__(self) -> str:
        return ' '.join(map(str, self))

    def get_highlight_text(self, surround) -> str:
        return high_light(self, surround)

class Content(BunchContainer):
    @property
    def doc_id(self) -> int:
        return int(self._bunch.docId)

    @property
    def uri(self) -> str:
        return self._bunch.uri

    @property
    def title(self) -> HighLightTextList:
        return HighLightTextList([HighLightText(bunch) for bunch in self._bunch.title])

    @property
    def abstract(self) -> HighLightTextList:
        return HighLightTextList([HighLightText(bunch) for bunch in self._bunch.abstract])

    @property
    def relevance(self) -> Bunch:
        return self._bunch.relevance

    @property
    def client_data(self) -> [Bunch]:
        return [bunch for bunch in self._bunch.clientData]

class ReplySetNode(BunchContainer):
    @property
    def key(self) -> object:
        return self._bunch.key

    @property
    def labels(self) -> [Bunch]:
        return [label for label in self._bunch.labels]

    @property
    def items(self) -> int:
        return int(self._bunch['items'])  # We can't use self._bunch.items function here, Bunch.items is a function

class ReplySetFacet(BunchContainer):
    @property
    def layout(self) -> str:
        return self._bunch.layout

    @property
    def type(self) -> str:
        return self._bunch.type

    @property
    def id(self) -> int:
        return self._bunch.id

    @property
    def nodes(self) -> [ReplySetNode]:
        return [ReplySetNode(node) for node in self._bunch.node]

    @property
    def pager(self) -> Pager:
        return Pager(self._bunch.pager)

    @property
    def labels(self) -> [Bunch]:
        return [label for label in self._bunch.labels]

    def node(self, key: object) -> ReplySetNode:
        for node_bunch in self._bunch.node:
            if node_bunch.key == key:
                return ReplySetNode(node_bunch)
        raise NotFoundException()

class ReplySet(BunchContainer):
    @property
    def contents(self) -> [Content]:
        return [Content(bunch) for bunch in self._bunch.content.reply]

    @property
    def meta_bunch(self) -> Bunch:
        return self._bunch.meta

    @property
    def pager(self) -> Pager:
        return Pager(self._bunch.pager)

    @property
    def facets(self) -> [ReplySetFacet]:
        return [ReplySetFacet(bunch) for bunch in self._bunch.facets.facet]

    def content(self, uri: object) -> Content:
        for content_bunch in self._bunch.content.reply:
            if content_bunch.uri == uri:
                return Content(content_bunch)
        raise NotFoundException()

    def facet(self, id: object) -> ReplySetFacet:
        for facet in self._bunch.facets.facet:
            if facet.id == id:
                return ReplySetFacet(facet)
        raise NotFoundException()

class Header(BunchContainer):
    @property
    def query_bunch(self) -> Bunch:
        return self._bunch.query

    @property
    def user_bunch(self) -> Bunch:
        return self._bunch.user

    @property
    def performance_bunch(self) -> Bunch:
        return self._bunch.performance

    @property
    def info_bunch(self) -> Bunch:
        return self._bunch.info

class Response(BunchContainer):
    @property
    def header(self) -> Header:
        return Header(self._bunch.header)

    @property
    def replies(self) -> list:
        return [ReplySet(bunch) for bunch in self._bunch.replySet]

    def reply(self, uri: object) -> ReplySet:
        for reply_set_bunch in self._bunch.replySet:
            if reply_set_bunch.meta.uri == uri:
                return ReplySet(reply_set_bunch)
        raise NotFoundException()
<|endoftext|>