<commit_before>#!/usr/bin/python
import appindicator
import gtk
import gtk.glade
import gconf
import sys
import pexpect
import os
from xdg.DesktopEntry import DesktopEntry

VERSION = "1.1.8"

class Fluxgui:
    def __init__(self):
        self.check_pid()
        self.indicator = Indicator(self)

        self.color = None;
        self.settings = Settings(self)
        self.start_xflux(self.settings.latitude, self.settings.longitude,
                         self.settings.zipcode, self.settings.color)

        if self.settings.latitude is "" and self.settings.zipcode is "":
            self.open_preferences("activate")

    def check_pid(self):
        pid = str(os.getpid())
        self.pidfile = "/tmp/fluxgui_%s.pid" % os.environ['USER']

        if os.path.isfile(self.pidfile):
          print "fluxgui is already running, exiting"
          sys.exit()
        else:
          file(self.pidfile, 'w').write(pid)

    def start_xflux(self, lat, lon, zipcode, color):
        args = []
        if zipcode:
            args = ["-z", zipcode, "-k", color, '-nofork']
        if lat:
            args = ["-l", lat, "-k", color, '-nofork']
            if lon:
                args = ["-l", lat, "-g", lon, "-k", color, '-nofork']

        if args:
            try:
                self.xflux = pexpect.spawn("/usr/bin/xflux", args)

                fout = file('/tmp/fluxlogstr.txt', 'w')
                self.xflux.logfile = fout

            except pexpect.ExceptionPexpect:
                print "\nError: Please install xflux in /usr/bin/ \n"
                os.unlink(self.pidfile)
                sys.exit(2)
        else:
            self.xflux = None

    def stop_xflux(self, item):
        self.indicator.item_turn_off.hide()
        self.indicator.item_turn_on.show()

        if self.xflux:
            # xflux is not running
            try:
                self.xflux.terminate(force=True)
            except Exception, e:
                # xflux has crashed in the meantime?
                pass

    def pause_xflux(self, item):
        self.indicator.item_turn_off.hide()
        self.indicator.item_turn_on.show()
        self.update_xflux("k=" + self.settings.get_color("4"))

    def unpause_xflux(self, item):
        self.indicator.item_turn_off.show()
        self.indicator.item_turn_on.hide()
        self.update_xflux("k=" + self.settings.color)

    def update_xflux(self, command):
        if self.xflux is None:
            self.start_xflux(self.settings.latitude, self.settings.longitude,
                         self.settings.zipcode, self.settings.color)
        else:
            self.xflux.sendline(command)

    def get_colortemp(self):
        if self.xflux:
            self.xflux.sendline("c")
            index = self.xflux.expect("Color.*")
            if index == 0:
                self.color = self.xflux.after[10:14]

    def preview_xflux(self, item):
      self.settings.set_colortemp(str(self.preferences.colsetting.get_active()))
      self.update_xflux("p")

    def open_preferences(self, item):
        self.get_colortemp()
        self.preferences = Preferences(self)

    #autostart code copied from AWN
    def get_autostart_file_path(self):
        autostart_dir = os.path.join(os.environ['HOME'], '.config',
                                     'autostart')
        return os.path.join(autostart_dir, 'fluxgui.desktop')

    def create_autostarter(self):
        autostart_file = self.get_autostart_file_path()
        autostart_dir = os.path.dirname(autostart_file)

        if not os.path.isdir(autostart_dir):
            #create autostart dir
            try:
                os.mkdir(autostart_dir)
            except Exeption, e:
                print "creation of autostart dir failed, please make it yourself: %s" % autostart_dir
                raise e

        if not os.path.isfile(autostart_file):
            #create autostart entry
            starter_item = DesktopEntry(autostart_file)
            starter_item.set('Name', 'f.lux indicator applet')
            starter_item.set('Exec', 'fluxgui')
            starter_item.set('Icon', 'fluxgui')
            starter_item.set('X-GNOME-Autostart-enabled', 'true')
            starter_item.write()
            self.settings.set_autostart(True)

    def delete_autostarter(self):
        autostart_file = self.get_autostart_file_path()
        if os.path.isfile(autostart_file):
            os.remove(autostart_file)
            self.settings.set_autostart(False)

    def run(self):
        gtk.main()

    def exit(self, widget, data=None):
        self.stop_xflux("activate")
        os.unlink(self.pidfile)
        gtk.main_quit()
        sys.exit(0)

class Indicator:

    def __init__(self, main):
        self.main = main
        self.setup_indicator()

    def setup_indicator(self):
        self.indicator = appindicator.Indicator(
          "fluxgui-indicator",
          "fluxgui",
          appindicator.CATEGORY_APPLICATION_STATUS)
        self.indicator.set_status(appindicator.STATUS_ACTIVE)

        # Check for special Ubuntu themes. copied from lookit

        try:
            theme = gtk.gdk.screen_get_default().get_setting('gtk-icon-theme-name')
        except:
            self.indicator.set_icon('fluxgui')
        else:
          if theme == 'ubuntu-mono-dark':
              self.indicator.set_icon('fluxgui-dark')
          elif theme == 'ubuntu-mono-light':
              self.indicator.set_icon('fluxgui-light')
          else:
              self.indicator.set_icon('fluxgui')

        self.indicator.set_menu(self.setup_menu())

    def setup_menu(self):
        menu = gtk.Menu()

        self.item_turn_off = gtk.MenuItem("_Pause f.lux")
        self.item_turn_off.connect("activate", self.main.pause_xflux)
        self.item_turn_off.show()
        menu.append(self.item_turn_off)

        self.item_turn_on = gtk.MenuItem("_Unpause f.lux")
        self.item_turn_on.connect("activate", self.main.unpause_xflux)
        self.item_turn_on.hide()
        menu.append(self.item_turn_on)

        item = gtk.MenuItem("_Preferences")
        item.connect("activate", self.main.open_preferences)
        item.show()
        menu.append(item)

        item = gtk.SeparatorMenuItem()
        item.show()
        menu.append(item)

        item = gtk.MenuItem("Quit")
        item.connect("activate", self.main.exit)
        item.show()
        menu.append(item)

        return menu

    def main(self):
        gtk.main()

class Preferences:

    def __init__(self, main):
        self.main = main
        self.gladefile = os.path.join(os.path.dirname(os.path.dirname(
          os.path.realpath(__file__))), "fluxgui/preferences.glade")
        self.wTree = gtk.glade.XML(self.gladefile)

        self.window = self.wTree.get_widget("window1")
        self.window.connect("destroy", self.delete_event)

        self.latsetting = self.wTree.get_widget("entry1")
        self.latsetting.set_text(self.main.settings.latitude)
        self.latsetting.connect("activate", self.delete_event)

        self.lonsetting = self.wTree.get_widget("entry3")
        self.lonsetting.set_text(self.main.settings.longitude)
        self.lonsetting.connect("activate", self.delete_event)

        self.zipsetting = self.wTree.get_widget("entry2")
        self.zipsetting.set_text(self.main.settings.zipcode)
        self.zipsetting.connect("activate", self.delete_event)

        self.colsetting = self.wTree.get_widget("combobox1")
        self.colsetting.set_active(int(self.main.settings.colortemp))

        self.colordisplay = self.wTree.get_widget("label6")
        if self.main.color:
            self.colordisplay.set_text("Current color temperature: "
                                       + self.main.color + "K")

        self.previewbutton = self.wTree.get_widget("button1")
        self.previewbutton.connect("clicked", self.main.preview_xflux)

        self.closebutton = self.wTree.get_widget("button2")
        self.closebutton.connect("clicked", self.delete_event)

        self.autostart = self.wTree.get_widget("checkbutton1")
        if self.main.settings.autostart is "1":
            self.autostart.set_active(True)
        else:
            self.autostart.set_active(False)

        if self.main.settings.latitude is ""\
           and self.main.settings.zipcode is "":
            md = gtk.MessageDialog(self.window,
                gtk.DIALOG_DESTROY_WITH_PARENT, gtk.MESSAGE_INFO,
                gtk.BUTTONS_OK, "The f.lux indicator applet needs to know " +
                "your latitude and longitude or zipcode to work correctly. " +
        "Please fill either of them in on the next screen and then hit enter.")
            md.set_title("f.lux indicator applet")
            md.run()
            md.destroy()
            self.window.show()
        else:
            self.window.show()

    def delete_event(self, widget, data=None):
        if self.main.settings.latitude != self.latsetting.get_text():
            self.main.settings.set_latitude(self.latsetting.get_text())

        if self.main.settings.longitude != self.lonsetting.get_text():
            self.main.settings.set_longitude(self.lonsetting.get_text())

        if self.main.settings.zipcode != self.zipsetting.get_text():
            self.main.settings.set_zipcode(self.zipsetting.get_text())

        if self.main.settings.colortemp != str(self.colsetting.get_active()):
            self.main.settings.set_colortemp(str(self.colsetting.get_active()))

        if self.autostart.get_active():
            self.main.create_autostarter()
        else:
            self.main.delete_autostarter()

        self.window.hide()
        return False

    def main(self):
        gtk.main()

