<commit_before>#!/usr/bin/env python3
#
# wtimegui - Working time class with GUI
#
import sys
import time
try:
    from Tkinter import *
    import ttk
    import tkMessageBox
    from threading import *
except ImportError:
    from tkinter import *
    from tkinter import ttk
    from threading import *
    from tkinter import messagebox as tkMessageBox

from wtime4 import wtime

__author__ = "Riccardo Bruno"
__copyright__ = "2017"
__license__ = "Apache"
__maintainer__ = "Riccardo Bruno"
__email__ = "riccardo.bruno@gmail.com"

class wtimeGUI:

    flag_ticket_reached = False
    flag_time_reached = False
    flag_thread_running = False
    interval_thread_waitcycles = 5
    check_time_thread=None
    wtime_out = {}

    winTITLE="wtime GUI"
    lblFONT=("Lucida Grande", 12)
    lblFGCOLOR='black'

    root = None

    GUI_data = (
        {"type": "text", "name": "t1", "title": "T1", "row": 0, "col": 0},
        {"type": "text", "name": "t2", "title": "T2", "row": 1, "col": 0},
        {"type": "text", "name": "t2t1", "title": "T2 - T1", "row": 1, "col": 2},
        {"type": "text", "name": "t3", "title": "T3", "row": 2, "col": 0},
        {"type": "text", "name": "t4", "title": "T4", "row": 3, "col": 0},
        {"type": "text", "name": "t4t3", "title": "T4 - T3", "row": 3, "col": 2},
        {"type": "text", "name": "pause time", "title": "Pause Time", "row": 4, "col": 0},
        {"type": "text", "name": "total time", "title": "Total Time", "row": 5, "col": 0},
        {"type": "text", "name": "overtime", "title": "Over Time", "row": 5, "col": 2},
        {"type": "text", "name": "time to reach", "title": "Time to reach", "row": 6, "col": 0},
        {"type": "text", "name": "time remaining", "title": "Time remain", "row": 7, "col": 0},
        {"type": "text", "name": "time remaining at", "title": "at", "row": 8, "col": 0},
        {"type": "text", "name": "ticket remaining", "title": "Ticket remain", "row": 9, "col": 0},
        {"type": "text", "name": "ticket remaining at", "title": "at", "row": 10, "col": 0},
        {"type": "text", "name": "ticket time", "title": "TicketTime", "row": 11, "col": 0},
        {"type": "text", "name": "time remaining perc", "title": "%", "row": 7, "col": 2},
        {"type": "text", "name": "ticket remaining perc", "title": "%", "row": 9, "col": 2},
        {"type": "progress", "name": "t1", "title": "Time progress", "row": 7, "col": 3},
        {"type": "progress", "name": "t1", "title": "Ticket progress", "row": 9, "col": 3},
        {"type": "button", "name": "Tx", "title": "T2", "row": 12, "col": 0},
        {"type": "button", "name": "Update", "title": "Update", "row": 12, "col": 1},
        {"type": "button", "name": "Exit", "title": "Exit", "row": 12, "col": 3},
    )

    def get_item(self, type, name):
        item_result = None
        for item in self.GUI_data:
            if item["type"] == type and item["name"] == name:
                item_result = item
                break
        return item_result

    def __init__(self):
        # wtime4
        self.t1, self.t2, self.t3, self.t4, self.ct = wtime.getTimes(sys.argv)
        self.wt = wtime(t1=self.t1, t2=self.t2, t3=self.t3, t4=self.t4, current_time=self.ct)
        # GUI
        self.root = Tk()
        self.root.title(self.winTITLE)
        self.gui_build()
        self.check_time()
        self.root.bind('<Return>',self.btnUpdate)
        self.root.bind('<space>',self.btnUpdate)
        self.root.bind('<Escape>',self.btnExit)
        self.root.lift()
        self.root.protocol("WM_DELETE_WINDOW", self.btnExit)
        self.root.call('wm', 'attributes', '.', '-topmost', True)
        self.root.after_idle(self.root.call, 'wm', 'attributes', '.', '-topmost', False)
        #Thread
        self.check_time_thread = Thread(target=self.check_time_thread, args=(self,))
        self.check_time_thread.start()
        # Main loop
        self.root.mainloop()

    def update_T_button(self):
        button = self.get_item("button","Tx")["button_ctl"]
        if self.t4 is not None:
            button["text"] = "T-"
            button["state"] = DISABLED
        elif self.t3 is not None:
            button["text"] = "T4"
        elif self.t2 is not None:
            button["text"] = "T3"
        else:
            pass

    def check_time(self):
        self.wtime_out = self.wt.calc2()
        self.gui_update()

    def btnTx(self, *args):
        ts = wtime.get_ts()
        if self.t2 is None:
            self.t2 = ts
        elif self.t3 is None:
            self.t3 = ts
        elif self.t4 is None:
            self.t4 = ts
        else:
            return
        self.wt = wtime(t1=self.t1, t2=self.t2, t3=self.t3, t4=self.t4, current_time=self.ct)
        self.update_T_button()
        self.btnUpdate()

    def btnExit(self, *args):
        self.btnUpdate()
        self.flag_thread_running = False
        self.root.destroy()
        sys.exit(0)

    def btnUpdate(self, *args):
        self.check_time()
        self.wt.printout(self.wtime_out)
        #print(self.wtime_out)

    def btnUnknown(self, *args):
        print("WARNING: Unknown button pressed")

    def show_message_box(self, message):
        self.root.attributes("-topmost", True)
        tkMessageBox.showinfo("wtimegui", message,parent=self.root)
        self.root.attributes("-topmost", False)

    def gui_build(self):
        for item in self.GUI_data:
            if item["type"] == "text":
                item["label_var"] = StringVar()
                item["value_var"] = StringVar()
                item["label_ctl"] = Label(self.root,
                                          textvariable=item["label_var"],
                                          text="None",
                                          font=self.lblFONT,
                                          fg=self.lblFGCOLOR).grid(row=item["row"],
                                                                   column=item["col"])
                item["value_ctl"] = Label(self.root,
                                         textvariable = item["value_var"],
                                         text="None",
                                         font=self.lblFONT,
                                         fg=self.lblFGCOLOR).grid(row=item["row"],
                                                                  column=item["col"]+1)
            elif item["type"] == "progress":
                item["progress_ctl"] = ttk.Progressbar(self.root,
                                                       orient=HORIZONTAL,
                                                       length=64,
                                                       mode='determinate')
                item["progress_ctl"].grid(row=item["row"],
                                          column=item["col"])
            elif item["type"] == "button":
                if item["title"] == "Exit":
                    callback = self.btnExit
                elif item["title"] == "Update":
                    callback = self.btnUpdate
                elif item["title"][0] == "T":
                    callback = self.btnTx
                else:
                    print("WARNING: Unhespected button named: %s" % item["title"])
                    callback = self.btnUnknown
                item['button_ctl'] = Button(self.root,
                                            text=item["title"],
                                            command=callback)
                item['button_ctl'].grid(row=item["row"],
                                        column=item["col"])
            else:
              print("WARNING: Skipping unknown type: '%s' for item '%s'"
                    % (item["type"], item["title"]))
        self.update_T_button()

    def gui_update(self):
        for item in self.GUI_data:
            if item["type"] == "text":
                if item["name"] == "time remaining perc" or item["name"] == "ticket remaining perc":
                    perc = max(0, self.wtime_out.get(item["name"],""))
                    item["label_var"].set("%2d%%: " % perc)
                elif item["name"] == "total time" and self.wtime_out.get("consume pause", None) is not None:
                    item["label_var"].set("Consuming pause: ")
                    item["value_var"].set(self.wtime_out["consume pause"])
                else:
                    value = self.wtime_out.get(item["name"],"")
                    if value != "":
                        item["label_var"].set(item["title"] + " :")
                        item["value_var"].set(value)
                    else:
                        pass
            elif item["type"] == "progress":
                item["progress_ctl"]["value"] = self.wtime_out["time remaining perc"]
                item["progress_ctl"]["value"] = self.wtime_out["ticket remaining perc"]
            elif item["type"] == "button":
                pass
            else:
                print("WARNING: Skipping unknown type: '%s' for item '%s'"
                        % (item["type"], item["title"]))

    def check_time_thread(self, *args):
        gui = args[0]
        if gui is None:
            print("ERROR: No GUI passed")
        time.sleep(.1)
        try:
            t = currentThread()
        except NameError as e:
            t = current_thread()
            print("wtime updating thread started")
            self.flag_thread_running = True
            while self.flag_thread_running:
                overtime = self.wtime_out.get("overtime", None)
                if overtime is not None and self.flag_time_reached == False:
                    print("You've DONE!!!")
                    self.flag_time_reached = True
                    self.flag_ticket_reached = True
                    gui.show_message_box("You've DONE!!!")
                elif overtime is not None and self.flag_time_reached == True:
                    self.TotalTimeText = "Overtime: "
                    self.gui_build()
                elif self.wtime_out["ticket remaining"] == "reached" and self.flag_ticket_reached == False:
                    print("Ticket reached!!!")
                    self.flag_ticket_reached = True
                    gui.show_message_box("Ticket reached!!!")
                self.check_time()
                for i in range(1,10 * self.interval_thread_waitcycles):
                    if self.flag_thread_running:
                        time.sleep(.1)
                    else:
                        break
        except Exception as e:
            print("Exception: %s" % type(e).__name__ )
            print("wtime updating thread not started: %s" % e)
        print("wtime updating thread terminated")

if __name__ == "__main__":
    gui = wtimeGUI()

<commit_msg>Ticket and Time info in bold, fix on progress bars<commit_after>#!/usr/bin/env python3
#
# wtimegui - Working time class with GUI
#
import sys
import time
try:
    from Tkinter import *
    import ttk
    import tkMessageBox
    from threading import *
except ImportError:
    from tkinter import *
    from tkinter import ttk
    from threading import *
    from tkinter import messagebox as tkMessageBox

from wtime4 import wtime