class Settings:

    def __init__(self, main):
        self.main = main
        self.client = gconf.client_get_default()
        self.prefs_key = "/apps/fluxgui"
        self.client.add_dir(self.prefs_key, gconf.CLIENT_PRELOAD_NONE)

        self.autostart = self.client.get_string(self.prefs_key + "/autostart")
        self.latitude = self.client.get_string(self.prefs_key + "/latitude")
        self.longitude = self.client.get_string(self.prefs_key + "/longitude")
        self.zipcode = self.client.get_string(self.prefs_key + "/zipcode")
        self.colortemp = self.client.get_string(self.prefs_key + "/colortemp")
        self.color = self.get_color(self.colortemp)

        if self.latitude is None:
            self.latitude = ""

        if self.longitude is None:
            self.longitude = ""

        if self.zipcode is None:
            self.zipcode = ""

        if not self.colortemp:
            self.colortemp = "1"

        if not self.autostart:
            self.autostart = "0"

    def set_latitude(self, latitude):
        self.client.set_string(self.prefs_key + "/latitude", latitude)
        self.latitude = latitude

        command = "l=" + latitude
        self.main.update_xflux(command)

    def set_longitude(self, longitude):
        self.client.set_string(self.prefs_key + "/longitude", longitude)
        self.longitude = longitude

        command = "g=" + longitude
        self.main.update_xflux(command)

    def set_zipcode(self, zipcode):
        self.client.set_string(self.prefs_key + "/zipcode", zipcode)
        self.zipcode = zipcode

        command = "z=" + zipcode
        self.main.update_xflux(command)

    def get_color(self, colortemp):
        color = "3400"
        if colortemp is "0":
            #tungsten
            color = "2700"
        elif colortemp is "1":
            #halogen
            color = "3400"
        elif colortemp is "2":
            #fluorescent
            color = "4200"
        elif colortemp is "3":
            #daylight
            color = "5000"
        elif colortemp is "4":
            #off
            color = "6500"

        return color

    def set_colortemp(self, colortemp):
        color = self.get_color(colortemp)

        self.client.set_string(self.prefs_key + "/colortemp", colortemp)
        self.colortemp = colortemp
        self.color = color

        command = "k=" + color
        self.main.update_xflux(command)

    def set_autostart(self, autostart):
        if autostart:
            self.client.set_string(self.prefs_key + "/autostart", "1")
            self.autostart = "1"
        else:
            self.client.set_string(self.prefs_key + "/autostart", "0")
            self.autostart = "0"

    def main(self):
        gtk.main()

if __name__ == "__main__":
    try:
        app = Fluxgui()
        app.run()
    except KeyboardInterrupt:
        app.stop_xflux("activate")
        os.unlink(app.pidfile)

<commit_msg>comment out logging<commit_after>#!/usr/bin/python
import appindicator
import gtk
import gtk.glade
import gconf
import sys
import pexpect
import os
from xdg.DesktopEntry import DesktopEntry

VERSION = "1.1.8"

class Fluxgui:
    def __init__(self):
        self.check_pid()
        self.indicator = Indicator(self)

        self.color = None;
        self.settings = Settings(self)
        self.start_xflux(self.settings.latitude, self.settings.longitude,
                         self.settings.zipcode, self.settings.color)

        if self.settings.latitude is "" and self.settings.zipcode is "":
            self.open_preferences("activate")

    def check_pid(self):
        pid = str(os.getpid())
        self.pidfile = "/tmp/fluxgui_%s.pid" % os.environ['USER']

        if os.path.isfile(self.pidfile):
          print "fluxgui is already running, exiting"
          sys.exit()
        else:
          file(self.pidfile, 'w').write(pid)

    def start_xflux(self, lat, lon, zipcode, color):
        args = []
        if zipcode:
            args = ["-z", zipcode, "-k", color, '-nofork']
        if lat:
            args = ["-l", lat, "-k", color, '-nofork']
            if lon:
                args = ["-l", lat, "-g", lon, "-k", color, '-nofork']

        if args:
            try:
                self.xflux = pexpect.spawn("/usr/bin/xflux", args)

                #fout = file('/tmp/fluxlogstr.txt', 'w')
                #self.xflux.logfile = fout

            except pexpect.ExceptionPexpect:
                print "\nError: Please install xflux in /usr/bin/ \n"
                os.unlink(self.pidfile)
                sys.exit(2)
        else:
            self.xflux = None

    def stop_xflux(self, item):
        self.indicator.item_turn_off.hide()
        self.indicator.item_turn_on.show()

        if self.xflux:
            # xflux is not running
            try:
                self.xflux.terminate(force=True)
            except Exception, e:
                # xflux has crashed in the meantime?
                pass

    def pause_xflux(self, item):
        self.indicator.item_turn_off.hide()
        self.indicator.item_turn_on.show()
        self.update_xflux("k=" + self.settings.get_color("4"))

    def unpause_xflux(self, item):
        self.indicator.item_turn_off.show()
        self.indicator.item_turn_on.hide()
        self.update_xflux("k=" + self.settings.color)

    def update_xflux(self, command):
        if self.xflux is None:
            self.start_xflux(self.settings.latitude, self.settings.longitude,
                         self.settings.zipcode, self.settings.color)
        else:
            self.xflux.sendline(command)

    def get_colortemp(self):
        if self.xflux:
            self.xflux.sendline("c")
            index = self.xflux.expect("Color.*")
            if index == 0:
                self.color = self.xflux.after[10:14]

    def preview_xflux(self, item):
      self.settings.set_colortemp(str(self.preferences.colsetting.get_active()))
      self.update_xflux("p")

    def open_preferences(self, item):
        self.get_colortemp()
        self.preferences = Preferences(self)

    #autostart code copied from AWN
    def get_autostart_file_path(self):
        autostart_dir = os.path.join(os.environ['HOME'], '.config',
                                     'autostart')
        return os.path.join(autostart_dir, 'fluxgui.desktop')

    def create_autostarter(self):
        autostart_file = self.get_autostart_file_path()
        autostart_dir = os.path.dirname(autostart_file)

        if not os.path.isdir(autostart_dir):
            #create autostart dir
            try:
                os.mkdir(autostart_dir)
            except Exeption, e:
                print "creation of autostart dir failed, please make it yourself: %s" % autostart_dir
                raise e

        if not os.path.isfile(autostart_file):
            #create autostart entry
            starter_item = DesktopEntry(autostart_file)
            starter_item.set('Name', 'f.lux indicator applet')
            starter_item.set('Exec', 'fluxgui')
            starter_item.set('Icon', 'fluxgui')
            starter_item.set('X-GNOME-Autostart-enabled', 'true')
            starter_item.write()
            self.settings.set_autostart(True)

    def delete_autostarter(self):
        autostart_file = self.get_autostart_file_path()
        if os.path.isfile(autostart_file):
            os.remove(autostart_file)
            self.settings.set_autostart(False)

    def run(self):
        gtk.main()

    def exit(self, widget, data=None):
        self.stop_xflux("activate")
        os.unlink(self.pidfile)
        gtk.main_quit()
        sys.exit(0)

class Indicator:

    def __init__(self, main):
        self.main = main
        self.setup_indicator()

    def setup_indicator(self):
        self.indicator = appindicator.Indicator(
          "fluxgui-indicator",
          "fluxgui",
          appindicator.CATEGORY_APPLICATION_STATUS)
        self.indicator.set_status(appindicator.STATUS_ACTIVE)

        # Check for special Ubuntu themes. copied from lookit

        try:
            theme = gtk.gdk.screen_get_default().get_setting('gtk-icon-theme-name')
        except:
            self.indicator.set_icon('fluxgui')
        else:
          if theme == 'ubuntu-mono-dark':
              self.indicator.set_icon('fluxgui-dark')
          elif theme == 'ubuntu-mono-light':
              self.indicator.set_icon('fluxgui-light')
          else:
              self.indicator.set_icon('fluxgui')

        self.indicator.set_menu(self.setup_menu())

    def setup_menu(self):
        menu = gtk.Menu()

        self.item_turn_off = gtk.MenuItem("_Pause f.lux")
        self.item_turn_off.connect("activate", self.main.pause_xflux)
        self.item_turn_off.show()
        menu.append(self.item_turn_off)

        self.item_turn_on = gtk.MenuItem("_Unpause f.lux")
        self.item_turn_on.connect("activate", self.main.unpause_xflux)
        self.item_turn_on.hide()
        menu.append(self.item_turn_on)

        item = gtk.MenuItem("_Preferences")
        item.connect("activate", self.main.open_preferences)
        item.show()
        menu.append(item)

        item = gtk.SeparatorMenuItem()
        item.show()
        menu.append(item)

        item = gtk.MenuItem("Quit")
        item.connect("activate", self.main.exit)
        item.show()
        menu.append(item)

        return menu

    def main(self):
        gtk.main()

class Preferences:

    def __init__(self, main):
        self.main = main
        self.gladefile = os.path.join(os.path.dirname(os.path.dirname(
          os.path.realpath(__file__))), "fluxgui/preferences.glade")
        self.wTree = gtk.glade.XML(self.gladefile)

        self.window = self.wTree.get_widget("window1")
        self.window.connect("destroy", self.delete_event)

        self.latsetting = self.wTree.get_widget("entry1")
        self.latsetting.set_text(self.main.settings.latitude)
        self.latsetting.connect("activate", self.delete_event)

        self.lonsetting = self.wTree.get_widget("entry3")
        self.lonsetting.set_text(self.main.settings.longitude)
        self.lonsetting.connect("activate", self.delete_event)

        self.zipsetting = self.wTree.get_widget("entry2")
        self.zipsetting.set_text(self.main.settings.zipcode)
        self.zipsetting.connect("activate", self.delete_event)

        self.colsetting = self.wTree.get_widget("combobox1")
        self.colsetting.set_active(int(self.main.settings.colortemp))

        self.colordisplay = self.wTree.get_widget("label6")
        if self.main.color:
            self.colordisplay.set_text("Current color temperature: "
                                       + self.main.color + "K")

        self.previewbutton = self.wTree.get_widget("button1")
        self.previewbutton.connect("clicked", self.main.preview_xflux)

        self.closebutton = self.wTree.get_widget("button2")
        self.closebutton.connect("clicked", self.delete_event)

        self.autostart = self.wTree.get_widget("checkbutton1")
        if self.main.settings.autostart is "1":
            self.autostart.set_active(True)
        else:
            self.autostart.set_active(False)

        if self.main.settings.latitude is ""\
           and self.main.settings.zipcode is "":
            md = gtk.MessageDialog(self.window,
                gtk.DIALOG_DESTROY_WITH_PARENT, gtk.MESSAGE_INFO,
                gtk.BUTTONS_OK, "The f.lux indicator applet needs to know " +
                "your latitude and longitude or zipcode to work correctly. " +
        "Please fill either of them in on the next screen and then hit enter.")
            md.set_title("f.lux indicator applet")
            md.run()
            md.destroy()
            self.window.show()
        else:
            self.window.show()

    def delete_event(self, widget, data=None):
        if self.main.settings.latitude != self.latsetting.get_text():
            self.main.settings.set_latitude(self.latsetting.get_text())

        if self.main.settings.longitude != self.lonsetting.get_text():
            self.main.settings.set_longitude(self.lonsetting.get_text())

        if self.main.settings.zipcode != self.zipsetting.get_text():
            self.main.settings.set_zipcode(self.zipsetting.get_text())

        if self.main.settings.colortemp != str(self.colsetting.get_active()):
            self.main.settings.set_colortemp(str(self.colsetting.get_active()))

        if self.autostart.get_active():
            self.main.create_autostarter()
        else:
            self.main.delete_autostarter()

        self.window.hide()
        return False

    def main(self):
        gtk.main()

class Settings:

    def __init__(self, main):
        self.main = main
        self.client = gconf.client_get_default()
        self.prefs_key = "/apps/fluxgui"
        self.client.add_dir(self.prefs_key, gconf.CLIENT_PRELOAD_NONE)

        self.autostart = self.client.get_string(self.prefs_key + "/autostart")
        self.latitude = self.client.get_string(self.prefs_key + "/latitude")
        self.longitude = self.client.get_string(self.prefs_key + "/longitude")
        self.zipcode = self.client.get_string(self.prefs_key + "/zipcode")
        self.colortemp = self.client.get_string(self.prefs_key + "/colortemp")
        self.color = self.get_color(self.colortemp)

        if self.latitude is None:
            self.latitude = ""

        if self.longitude is None:
            self.longitude = ""

        if self.zipcode is None:
            self.zipcode = ""

        if not self.colortemp:
            self.colortemp = "1"

        if not self.autostart:
            self.autostart = "0"

    def set_latitude(self, latitude):
        self.client.set_string(self.prefs_key + "/latitude", latitude)
        self.latitude = latitude

        command = "l=" + latitude
        self.main.update_xflux(command)

    def set_longitude(self, longitude):
        self.client.set_string(self.prefs_key + "/longitude", longitude)
        self.longitude = longitude

        command = "g=" + longitude
        self.main.update_xflux(command)

    def set_zipcode(self, zipcode):
        self.client.set_string(self.prefs_key + "/zipcode", zipcode)
        self.zipcode = zipcode

        command = "z=" + zipcode
        self.main.update_xflux(command)

    def get_color(self, colortemp):
        color = "3400"
        if colortemp is "0":
            #tungsten
            color = "2700"
        elif colortemp is "1":
            #halogen
            color = "3400"
        elif colortemp is "2":
            #fluorescent
            color = "4200"
        elif colortemp is "3":
            #daylight
            color = "5000"
        elif colortemp is "4":
            #off
            color = "6500"

        return color

    def set_colortemp(self, colortemp):
        color = self.get_color(colortemp)

        self.client.set_string(self.prefs_key + "/colortemp", colortemp)
        self.colortemp = colortemp
        self.color = color

        command = "k=" + color
        self.main.update_xflux(command)

    def set_autostart(self, autostart):
        if autostart:
            self.client.set_string(self.prefs_key + "/autostart", "1")
            self.autostart = "1"
        else:
            self.client.set_string(self.prefs_key + "/autostart", "0")
            self.autostart = "0"

    def main(self):
        gtk.main()

if __name__ == "__main__":
    try:
        app = Fluxgui()
        app.run()
    except KeyboardInterrupt:
        app.stop_xflux("activate")
        os.unlink(app.pidfile)

<|endoftext|>


<commit_before># TODO: write tests
<commit_msg>add asyncio exception handler test<commit_after>import asyncio

from mitmproxy.test.taddons import RecordingMaster

async def err():
    raise RuntimeError

async def test_exception_handler():
    m = RecordingMaster(None)
    running = asyncio.create_task(m.run())
    asyncio.create_task(err())
    await m.await_log("Traceback", level="error")
    m.shutdown()
    await running
<|endoftext|>


<commit_before>#   Foremast - Pipeline Tooling
#
#   Copyright 2018 Gogo, LLC
#
#   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.
"""Load base config and export package constants.

The ``foremast`` configuration file is read from the following locations in
descending order. First found wins.

* ``./.foremast/foremast.cfg``
* ``~/.foremast/foremast.cfg``
* ``/etc/foremast/foremast.cfg``

.. literalinclude:: ../../src/foremast/templates/configs/foremast.cfg.example
   :language: ini
"""
import ast
import json
import logging
import sys
from configparser import ConfigParser
from os import getcwd, path
from os.path import exists, expanduser, expandvars

LOG = logging.getLogger(__name__)
LOGGING_FORMAT = '%(asctime)s [%(levelname)s] %(name)s:%(funcName)s:%(lineno)d - %(message)s'
SHORT_LOGGING_FORMAT = '[%(levelname)s] %(message)s'

logging.basicConfig(format=LOGGING_FORMAT)
logging.getLogger(__package__.split('.')[0]).setLevel(logging.INFO)

GOOD_STATUSES = frozenset(('SUCCEEDED', ))
SKIP_STATUSES = frozenset(('NOT_STARTED', ))

def validate_key_values(config_handle, section, key, default=None):
    """Warn when *key* is missing from configuration *section*.

    Args:
        config_handle (configparser.ConfigParser): Instance of configurations.
        section (str): Name of configuration section to retrieve.
        key (str): Configuration key to look up.
        default (object): Default object to use when *key* is not found.

    Returns:
        object: ``str`` when *key* exists, otherwise *default* object.
    """
    if section not in config_handle:
        LOG.info('Section missing from configurations: [%s]', section)

    try:
        value = config_handle[section][key]
    except KeyError:
        LOG.warning('[%s] missing key "%s", using %r.', section, key, default)
        value = default

    return value

def extract_formats(config_handle):
    """Get application formats.

    See :class:`gogoutils.Formats` for available options.

    Args:
        config_handle (configparser.ConfigParser): Instance of configurations.

    Returns:
        dict: Formats in ``{$format_type: $format_pattern}``.

    """
    configurations = dict(config_handle)
    formats = dict(configurations.get('formats', {}))
    return formats

def load_dynamic_config(config_dir=getcwd()):
    """Load and parse dynamic config"""
    dynamic_configurations = {}

    # Create full path of config
    config_file = '{path}/config.py'.format(path=config_dir)

    # Insert config path so we can import it
    sys.path.insert(0, path.dirname(path.abspath(config_file)))
    try:
        config_module = __import__('config')

        dynamic_configurations = config_module.CONFIG
    except ImportError:
        # Provide a default if config not found
        LOG.error('ImportError: Unable to load dynamic config. Check config.py file imports!')

    return dynamic_configurations

def find_config():
    """Look for **foremast.cfg** in config_locations or ``./config.py``.

    Raises:
        SystemExit: No configuration file found.

    Returns:
        dict: Found dynamic or static configuration.

    """
    config_locations = [
        '/etc/foremast/foremast.cfg',
        expanduser('~/.foremast/foremast.cfg'),
        './.foremast/foremast.cfg',
    ]
    configurations = ConfigParser()

    cfg_file = configurations.read(config_locations)
    dynamic_config_file = '{path}/config.py'.format(path=getcwd())

    if cfg_file:
        LOG.info('Loading static configuration file.')
    elif exists(dynamic_config_file):
        LOG.info('Loading dynamic configuration file.')
        configurations = load_dynamic_config()
    else:
        config_locations.append(dynamic_config_file)
        LOG.warning('No configuration found in the following locations:\n%s', '\n'.join(config_locations))
        LOG.warning('Using defaults...')

    return dict(configurations)

def _remove_empty_entries(entries):
    """Remove empty entries in a list"""
    valid_entries = []
    for entry in set(entries):
        if entry:
            valid_entries.append(entry)
    return sorted(valid_entries)

def _convert_string_to_native(value):
    """Convert a string to its native python type"""
    result = None

    try:
        result = ast.literal_eval(value)
    except (SyntaxError, ValueError):
        # Likely a string
        result = value.split(',')
    return result

def _generate_security_groups(config_key):
    """Read config file and generate security group dict by environment.

    Args:
        config_key (str): Configuration file key

    Returns:
        dict: of environments in {'env1': ['group1', 'group2']} format
    """
    raw_default_groups = validate_key_values(CONFIG, 'base', config_key, default='')
    default_groups = _convert_string_to_native(raw_default_groups)
    LOG.debug('Default security group for %s is %s', config_key, default_groups)

    entries = {}
    for env in ENVS:
        entries[env] = []

    if isinstance(default_groups, (list)):
        groups = _remove_empty_entries(default_groups)
        for env in entries:
            entries[env] = groups
    elif isinstance(default_groups, (dict)):
        entries.update(default_groups)

    LOG.debug('Generated security group: %s', entries)
    return entries

CONFIG = find_config()

API_URL = validate_key_values(CONFIG, 'base', 'gate_api_url')
GIT_URL = validate_key_values(CONFIG, 'base', 'git_url')
DOMAIN = validate_key_values(CONFIG, 'base', 'domain', default='example.com')
ENVS = set(validate_key_values(CONFIG, 'base', 'envs', default='').split(','))
REGIONS = set(validate_key_values(CONFIG, 'base', 'regions', default='').split(','))
ALLOWED_TYPES = set(
    validate_key_values(CONFIG, 'base', 'types', default='ec2,lambda,s3,datapipeline,rolling').split(','))
TEMPLATES_PATH = validate_key_values(CONFIG, 'base', 'templates_path')
AMI_JSON_URL = validate_key_values(CONFIG, 'base', 'ami_json_url')
DEFAULT_RUN_AS_USER = validate_key_values(CONFIG, 'base', 'default_run_as_user', default=None)
DEFAULT_SECURITYGROUP_RULES = _generate_security_groups('default_securitygroup_rules')
DEFAULT_EC2_SECURITYGROUPS = _generate_security_groups('default_ec2_securitygroups')
DEFAULT_ELB_SECURITYGROUPS = _generate_security_groups('default_elb_securitygroups')

EC2_PIPELINE_TYPES = tuple(validate_key_values(CONFIG, 'base', 'ec2_pipeline_types', default='ec2,rolling').split(','))
"""Comma separated list of Pipeline Types to treat as EC2 deployments.

This is useful when defining custom Pipeline Types. When Pipeline Type matches,
EC2 specific data is used in deployment, such as Auto Scaling Groups and
Availability Zones.

    | *Default*: ``ec2,rolling``
    | *Required*: No
    | *Example*: ``ec2,infrastructure,propeller``
"""