__author__ = "Riccardo Bruno"
__copyright__ = "2017"
__license__ = "Apache"
__maintainer__ = "Riccardo Bruno"
__email__ = "riccardo.bruno@gmail.com"

class wtimeGUI:

    flag_ticket_reached = False
    flag_time_reached = False
    flag_thread_running = False
    interval_thread_waitcycles = 5
    check_time_thread=None
    wtime_out = {}

    winTITLE="wtime GUI"
    lblFONT=("Lucida Grande", 12)
    lblFGCOLOR='black'

    root = None

    GUI_data = (
        {"type": "text", "name": "t1", "title": "T1", "row": 0, "col": 0},
        {"type": "text", "name": "t2", "title": "T2", "row": 1, "col": 0},
        {"type": "text", "name": "t2t1", "title": "T2 - T1", "row": 1, "col": 2},
        {"type": "text", "name": "t3", "title": "T3", "row": 2, "col": 0},
        {"type": "text", "name": "t4", "title": "T4", "row": 3, "col": 0},
        {"type": "text", "name": "t4t3", "title": "T4 - T3", "row": 3, "col": 2},
        {"type": "text", "name": "pause time", "title": "Pause Time", "row": 4, "col": 0},
        {"type": "text", "name": "total time", "title": "Total Time", "row": 5, "col": 0},
        {"type": "text", "name": "overtime", "title": "Over Time", "row": 5, "col": 2},
        {"type": "text", "name": "time to reach", "title": "Time to reach", "row": 6, "col": 0},
        {"type": "text", "name": "ticket remaining", "title": "Ticket remain", "row": 7, "col": 0},
        {"type": "text", "name": "ticket remaining at", "title": "at", "row": 8, "col": 0},
        {"type": "text", "name": "ticket time", "title": "TicketTime", "row": 9, "col": 0},
        {"type": "text", "name": "ticket remaining perc", "title": "%", "row": 8, "col": 2},
        {"type": "progress", "name": "ticket progress", "title": "Ticket progress", "row": 8, "col": 3},
        {"type": "text", "name": "time remaining", "title": "Time remain", "row": 10, "col": 0},
        {"type": "text", "name": "time remaining at", "title": "at", "row": 11, "col": 0},
        {"type": "text", "name": "time remaining perc", "title": "%", "row": 11, "col": 2},
        {"type": "progress", "name": "time progress", "title": "Time progress", "row": 11, "col": 3},
        {"type": "button", "name": "Tx", "title": "T2", "row": 12, "col": 0},
        {"type": "button", "name": "Update", "title": "Update", "row": 12, "col": 1},
        {"type": "button", "name": "Exit", "title": "Exit", "row": 12, "col": 3},
    )

    def get_item(self, type, name):
        item_result = None
        for item in self.GUI_data:
            if item["type"] == type and item["name"] == name:
                item_result = item
                break
        return item_result

    def __init__(self):
        # wtime4
        self.t1, self.t2, self.t3, self.t4, self.ct = wtime.getTimes(sys.argv)
        self.wt = wtime(t1=self.t1, t2=self.t2, t3=self.t3, t4=self.t4, current_time=self.ct)
        # GUI
        self.root = Tk()
        self.root.title(self.winTITLE)
        self.gui_build()
        self.check_time()
        self.root.bind('<Return>',self.btnUpdate)
        self.root.bind('<space>',self.btnUpdate)
        self.root.bind('<Escape>',self.btnExit)
        self.root.lift()
        self.root.protocol("WM_DELETE_WINDOW", self.btnExit)
        self.root.call('wm', 'attributes', '.', '-topmost', True)
        self.root.after_idle(self.root.call, 'wm', 'attributes', '.', '-topmost', False)
        #Thread
        self.check_time_thread = Thread(target=self.check_time_thread, args=(self,))
        self.check_time_thread.start()
        # Main loop
        self.root.mainloop()

    def update_T_button(self):
        button = self.get_item("button","Tx")["button_ctl"]
        if self.t4 is not None:
            button["text"] = "T-"
            button["state"] = DISABLED
        elif self.t3 is not None:
            button["text"] = "T4"
        elif self.t2 is not None:
            button["text"] = "T3"
        else:
            pass

    def check_time(self):
        self.wtime_out = self.wt.calc2()
        self.gui_update()

    def btnTx(self, *args):
        ts = wtime.get_ts()
        if self.t2 is None:
            self.t2 = ts
        elif self.t3 is None:
            self.t3 = ts
        elif self.t4 is None:
            self.t4 = ts
        else:
            return
        self.wt = wtime(t1=self.t1, t2=self.t2, t3=self.t3, t4=self.t4, current_time=self.ct)
        self.update_T_button()
        self.btnUpdate()

    def btnExit(self, *args):
        self.btnUpdate()
        self.flag_thread_running = False
        self.root.destroy()
        sys.exit(0)

    def btnUpdate(self, *args):
        self.check_time()
        self.wt.printout(self.wtime_out)
        #print(self.wtime_out)

    def btnUnknown(self, *args):
        print("WARNING: Unknown button pressed")

    def show_message_box(self, message):
        self.root.attributes("-topmost", True)
        tkMessageBox.showinfo("wtimegui", message,parent=self.root)
        self.root.attributes("-topmost", False)

    def gui_build(self):
        for item in self.GUI_data:
            if item["type"] == "text":
                if item["name"] in ("ticket remaining",
                                    "ticket remaining at",
                                    "time remaining",
                                    "time remaining at",
                                    "overtime"):
                    lblFONT_val_style = ("bold",)
                else:
                    lblFONT_val_style = ()
                if item["name"] in ("ticket remaining",
                                    "ticket remaining perc",
                                    "time remaining",
                                    "time remaining perc",
                                    "overtime"):
                    lblFONT_lbl_style = ("bold",)
                else:
                    lblFONT_lbl_style = ()
                item["label_var"] = StringVar()
                item["value_var"] = StringVar()
                item["label_ctl"] = Label(self.root,
                                          textvariable=item["label_var"],
                                          text="None",
                                          font=self.lblFONT + lblFONT_lbl_style,
                                          fg=self.lblFGCOLOR).grid(row=item["row"],
                                                                   column=item["col"])
                item["value_ctl"] = Label(self.root,
                                         textvariable = item["value_var"],
                                         text="None",
                                         font=self.lblFONT + lblFONT_val_style,
                                         fg=self.lblFGCOLOR).grid(row=item["row"],
                                                                  column=item["col"]+1)
            elif item["type"] == "progress":
                item["progress_ctl"] = ttk.Progressbar(self.root,
                                                       orient=HORIZONTAL,
                                                       length=64,
                                                       mode='determinate')
                item["progress_ctl"].grid(row=item["row"],
                                          column=item["col"])
            elif item["type"] == "button":
                if item["title"] == "Exit":
                    callback = self.btnExit
                elif item["title"] == "Update":
                    callback = self.btnUpdate
                elif item["title"][0] == "T":
                    callback = self.btnTx
                else:
                    print("WARNING: Unhespected button named: %s" % item["title"])
                    callback = self.btnUnknown
                item['button_ctl'] = Button(self.root,
                                            text=item["title"],
                                            command=callback)
                item['button_ctl'].grid(row=item["row"],
                                        column=item["col"])
            else:
              print("WARNING: Skipping unknown type: '%s' for item '%s'"
                    % (item["type"], item["title"]))
        self.update_T_button()

    def gui_update(self):
        for item in self.GUI_data:
            if item["type"] == "text":
                if item["name"] == "time remaining perc" or item["name"] == "ticket remaining perc":
                    perc = max(0, self.wtime_out.get(item["name"],""))
                    item["label_var"].set("%2d%%: " % perc)
                elif item["name"] == "total time" and self.wtime_out.get("consume pause", None) is not None:
                    item["label_var"].set("Consuming pause: ")
                    item["value_var"].set(self.wtime_out["consume pause"])
                else:
                    value = self.wtime_out.get(item["name"],"")
                    if value != "":
                        item["label_var"].set(item["title"] + " :")
                        item["value_var"].set(value)
                    else:
                        pass
            elif item["type"] == "progress":
                if item["name"] == "time progress":
                    item["progress_ctl"]["value"] = self.wtime_out["time remaining perc"]
                elif item["name"] == "ticket progress":
                    item["progress_ctl"]["value"] = self.wtime_out["ticket remaining perc"]
                else:
                    pass
            elif item["type"] == "button":
                pass
            else:
                print("WARNING: Skipping unknown type: '%s' for item '%s'"
                        % (item["type"], item["title"]))

    def check_time_thread(self, *args):
        gui = args[0]
        if gui is None:
            print("ERROR: No GUI passed")
        time.sleep(.1)
        try:
            t = currentThread()
        except NameError as e:
            t = current_thread()
            print("wtime updating thread started")
            self.flag_thread_running = True
            while self.flag_thread_running:
                overtime = self.wtime_out.get("overtime", None)
                if overtime is not None and self.flag_time_reached == False:
                    print("You've DONE!!!")
                    self.flag_time_reached = True
                    self.flag_ticket_reached = True
                    gui.show_message_box("You've DONE!!!")
                elif overtime is not None and self.flag_time_reached == True:
                    self.TotalTimeText = "Overtime: "
                    self.gui_build()
                elif self.wtime_out["ticket remaining"] == "reached" and self.flag_ticket_reached == False:
                    print("Ticket reached!!!")
                    self.flag_ticket_reached = True
                    gui.show_message_box("Ticket reached!!!")
                self.check_time()
                for i in range(1,10 * self.interval_thread_waitcycles):
                    if self.flag_thread_running:
                        time.sleep(.1)
                    else:
                        break
        except Exception as e:
            print("Exception: %s" % type(e).__name__ )
            print("wtime updating thread not started: %s" % e)
        print("wtime updating thread terminated")

if __name__ == "__main__":
    gui = wtimeGUI()

<|endoftext|>


<commit_before>from __future__ import absolute_import

import logging

from datetime import datetime, timedelta
from functools import wraps
from threading import local, Lock
from uuid import uuid4

from changes.config import db, queue
from changes.constants import Result, Status
from changes.db.utils import get_or_create
from changes.experimental.stats import incr, decr
from changes.models import Task
from changes.utils.locking import lock

BASE_RETRY_COUNTDOWN = 60
CONTINUE_COUNTDOWN = 5

RUN_TIMEOUT = timedelta(minutes=60)
EXPIRE_TIMEOUT = timedelta(minutes=120)
HARD_TIMEOUT = timedelta(hours=12)

MAX_RETRIES = 10

class NotFinished(Exception):
    pass

class TooManyRetries(Exception):
    pass

class TrackedTask(local):
    """
    Tracks the state of the given Task and it's children.

    Tracked tasks **never** return a result.

    >>> @tracked_task
    >>> def func(foo):
    >>>    if random.randint(0, 1) == 1:
    >>>        # re-queue for further results
    >>>        raise func.NotFinished
    >>>
    >>>    elif random.randint(0, 1) == 1:
    >>>        # cause an exception to retry
    >>>        raise Exception
    >>>
    >>>    # finish normally to update Status
    >>>    print "func", foo
    >>> foo.delay(foo='bar', task_id='bar')
    """
    NotFinished = NotFinished

    def __init__(self, func, max_retries=MAX_RETRIES, on_abort=None):
        incr('task_construct')
        self.func = lock(func)
        self.task_name = func.__name__
        self.parent_id = None
        self.task_id = None
        self.lock = Lock()
        self.logger = logging.getLogger('jobs.{0}'.format(self.task_name))

        self.max_retries = max_retries
        self.on_abort = on_abort

        self.__name__ = func.__name__
        self.__doc__ = func.__doc__
        self.__wraps__ = getattr(func, '__wraps__', func)
        self.__code__ = getattr(func, '__code__', None)

    def __call__(self, **kwargs):
        incr('task_lock_wait')
        with self.lock:
            decr('task_lock_wait')
            incr('task_run')
            self._run(kwargs)
            decr('task_run')

    def __repr__(self):
        return '<%s: task_name=%s>' % (type(self), self.task_name)

    def _run(self, kwargs):
        self.task_id = kwargs.pop('task_id', None)
        if self.task_id is None:
            self.task_id = uuid4().hex

        self.parent_id = kwargs.pop('parent_task_id', None)
        self.kwargs = kwargs

        date_started = datetime.utcnow()

        self._update({
            Task.date_modified: datetime.utcnow(),
        })

        try:
            self.func(**kwargs)

        except NotFinished:
            self.logger.info(
                'Task marked as not finished: %s %s', self.task_name, self.task_id)

            self._continue(kwargs)

        except Exception as exc:
            db.session.rollback()

            self.logger.exception(unicode(exc))

            try:
                self._retry()
            except TooManyRetries as exc:
                date_finished = datetime.utcnow()

                self._update({
                    Task.date_finished: date_finished,
                    Task.date_modified: date_finished,
                    Task.status: Status.finished,
                    Task.result: Result.failed,
                })
                self.logger.exception(unicode(exc))

                if self.on_abort:
                    self.on_abort(self)
            except Exception as exc:
                self.logger.exception(unicode(exc))
                raise

        else:
            date_finished = datetime.utcnow()

            try:
                self._update({
                    Task.date_started: date_started,
                    Task.date_finished: date_finished,
                    Task.date_modified: date_finished,
                    Task.status: Status.finished,
                    Task.result: Result.passed,
                })
            except Exception as exc:
                self.logger.exception(unicode(exc))
                raise

            db.session.commit()
        finally:
            db.session.expire_all()

            self.task_id = None
            self.parent_id = None
            self.kwargs = kwargs

    def _update(self, kwargs):
        """
        Update's the state of this Task.

        >>> task._update(status=Status.finished)
        """
        assert self.task_id

        Task.query.filter(
            Task.task_name == self.task_name,
            Task.task_id == self.task_id,
            Task.parent_id == self.parent_id,
        ).update(kwargs, synchronize_session=False)

    def _continue(self, kwargs):
        kwargs['task_id'] = self.task_id
        kwargs['parent_task_id'] = self.parent_id

        self._update({
            Task.date_modified: datetime.utcnow(),
            Task.status: Status.in_progress,
        })

        db.session.commit()

        queue.delay(
            self.task_name,
            kwargs=kwargs,
            countdown=CONTINUE_COUNTDOWN,
        )

    def _retry(self):
        """
        Retry this task and update it's state.

        >>> task.retry()
        """
        # TODO(dcramer): this needs to handle too-many-retries itself
        assert self.task_id

        task = Task.query.filter(
            Task.task_name == self.task_name,
            Task.task_id == self.task_id,
            Task.parent_id == self.parent_id,
        ).first()
        if task and self.max_retries and task.num_retries > self.max_retries:
            date_finished = datetime.utcnow()
            self._update({
                Task.date_finished: date_finished,
                Task.date_modified: date_finished,
                Task.status: Status.finished,
                Task.result: Result.failed,
            })
            db.session.commit()

            raise TooManyRetries('%s failed after %d retries' % (self.task_name, task.num_retries))

        self._update({
            Task.date_modified: datetime.utcnow(),
            Task.status: Status.in_progress,
            Task.num_retries: Task.num_retries + 1,
        })

        db.session.commit()

        kwargs = self.kwargs.copy()
        kwargs['task_id'] = self.task_id
        kwargs['parent_task_id'] = self.parent_id

        retry_number = db.session.query(Task.num_retries).filter(
            Task.task_name == self.task_name,
            Task.task_id == self.task_id,
            Task.parent_id == self.parent_id,
        ).scalar() or 0

        retry_countdown = min(BASE_RETRY_COUNTDOWN + (retry_number ** 2), 300)

        queue.delay(
            self.task_name,
            kwargs=kwargs,
            countdown=retry_countdown,
        )

    def needs_requeued(self, task):
        if self.max_retries and task.num_retries >= self.max_retries:
            return False

        current_datetime = datetime.utcnow()

        timeout_datetime = current_datetime - HARD_TIMEOUT
        if task.date_created < timeout_datetime:
            return False

        run_datetime = current_datetime - RUN_TIMEOUT
        return task.date_modified < run_datetime

    def needs_expired(self, task):
        if self.max_retries and task.num_retries >= self.max_retries:
            return True

        current_datetime = datetime.utcnow()

        timeout_datetime = current_datetime - HARD_TIMEOUT
        if task.date_created < timeout_datetime:
            return True

        expire_datetime = current_datetime - EXPIRE_TIMEOUT
        if task.date_modified < expire_datetime:
            return True

        return False

    def delay_if_needed(self, **kwargs):
        """
        Enqueue this task if it's new or hasn't checked in in a reasonable
        amount of time.

        >>> task.delay_if_needed(
        >>>     task_id='33846695b2774b29a71795a009e8168a',
        >>>     parent_task_id='659974858dcf4aa08e73a940e1066328',
        >>> )
        """
        kwargs.setdefault('task_id', uuid4().hex)

        fn_kwargs = dict(
            (k, v) for k, v in kwargs.iteritems()
            if k not in ('task_id', 'parent_task_id')
        )

        task, created = get_or_create(Task, where={
            'task_name': self.task_name,
            'task_id': kwargs['task_id'],
        }, defaults={
            'parent_id': kwargs.get('parent_task_id'),
            'data': {
                'kwargs': fn_kwargs,
            },
            'status': Status.queued,
        })

        if created or self.needs_requeued(task):
            if not created:
                task.date_modified = datetime.utcnow()
                db.session.add(task)

            db.session.commit()

            queue.delay(
                self.task_name,
                kwargs=kwargs,
                countdown=CONTINUE_COUNTDOWN,
            )

    def delay(self, **kwargs):
        """
        Enqueue this task.

        >>> task.delay(
        >>>     task_id='33846695b2774b29a71795a009e8168a',
        >>>     parent_task_id='659974858dcf4aa08e73a940e1066328',
        >>> )
        """
        kwargs.setdefault('task_id', uuid4().hex)

        fn_kwargs = dict(
            (k, v) for k, v in kwargs.iteritems()
            if k not in ('task_id', 'parent_task_id')
        )

        task, created = get_or_create(Task, where={
            'task_name': self.task_name,
            'task_id': kwargs['task_id'],
        }, defaults={
            'parent_id': kwargs.get('parent_task_id'),
            'data': {
                'kwargs': fn_kwargs,
            },
            'status': Status.queued,
        })

        if not created:
            task.date_modified = datetime.utcnow()
            db.session.add(task)

        db.session.commit()

        queue.delay(
            self.task_name,
            kwargs=kwargs,
            countdown=CONTINUE_COUNTDOWN,
        )

    def verify_all_children(self):
        task_list = list(Task.query.filter(
            Task.parent_id == self.task_id,
            Task.status != Status.finished,
        ))

        if not task_list:
            return Status.finished

        current_datetime = datetime.utcnow()

        need_expire = set()
        need_run = set()

        has_pending = False

        for task in task_list:
            if self.needs_expired(task):
                need_expire.add(task)
                continue

            has_pending = True

            if self.needs_requeued(task) and 'kwargs' in task.data:
                need_run.add(task)

        if need_expire:
            Task.query.filter(
                Task.id.in_([n.id for n in need_expire]),
            ).update({
                Task.date_modified: current_datetime,
                Task.date_finished: current_datetime,
                Task.status: Status.finished,
                Task.result: Result.aborted,
            }, synchronize_session=False)
            db.session.commit()

        if need_run:
            for task in need_run:
                child_kwargs = task.data['kwargs'].copy()
                child_kwargs['parent_task_id'] = task.parent_id.hex
                child_kwargs['task_id'] = task.task_id.hex
                queue.delay(task.task_name, kwargs=child_kwargs)

            Task.query.filter(
                Task.id.in_([n.id for n in need_run]),
            ).update({
                Task.date_modified: current_datetime,
            }, synchronize_session=False)
            db.session.commit()

        if has_pending:
            status = Status.in_progress

        else:
            status = Status.finished

        return status