SECURITYGROUP_REPLACEMENTS = _convert_string_to_native(
    validate_key_values(
        CONFIG,
        'base',
        'securitygroup_replacements',
        default='{}',
    ))
GITLAB_TOKEN = validate_key_values(CONFIG, 'credentials', 'gitlab_token')
SLACK_TOKEN = validate_key_values(CONFIG, 'credentials', 'slack_token')
DEFAULT_TASK_TIMEOUT = validate_key_values(CONFIG, 'task_timeouts', 'default', default=120)
TASK_TIMEOUTS = json.loads(validate_key_values(CONFIG, 'task_timeouts', 'envs', default="{}"))
ASG_WHITELIST = set(validate_key_values(CONFIG, 'whitelists', 'asg_whitelist', default='').split(','))
APP_FORMATS = extract_formats(CONFIG)
GATE_CLIENT_CERT = expandvars(expanduser(validate_key_values(CONFIG, 'base', 'gate_client_cert', default='')))
GATE_CA_BUNDLE = expandvars(expanduser(validate_key_values(CONFIG, 'base', 'gate_ca_bundle', default='')))
LINKS = _convert_string_to_native(validate_key_values(CONFIG, 'links', 'default', default='{}'))

HEADERS = {
    'accept': '*/*',
    'content-type': 'application/json',
    'user-agent': 'foremast',
}
<commit_msg>feat: Environment variable FOREMAST_CONFIG_FILE<commit_after>#   Foremast - Pipeline Tooling
#
#   Copyright 2018 Gogo, LLC
#
#   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.
"""Load base config and export package constants.

The ``foremast`` configuration file is read from the following locations in
descending order. First found wins.

* ``./.foremast/foremast.cfg``
* ``~/.foremast/foremast.cfg``
* ``/etc/foremast/foremast.cfg``

.. literalinclude:: ../../src/foremast/templates/configs/foremast.cfg.example
   :language: ini
"""
import ast
import json
import logging
import sys
from configparser import ConfigParser
from os import getcwd, getenv, path
from os.path import exists, expanduser, expandvars

LOG = logging.getLogger(__name__)
LOGGING_FORMAT = '%(asctime)s [%(levelname)s] %(name)s:%(funcName)s:%(lineno)d - %(message)s'
SHORT_LOGGING_FORMAT = '[%(levelname)s] %(message)s'

logging.basicConfig(format=LOGGING_FORMAT)
logging.getLogger(__package__.split('.')[0]).setLevel(logging.INFO)

GOOD_STATUSES = frozenset(('SUCCEEDED', ))
SKIP_STATUSES = frozenset(('NOT_STARTED', ))

DEFAULT_DYNAMIC_CONFIG_FILE = '{path}/config.py'.format(path=getcwd())
"""Default `config.py` file path is in the current directory.

To override, use the `FOREMAST_CONFIG_FILE` environment variable.
"""

def validate_key_values(config_handle, section, key, default=None):
    """Warn when *key* is missing from configuration *section*.

    Args:
        config_handle (configparser.ConfigParser): Instance of configurations.
        section (str): Name of configuration section to retrieve.
        key (str): Configuration key to look up.
        default (object): Default object to use when *key* is not found.

    Returns:
        object: ``str`` when *key* exists, otherwise *default* object.
    """
    if section not in config_handle:
        LOG.info('Section missing from configurations: [%s]', section)

    try:
        value = config_handle[section][key]
    except KeyError:
        LOG.warning('[%s] missing key "%s", using %r.', section, key, default)
        value = default

    return value

def extract_formats(config_handle):
    """Get application formats.

    See :class:`gogoutils.Formats` for available options.

    Args:
        config_handle (configparser.ConfigParser): Instance of configurations.

    Returns:
        dict: Formats in ``{$format_type: $format_pattern}``.

    """
    configurations = dict(config_handle)
    formats = dict(configurations.get('formats', {}))
    return formats

def load_dynamic_config(config_file=DEFAULT_DYNAMIC_CONFIG_FILE):
    """Load and parse dynamic config"""
    dynamic_configurations = {}

    # Insert config path so we can import it
    sys.path.insert(0, path.dirname(path.abspath(config_file)))
    try:
        config_module = __import__('config')

        dynamic_configurations = config_module.CONFIG
    except ImportError:
        # Provide a default if config not found
        LOG.error('ImportError: Unable to load dynamic config. Check config.py file imports!')

    return dynamic_configurations

def find_config():
    """Look for **foremast.cfg** in config_locations or ``./config.py``.

    Raises:
        SystemExit: No configuration file found.

    Returns:
        dict: Found dynamic or static configuration.

    """
    config_locations = [
        '/etc/foremast/foremast.cfg',
        expanduser('~/.foremast/foremast.cfg'),
        './.foremast/foremast.cfg',
    ]
    configurations = ConfigParser()

    cfg_file = configurations.read(config_locations)
    dynamic_config_file = getenv('FOREMAST_CONFIG_FILE', DEFAULT_DYNAMIC_CONFIG_FILE)

    if cfg_file:
        LOG.info('Loading static configuration file.')
    elif exists(dynamic_config_file):
        LOG.info('Loading dynamic configuration file.')
        configurations = load_dynamic_config(config_file=dynamic_config_file)
    else:
        config_locations.append(dynamic_config_file)
        LOG.warning('No configuration found in the following locations:\n%s', '\n'.join(config_locations))
        LOG.warning('Using defaults...')

    return dict(configurations)

def _remove_empty_entries(entries):
    """Remove empty entries in a list"""
    valid_entries = []
    for entry in set(entries):
        if entry:
            valid_entries.append(entry)
    return sorted(valid_entries)

def _convert_string_to_native(value):
    """Convert a string to its native python type"""
    result = None

    try:
        result = ast.literal_eval(value)
    except (SyntaxError, ValueError):
        # Likely a string
        result = value.split(',')
    return result

def _generate_security_groups(config_key):
    """Read config file and generate security group dict by environment.

    Args:
        config_key (str): Configuration file key

    Returns:
        dict: of environments in {'env1': ['group1', 'group2']} format
    """
    raw_default_groups = validate_key_values(CONFIG, 'base', config_key, default='')
    default_groups = _convert_string_to_native(raw_default_groups)
    LOG.debug('Default security group for %s is %s', config_key, default_groups)

    entries = {}
    for env in ENVS:
        entries[env] = []

    if isinstance(default_groups, (list)):
        groups = _remove_empty_entries(default_groups)
        for env in entries:
            entries[env] = groups
    elif isinstance(default_groups, (dict)):
        entries.update(default_groups)

    LOG.debug('Generated security group: %s', entries)
    return entries

CONFIG = find_config()

API_URL = validate_key_values(CONFIG, 'base', 'gate_api_url')
GIT_URL = validate_key_values(CONFIG, 'base', 'git_url')
DOMAIN = validate_key_values(CONFIG, 'base', 'domain', default='example.com')
ENVS = set(validate_key_values(CONFIG, 'base', 'envs', default='').split(','))
REGIONS = set(validate_key_values(CONFIG, 'base', 'regions', default='').split(','))
ALLOWED_TYPES = set(
    validate_key_values(CONFIG, 'base', 'types', default='ec2,lambda,s3,datapipeline,rolling').split(','))
TEMPLATES_PATH = validate_key_values(CONFIG, 'base', 'templates_path')
AMI_JSON_URL = validate_key_values(CONFIG, 'base', 'ami_json_url')
DEFAULT_RUN_AS_USER = validate_key_values(CONFIG, 'base', 'default_run_as_user', default=None)
DEFAULT_SECURITYGROUP_RULES = _generate_security_groups('default_securitygroup_rules')
DEFAULT_EC2_SECURITYGROUPS = _generate_security_groups('default_ec2_securitygroups')
DEFAULT_ELB_SECURITYGROUPS = _generate_security_groups('default_elb_securitygroups')

EC2_PIPELINE_TYPES = tuple(validate_key_values(CONFIG, 'base', 'ec2_pipeline_types', default='ec2,rolling').split(','))
"""Comma separated list of Pipeline Types to treat as EC2 deployments.

This is useful when defining custom Pipeline Types. When Pipeline Type matches,
EC2 specific data is used in deployment, such as Auto Scaling Groups and
Availability Zones.

    | *Default*: ``ec2,rolling``
    | *Required*: No
    | *Example*: ``ec2,infrastructure,propeller``
"""

SECURITYGROUP_REPLACEMENTS = _convert_string_to_native(
    validate_key_values(
        CONFIG,
        'base',
        'securitygroup_replacements',
        default='{}',
    ))
GITLAB_TOKEN = validate_key_values(CONFIG, 'credentials', 'gitlab_token')
SLACK_TOKEN = validate_key_values(CONFIG, 'credentials', 'slack_token')
DEFAULT_TASK_TIMEOUT = validate_key_values(CONFIG, 'task_timeouts', 'default', default=120)
TASK_TIMEOUTS = json.loads(validate_key_values(CONFIG, 'task_timeouts', 'envs', default="{}"))
ASG_WHITELIST = set(validate_key_values(CONFIG, 'whitelists', 'asg_whitelist', default='').split(','))
APP_FORMATS = extract_formats(CONFIG)
GATE_CLIENT_CERT = expandvars(expanduser(validate_key_values(CONFIG, 'base', 'gate_client_cert', default='')))
GATE_CA_BUNDLE = expandvars(expanduser(validate_key_values(CONFIG, 'base', 'gate_ca_bundle', default='')))
LINKS = _convert_string_to_native(validate_key_values(CONFIG, 'links', 'default', default='{}'))

HEADERS = {
    'accept': '*/*',
    'content-type': 'application/json',
    'user-agent': 'foremast',
}
<|endoftext|>


<commit_before>import sys
from math import ceil

from OpenGL.GL import GL_LESS, GL_TRUE, GL_DEPTH_TEST, GL_STENCIL_TEST
from OpenGL.GL import GL_COLOR_ARRAY, GL_VERTEX_ARRAY, GL_TRIANGLES
from OpenGL.GL import GL_COLOR_BUFFER_BIT, GL_DEPTH_BUFFER_BIT
from OpenGL.GL import GL_UNSIGNED_SHORT, GL_FLOAT
from OpenGL.GL import GL_CULL_FACE, GL_FRONT, GL_MODELVIEW_MATRIX