# bind to a decorator-like naming scheme
def tracked_task(func=None, **kwargs):
    def wrapped(func):
        return wraps(func)(TrackedTask(func, **kwargs))

    if func:
        return wraps(func)(wrapped(func))
    return wrapped
<commit_msg>Revert "Track tasks only"<commit_after>from __future__ import absolute_import

import logging

from datetime import datetime, timedelta
from functools import wraps
from threading import local, Lock
from uuid import uuid4

from changes.config import db, queue
from changes.constants import Result, Status
from changes.db.utils import get_or_create
from changes.models import Task
from changes.utils.locking import lock

BASE_RETRY_COUNTDOWN = 60
CONTINUE_COUNTDOWN = 5

RUN_TIMEOUT = timedelta(minutes=60)
EXPIRE_TIMEOUT = timedelta(minutes=120)
HARD_TIMEOUT = timedelta(hours=12)

MAX_RETRIES = 10

class NotFinished(Exception):
    pass

class TooManyRetries(Exception):
    pass

class TrackedTask(local):
    """
    Tracks the state of the given Task and it's children.

    Tracked tasks **never** return a result.

    >>> @tracked_task
    >>> def func(foo):
    >>>    if random.randint(0, 1) == 1:
    >>>        # re-queue for further results
    >>>        raise func.NotFinished
    >>>
    >>>    elif random.randint(0, 1) == 1:
    >>>        # cause an exception to retry
    >>>        raise Exception
    >>>
    >>>    # finish normally to update Status
    >>>    print "func", foo
    >>> foo.delay(foo='bar', task_id='bar')
    """
    NotFinished = NotFinished

    def __init__(self, func, max_retries=MAX_RETRIES, on_abort=None):
        self.func = lock(func)
        self.task_name = func.__name__
        self.parent_id = None
        self.task_id = None
        self.lock = Lock()
        self.logger = logging.getLogger('jobs.{0}'.format(self.task_name))

        self.max_retries = max_retries
        self.on_abort = on_abort

        self.__name__ = func.__name__
        self.__doc__ = func.__doc__
        self.__wraps__ = getattr(func, '__wraps__', func)
        self.__code__ = getattr(func, '__code__', None)

    def __call__(self, **kwargs):
        with self.lock:
            self._run(kwargs)

    def __repr__(self):
        return '<%s: task_name=%s>' % (type(self), self.task_name)

    def _run(self, kwargs):
        self.task_id = kwargs.pop('task_id', None)
        if self.task_id is None:
            self.task_id = uuid4().hex

        self.parent_id = kwargs.pop('parent_task_id', None)
        self.kwargs = kwargs

        date_started = datetime.utcnow()

        self._update({
            Task.date_modified: datetime.utcnow(),
        })

        try:
            self.func(**kwargs)

        except NotFinished:
            self.logger.info(
                'Task marked as not finished: %s %s', self.task_name, self.task_id)

            self._continue(kwargs)

        except Exception as exc:
            db.session.rollback()

            self.logger.exception(unicode(exc))

            try:
                self._retry()
            except TooManyRetries as exc:
                date_finished = datetime.utcnow()

                self._update({
                    Task.date_finished: date_finished,
                    Task.date_modified: date_finished,
                    Task.status: Status.finished,
                    Task.result: Result.failed,
                })
                self.logger.exception(unicode(exc))

                if self.on_abort:
                    self.on_abort(self)
            except Exception as exc:
                self.logger.exception(unicode(exc))
                raise

        else:
            date_finished = datetime.utcnow()

            try:
                self._update({
                    Task.date_started: date_started,
                    Task.date_finished: date_finished,
                    Task.date_modified: date_finished,
                    Task.status: Status.finished,
                    Task.result: Result.passed,
                })
            except Exception as exc:
                self.logger.exception(unicode(exc))
                raise

            db.session.commit()
        finally:
            db.session.expire_all()

            self.task_id = None
            self.parent_id = None
            self.kwargs = kwargs

    def _update(self, kwargs):
        """
        Update's the state of this Task.

        >>> task._update(status=Status.finished)
        """
        assert self.task_id

        Task.query.filter(
            Task.task_name == self.task_name,
            Task.task_id == self.task_id,
            Task.parent_id == self.parent_id,
        ).update(kwargs, synchronize_session=False)

    def _continue(self, kwargs):
        kwargs['task_id'] = self.task_id
        kwargs['parent_task_id'] = self.parent_id

        self._update({
            Task.date_modified: datetime.utcnow(),
            Task.status: Status.in_progress,
        })

        db.session.commit()

        queue.delay(
            self.task_name,
            kwargs=kwargs,
            countdown=CONTINUE_COUNTDOWN,
        )

    def _retry(self):
        """
        Retry this task and update it's state.

        >>> task.retry()
        """
        # TODO(dcramer): this needs to handle too-many-retries itself
        assert self.task_id

        task = Task.query.filter(
            Task.task_name == self.task_name,
            Task.task_id == self.task_id,
            Task.parent_id == self.parent_id,
        ).first()
        if task and self.max_retries and task.num_retries > self.max_retries:
            date_finished = datetime.utcnow()
            self._update({
                Task.date_finished: date_finished,
                Task.date_modified: date_finished,
                Task.status: Status.finished,
                Task.result: Result.failed,
            })
            db.session.commit()

            raise TooManyRetries('%s failed after %d retries' % (self.task_name, task.num_retries))

        self._update({
            Task.date_modified: datetime.utcnow(),
            Task.status: Status.in_progress,
            Task.num_retries: Task.num_retries + 1,
        })

        db.session.commit()

        kwargs = self.kwargs.copy()
        kwargs['task_id'] = self.task_id
        kwargs['parent_task_id'] = self.parent_id

        retry_number = db.session.query(Task.num_retries).filter(
            Task.task_name == self.task_name,
            Task.task_id == self.task_id,
            Task.parent_id == self.parent_id,
        ).scalar() or 0

        retry_countdown = min(BASE_RETRY_COUNTDOWN + (retry_number ** 2), 300)

        queue.delay(
            self.task_name,
            kwargs=kwargs,
            countdown=retry_countdown,
        )

    def needs_requeued(self, task):
        if self.max_retries and task.num_retries >= self.max_retries:
            return False

        current_datetime = datetime.utcnow()

        timeout_datetime = current_datetime - HARD_TIMEOUT
        if task.date_created < timeout_datetime:
            return False

        run_datetime = current_datetime - RUN_TIMEOUT
        return task.date_modified < run_datetime

    def needs_expired(self, task):
        if self.max_retries and task.num_retries >= self.max_retries:
            return True

        current_datetime = datetime.utcnow()

        timeout_datetime = current_datetime - HARD_TIMEOUT
        if task.date_created < timeout_datetime:
            return True

        expire_datetime = current_datetime - EXPIRE_TIMEOUT
        if task.date_modified < expire_datetime:
            return True

        return False

    def delay_if_needed(self, **kwargs):
        """
        Enqueue this task if it's new or hasn't checked in in a reasonable
        amount of time.

        >>> task.delay_if_needed(
        >>>     task_id='33846695b2774b29a71795a009e8168a',
        >>>     parent_task_id='659974858dcf4aa08e73a940e1066328',
        >>> )
        """
        kwargs.setdefault('task_id', uuid4().hex)

        fn_kwargs = dict(
            (k, v) for k, v in kwargs.iteritems()
            if k not in ('task_id', 'parent_task_id')
        )

        task, created = get_or_create(Task, where={
            'task_name': self.task_name,
            'task_id': kwargs['task_id'],
        }, defaults={
            'parent_id': kwargs.get('parent_task_id'),
            'data': {
                'kwargs': fn_kwargs,
            },
            'status': Status.queued,
        })

        if created or self.needs_requeued(task):
            if not created:
                task.date_modified = datetime.utcnow()
                db.session.add(task)

            db.session.commit()

            queue.delay(
                self.task_name,
                kwargs=kwargs,
                countdown=CONTINUE_COUNTDOWN,
            )

    def delay(self, **kwargs):
        """
        Enqueue this task.

        >>> task.delay(
        >>>     task_id='33846695b2774b29a71795a009e8168a',
        >>>     parent_task_id='659974858dcf4aa08e73a940e1066328',
        >>> )
        """
        kwargs.setdefault('task_id', uuid4().hex)

        fn_kwargs = dict(
            (k, v) for k, v in kwargs.iteritems()
            if k not in ('task_id', 'parent_task_id')
        )

        task, created = get_or_create(Task, where={
            'task_name': self.task_name,
            'task_id': kwargs['task_id'],
        }, defaults={
            'parent_id': kwargs.get('parent_task_id'),
            'data': {
                'kwargs': fn_kwargs,
            },
            'status': Status.queued,
        })

        if not created:
            task.date_modified = datetime.utcnow()
            db.session.add(task)

        db.session.commit()

        queue.delay(
            self.task_name,
            kwargs=kwargs,
            countdown=CONTINUE_COUNTDOWN,
        )

    def verify_all_children(self):
        task_list = list(Task.query.filter(
            Task.parent_id == self.task_id,
            Task.status != Status.finished,
        ))

        if not task_list:
            return Status.finished

        current_datetime = datetime.utcnow()

        need_expire = set()
        need_run = set()

        has_pending = False

        for task in task_list:
            if self.needs_expired(task):
                need_expire.add(task)
                continue

            has_pending = True

            if self.needs_requeued(task) and 'kwargs' in task.data:
                need_run.add(task)

        if need_expire:
            Task.query.filter(
                Task.id.in_([n.id for n in need_expire]),
            ).update({
                Task.date_modified: current_datetime,
                Task.date_finished: current_datetime,
                Task.status: Status.finished,
                Task.result: Result.aborted,
            }, synchronize_session=False)
            db.session.commit()

        if need_run:
            for task in need_run:
                child_kwargs = task.data['kwargs'].copy()
                child_kwargs['parent_task_id'] = task.parent_id.hex
                child_kwargs['task_id'] = task.task_id.hex
                queue.delay(task.task_name, kwargs=child_kwargs)

            Task.query.filter(
                Task.id.in_([n.id for n in need_run]),
            ).update({
                Task.date_modified: current_datetime,
            }, synchronize_session=False)
            db.session.commit()

        if has_pending:
            status = Status.in_progress

        else:
            status = Status.finished

        return status