from OpenGL.GL import glLoadIdentity, glOrtho
from OpenGL.GL import glDepthMask, glDepthFunc, glCullFace, glDisable
from OpenGL.GL import glEnable, glClearColor, glEnableClientState, glClear
from OpenGL.GL import glDrawElements, glGetFloatv, glFinish
from OpenGL.GL import glActiveTexture, glBindFramebuffer
from OpenGL.GL import GL_TEXTURE0, GL_TEXTURE1, GL_FRAMEBUFFER

from OpenGL.GL import glReadBuffer, glReadPixels, GL_RGBA, glPolygonOffset
from OpenGL.GL import GL_POLYGON_OFFSET_FILL

from OpenGL.GLUT import GLUT_DEPTH, GLUT_RGB, GLUT_ALPHA, GLUT_DOUBLE

from OpenGL.GLUT import glutSwapBuffers
from OpenGL.GLUT import glutInitWindowSize, glutPostRedisplay
from OpenGL.GLUT import glutCreateWindow, glutInit, glutInitWindowPosition
from OpenGL.GLUT import glutInitDisplayMode, glutLeaveMainLoop, glutDisplayFunc

from OpenGL.GLU import gluLookAt

from numpy import zeros, array

from .ShadersHelper import ShadersHelper
from .normals import get_normals
from .face import calculate_face, init_face_calculator

class View:
    """Viewport for Faces."""

    __triangles = None
    __principal_components = None
    __principal_components_flattened = None
    __deviations = None
    __mean_face = None

    def __init__(self, size):
        """Initialize viewport with initial Face rotation and position."""
        self.__size = size

        self.__height, self.__width = self.__size
        self.__output_image = zeros(self.__width * self.__height * 4,
                                    dtype='f')

        self.__light = None
        self.__face = None
        self.__face_vertices = None
        self.__model_matrix = zeros((4, 4), dtype='f')
        self.__light_matrix = zeros((4, 4), dtype='f')

        self.__init_display()
        self.__enable_depth_test()

        glEnableClientState(GL_COLOR_ARRAY)
        glEnableClientState(GL_VERTEX_ARRAY)

        glClearColor(1., 1., 1., 0.)

        self.__sh = ShadersHelper(['face.vert', 'depth.vert'],
                                  ['face.frag', 'depth.frag'], 2, 1)

        glutDisplayFunc(self.__display)
        self.__callback = None

        self.__sh.add_attribute(0, array([]), 'face_vertices')
        self.__sh.add_attribute(1, array([]), 'normal_vector')
        self.__sh.bind_buffer()
        self.__sh.use_shaders()
        self.__sh.link_texture('depth_map', 0)
        self.__sh.bind_depth_texture(self.__size)

    def get_size(self):
        """Get size of the viewport."""
        return self.__size

    @property
    def light(self):
        """Get light direction."""
        return self.__light

    @light.setter
    def light(self, light):
        """Set light direction."""
        self.__light = light

    @property
    def face(self):
        """Get current Face."""
        return self.__face

    @face.setter
    def face(self, face):
        """Set current Face."""
        self.__face = face
        self.__face_vertices = calculate_face(self.__face.coefficients)

    def redraw(self, callback=None):
        """Trigger redisplay and trigger callback after render."""
        self.__callback = callback
        glutPostRedisplay()

    def get_image(self):
        """Copy RGBA data from the viewport to NumPy Matrix of float."""
        glReadBuffer(GL_FRONT)
        glReadPixels(0, 0, self.__width, self.__height, GL_RGBA, GL_FLOAT,
                     self.__output_image)
        return self.__output_image

    def close(self):
        """Close the viewport."""
        glutLeaveMainLoop()

    @staticmethod
    def set_triangles(triangles):
        """Set triangles of the model.

        Allows to not allocate additional memory for each point,
        which is used in multiple triangles.
        """
        View.__triangles = triangles

    @staticmethod
    def set_principal_components(principal_components):
        """Set principal components for Face calculation."""
        View.__principal_components = principal_components
        View.__principal_components_flattened = principal_components.flatten()

    @staticmethod
    def set_deviations(deviations):
        """Get principal components deviations."""
        View.__deviations = deviations

    @staticmethod
    def set_mean_face(mean_face):
        """Set mean Face for modelling."""
        View.__mean_face = mean_face

    @staticmethod
    def finalize_initialization():
        init_face_calculator(
            View.__mean_face,
            View.__principal_components_flattened,
            View.__deviations)

    @staticmethod
    def __get_rotation_matrix(coordinates, side_length):
        """Get rotation matrix from specific point of view and scale."""
        assert len(coordinates) == 3
        glLoadIdentity()
        glOrtho(-side_length, side_length, -side_length, side_length,
                -4 * side_length, 4 * side_length)
        x, y, z = coordinates
        gluLookAt(x, y, z, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0)
        return array(glGetFloatv(GL_MODELVIEW_MATRIX), dtype='f')

    def __display(self):
        """Render the model by existent vertices, colors and triangles."""
        self.__rotate_model()
        self.__generate_shadows()
        self.__generate_model()

        glutSwapBuffers()
        if self.__callback is not None:
            self.__callback()

    def __rotate_model(self):
        """Update model rotation matrix."""
        self.__model_matrix = self.__get_rotation_matrix(
            self.__face.position_cartesian,
            (1 + self.__face.position[2]) * 0.5)

        light = self.__face.directed_light_cartesian
        self.__light_matrix = self.__get_rotation_matrix(
            (light[0], light[1], light[2]), 1.0)

    def __generate_shadows(self):
        """Generate shadow matrix for rotated model."""
        glDisable(GL_CULL_FACE)
        glEnable(GL_POLYGON_OFFSET_FILL)
        glPolygonOffset(3, 0)

        self.__sh.change_shader(vertex=1, fragment=1)
        self.__prepare_shaders(self.__model_matrix, self.__light_matrix, True)
        self.__sh.bind_fbo()
        glClear(GL_DEPTH_BUFFER_BIT)
        glDrawElements(GL_TRIANGLES, View.__triangles.size,
                       GL_UNSIGNED_SHORT, View.__triangles)
        glFinish()
        glBindFramebuffer(GL_FRAMEBUFFER, 0)
        self.__sh.clear()

    def __generate_model(self):
        """Generate rotated model with shadows."""
        glEnable(GL_CULL_FACE)

        self.__sh.change_shader(vertex=0, fragment=0)
        self.__prepare_shaders(self.__model_matrix, self.__light_matrix, False)
        self.__sh.bind_buffer()
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glDrawElements(GL_TRIANGLES, View.__triangles.size,
                       GL_UNSIGNED_SHORT, View.__triangles)
        self.__sh.clear()

    def __prepare_shaders(self, rotation_matrix=None, light_matrix=None,
                          depth=True):
        """Generic shaders preparation method for depth map and final scene."""
        self.__sh.add_attribute(0, self.__face_vertices, 'face_vertices')
        self.__sh.bind_buffer()

        self.__sh.use_shaders()

        self.__sh.bind_uniform_matrix(light_matrix, 'light_matrix')
        if not depth:
            face = self.__face_vertices.reshape(self.__mean_face.size // 3, 3)
            normals = get_normals(face, View.__triangles.flatten())
            self.__sh.add_attribute(1, normals, 'normal_vector')

            self.__sh.bind_uniform_matrix(rotation_matrix, 'rotation_matrix')
            self.__sh.bind_uniform_vector(self.__face.light_cartesian,
                                          'light_vector')

        if not depth:
            glActiveTexture(GL_TEXTURE0)
            self.__sh.bind_texture(0)

    def __init_display(self):
        """Initialize the viewport with specified size."""
        glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_ALPHA | GLUT_DEPTH)
        glutInitWindowSize(*self.__size)

        glutInitWindowPosition(0, 0)
        glutInit(sys.argv)
        glutCreateWindow(b"Morphable face model")

    def __enable_depth_test(self):
        """Enable depth test and faces culling.

        Needed to not render invisible vertices and triangles.
        """
        glDepthMask(GL_TRUE)
        glDepthFunc(GL_LESS)
        glCullFace(GL_FRONT)

        glEnable(GL_DEPTH_TEST)
        glEnable(GL_STENCIL_TEST)
<commit_msg>Added doc comment to C modules initialization in View<commit_after>import sys
from math import ceil

from OpenGL.GL import GL_LESS, GL_TRUE, GL_DEPTH_TEST, GL_STENCIL_TEST
from OpenGL.GL import GL_COLOR_ARRAY, GL_VERTEX_ARRAY, GL_TRIANGLES
from OpenGL.GL import GL_COLOR_BUFFER_BIT, GL_DEPTH_BUFFER_BIT
from OpenGL.GL import GL_UNSIGNED_SHORT, GL_FLOAT
from OpenGL.GL import GL_CULL_FACE, GL_FRONT, GL_MODELVIEW_MATRIX

from OpenGL.GL import glLoadIdentity, glOrtho
from OpenGL.GL import glDepthMask, glDepthFunc, glCullFace, glDisable
from OpenGL.GL import glEnable, glClearColor, glEnableClientState, glClear
from OpenGL.GL import glDrawElements, glGetFloatv, glFinish
from OpenGL.GL import glActiveTexture, glBindFramebuffer
from OpenGL.GL import GL_TEXTURE0, GL_TEXTURE1, GL_FRAMEBUFFER

from OpenGL.GL import glReadBuffer, glReadPixels, GL_RGBA, glPolygonOffset
from OpenGL.GL import GL_POLYGON_OFFSET_FILL

from OpenGL.GLUT import GLUT_DEPTH, GLUT_RGB, GLUT_ALPHA, GLUT_DOUBLE

from OpenGL.GLUT import glutSwapBuffers
from OpenGL.GLUT import glutInitWindowSize, glutPostRedisplay
from OpenGL.GLUT import glutCreateWindow, glutInit, glutInitWindowPosition
from OpenGL.GLUT import glutInitDisplayMode, glutLeaveMainLoop, glutDisplayFunc

from OpenGL.GLU import gluLookAt

from numpy import zeros, array

from .ShadersHelper import ShadersHelper
from .normals import get_normals
from .face import calculate_face, init_face_calculator

class View:
    """Viewport for Faces."""

    __triangles = None
    __principal_components = None
    __principal_components_flattened = None
    __deviations = None
    __mean_face = None

    def __init__(self, size):
        """Initialize viewport with initial Face rotation and position."""
        self.__size = size

        self.__height, self.__width = self.__size
        self.__output_image = zeros(self.__width * self.__height * 4,
                                    dtype='f')

        self.__light = None
        self.__face = None
        self.__face_vertices = None
        self.__model_matrix = zeros((4, 4), dtype='f')
        self.__light_matrix = zeros((4, 4), dtype='f')

        self.__init_display()
        self.__enable_depth_test()

        glEnableClientState(GL_COLOR_ARRAY)
        glEnableClientState(GL_VERTEX_ARRAY)

        glClearColor(1., 1., 1., 0.)

        self.__sh = ShadersHelper(['face.vert', 'depth.vert'],
                                  ['face.frag', 'depth.frag'], 2, 1)

        glutDisplayFunc(self.__display)
        self.__callback = None

        self.__sh.add_attribute(0, array([]), 'face_vertices')
        self.__sh.add_attribute(1, array([]), 'normal_vector')
        self.__sh.bind_buffer()
        self.__sh.use_shaders()
        self.__sh.link_texture('depth_map', 0)
        self.__sh.bind_depth_texture(self.__size)

    def get_size(self):
        """Get size of the viewport."""
        return self.__size

    @property
    def light(self):
        """Get light direction."""
        return self.__light

    @light.setter
    def light(self, light):
        """Set light direction."""
        self.__light = light

    @property
    def face(self):
        """Get current Face."""
        return self.__face

    @face.setter
    def face(self, face):
        """Set current Face."""
        self.__face = face
        self.__face_vertices = calculate_face(self.__face.coefficients)

    def redraw(self, callback=None):
        """Trigger redisplay and trigger callback after render."""
        self.__callback = callback
        glutPostRedisplay()

    def get_image(self):
        """Copy RGBA data from the viewport to NumPy Matrix of float."""
        glReadBuffer(GL_FRONT)
        glReadPixels(0, 0, self.__width, self.__height, GL_RGBA, GL_FLOAT,
                     self.__output_image)
        return self.__output_image

    def close(self):
        """Close the viewport."""
        glutLeaveMainLoop()

    @staticmethod
    def set_triangles(triangles):
        """Set triangles of the model.

        Allows to not allocate additional memory for each point,
        which is used in multiple triangles.
        """
        View.__triangles = triangles

    @staticmethod
    def set_principal_components(principal_components):
        """Set principal components for Face calculation."""
        View.__principal_components = principal_components
        View.__principal_components_flattened = principal_components.flatten()

    @staticmethod
    def set_deviations(deviations):
        """Get principal components deviations."""
        View.__deviations = deviations

    @staticmethod
    def set_mean_face(mean_face):
        """Set mean Face for modelling."""
        View.__mean_face = mean_face

    @staticmethod
    def finalize_initialization():
        """Initialize C methods."""
        init_face_calculator(
            View.__mean_face,
            View.__principal_components_flattened,
            View.__deviations)

    @staticmethod
    def __get_rotation_matrix(coordinates, side_length):
        """Get rotation matrix from specific point of view and scale."""
        assert len(coordinates) == 3
        glLoadIdentity()
        glOrtho(-side_length, side_length, -side_length, side_length,
                -4 * side_length, 4 * side_length)
        x, y, z = coordinates
        gluLookAt(x, y, z, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0)
        return array(glGetFloatv(GL_MODELVIEW_MATRIX), dtype='f')

    def __display(self):
        """Render the model by existent vertices, colors and triangles."""
        self.__rotate_model()
        self.__generate_shadows()
        self.__generate_model()

        glutSwapBuffers()
        if self.__callback is not None:
            self.__callback()

    def __rotate_model(self):
        """Update model rotation matrix."""
        self.__model_matrix = self.__get_rotation_matrix(
            self.__face.position_cartesian,
            (1 + self.__face.position[2]) * 0.5)

        light = self.__face.directed_light_cartesian
        self.__light_matrix = self.__get_rotation_matrix(
            (light[0], light[1], light[2]), 1.0)

    def __generate_shadows(self):
        """Generate shadow matrix for rotated model."""
        glDisable(GL_CULL_FACE)
        glEnable(GL_POLYGON_OFFSET_FILL)
        glPolygonOffset(3, 0)

        self.__sh.change_shader(vertex=1, fragment=1)
        self.__prepare_shaders(self.__model_matrix, self.__light_matrix, True)
        self.__sh.bind_fbo()
        glClear(GL_DEPTH_BUFFER_BIT)
        glDrawElements(GL_TRIANGLES, View.__triangles.size,
                       GL_UNSIGNED_SHORT, View.__triangles)
        glFinish()
        glBindFramebuffer(GL_FRAMEBUFFER, 0)
        self.__sh.clear()

    def __generate_model(self):
        """Generate rotated model with shadows."""
        glEnable(GL_CULL_FACE)

        self.__sh.change_shader(vertex=0, fragment=0)
        self.__prepare_shaders(self.__model_matrix, self.__light_matrix, False)
        self.__sh.bind_buffer()
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glDrawElements(GL_TRIANGLES, View.__triangles.size,
                       GL_UNSIGNED_SHORT, View.__triangles)
        self.__sh.clear()

    def __prepare_shaders(self, rotation_matrix=None, light_matrix=None,
                          depth=True):
        """Generic shaders preparation method for depth map and final scene."""
        self.__sh.add_attribute(0, self.__face_vertices, 'face_vertices')
        self.__sh.bind_buffer()

        self.__sh.use_shaders()

        self.__sh.bind_uniform_matrix(light_matrix, 'light_matrix')
        if not depth:
            face = self.__face_vertices.reshape(self.__mean_face.size // 3, 3)
            normals = get_normals(face, View.__triangles.flatten())
            self.__sh.add_attribute(1, normals, 'normal_vector')

            self.__sh.bind_uniform_matrix(rotation_matrix, 'rotation_matrix')
            self.__sh.bind_uniform_vector(self.__face.light_cartesian,
                                          'light_vector')

        if not depth:
            glActiveTexture(GL_TEXTURE0)
            self.__sh.bind_texture(0)

    def __init_display(self):
        """Initialize the viewport with specified size."""
        glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_ALPHA | GLUT_DEPTH)
        glutInitWindowSize(*self.__size)

        glutInitWindowPosition(0, 0)
        glutInit(sys.argv)
        glutCreateWindow(b"Morphable face model")

    def __enable_depth_test(self):
        """Enable depth test and faces culling.

        Needed to not render invisible vertices and triangles.
        """
        glDepthMask(GL_TRUE)
        glDepthFunc(GL_LESS)
        glCullFace(GL_FRONT)

        glEnable(GL_DEPTH_TEST)
        glEnable(GL_STENCIL_TEST)
<|endoftext|>


<commit_before>"""Handle database functions."""

import sqlite3
import sys
import os
from os.path import basename, join, exists

ATRAM_VERSION = 'v2.1.0'
DB_VERSION = '2.0'

BATCH_SIZE = 1e6  # How many sequence records to insert at a time

def connect(blast_db, check_version=False, clean=False):
    """Create DB connection."""
    db_name = get_db_name(blast_db)

    if clean and exists(db_name):
        os.remove(db_name)

    if check_version and not exists(db_name):
        err = 'Could not find the database file "{}".'.format(db_name)
        sys.exit(err)

    cxn = sqlite3.connect(db_name)

    cxn.execute("PRAGMA page_size = {}".format(2**16))
    cxn.execute("PRAGMA busy_timeout = 10000")
    cxn.execute("PRAGMA journal_mode = WAL")
    # cxn.execute("PRAGMA synchronous = OFF")

    if check_version:
        check_versions(cxn)

    return cxn

def get_db_name(blast_db):
    """Build the SQLite DB name from the blast DB argument."""
    return '{}.sqlite.db'.format(blast_db)

def aux_db(cxn, temp_dir, blast_db, query_name):
    """Create & attach an temporary database to the current DB connection."""
    db_dir = join(temp_dir, 'db')
    os.makedirs(db_dir, exist_ok=True)

    db_name = '{}_{}_temp.sqlite.db'.format(
        basename(blast_db), basename(query_name))
    db_name = join(db_dir, db_name)

    sql = """ATTACH DATABASE '{}' AS aux""".format(db_name)
    cxn.execute(sql)

def aux_detach(cxn):
    """Detach the temporary database."""
    cxn.execute('DETACH DATABASE aux')

# ########################### misc functions #################################
# DB_VERSION != version in DB. Don't force DB changes until required. So
# this version will tend to lag ATRAM_VERSION.

def check_versions(cxn):
    """Make sure the database version matches what we built it with."""
    version = get_version(cxn)
    if version != DB_VERSION:
        err = ('The database was built with version {} but you are running '
               'version {}. You need to rebuild the atram database by '
               'running atram_preprocessor.py again.').format(
                   version, DB_VERSION)
        sys.exit(err)

# ########################## metadata table ##################################

def create_metadata_table(cxn):
    """
    Create the metadata table.

    A single record used to tell if we are running atram.py against the
    schema version we built with atram_preprocessor.py.
    """
    cxn.execute("""DROP TABLE IF EXISTS metadata""")
    sql = 'CREATE TABLE metadata (label TEXT, value TEXT)'
    cxn.execute(sql)

    with cxn:
        sql = """INSERT INTO metadata (label, value) VALUES (?, ?)"""
        cxn.execute(sql, ('version', DB_VERSION))
        cxn.commit()

def get_version(cxn):
    """Get the current database version."""
    sql = """SELECT value FROM metadata WHERE label = ?"""
    try:
        result = cxn.execute(sql, ('version', ))
        return result.fetchone()[0]
    except sqlite3.OperationalError:
        return '1.0'

# ########################## sequences table ##################################

def create_sequences_table(cxn):
    """Create the sequence table."""
    cxn.execute("""DROP TABLE IF EXISTS sequences""")
    sql = 'CREATE TABLE sequences (seq_name TEXT, seq_end TEXT, seq TEXT)'
    cxn.execute(sql)

def create_sequences_index(cxn):
    """
    Create the sequences index after we build the table.

    This speeds up the program significantly.
    """
    sql = 'CREATE INDEX sequences_index ON sequences (seq_name, seq_end)'
    cxn.execute(sql)