# bind to a decorator-like naming scheme
def tracked_task(func=None, **kwargs):
    def wrapped(func):
        return wraps(func)(TrackedTask(func, **kwargs))

    if func:
        return wraps(func)(wrapped(func))
    return wrapped
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding:utf-8 -*-
from utils.mysql_util import MysqlUtil
from sqlalchemy import desc
from models import Article

import datetime

class ArticleDao:
    def __init__(self):
        self.session = MysqlUtil().get_session()

    def get_articles(self, page=0, size=20):
        article_num = self.get_article_num()
        if page * size > article_num:
            return None, u'分页数超出最大值'
        info = u'more' if (page + 1) * size < article_num else u'nomore'
        return self.session.query(Article).order_by(desc(Article.modified_time)).offset(size * page).limit(size).all(), info

    def get_article_num(self):
        return self.session.query(Article).count()

    def get_article_by_title(self, title):
        return self.session.query(Article).filter_by(title=title).first()

    def get_article_by_id(self, id):
        return self.session.query(Article).filter_by(id=id).first()

    def new_article(self, title, author_id, author_name, cate_id, cate_name, intro, filepath, tags):
        article = Article(
            title=title,
            intro=intro,
            is_public=1,
            auth_id=author_id,
            auth_name=author_name,
            cate_id=cate_id,
            cate_name=cate_name,
            file_path=filepath,
            tags=tags,
            create_time=datetime.datetime.now(),
            modified_time=datetime.datetime.now()
        )

        if self.get_article_by_title(title) is None:
            self.session.add(article)
            self.session.commit()
            return True, u'成功添加新文章!'
        else:
            return False, u'文章标题已经存在!'

    def delete_article(self, article):
        self.session.delete(article)
        self.session.commit()
        return True, u'删除成功!'

    def get_articles_by_cate(self, cate, page=0, size=20):
        article_num = self.get_article_num()
        if page * size > article_num:
            return None, u'分页数超出最大值'
        info = u'more' if (page + 1) * size < article_num else u'nomore'
        return self.session.query(Article).filter_by(cate_id=cate).order_by(desc(Article.modified_time)).offset(size * page).limit(size).all(), info

    def get_article_by_tag(self, tag, page=0, size=20):
        article_num = self.get_article_num()
        if page * size > article_num:
            return None, u'分页数超出最大值'
        info = u'more' if (page + 1) * size < article_num else u'nomore'
        return self.session.query(Article)\
                   .filter(Article.tags.like('%{tag}%'.format(tag=tag)))\
                   .order_by(desc(Article.modified_time)).offset(size * page).limit(size).all(), info

if __name__ == '__main__':
    print ArticleDao().get_article_by_tag('hive')<commit_msg>no message<commit_after>#!/usr/bin/python
# -*- coding:utf-8 -*-
from utils.mysql_util import MysqlUtil
from sqlalchemy import desc
from models import Article

import datetime

class ArticleDao:
    def __init__(self):
        self.session = MysqlUtil().get_session()

    def get_articles(self, page=0, size=20):
        article_num = self.get_article_num()
        if page * size > article_num:
            return None, u'分页数超出最大值'
        info = u'more' if (page + 1) * size < article_num else u'nomore'
        return self.session.query(Article).order_by(desc(Article.modified_time)).offset(size * page).limit(size).all(), info

    def get_article_num(self):
        return self.session.query(Article).count()

    def get_article_by_title(self, title):
        return self.session.query(Article).filter_by(title=title).first()

    def get_article_by_id(self, id):
        return self.session.query(Article).filter_by(id=id).first()

    def new_article(self, title, author_id, author_name, cate_id, cate_name, intro, filepath, tags):
        article = Article(
            title=title,
            intro=intro,
            is_public=1,
            auth_id=author_id,
            auth_name=author_name,
            cate_id=cate_id,
            cate_name=cate_name,
            file_path=filepath,
            tags=tags,
            create_time=datetime.datetime.now(),
            modified_time=datetime.datetime.now()
        )

        if self.get_article_by_title(title) is None:
            self.session.add(article)
            self.session.commit()
            return True, u'成功添加新文章!'
        else:
            return False, u'文章标题已经存在!'

    def delete_article(self, article):
        self.session.delete(article)
        self.session.commit()
        return True, u'删除成功!'

    def get_articles_by_cate(self, cate, page=0, size=20):
        article_num = self.get_article_num()
        if page * size > article_num:
            return None, u'分页数超出最大值'
        info = u'more' if (page + 1) * size < article_num else u'nomore'
        return self.session.query(Article).filter_by(cate_id=cate).order_by(desc(Article.modified_time)).offset(size * page).limit(size).all(), info

    def get_article_by_tag(self, tag, page=0, size=20):
        article_num = self.get_article_num()
        if page * size > article_num:
            return None, u'分页数超出最大值'
        info = u'more' if (page + 1) * size < article_num else u'nomore'
        return self.session.query(Article)\
                   .filter(Article.tags.like(u'%{tag}%'.format(tag=tag)))\
                   .order_by(desc(Article.modified_time)).offset(size * page).limit(size).all(), info

if __name__ == '__main__':
    print ArticleDao().get_article_by_tag('hive')<|endoftext|>


<commit_before>"""
This migration will add original_node and node associated with the log to nodelogs. It will then make
copies of each nodelog for the remaining nodes in the backref (registrations and forks),
changing the node to the current node.
"""

import sys
import logging
from modularodm import Q
from website.app import init_app
from website import models
from scripts import utils as script_utils
from framework.transactions.context import TokuTransaction

logger = logging.getLogger(__name__)

def get_original_node(log):
    """
    Retrieves original node on nodelog
    """
    return models.Node.find(Q('_id', 'eq', log._backrefs['logged']['node']['logs'][0]))[0]

def main(dry=True):
    init_app(routes=False)
    node_logs = models.NodeLog.find()
    log_count = node_logs.count()
    errored_logs = []
    for count, log in enumerate(node_logs):
        with TokuTransaction():
            log.original_node = get_original_node(log)
            log.node = get_original_node(log)
            if not dry:
                try:
                    log.save()
                except KeyError as error:
                    logger.error('Could not migrate log due to error')
                    logger.exception(error)
                    errored_logs.append(log)
                else:
                    logger.info('{}/{} Log {} "original_node and node" added'.format(count, log_count, log._id))

    if errored_logs:
        logger.error('{} errored'.format(len(errored_logs)))
        logger.error('\n'.join([each._id for each in errored_logs]))

    for count, log in enumerate(node_logs): # for every log in node_logs
        with TokuTransaction():
            errored_clones = []
            for node in log._backrefs['logged']['node']['logs']:
                if node != log.original_node._id:
                    clone = log.clone_node_log(node)
                    clone.original_node = get_original_node(log)
                    try:
                        clone.save()
                        logger.info('Log {} cloned for node {}. New log is {}'.format(log._id, node, clone._id))
                    except KeyError as error:
                        logger.error('Could not copy node log due to error')
                        logger.exception(error)
                        errored_clones.append(log, clone)

if __name__ == '__main__':

    dry_run = 'dry' in sys.argv
    if not dry_run:
        script_utils.add_file_logger(logger, __file__)
    main(dry=dry_run)
<commit_msg>Modify script to clone logs for other nodes in backref as looping through original logs.<commit_after>"""
This migration will add original_node and node associated with the log to nodelogs. It will then make
copies of each nodelog for the remaining nodes in the backref (registrations and forks),
changing the node to the current node.
"""

import sys
import logging
from modularodm import Q
from website.app import init_app
from website import models
from scripts import utils as script_utils
from framework.transactions.context import TokuTransaction

logger = logging.getLogger(__name__)

def get_original_node(log):
    """
    Retrieves original node on nodelog
    """
    return models.Node.find(Q('_id', 'eq', log._backrefs['logged']['node']['logs'][0]))[0]

def main(dry=True):
    init_app(routes=False)
    node_logs = models.NodeLog.find(Q('original_node', 'eq', None))
    total_log_count = node_logs.count()
    count = 0
    errored_logs = []
    for log in node_logs:
        count += 1
        with TokuTransaction():
            log.original_node = get_original_node(log)
            log.node = get_original_node(log)
            if not dry:
                try:
                    log.save()
                    logger.info('{}/{} Log {} "original_node = {} and node = {}" added'.format(count, total_log_count, log._id, log.original_node._id, log.node._id))
                except KeyError as error:
                    logger.error('Could not migrate log due to error')
                    logger.exception(error)
                    errored_logs.append(log)

            errored_clones = []
            for node in log._backrefs['logged']['node']['logs']:
                if node != log.original_node._id:
                    clone = log.clone_node_log(node)
                    print clone.original_node
                    print clone.node
                    clone.original_node = get_original_node(log)
                    try:
                        clone.save()
                        logger.info('Log {} cloned for node {}. New log is {}'.format(log._id, node, clone._id))
                    except KeyError as error:
                        logger.error('Could not copy node log due to error')
                        logger.exception(error)
                        errored_clones.append([log, clone])

if __name__ == '__main__':

    dry_run = 'dry' in sys.argv
    if not dry_run:
        script_utils.add_file_logger(logger, __file__)
    main(dry=dry_run)
<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst

from .pickle_helpers import *<commit_msg>Whitespace<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst

from .pickle_helpers import *
<|endoftext|>


<commit_before># Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Dynamic Linear Regression model."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python import distributions as tfd
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import prefer_static

from tensorflow_probability.python.sts.internal import util as sts_util
from tensorflow_probability.python.sts.structural_time_series import Parameter
from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries

class DynamicLinearRegressionStateSpaceModel(tfd.LinearGaussianStateSpaceModel):
  """State space model for a dynamic linear regression from provided covariates.

  A state space model (SSM) posits a set of latent (unobserved) variables that
  evolve over time with dynamics specified by a probabilistic transition model
  `p(z[t+1] | z[t])`. At each timestep, we observe a value sampled from an
  observation model conditioned on the current state, `p(x[t] | z[t])`. The
  special case where both the transition and observation models are Gaussians
  with mean specified as a linear function of the inputs, is known as a linear
  Gaussian state space model and supports tractable exact probabilistic
  calculations; see `tfp.distributions.LinearGaussianStateSpaceModel` for
  details.

  The dynamic linear regression model is a special case of a linear Gaussian SSM
  and a generalization of typical (static) linear regression. The model
  represents regression `weights` with a latent state which evolves via a
  Gaussian random walk:

  ```
  weights[t] ~ Normal(weights[t-1], drift_scale)
  ```

  The latent state (the weights) has dimension `num_features`, while the
  parameters `drift_scale` and `observation_noise_scale` are each (a batch of)
  scalars. The batch shape of this `Distribution` is the broadcast batch shape
  of these parameters, the `initial_state_prior`, and the
  `design_matrix`. `num_features` is determined from the last dimension of
  `design_matrix` (equivalent to the number of columns in the design matrix in
  linear regression).

  #### Mathematical Details

  The local level model implements a
  `tfp.distributions.LinearGaussianStateSpaceModel` with `latent_size =
  num_features` and `observation_size = 1` following the transition model:

  ```
  transition_matrix = eye(num_features)
  transition_noise ~ Normal(0, diag([drift_scale]))
  ```

  which implements the evolution of `weights` described above. The observation
  model is:

  ```
  observation_matrix[t] = design_matrix[t]
  observation_noise ~ Normal(0, observation_noise_scale)
  ```

  #### Examples

  Given `series1`, `series2` as `Tensors` each of shape `[num_timesteps] = [42]`
  representing covariate time series, we create a dynamic regression model which
  conditions on these via the following:

  ```python
  dynamic_regression_ssm = DynamicLinearRegressionStateSpaceModel(
      num_timesteps=42,
      design_matrix=tf.stack([series1, series2], axis=-1),
      drift_scale=3.14,
      initial_state_prior=tfd.MultivariateNormalDiag(scale_diag=[1., 2.]),
      observation_noise_scale=1.)

  y = dynamic_regression_ssm.sample()  # shape [42, 1]
  lp = dynamic_regression_ssm.log_prob(y)  # scalar
  ```

  Passing additional parameter and `initial_state_prior` dimensions constructs a
  batch of models, consider the following:

  ```python
  dynamic_regression_ssm = DynamicLinearRegressionStateSpaceModel(
      num_timesteps=42,
      design_matrix=tf.stack([series1, series2], axis=-1),
      drift_scale=[3.14, 1.],
      initial_state_prior=tfd.MultivariateNormalDiag(scale_diag=[1., 2.]),
      observation_noise_scale=[1., 2.])

  y = dynamic_regression_ssm.sample(3)  # shape [3, 2, 42, 1]
  lp = dynamic_regression_ssm.log_prob(y)  # shape [3, 2]
  ```

  Which (effectively) constructs two independent state space models; the first
  with `drift_scale = 3.14` and `observation_noise_scale = 1.`, the second with
  `drift_scale = 1.` and `observation_noise_scale = 2.`. We then sample from
  each of the models three times and calculate the log probability of each of
  the samples under each of the models.

  Similarly, it is also possible to add batch dimensions via the
  `design_matrix`.

  """

  def __init__(self,
               num_timesteps,
               design_matrix,
               drift_scale,
               initial_state_prior,
               observation_noise_scale=0.,
               initial_step=0,
               validate_args=False,
               allow_nan_stats=True,
               name=None):
  """State space model for a dynamic linear regression.

  Args:
    num_timesteps: Scalar `int` `Tensor` number of timesteps to model
      with this distribution.
    design_matrix: float `Tensor` of shape `concat([batch_shape,
      [num_timesteps, num_features]])`.
    drift_scale: Scalar (any additional dimensions are treated as batch
      dimensions) `float` `Tensor` indicating the standard deviation of the
      latent state transitions.
    initial_state_prior: instance of `tfd.MultivariateNormal`
      representing the prior distribution on latent states.  Must have
      event shape `[num_features]`.
    observation_noise_scale: Scalar (any additional dimensions are
      treated as batch dimensions) `float` `Tensor` indicating the standard
      deviation of the observation noise.
      Default value: `0.`.
    initial_step: scalar `int` `Tensor` specifying the starting timestep.
      Default value: `0`.
    validate_args: Python `bool`. Whether to validate input with asserts. If
      `validate_args` is `False`, and the inputs are invalid, correct behavior
      is not guaranteed.
      Default value: `False`.
    allow_nan_stats: Python `bool`. If `False`, raise an
      exception if a statistic (e.g. mean/mode/etc...) is undefined for any
      batch member. If `True`, batch members with valid parameters leading to
      undefined statistics will return NaN for this statistic.
      Default value: `True`.
    name: Python `str` name prefixed to ops created by this class.
      Default value: "DynamicLinearRegressionStateSpaceModel".

  """

    with tf.compat.v1.name_scope(name, 'DynamicLinearRegressionStateSpaceModel',
                                 values=[drift_scale]) as name:

      dtype = dtype_util.common_dtype(
          [design_matrix, drift_scale, initial_state_prior])

      design_matrix = tf.convert_to_tensor(
          value=design_matrix, name='design_matrix', dtype=dtype)

      drift_scale = tf.convert_to_tensor(
          value=drift_scale, name='drift_scale', dtype=dtype)

      observation_noise_scale = tf.convert_to_tensor(
          value=observation_noise_scale,
          name='observation_noise_scale',
          dtype=dtype)

      num_features = prefer_static.shape(design_matrix)[-1]

      def observation_matrix_fn(t):
        observation_matrix = tf.linalg.LinearOperatorFullMatrix(
            design_matrix[..., t, tf.newaxis, :], name='observation_matrix')
        return observation_matrix

      self._drift_scale = drift_scale
      self._observation_noise_scale = observation_noise_scale

      super(DynamicLinearRegressionStateSpaceModel, self).__init__(
          num_timesteps=num_timesteps,
          transition_matrix=tf.linalg.LinearOperatorIdentity(
              num_rows=num_features,
              dtype=dtype,
              name='transition_matrix'),
          transition_noise=tfd.MultivariateNormalDiag(
              scale_diag=tf.stack([drift_scale] * num_features, axis=-1),
              name='transition_noise'),
          observation_matrix=observation_matrix_fn,
          observation_noise=tfd.MultivariateNormalDiag(
              scale_diag=observation_noise_scale[..., tf.newaxis],
              name='observation_noise'),
          initial_state_prior=initial_state_prior,
          initial_step=initial_step,
          allow_nan_stats=allow_nan_stats,
          validate_args=validate_args,
          name=name)

  @property
  def drift_scale(self):
    """Standard deviation of the drift in weights at each timestep."""
    return self._drift_scale

  @property
  def observation_noise_scale(self):
    """Standard deviation of the observation noise."""
    return self._observation_noise_scale

class DynamicLinearRegression(StructuralTimeSeries):
  """Formal representation of a dynamic linear regresson model.

  The dynamic linear regression model is a special case of a linear Gaussian SSM
  and a generalization of typical (static) linear regression. The model
  represents regression `weights` with a latent state which evolves via a
  Gaussian random walk:

  ```
  weights[t] ~ Normal(weights[t-1], drift_scale)
  ```

  The latent state has dimension `num_features`, while the parameters
  `drift_scale` and `observation_noise_scale` are each (a batch of) scalars. The
  batch shape of this `Distribution` is the broadcast batch shape of these
  parameters, the `initial_state_prior`, and the `design_matrix`. `num_features`
  is determined from the last dimension of `design_matrix` (equivalent to the
  number of columns in the design matrix in linear regression).

  """

  def __init__(self,
               design_matrix,
               drift_scale_prior=None,
               initial_weights_prior=None,
               observed_time_series=None,
               name=None):
  """Specify a dynamic linear regression.

  Args:
    design_matrix: float `Tensor` of shape `concat([batch_shape,
      [num_timesteps, num_features]])`.
    drift_scale_prior: instance of `tfd.Distribution` specifying a prior on the
      `drift_scale` parameter. If `None`, a heuristic default prior is
      constructed based on the provided `observed_time_series`.
      Default value: `None`.
    initial_weights_prior: instance of `tfd.MultivariateNormal` representing the
      prior distribution on the latent states (the regression weights). Must
      have event shape `[num_features]`. If `None`, a weakly-informative
      Normal(0., 10.) prior is used.
      Default value: `None`.
    observed_time_series: `float` `Tensor` of shape `batch_shape + [T, 1]`
      (omitting the trailing unit dimension is also supported when `T > 1`),
      specifying an observed time series. Any priors not explicitly set will be
      given default values according to the scale of the observed time series
      (or batch of time series). May optionally be an instance of
      `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify
      timesteps with missing observations.
      Default value: `None`.
    name: Python `str` for the name of this component.
      Default value: "DynamicLinearRegression".

  """

    with tf.compat.v1.name_scope(
        name, 'DynamicLinearRegression', values=[observed_time_series]) as name:

      dtype = dtype_util.common_dtype(
          [design_matrix, drift_scale_prior, initial_weights_prior])

      num_features = prefer_static.shape(design_matrix)[-1]

      # Default to a weakly-informative Normal(0., 10.) for the initital state
      if initial_weights_prior is None:
        initial_weights_prior = tfd.MultivariateNormalDiag(
            scale_diag=10. * tf.ones([num_features], dtype=dtype))

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if drift_scale_prior is None:
        if observed_time_series is None:
          observed_stddev = tf.constant(1.0, dtype=dtype)
        else:
          _, observed_stddev, _ = sts_util.empirical_statistics(
              observed_time_series)

        drift_scale_prior = tfd.LogNormal(
            loc=tf.math.log(.05 * observed_stddev),
            scale=3.,
            name='drift_scale_prior')

      self._initial_state_prior = initial_weights_prior
      self._design_matrix = design_matrix

      super(DynamicLinearRegression, self).__init__(
          parameters=[
              Parameter('drift_scale', drift_scale_prior, tfb.Softplus())
          ],
          latent_size=num_features,
          name=name)

  @property
  def initial_state_prior(self):
    """Prior distribution on the initial latent state (level and scale)."""
    return self._initial_state_prior

  @property
  def design_matrix(self):
    """Tensor representing the design matrix."""
    return self._design_matrix

  def _make_state_space_model(self,
                              num_timesteps,
                              param_map,
                              initial_state_prior=None,
                              initial_step=0):

    if initial_state_prior is None:
      initial_state_prior = self.initial_state_prior

    return DynamicLinearRegressionStateSpaceModel(
        num_timesteps=num_timesteps,
        design_matrix=self.design_matrix,
        initial_state_prior=initial_state_prior,
        initial_step=initial_step,
        **param_map)