def insert_sequences_batch(cxn, batch):
    """Insert a batch of sequence records into the database."""
    if batch:
        sql = """INSERT INTO sequences (seq_name, seq_end, seq)
                      VALUES (?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def get_sequence_count(cxn):
    """Get the number of sequences in the table."""
    result = cxn.execute('SELECT COUNT(*) FROM sequences')
    return result.fetchone()[0]

def get_shard_cut(cxn, offset):
    """Get the sequence name at the given offset."""
    sql = 'SELECT seq_name FROM sequences ORDER BY seq_name LIMIT 1 OFFSET {}'
    result = cxn.execute(sql.format(offset))
    cut = result.fetchone()[0]
    return cut

def get_sequences_in_shard(cxn, start, end):
    """Get all sequences in a shard."""
    sql = """
        SELECT seq_name, seq_end, seq
          FROM sequences
         WHERE seq_name >= ?
           AND seq_name < ?
        """
    return cxn.execute(sql, (start, end))

def get_sequence_ends(cxn):
    """Get a list of all seq_ends in the database."""
    return cxn.execute('SELECT DISTINCT seq_end FROM sequences')

def get_all_sequences(cxn):
    """Get a list of all seq_ends in the database."""
    return cxn.execute('SELECT * FROM sequences')

# ######################## sra_blast_hits table ###############################

def create_sra_blast_hits_table(cxn):
    """
    Reset the DB.

    Delete the tables and recreate them.
    """
    cxn.execute("""DROP TABLE IF EXISTS aux.sra_blast_hits""")
    sql = """
        CREATE TABLE aux.sra_blast_hits
                   (iteration INTEGER, seq_name TEXT, seq_end TEXT, shard TEXT)
        """
    cxn.execute(sql)

    sql = """
        CREATE INDEX aux.sra_blast_hits_index
                  ON sra_blast_hits (iteration, seq_name, seq_end)
        """
    cxn.execute(sql)

def insert_blast_hit_batch(cxn, batch):
    """Insert a batch of blast hit records into the database."""
    if batch:
        sql = """
            INSERT INTO aux.sra_blast_hits
                        (iteration, seq_end, seq_name, shard)
                        VALUES (?, ?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def sra_blast_hits_count(cxn, iteration):
    """Count the blast hist for select the iteration."""
    sql = """
        SELECT COUNT(*) AS count
          FROM aux.sra_blast_hits
         WHERE iteration = ?
        """

    result = cxn.execute(sql, (iteration, ))
    return result.fetchone()[0]

def get_sra_blast_hits(cxn, iteration):
    """Get all blast hits for the iteration."""
    sql = """
        SELECT seq_name, seq_end, seq
          FROM sequences
         WHERE seq_name IN (SELECT DISTINCT seq_name
                              FROM aux.sra_blast_hits
                             WHERE iteration = ?)
      ORDER BY seq_name, seq_end
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (iteration, ))

def get_blast_hits_by_end_count(cxn, iteration, end_count):
    """Get all blast hits for the iteration."""
    sql = """
        SELECT seq_name, seq_end, seq
          FROM sequences
         WHERE seq_name IN (SELECT seq_name
                              FROM sequences
                             WHERE seq_name IN (SELECT DISTINCT seq_name
                                                  FROM aux.sra_blast_hits
                                                 WHERE iteration = ?)
                          GROUP BY seq_name
                            HAVING COUNT(*) = ?)
      ORDER BY seq_name, seq_end
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (iteration, end_count))

# ####################### contig_blast_hits table #############################

def create_contig_blast_hits_table(cxn):
    """Reset the database. Delete the tables and recreate them."""
    cxn.execute("""DROP TABLE IF EXISTS aux.contig_blast_hits""")
    sql = """
        CREATE TABLE aux.contig_blast_hits
                     (iteration INTEGER, contig_id TEXT, description TEXT,
                      bit_score NUMERIC, len INTEGER,
                      query_from INTEGER, query_to INTEGER, query_strand TEXT,
                      hit_from INTEGER, hit_to INTEGER, hit_strand TEXT)
        """
    cxn.execute(sql)

    sql = """
        CREATE INDEX aux.contig_blast_hits_index
                  ON contig_blast_hits (iteration, bit_score, len)
        """
    cxn.execute(sql)

def insert_contig_hit_batch(cxn, batch):
    """Insert a batch of blast hit records into the database."""
    if batch:
        sql = """
            INSERT INTO aux.contig_blast_hits
                        (iteration, contig_id, description, bit_score, len,
                         query_from, query_to, query_strand,
                         hit_from, hit_to, hit_strand)
                 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def get_contig_blast_hits(cxn, iteration):
    """Get all blast hits for the iteration."""
    sql = """
        SELECT iteration, contig_id, description, bit_score, len,
               query_from, query_to, query_strand,
               hit_from, hit_to, hit_strand
          FROM aux.contig_blast_hits
         WHERE iteration = ?
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (iteration, ))

# ####################### assembled_contigs table #############################

def create_assembled_contigs_table(cxn):
    """Reset the database. Delete the tables and recreate them."""
    cxn.execute("""DROP TABLE IF EXISTS aux.assembled_contigs""")
    sql = """
        CREATE TABLE aux.assembled_contigs
                     (iteration INTEGER, contig_id TEXT, seq TEXT,
                      description TEXT, bit_score NUMERIC, len INTEGER,
                      query_from INTEGER, query_to INTEGER, query_strand TEXT,
                      hit_from INTEGER, hit_to INTEGER, hit_strand TEXT)
        """
    cxn.execute(sql)

    sql = """
        CREATE INDEX aux.assembled_contigs_index
                  ON assembled_contigs (iteration, contig_id)
        """
    cxn.execute(sql)

def assembled_contigs_count(cxn, iteration, bit_score, length):
    """Count the blast hist for the iteration."""
    sql = """
        SELECT COUNT(*) AS count
          FROM aux.assembled_contigs
         WHERE iteration = ?
           AND bit_score >= ?
           AND len >= ?
        """

    result = cxn.execute(sql, (iteration, bit_score, length))
    return result.fetchone()[0]

def iteration_overlap_count(cxn, iteration, bit_score, length):
    """Count how many assembled contigs match what's in the last iteration."""
    sql = """
        SELECT COUNT(*) AS overlap
          FROM aux.assembled_contigs AS curr_iter
          JOIN aux.assembled_contigs AS prev_iter
            ON (     curr_iter.contig_id = prev_iter.contig_id
                 AND curr_iter.iteration = prev_iter.iteration + 1)
         WHERE curr_iter.iteration = ?
           AND curr_iter.seq = prev_iter.seq
           AND curr_iter.bit_score >= ?
           AND prev_iter.bit_score >= ?
           AND curr_iter.len >= ?
        """
    result = cxn.execute(
        sql, (iteration, bit_score, bit_score, length))
    return result.fetchone()[0]

def insert_assembled_contigs_batch(cxn, batch):
    """Insert a batch of blast hit records into the database."""
    if batch:
        sql = """
            INSERT INTO aux.assembled_contigs
                        (iteration, contig_id, seq, description,
                         bit_score, len,
                         query_from, query_to, query_strand,
                         hit_from, hit_to, hit_strand)
                 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def get_assembled_contigs(cxn, iteration, bit_score, length):
    """
    Get all assembled contigs for the iteration.

    We will use them as the queries in the next atram iteration.
    """
    sql = """
        SELECT contig_id, seq
          FROM aux.assembled_contigs
         WHERE iteration = ?
           AND bit_score >= ?
           AND len >= ?
         """
    return cxn.execute(sql, (iteration, bit_score, length))

def get_all_assembled_contigs(cxn, bit_score=0, length=0):
    """Get all assembled contigs."""
    sql = """
        SELECT iteration, contig_id, seq, description, bit_score, len,
               query_from, query_to, query_strand,
               hit_from, hit_to, hit_strand
          FROM aux.assembled_contigs
         WHERE bit_score >= ?
           AND len >= ?
      ORDER BY bit_score DESC, iteration
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (bit_score, length))

def all_assembled_contigs_count(cxn, bit_score=0, length=0):
    """Count all assembed contigs."""
    sql = """
        SELECT COUNT(*) AS count
          FROM aux.assembled_contigs
         WHERE bit_score >= ?
           AND len >= ?
        """

    result = cxn.execute(sql, (bit_score, length))
    return result.fetchone()[0]
<commit_msg>Remove dead code<commit_after>"""Handle database functions."""

import sqlite3
import sys
import os
from os.path import basename, join, exists

ATRAM_VERSION = 'v2.1.0'
DB_VERSION = '2.0'

BATCH_SIZE = 1e6  # How many sequence records to insert at a time

def connect(blast_db, check_version=False, clean=False):
    """Create DB connection."""
    db_name = get_db_name(blast_db)

    if clean and exists(db_name):
        os.remove(db_name)

    if check_version and not exists(db_name):
        err = 'Could not find the database file "{}".'.format(db_name)
        sys.exit(err)

    cxn = sqlite3.connect(db_name)

    cxn.execute("PRAGMA page_size = {}".format(2**16))
    cxn.execute("PRAGMA busy_timeout = 10000")
    cxn.execute("PRAGMA journal_mode = WAL")

    if check_version:
        check_versions(cxn)

    return cxn

def get_db_name(blast_db):
    """Build the SQLite DB name from the blast DB argument."""
    return '{}.sqlite.db'.format(blast_db)

def aux_db(cxn, temp_dir, blast_db, query_name):
    """Create & attach an temporary database to the current DB connection."""
    db_dir = join(temp_dir, 'db')
    os.makedirs(db_dir, exist_ok=True)

    db_name = '{}_{}_temp.sqlite.db'.format(
        basename(blast_db), basename(query_name))
    db_name = join(db_dir, db_name)

    sql = """ATTACH DATABASE '{}' AS aux""".format(db_name)
    cxn.execute(sql)

def aux_detach(cxn):
    """Detach the temporary database."""
    cxn.execute('DETACH DATABASE aux')

# ########################### misc functions #################################
# DB_VERSION != version in DB. Don't force DB changes until required. So
# this version will tend to lag ATRAM_VERSION.