<commit_msg>fix indentation<commit_after># Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Dynamic Linear Regression model."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python import distributions as tfd
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import prefer_static

from tensorflow_probability.python.sts.internal import util as sts_util
from tensorflow_probability.python.sts.structural_time_series import Parameter
from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries

class DynamicLinearRegressionStateSpaceModel(tfd.LinearGaussianStateSpaceModel):
  """State space model for a dynamic linear regression from provided covariates.

  A state space model (SSM) posits a set of latent (unobserved) variables that
  evolve over time with dynamics specified by a probabilistic transition model
  `p(z[t+1] | z[t])`. At each timestep, we observe a value sampled from an
  observation model conditioned on the current state, `p(x[t] | z[t])`. The
  special case where both the transition and observation models are Gaussians
  with mean specified as a linear function of the inputs, is known as a linear
  Gaussian state space model and supports tractable exact probabilistic
  calculations; see `tfp.distributions.LinearGaussianStateSpaceModel` for
  details.

  The dynamic linear regression model is a special case of a linear Gaussian SSM
  and a generalization of typical (static) linear regression. The model
  represents regression `weights` with a latent state which evolves via a
  Gaussian random walk:

  ```
  weights[t] ~ Normal(weights[t-1], drift_scale)
  ```

  The latent state (the weights) has dimension `num_features`, while the
  parameters `drift_scale` and `observation_noise_scale` are each (a batch of)
  scalars. The batch shape of this `Distribution` is the broadcast batch shape
  of these parameters, the `initial_state_prior`, and the
  `design_matrix`. `num_features` is determined from the last dimension of
  `design_matrix` (equivalent to the number of columns in the design matrix in
  linear regression).

  #### Mathematical Details

  The local level model implements a
  `tfp.distributions.LinearGaussianStateSpaceModel` with `latent_size =
  num_features` and `observation_size = 1` following the transition model:

  ```
  transition_matrix = eye(num_features)
  transition_noise ~ Normal(0, diag([drift_scale]))
  ```

  which implements the evolution of `weights` described above. The observation
  model is:

  ```
  observation_matrix[t] = design_matrix[t]
  observation_noise ~ Normal(0, observation_noise_scale)
  ```

  #### Examples

  Given `series1`, `series2` as `Tensors` each of shape `[num_timesteps] = [42]`
  representing covariate time series, we create a dynamic regression model which
  conditions on these via the following:

  ```python
  dynamic_regression_ssm = DynamicLinearRegressionStateSpaceModel(
      num_timesteps=42,
      design_matrix=tf.stack([series1, series2], axis=-1),
      drift_scale=3.14,
      initial_state_prior=tfd.MultivariateNormalDiag(scale_diag=[1., 2.]),
      observation_noise_scale=1.)

  y = dynamic_regression_ssm.sample()  # shape [42, 1]
  lp = dynamic_regression_ssm.log_prob(y)  # scalar
  ```

  Passing additional parameter and `initial_state_prior` dimensions constructs a
  batch of models, consider the following:

  ```python
  dynamic_regression_ssm = DynamicLinearRegressionStateSpaceModel(
      num_timesteps=42,
      design_matrix=tf.stack([series1, series2], axis=-1),
      drift_scale=[3.14, 1.],
      initial_state_prior=tfd.MultivariateNormalDiag(scale_diag=[1., 2.]),
      observation_noise_scale=[1., 2.])

  y = dynamic_regression_ssm.sample(3)  # shape [3, 2, 42, 1]
  lp = dynamic_regression_ssm.log_prob(y)  # shape [3, 2]
  ```

  Which (effectively) constructs two independent state space models; the first
  with `drift_scale = 3.14` and `observation_noise_scale = 1.`, the second with
  `drift_scale = 1.` and `observation_noise_scale = 2.`. We then sample from
  each of the models three times and calculate the log probability of each of
  the samples under each of the models.

  Similarly, it is also possible to add batch dimensions via the
  `design_matrix`.

  """

  def __init__(self,
               num_timesteps,
               design_matrix,
               drift_scale,
               initial_state_prior,
               observation_noise_scale=0.,
               initial_step=0,
               validate_args=False,
               allow_nan_stats=True,
               name=None):
    """State space model for a dynamic linear regression.

    Args:
      num_timesteps: Scalar `int` `Tensor` number of timesteps to model
        with this distribution.
      design_matrix: float `Tensor` of shape `concat([batch_shape,
        [num_timesteps, num_features]])`.
      drift_scale: Scalar (any additional dimensions are treated as batch
        dimensions) `float` `Tensor` indicating the standard deviation of the
        latent state transitions.
      initial_state_prior: instance of `tfd.MultivariateNormal`
        representing the prior distribution on latent states.  Must have
        event shape `[num_features]`.
      observation_noise_scale: Scalar (any additional dimensions are
        treated as batch dimensions) `float` `Tensor` indicating the standard
        deviation of the observation noise.
        Default value: `0.`.
      initial_step: scalar `int` `Tensor` specifying the starting timestep.
        Default value: `0`.
      validate_args: Python `bool`. Whether to validate input with asserts. If
        `validate_args` is `False`, and the inputs are invalid, correct behavior
        is not guaranteed.
        Default value: `False`.
      allow_nan_stats: Python `bool`. If `False`, raise an
        exception if a statistic (e.g. mean/mode/etc...) is undefined for any
        batch member. If `True`, batch members with valid parameters leading to
        undefined statistics will return NaN for this statistic.
        Default value: `True`.
      name: Python `str` name prefixed to ops created by this class.
        Default value: 'DynamicLinearRegressionStateSpaceModel'.

      """

    with tf.compat.v1.name_scope(name, 'DynamicLinearRegressionStateSpaceModel',
                                 values=[drift_scale]) as name:

      dtype = dtype_util.common_dtype(
          [design_matrix, drift_scale, initial_state_prior])

      design_matrix = tf.convert_to_tensor(
          value=design_matrix, name='design_matrix', dtype=dtype)

      drift_scale = tf.convert_to_tensor(
          value=drift_scale, name='drift_scale', dtype=dtype)

      observation_noise_scale = tf.convert_to_tensor(
          value=observation_noise_scale,
          name='observation_noise_scale',
          dtype=dtype)

      num_features = prefer_static.shape(design_matrix)[-1]

      def observation_matrix_fn(t):
        observation_matrix = tf.linalg.LinearOperatorFullMatrix(
            design_matrix[..., t, tf.newaxis, :], name='observation_matrix')
        return observation_matrix

      self._drift_scale = drift_scale
      self._observation_noise_scale = observation_noise_scale

      super(DynamicLinearRegressionStateSpaceModel, self).__init__(
          num_timesteps=num_timesteps,
          transition_matrix=tf.linalg.LinearOperatorIdentity(
              num_rows=num_features,
              dtype=dtype,
              name='transition_matrix'),
          transition_noise=tfd.MultivariateNormalDiag(
              scale_diag=tf.stack([drift_scale] * num_features, axis=-1),
              name='transition_noise'),
          observation_matrix=observation_matrix_fn,
          observation_noise=tfd.MultivariateNormalDiag(
              scale_diag=observation_noise_scale[..., tf.newaxis],
              name='observation_noise'),
          initial_state_prior=initial_state_prior,
          initial_step=initial_step,
          allow_nan_stats=allow_nan_stats,
          validate_args=validate_args,
          name=name)

  @property
  def drift_scale(self):
    """Standard deviation of the drift in weights at each timestep."""
    return self._drift_scale

  @property
  def observation_noise_scale(self):
    """Standard deviation of the observation noise."""
    return self._observation_noise_scale

class DynamicLinearRegression(StructuralTimeSeries):
  """Formal representation of a dynamic linear regresson model.

  The dynamic linear regression model is a special case of a linear Gaussian SSM
  and a generalization of typical (static) linear regression. The model
  represents regression `weights` with a latent state which evolves via a
  Gaussian random walk:

  ```
  weights[t] ~ Normal(weights[t-1], drift_scale)
  ```

  The latent state has dimension `num_features`, while the parameters
  `drift_scale` and `observation_noise_scale` are each (a batch of) scalars. The
  batch shape of this `Distribution` is the broadcast batch shape of these
  parameters, the `initial_state_prior`, and the `design_matrix`. `num_features`
  is determined from the last dimension of `design_matrix` (equivalent to the
  number of columns in the design matrix in linear regression).

  """

  def __init__(self,
               design_matrix,
               drift_scale_prior=None,
               initial_weights_prior=None,
               observed_time_series=None,
               name=None):
    """Specify a dynamic linear regression.

    Args:
      design_matrix: float `Tensor` of shape `concat([batch_shape,
        [num_timesteps, num_features]])`.
      drift_scale_prior: instance of `tfd.Distribution` specifying a prior on
        the `drift_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_weights_prior: instance of `tfd.MultivariateNormal` representing
        the prior distribution on the latent states (the regression weights).
        Must have event shape `[num_features]`. If `None`, a weakly-informative
        Normal(0., 10.) prior is used.
        Default value: `None`.
      observed_time_series: `float` `Tensor` of shape `batch_shape + [T, 1]`
        (omitting the trailing unit dimension is also supported when `T > 1`),
        specifying an observed time series. Any priors not explicitly set will
        be given default values according to the scale of the observed time
        series (or batch of time series). May optionally be an instance of
        `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify
        timesteps with missing observations.
        Default value: `None`.
      name: Python `str` for the name of this component.
        Default value: 'DynamicLinearRegression'.

    """

    with tf.compat.v1.name_scope(
        name, 'DynamicLinearRegression', values=[observed_time_series]) as name:

      dtype = dtype_util.common_dtype(
          [design_matrix, drift_scale_prior, initial_weights_prior])

      num_features = prefer_static.shape(design_matrix)[-1]

      # Default to a weakly-informative Normal(0., 10.) for the initital state
      if initial_weights_prior is None:
        initial_weights_prior = tfd.MultivariateNormalDiag(
            scale_diag=10. * tf.ones([num_features], dtype=dtype))

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if drift_scale_prior is None:
        if observed_time_series is None:
          observed_stddev = tf.constant(1.0, dtype=dtype)
        else:
          _, observed_stddev, _ = sts_util.empirical_statistics(
              observed_time_series)

        drift_scale_prior = tfd.LogNormal(
            loc=tf.math.log(.05 * observed_stddev),
            scale=3.,
            name='drift_scale_prior')

      self._initial_state_prior = initial_weights_prior
      self._design_matrix = design_matrix

      super(DynamicLinearRegression, self).__init__(
          parameters=[
              Parameter('drift_scale', drift_scale_prior, tfb.Softplus())
          ],
          latent_size=num_features,
          name=name)

  @property
  def initial_state_prior(self):
    """Prior distribution on the initial latent state (level and scale)."""
    return self._initial_state_prior

  @property
  def design_matrix(self):
    """Tensor representing the design matrix."""
    return self._design_matrix

  def _make_state_space_model(self,
                              num_timesteps,
                              param_map,
                              initial_state_prior=None,
                              initial_step=0):

    if initial_state_prior is None:
      initial_state_prior = self.initial_state_prior

    return DynamicLinearRegressionStateSpaceModel(
        num_timesteps=num_timesteps,
        design_matrix=self.design_matrix,
        initial_state_prior=initial_state_prior,
        initial_step=initial_step,
        **param_map)
<|endoftext|>


<commit_before># Tai Sakuma <tai.sakuma@gmail.com>

import math
import collections
import logging

from .ReturnTrue import ReturnTrue

##__________________________________________________________________||
class Round(object):
    """Equal width binning

    Parameters
    ----------
    width : float or int, default 1
        The width.
    aboundary : float or int, optional
        A boundary. If not given, `width`/2 will be used.
    min : float or int, optional
        The lowest bin will be the bin that `min` falls in. If given,
        `__call__(val)` returns `underflow_bin` if the `val` is less than the
        lower edge of the lowest bin.
    underflow_bin : optional
        The underflow bin. When `min` is given, the `__call__(val)` returns
        `underflow_bin` if the `val` is less than the lower edge of the lowest
        bin.
    max : float or int, optional
        The highest bin will be the bin that `max` falls in except when `max`
        is one of boundaries. When `max` is one of boundaries, the highest bin
        is the bin whose upper edge is `max`. If given, `__call__(val)` returns
        the overflow bin if the `val` is greater than or equal to the upper
        edge of the highest bin.
    overflow_bin : optional

        The overflow bin if `overflow_bin` is any value other than `True`. If
        `overflow_bin` is `True`, the overflow bin will be the upper edge of
        the highest bin. When `max` is given, the `__call__(val)` returns the
        overflow bin if the `val` is greater than or equal to the upper edge of
        the highest bin.
    valid : function, optional
        Boolean function to test if value is valid

    """
    def __init__(self, width=1, aboundary=None,
                 min=None, underflow_bin=None,
                 max=None, overflow_bin=None,
                 valid=ReturnTrue()):

        self.width = width
        self.aboundary = aboundary
        if aboundary is None:
            halfWidth = self.width//2 if self.width % 2 == 0 else float(self.width)/2
            aboundary = halfWidth
        self.boundaries = collections.deque([aboundary])
        self.min = min
        self.underflow_bin = underflow_bin
        self.max = max
        self.overflow_bin = overflow_bin
        self.valid = valid

        if self.min is not None:
            self._update_boundaries(self.min)

        if self.max is not None:
            self._update_boundaries(self.max)
            if self.overflow_bin is True:
                self.overflow_bin = self.boundaries[-1]

    def __repr__(self):
        return '{}(width={!r}, aboundary={!r}, min={!r}, underflow_bin={!r}, max={!r}, overflow_bin={!r}, valid={!r})'.format(
            self.__class__.__name__,
            self.width,
            self.aboundary,
            self.min,
            self.underflow_bin,
            self.max,
            self.overflow_bin,
            self.valid
        )

    def __call__(self, val):
        return self._lower_boundary(val)

    def _lower_boundary(self, val):

        if not self.valid(val):
            return None

        if self.min is not None:
            if not self.boundaries[0] <= val:
                return self.underflow_bin

        if self.max is not None:
            if not val < self.boundaries[-1]:
                return self.overflow_bin

        if math.isinf(val):
            logger = logging.getLogger(__name__)
            logger.warning('val={}. will return {}'.format(val, None))
            return None

        self._update_boundaries(val)

        bin = self.boundaries[0]
        for b in self.boundaries:
            if b <= val:
                bin = b
            else:
                break

        return bin

    def _update_boundaries(self, val):

        while val < self.boundaries[0]:
            self.boundaries.appendleft(self.boundaries[0] - self.width)

        while val > self.boundaries[-1]:
            self.boundaries.append(self.boundaries[-1] + self.width)

    def next(self, bin):
        return self._next_lower_boundary(bin)

    def _next_lower_boundary(self, bin):

        bin = self._lower_boundary(bin)

        if bin is None:
            return None

        if bin == self.underflow_bin:
            return self._lower_boundary(self.min)

        if bin == self.overflow_bin:
            return self.overflow_bin

        self._update_boundaries(bin)

        return self._lower_boundary(bin + self.width*1.001)

##__________________________________________________________________||
<commit_msg>check for None in Round<commit_after># Tai Sakuma <tai.sakuma@gmail.com>

import math
import collections
import logging

from .ReturnTrue import ReturnTrue

##__________________________________________________________________||
class Round(object):
    """Equal width binning

    Parameters
    ----------
    width : float or int, default 1
        The width.
    aboundary : float or int, optional
        A boundary. If not given, `width`/2 will be used.
    min : float or int, optional
        The lowest bin will be the bin that `min` falls in. If given,
        `__call__(val)` returns `underflow_bin` if the `val` is less than the
        lower edge of the lowest bin.
    underflow_bin : optional
        The underflow bin. When `min` is given, the `__call__(val)` returns
        `underflow_bin` if the `val` is less than the lower edge of the lowest
        bin.
    max : float or int, optional
        The highest bin will be the bin that `max` falls in except when `max`
        is one of boundaries. When `max` is one of boundaries, the highest bin
        is the bin whose upper edge is `max`. If given, `__call__(val)` returns
        the overflow bin if the `val` is greater than or equal to the upper
        edge of the highest bin.
    overflow_bin : optional

        The overflow bin if `overflow_bin` is any value other than `True`. If
        `overflow_bin` is `True`, the overflow bin will be the upper edge of
        the highest bin. When `max` is given, the `__call__(val)` returns the
        overflow bin if the `val` is greater than or equal to the upper edge of
        the highest bin.
    valid : function, optional
        Boolean function to test if value is valid

    """
    def __init__(self, width=1, aboundary=None,
                 min=None, underflow_bin=None,
                 max=None, overflow_bin=None,
                 valid=ReturnTrue()):

        self.width = width
        self.aboundary = aboundary
        if aboundary is None:
            halfWidth = self.width//2 if self.width % 2 == 0 else float(self.width)/2
            aboundary = halfWidth
        self.boundaries = collections.deque([aboundary])
        self.min = min
        self.underflow_bin = underflow_bin
        self.max = max
        self.overflow_bin = overflow_bin
        self.valid = valid

        if self.min is not None:
            self._update_boundaries(self.min)

        if self.max is not None:
            self._update_boundaries(self.max)
            if self.overflow_bin is True:
                self.overflow_bin = self.boundaries[-1]

    def __repr__(self):
        return '{}(width={!r}, aboundary={!r}, min={!r}, underflow_bin={!r}, max={!r}, overflow_bin={!r}, valid={!r})'.format(
            self.__class__.__name__,
            self.width,
            self.aboundary,
            self.min,
            self.underflow_bin,
            self.max,
            self.overflow_bin,
            self.valid
        )

    def __call__(self, val):
        return self._lower_boundary(val)

    def _lower_boundary(self, val):

        if val is None:
            return None

        if not self.valid(val):
            return None

        if self.min is not None:
            if not self.boundaries[0] <= val:
                return self.underflow_bin

        if self.max is not None:
            if not val < self.boundaries[-1]:
                return self.overflow_bin

        if math.isinf(val):
            logger = logging.getLogger(__name__)
            logger.warning('val={}. will return {}'.format(val, None))
            return None

        self._update_boundaries(val)

        bin = self.boundaries[0]
        for b in self.boundaries:
            if b <= val:
                bin = b
            else:
                break

        return bin

    def _update_boundaries(self, val):

        while val < self.boundaries[0]:
            self.boundaries.appendleft(self.boundaries[0] - self.width)

        while val > self.boundaries[-1]:
            self.boundaries.append(self.boundaries[-1] + self.width)

    def next(self, bin):
        return self._next_lower_boundary(bin)

    def _next_lower_boundary(self, bin):

        bin = self._lower_boundary(bin)

        if bin is None:
            return None

        if bin == self.underflow_bin:
            return self._lower_boundary(self.min)

        if bin == self.overflow_bin:
            return self.overflow_bin

        self._update_boundaries(bin)

        return self._lower_boundary(bin + self.width*1.001)

##__________________________________________________________________||
<|endoftext|>