def check_versions(cxn):
    """Make sure the database version matches what we built it with."""
    version = get_version(cxn)
    if version != DB_VERSION:
        err = ('The database was built with version {} but you are running '
               'version {}. You need to rebuild the atram database by '
               'running atram_preprocessor.py again.').format(
                   version, DB_VERSION)
        sys.exit(err)

# ########################## metadata table ##################################

def create_metadata_table(cxn):
    """
    Create the metadata table.

    A single record used to tell if we are running atram.py against the
    schema version we built with atram_preprocessor.py.
    """
    cxn.execute("""DROP TABLE IF EXISTS metadata""")
    sql = 'CREATE TABLE metadata (label TEXT, value TEXT)'
    cxn.execute(sql)

    with cxn:
        sql = """INSERT INTO metadata (label, value) VALUES (?, ?)"""
        cxn.execute(sql, ('version', DB_VERSION))
        cxn.commit()

def get_version(cxn):
    """Get the current database version."""
    sql = """SELECT value FROM metadata WHERE label = ?"""
    try:
        result = cxn.execute(sql, ('version', ))
        return result.fetchone()[0]
    except sqlite3.OperationalError:
        return '1.0'

# ########################## sequences table ##################################

def create_sequences_table(cxn):
    """Create the sequence table."""
    cxn.execute("""DROP TABLE IF EXISTS sequences""")
    sql = 'CREATE TABLE sequences (seq_name TEXT, seq_end TEXT, seq TEXT)'
    cxn.execute(sql)

def create_sequences_index(cxn):
    """
    Create the sequences index after we build the table.

    This speeds up the program significantly.
    """
    sql = 'CREATE INDEX sequences_index ON sequences (seq_name, seq_end)'
    cxn.execute(sql)

def insert_sequences_batch(cxn, batch):
    """Insert a batch of sequence records into the database."""
    if batch:
        sql = """INSERT INTO sequences (seq_name, seq_end, seq)
                      VALUES (?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def get_sequence_count(cxn):
    """Get the number of sequences in the table."""
    result = cxn.execute('SELECT COUNT(*) FROM sequences')
    return result.fetchone()[0]

def get_shard_cut(cxn, offset):
    """Get the sequence name at the given offset."""
    sql = 'SELECT seq_name FROM sequences ORDER BY seq_name LIMIT 1 OFFSET {}'
    result = cxn.execute(sql.format(offset))
    cut = result.fetchone()[0]
    return cut

def get_sequences_in_shard(cxn, start, end):
    """Get all sequences in a shard."""
    sql = """
        SELECT seq_name, seq_end, seq
          FROM sequences
         WHERE seq_name >= ?
           AND seq_name < ?
        """
    return cxn.execute(sql, (start, end))

def get_sequence_ends(cxn):
    """Get a list of all seq_ends in the database."""
    return cxn.execute('SELECT DISTINCT seq_end FROM sequences')

def get_all_sequences(cxn):
    """Get a list of all seq_ends in the database."""
    return cxn.execute('SELECT * FROM sequences')

# ######################## sra_blast_hits table ###############################

def create_sra_blast_hits_table(cxn):
    """
    Reset the DB.

    Delete the tables and recreate them.
    """
    cxn.execute("""DROP TABLE IF EXISTS aux.sra_blast_hits""")
    sql = """
        CREATE TABLE aux.sra_blast_hits
                   (iteration INTEGER, seq_name TEXT, seq_end TEXT, shard TEXT)
        """
    cxn.execute(sql)

    sql = """
        CREATE INDEX aux.sra_blast_hits_index
                  ON sra_blast_hits (iteration, seq_name, seq_end)
        """
    cxn.execute(sql)

def insert_blast_hit_batch(cxn, batch):
    """Insert a batch of blast hit records into the database."""
    if batch:
        sql = """
            INSERT INTO aux.sra_blast_hits
                        (iteration, seq_end, seq_name, shard)
                        VALUES (?, ?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def sra_blast_hits_count(cxn, iteration):
    """Count the blast hist for select the iteration."""
    sql = """
        SELECT COUNT(*) AS count
          FROM aux.sra_blast_hits
         WHERE iteration = ?
        """

    result = cxn.execute(sql, (iteration, ))
    return result.fetchone()[0]

def get_sra_blast_hits(cxn, iteration):
    """Get all blast hits for the iteration."""
    sql = """
        SELECT seq_name, seq_end, seq
          FROM sequences
         WHERE seq_name IN (SELECT DISTINCT seq_name
                              FROM aux.sra_blast_hits
                             WHERE iteration = ?)
      ORDER BY seq_name, seq_end
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (iteration, ))

def get_blast_hits_by_end_count(cxn, iteration, end_count):
    """Get all blast hits for the iteration."""
    sql = """
        SELECT seq_name, seq_end, seq
          FROM sequences
         WHERE seq_name IN (SELECT seq_name
                              FROM sequences
                             WHERE seq_name IN (SELECT DISTINCT seq_name
                                                  FROM aux.sra_blast_hits
                                                 WHERE iteration = ?)
                          GROUP BY seq_name
                            HAVING COUNT(*) = ?)
      ORDER BY seq_name, seq_end
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (iteration, end_count))

# ####################### contig_blast_hits table #############################

def create_contig_blast_hits_table(cxn):
    """Reset the database. Delete the tables and recreate them."""
    cxn.execute("""DROP TABLE IF EXISTS aux.contig_blast_hits""")
    sql = """
        CREATE TABLE aux.contig_blast_hits
                     (iteration INTEGER, contig_id TEXT, description TEXT,
                      bit_score NUMERIC, len INTEGER,
                      query_from INTEGER, query_to INTEGER, query_strand TEXT,
                      hit_from INTEGER, hit_to INTEGER, hit_strand TEXT)
        """
    cxn.execute(sql)

    sql = """
        CREATE INDEX aux.contig_blast_hits_index
                  ON contig_blast_hits (iteration, bit_score, len)
        """
    cxn.execute(sql)

def insert_contig_hit_batch(cxn, batch):
    """Insert a batch of blast hit records into the database."""
    if batch:
        sql = """
            INSERT INTO aux.contig_blast_hits
                        (iteration, contig_id, description, bit_score, len,
                         query_from, query_to, query_strand,
                         hit_from, hit_to, hit_strand)
                 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def get_contig_blast_hits(cxn, iteration):
    """Get all blast hits for the iteration."""
    sql = """
        SELECT iteration, contig_id, description, bit_score, len,
               query_from, query_to, query_strand,
               hit_from, hit_to, hit_strand
          FROM aux.contig_blast_hits
         WHERE iteration = ?
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (iteration, ))

# ####################### assembled_contigs table #############################

def create_assembled_contigs_table(cxn):
    """Reset the database. Delete the tables and recreate them."""
    cxn.execute("""DROP TABLE IF EXISTS aux.assembled_contigs""")
    sql = """
        CREATE TABLE aux.assembled_contigs
                     (iteration INTEGER, contig_id TEXT, seq TEXT,
                      description TEXT, bit_score NUMERIC, len INTEGER,
                      query_from INTEGER, query_to INTEGER, query_strand TEXT,
                      hit_from INTEGER, hit_to INTEGER, hit_strand TEXT)
        """
    cxn.execute(sql)

    sql = """
        CREATE INDEX aux.assembled_contigs_index
                  ON assembled_contigs (iteration, contig_id)
        """
    cxn.execute(sql)

def assembled_contigs_count(cxn, iteration, bit_score, length):
    """Count the blast hist for the iteration."""
    sql = """
        SELECT COUNT(*) AS count
          FROM aux.assembled_contigs
         WHERE iteration = ?
           AND bit_score >= ?
           AND len >= ?
        """

    result = cxn.execute(sql, (iteration, bit_score, length))
    return result.fetchone()[0]

def iteration_overlap_count(cxn, iteration, bit_score, length):
    """Count how many assembled contigs match what's in the last iteration."""
    sql = """
        SELECT COUNT(*) AS overlap
          FROM aux.assembled_contigs AS curr_iter
          JOIN aux.assembled_contigs AS prev_iter
            ON (     curr_iter.contig_id = prev_iter.contig_id
                 AND curr_iter.iteration = prev_iter.iteration + 1)
         WHERE curr_iter.iteration = ?
           AND curr_iter.seq = prev_iter.seq
           AND curr_iter.bit_score >= ?
           AND prev_iter.bit_score >= ?
           AND curr_iter.len >= ?
        """
    result = cxn.execute(
        sql, (iteration, bit_score, bit_score, length))
    return result.fetchone()[0]

def insert_assembled_contigs_batch(cxn, batch):
    """Insert a batch of blast hit records into the database."""
    if batch:
        sql = """
            INSERT INTO aux.assembled_contigs
                        (iteration, contig_id, seq, description,
                         bit_score, len,
                         query_from, query_to, query_strand,
                         hit_from, hit_to, hit_strand)
                 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """
        cxn.executemany(sql, batch)
        cxn.commit()

def get_assembled_contigs(cxn, iteration, bit_score, length):
    """
    Get all assembled contigs for the iteration.

    We will use them as the queries in the next atram iteration.
    """
    sql = """
        SELECT contig_id, seq
          FROM aux.assembled_contigs
         WHERE iteration = ?
           AND bit_score >= ?
           AND len >= ?
         """
    return cxn.execute(sql, (iteration, bit_score, length))

def get_all_assembled_contigs(cxn, bit_score=0, length=0):
    """Get all assembled contigs."""
    sql = """
        SELECT iteration, contig_id, seq, description, bit_score, len,
               query_from, query_to, query_strand,
               hit_from, hit_to, hit_strand
          FROM aux.assembled_contigs
         WHERE bit_score >= ?
           AND len >= ?
      ORDER BY bit_score DESC, iteration
        """

    cxn.row_factory = sqlite3.Row
    return cxn.execute(sql, (bit_score, length))

def all_assembled_contigs_count(cxn, bit_score=0, length=0):
    """Count all assembed contigs."""
    sql = """
        SELECT COUNT(*) AS count
          FROM aux.assembled_contigs
         WHERE bit_score >= ?
           AND len >= ?
        """

    result = cxn.execute(sql, (bit_score, length))
    return result.fetchone()[0]
<|endoftext|>