<commit_before>from bigbang.parse import get_date
import urllib2
import urllib
import gzip
import re
import os
import mailbox
import parse
import pandas as pd
from pprint import pprint as pp
import w3crawl
import warnings

ml_exp = re.compile('/([\w-]*)/$')

gz_exp = re.compile('href="(\d\d\d\d-\w*\.txt\.gz)"')
ietf_ml_exp = re.compile('href="([\d-]+.mail)"')
w3c_archives_exp = re.compile('lists\.w3\.org')

mailing_list_path_expressions = [gz_exp, ietf_ml_exp]

class InvalidURLException(Exception):

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

class MissingDataException(Exception):

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

def load_data(name,archive_dir="archives",mbox=False):
    """
    Loads the data associated with an archive name, given
    as a string.

    Attempts to open {archives-directory}/NAME.csv as data.

    Failing that, if the the name is a URL, it will try to derive
    the list name from that URL and load the .csv again.

    Failing that, it will collect the data from the web and create the CSV archive.
    """

    if mbox:
        return open_list_archives(name, archive_dir=archive_dir, mbox=True)

    # a first pass at detecting if the string is a URL...
    if not (name.startswith("http://") or name.startswith("https://")):
        path = os.path.join(archive_dir,name + ".csv")

        if os.path.exists(path):
            data = pd.read_csv(path)
            return data
        else:
            print "No data available at %s" % (path)
    else:
        path = os.path.join(archive_dir,get_list_name(name) + ".csv")

        if os.path.exists(path):
            data = pd.read_csv(path)
            return data
        else:
            print "No data found at %s. Attempting to collect data from URL." % (name)
            print "This could take a while."
            return collect_from_url(name,archive_dir=archive_dir)

def collect_from_url(url,archive_dir="archives"):
    url = url.rstrip()
    try:
        has_archives = collect_archive_from_url(url)
    except urllib2.HTTPError as e:
        print "HTTP 404 Error: %s" % (url)
        return None

    if has_archives:
        unzip_archive(url)
        data = open_list_archives(url)

        # hard coding the archives directory in too many places
        # need to push this default to a configuration file
        path = os.path.join(archive_dir, get_list_name(url) + ".csv")

        try:
            data.to_csv(path, ",", encoding="utf-8")
        except:
            # if encoding doesn't work...don't encode?
            data.to_csv(path, ",")

        return data
    else:
        return None

def collect_from_file(urls_file):
    for url in open(urls_file):
        collect_from_url(url)

# gets the 'list name' from a canonical mailman archive url
# does nothing if it's not this kind of url
# it would be better to catch these non-url cases earlier
def get_list_name(url):
    url = url.rstrip()

    if ml_exp.search(url) is not None:
        return ml_exp.search(url).groups()[0]
    else:
        warnings.warn("No mailing list name found at %s" % url)
        return url

def archive_directory(base_dir, list_name):
    arc_dir = os.path.join(base_dir, list_name)
    if not os.path.exists(arc_dir):
        os.makedirs(arc_dir)
    return arc_dir

def collect_archive_from_url(url, archive_dir="archives"):
    """
    Collects archives (generally tar.gz) files from mailmain
    archive page.

    Returns True if archives were downloaded, False otherwise
    (for example if the page lists no accessible archive files).
    """
    list_name = get_list_name(url)
    pp("Getting archive page for %s" % list_name)

    if w3c_archives_exp.search(url):
        return w3crawl.collect_from_url(url, archive_dir)

    response = urllib2.urlopen(url)
    html = response.read()

    results = []
    for exp in mailing_list_path_expressions:
        results.extend(exp.findall(html))

    pp(results)

    # directory for downloaded files
    arc_dir = archive_directory(archive_dir, list_name)

    # download monthly archives
    for res in results:
        result_path = os.path.join(arc_dir, res)
        # this check is redundant with urlretrieve
        if not os.path.isfile(result_path):
            gz_url = "/".join([url.strip("/"),res])
            pp('retrieving %s' % gz_url)
            resp = urllib2.urlopen(gz_url)
            if resp.getcode() == 200:
                print("200 - writing file to %s" % (result_path))
                output = open(result_path, 'wb')
                output.write(resp.read())
                output.close()
            else:
                print("%s error code trying to retrieve %s" %
                      (str(resp.getcode(), gz_url)))

    # return True if any archives collected, false otherwise
    return len(results) > 0

def unzip_archive(url, archive_dir="archives"):
    arc_dir = archive_directory(archive_dir, get_list_name(url))

    gzs = [os.path.join(arc_dir, fn) for fn
           in os.listdir(arc_dir)
           if fn.endswith('.txt.gz')]

    print 'unzipping %d archive files' % (len(gzs))

    for gz in gzs:
        try:
            f = gzip.open(gz, 'rb')
            content = f.read()
            f.close()

            txt_fn = str(gz[:-3])

            f2 = open(txt_fn, 'w')
            f2.write(content)
            f2.close()
        except Exception as e:
            print e

# This works for the names of the files. Order them.
# datetime.datetime.strptime('2000-November',"%Y-%B")

# This doesn't yet work for parsing the dates. Because of %z Bullshit
# datetime.datetime.strptime(arch[0][0].get('Date'),"%a, %d %b %Y %H:%M:%S %z")

# x is a String, a Message, or a list of Messages
# The payload of a Message may be a String, a Message, or a list of Messages.
# OR maybe it's never just a Message, but always a list of them.
def recursive_get_payload(x):
    if isinstance(x,str):
        return x
    elif isinstance(x,list):
        #return [recursive_get_payload(y) for y in x]
        return recursive_get_payload(x[0])
    elif isinstance(x,email.message.Message):
        return recursive_get_payload(x.get_payload())
    else:
        print x
        return None

def open_list_archives(url, archive_dir="archives", mbox=False):
    """
    Returns a list of all email messages contained in the specified directory.

    The argument *url* here is taken to be the name of a subdirectory
    of the directory specified in argument *archive_dir*.

    This directory is expected to contain files with extensions .txt,
    .mail, or .mbox. These files are all expected to be in mbox format--
    i.e. a series of blocks of text starting with headers (colon-separated
    key-value pairs) followed by an email body.
    """

    messages = None

    if mbox and (os.path.isfile(os.path.join(archive_dir, url))):
        # treat string as the path to a file that is an mbox
        box = mailbox.mbox(os.path.join(archive_dir, url), create=False)
        messages = box.values()
    else:
        # assume string is the path to a directory with many

        list_name = get_list_name(url)
        arc_dir = archive_directory(archive_dir, list_name)

        file_extensions = [".txt", ".mail", ".mbox"]

        txts = [os.path.join(arc_dir, fn) for fn
                in os.listdir(arc_dir)
                if any([fn.endswith(extension) for extension in file_extensions])]

        print 'Opening %d archive files' % (len(txts))
        arch = [mailbox.mbox(txt, create=False).values() for txt in txts]

        messages = [item for sublist in arch for item in sublist]

        if len(messages) == 0:
            raise MissingDataException(
                ("No messages in %s under %s. Did you run the "
                 "collect_mail.py script?") %
                (archive_dir, list_name))

    return messages_to_dataframe(messages)

def get_text(msg):
    import chardet
    text = u""
    if msg.is_multipart():
        html = None
        for part in msg.walk():
            if part.get_content_charset() is None:
                charset = chardet.detect(str(part))['encoding']
            else:
                charset = part.get_content_charset()
            if part.get_content_type() == 'text/plain':
                try:
                    text = unicode(part.get_payload(decode=True), str(charset), "ignore")
                except LookupError as e:
                    print "%s unknown encoding in message %s, using UTF-8 instead" % (charset,msg['Message-ID'])
                    charset = "utf-8"
                    text = unicode(part.get_payload(decode=True), str(charset), "ignore")

            if part.get_content_type() == 'text/html':
                html = unicode(part.get_payload(decode=True), str(charset), "ignore")
        if text is not None:
            return text.strip()
        else:
            import html2text
            h = html2text.HTML2Text()
            h.encoding = 'utf-8'
            return unicode(h.handle(html))
    else:
        charset = msg.get_content_charset() or 'utf-8'
        text = unicode(msg.get_payload(), encoding=charset, errors='ignore')
        return text.strip()

def messages_to_dataframe(messages):
    """
    Turn a list of parsed messages into a dataframe of message data,
    indexed by message-id, with column-names from headers.

    """
    def safe_unicode(t):
        return t and unicode(t, 'utf-8', 'ignore')
    # extract data into a list of tuples -- records -- with
    # the Message-ID separated out as an index
    pm = [(m.get('Message-ID'),
           safe_unicode(m.get('From')),
           safe_unicode(m.get('Subject')),
           get_date(m),
           safe_unicode(m.get('In-Reply-To')),
           safe_unicode(m.get('References')),
           get_text(m))
          for m in messages if m.get('Message-ID')]

    mdf = pd.DataFrame.from_records(list(pm),
                                    index='Message-ID',
                                    columns=['Message-ID', 'From',
                                             'Subject',
                                             'Date',
                                             'In-Reply-To',
                                             'References',
                                             'Body'])
    mdf.index.name = 'Message-ID'

    return mdf
<commit_msg>fail gracefully if mail archive data cannot be exported to csv<commit_after>from bigbang.parse import get_date
import urllib2
import urllib
import gzip
import re
import os
import mailbox
import parse
import pandas as pd
from pprint import pprint as pp
import w3crawl
import warnings

ml_exp = re.compile('/([\w-]*)/$')

gz_exp = re.compile('href="(\d\d\d\d-\w*\.txt\.gz)"')
ietf_ml_exp = re.compile('href="([\d-]+.mail)"')
w3c_archives_exp = re.compile('lists\.w3\.org')

mailing_list_path_expressions = [gz_exp, ietf_ml_exp]

class InvalidURLException(Exception):

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

class MissingDataException(Exception):

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

def load_data(name,archive_dir="archives",mbox=False):
    """
    Loads the data associated with an archive name, given
    as a string.

    Attempts to open {archives-directory}/NAME.csv as data.

    Failing that, if the the name is a URL, it will try to derive
    the list name from that URL and load the .csv again.

    Failing that, it will collect the data from the web and create the CSV archive.
    """

    if mbox:
        return open_list_archives(name, archive_dir=archive_dir, mbox=True)

    # a first pass at detecting if the string is a URL...
    if not (name.startswith("http://") or name.startswith("https://")):
        path = os.path.join(archive_dir,name + ".csv")

        if os.path.exists(path):
            data = pd.read_csv(path)
            return data
        else:
            print "No data available at %s" % (path)
    else:
        path = os.path.join(archive_dir,get_list_name(name) + ".csv")

        if os.path.exists(path):
            data = pd.read_csv(path)
            return data
        else:
            print "No data found at %s. Attempting to collect data from URL." % (name)
            print "This could take a while."
            return collect_from_url(name,archive_dir=archive_dir)

def collect_from_url(url,archive_dir="archives"):
    url = url.rstrip()
    try:
        has_archives = collect_archive_from_url(url)
    except urllib2.HTTPError as e:
        print "HTTP 404 Error: %s" % (url)
        return None

    if has_archives:
        unzip_archive(url)
        data = open_list_archives(url)

        # hard coding the archives directory in too many places
        # need to push this default to a configuration file
        path = os.path.join(archive_dir, get_list_name(url) + ".csv")

        try:
            data.to_csv(path, ",", encoding="utf-8")
        except Exception as e:
            print e
            # if encoding doesn't work...don't encode?
            try:
                data.to_csv(path, ",")
            except Exception as e:
                print e
                print "Can't export data. Aborting."
                return None

        return data
    else:
        return None

def collect_from_file(urls_file):
    for url in open(urls_file):
        collect_from_url(url)

# gets the 'list name' from a canonical mailman archive url
# does nothing if it's not this kind of url
# it would be better to catch these non-url cases earlier
def get_list_name(url):
    url = url.rstrip()

    if ml_exp.search(url) is not None:
        return ml_exp.search(url).groups()[0]
    else:
        warnings.warn("No mailing list name found at %s" % url)
        return url

def archive_directory(base_dir, list_name):
    arc_dir = os.path.join(base_dir, list_name)
    if not os.path.exists(arc_dir):
        os.makedirs(arc_dir)
    return arc_dir

def collect_archive_from_url(url, archive_dir="archives"):
    """
    Collects archives (generally tar.gz) files from mailmain
    archive page.

    Returns True if archives were downloaded, False otherwise
    (for example if the page lists no accessible archive files).
    """
    list_name = get_list_name(url)
    pp("Getting archive page for %s" % list_name)

    if w3c_archives_exp.search(url):
        return w3crawl.collect_from_url(url, archive_dir)

    response = urllib2.urlopen(url)
    html = response.read()

    results = []
    for exp in mailing_list_path_expressions:
        results.extend(exp.findall(html))

    pp(results)

    # directory for downloaded files
    arc_dir = archive_directory(archive_dir, list_name)

    # download monthly archives
    for res in results:
        result_path = os.path.join(arc_dir, res)
        # this check is redundant with urlretrieve
        if not os.path.isfile(result_path):
            gz_url = "/".join([url.strip("/"),res])
            pp('retrieving %s' % gz_url)
            resp = urllib2.urlopen(gz_url)
            if resp.getcode() == 200:
                print("200 - writing file to %s" % (result_path))
                output = open(result_path, 'wb')
                output.write(resp.read())
                output.close()
            else:
                print("%s error code trying to retrieve %s" %
                      (str(resp.getcode(), gz_url)))

    # return True if any archives collected, false otherwise
    return len(results) > 0

def unzip_archive(url, archive_dir="archives"):
    arc_dir = archive_directory(archive_dir, get_list_name(url))

    gzs = [os.path.join(arc_dir, fn) for fn
           in os.listdir(arc_dir)
           if fn.endswith('.txt.gz')]

    print 'unzipping %d archive files' % (len(gzs))

    for gz in gzs:
        try:
            f = gzip.open(gz, 'rb')
            content = f.read()
            f.close()

            txt_fn = str(gz[:-3])

            f2 = open(txt_fn, 'w')
            f2.write(content)
            f2.close()
        except Exception as e:
            print e

# This works for the names of the files. Order them.
# datetime.datetime.strptime('2000-November',"%Y-%B")

# This doesn't yet work for parsing the dates. Because of %z Bullshit
# datetime.datetime.strptime(arch[0][0].get('Date'),"%a, %d %b %Y %H:%M:%S %z")

# x is a String, a Message, or a list of Messages
# The payload of a Message may be a String, a Message, or a list of Messages.
# OR maybe it's never just a Message, but always a list of them.
def recursive_get_payload(x):
    if isinstance(x,str):
        return x
    elif isinstance(x,list):
        #return [recursive_get_payload(y) for y in x]
        return recursive_get_payload(x[0])
    elif isinstance(x,email.message.Message):
        return recursive_get_payload(x.get_payload())
    else:
        print x
        return None

def open_list_archives(url, archive_dir="archives", mbox=False):
    """
    Returns a list of all email messages contained in the specified directory.

    The argument *url* here is taken to be the name of a subdirectory
    of the directory specified in argument *archive_dir*.

    This directory is expected to contain files with extensions .txt,
    .mail, or .mbox. These files are all expected to be in mbox format--
    i.e. a series of blocks of text starting with headers (colon-separated
    key-value pairs) followed by an email body.
    """

    messages = None

    if mbox and (os.path.isfile(os.path.join(archive_dir, url))):
        # treat string as the path to a file that is an mbox
        box = mailbox.mbox(os.path.join(archive_dir, url), create=False)
        messages = box.values()
    else:
        # assume string is the path to a directory with many

        list_name = get_list_name(url)
        arc_dir = archive_directory(archive_dir, list_name)

        file_extensions = [".txt", ".mail", ".mbox"]

        txts = [os.path.join(arc_dir, fn) for fn
                in os.listdir(arc_dir)
                if any([fn.endswith(extension) for extension in file_extensions])]

        print 'Opening %d archive files' % (len(txts))
        arch = [mailbox.mbox(txt, create=False).values() for txt in txts]

        messages = [item for sublist in arch for item in sublist]

        if len(messages) == 0:
            raise MissingDataException(
                ("No messages in %s under %s. Did you run the "
                 "collect_mail.py script?") %
                (archive_dir, list_name))

    return messages_to_dataframe(messages)

def get_text(msg):
    import chardet
    text = u""
    if msg.is_multipart():
        html = None
        for part in msg.walk():
            if part.get_content_charset() is None:
                charset = chardet.detect(str(part))['encoding']
            else:
                charset = part.get_content_charset()
            if part.get_content_type() == 'text/plain':
                try:
                    text = unicode(part.get_payload(decode=True), str(charset), "ignore")
                except LookupError as e:
                    print "%s unknown encoding in message %s, using UTF-8 instead" % (charset,msg['Message-ID'])
                    charset = "utf-8"
                    text = unicode(part.get_payload(decode=True), str(charset), "ignore")

            if part.get_content_type() == 'text/html':
                html = unicode(part.get_payload(decode=True), str(charset), "ignore")
        if text is not None:
            return text.strip()
        else:
            import html2text
            h = html2text.HTML2Text()
            h.encoding = 'utf-8'
            return unicode(h.handle(html))
    else:
        charset = msg.get_content_charset() or 'utf-8'
        text = unicode(msg.get_payload(), encoding=charset, errors='ignore')
        return text.strip()

def messages_to_dataframe(messages):
    """
    Turn a list of parsed messages into a dataframe of message data,
    indexed by message-id, with column-names from headers.

    """
    def safe_unicode(t):
        return t and unicode(t, 'utf-8', 'ignore')
    # extract data into a list of tuples -- records -- with
    # the Message-ID separated out as an index
    pm = [(m.get('Message-ID'),
           safe_unicode(m.get('From')),
           safe_unicode(m.get('Subject')),
           get_date(m),
           safe_unicode(m.get('In-Reply-To')),
           safe_unicode(m.get('References')),
           get_text(m))
          for m in messages if m.get('Message-ID')]

    mdf = pd.DataFrame.from_records(list(pm),
                                    index='Message-ID',
                                    columns=['Message-ID', 'From',
                                             'Subject',
                                             'Date',
                                             'In-Reply-To',
                                             'References',
                                             'Body'])
    mdf.index.name = 'Message-ID'

    return mdf
<|endoftext|>


<commit_before>#!/usr/bin/env python

import cPickle
import sys
from matplotlib import pyplot as plt
from matplotlib import rcParams
import matplotlib.gridspec as gridspec
import numpy as np
import math
import argparse
import copy
import random
import wifi_im

def plot(data, time, args):
  xyr = [point for point in data["xyr"] if point[0] < time]

  states = [state for state in data["state_log"] if state[0] < time]
  fit_model = [state for state in states if state[1] == "fit_model"]

  if args.model:
    xyr_s = []
    if len(fit_model) > 0:
      xyr_s = copy.deepcopy([point for point in xyr if point[0] <= fit_model[-1][0]])
      random.shuffle(xyr_s)
    xx = np.linspace(-10, 0, 20)
    yy = np.linspace(-15, 20, 70)
    rssi_plot = np.zeros((len(xx), len(yy)))
    if len(xyr_s) > 0:
      model = wifi_im.ScaledModel(wifi_im.FNN([10, 5]))
      model.fit([point[1:3] for point in xyr_s], [point[3] for point in xyr_s])
      for i in range(len(xx)):
        for j in range(len(yy)):
          rssi_plot[i, j] = model.predict([[xx[i], yy[j]]])

  gs = gridspec.GridSpec(2, 3)
  plt.subplot(gs[:, 0])
  plt.title("hexbin")
  plt.hexbin([x for t, x, y, r in xyr], [y for t, x, y, r in xyr], [r for t, x, y, r in xyr], gridsize=20, cmap=plt.get_cmap("gnuplot2"), vmin=-80, vmax=-20, extent=(-10, 0, -15, 20))
  plt.plot([x for t, x, y, r in xyr], [y for t, x, y, r in xyr], "g-")
  plt.gca().set_xlim((-10, 0))
  plt.gca().set_xlabel("x [m]")
  plt.gca().set_ylim((-15, 20))
  plt.gca().set_ylabel("y [m]")
  cbar = plt.colorbar()
  cbar.set_label("mean rssi [dB]")
  if args.model:
    plt.subplot(gs[:, 1])
    plt.title("model")
    plt.pcolormesh(xx, yy, rssi_plot.T, cmap=plt.get_cmap("gnuplot2"), vmin=-80, vmax=-20)
    plt.gca().set_xlim((-10, 0))
    plt.gca().set_xlabel("x [m]")
    plt.gca().set_ylim((-15, 20))
    plt.gca().set_ylabel("y [m]")
    cbar = plt.colorbar()
    cbar.set_label("mean rssi [dB]")
  if args.model:
    plt.subplot(gs[0, 2:])
  else:
    plt.subplot(gs[0, 1:])
  plt.title("log by time")
  plt.plot([rssi[0]-data["xyr"][0][0] for rssi in xyr], [rssi[3] for rssi in xyr], "b.")
  plt.gca().set_xlabel("time [s]")
  plt.gca().set_ylabel("rssi [dB]")
  plt.gca().set_ylim((-80, -20))
  if args.model:
    plt.subplot(gs[1, 2:])
  else:
    plt.subplot(gs[1, 1:])
  plt.title("log by distance")
  trajectory = [xyr[i][1:3] for i in range(len(xyr))]
  lengths = [math.hypot(trajectory[i][0] - trajectory[i-1][0], trajectory[i][1]-trajectory[i][1]) for i in range(1, len(trajectory))]
  dist = [0.0]
  for length in lengths:
    dist.append(dist[-1]+length)
  plt.plot(dist, [rssi[3] for rssi in xyr], "b.")
  plt.gca().set_xlabel("distance travelled [m]")
  plt.gca().set_ylabel("rssi [dB]")
  plt.gca().set_ylim((-80, -20))

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="analyze pickle files generated by wifi_im")
  parser.add_argument("-v", "--video", help="generate a video", action="store_true")
  parser.add_argument("-m", "--model", help="also show the model", action="store_true")
  parser.add_argument("filename", help="pickle file")

  args = parser.parse_args()

  data = cPickle.load(open(args.filename, "rb"))

  time_start = data["xyr"][0][0]
  time_end = data["xyr"][-1][0]
  times = [time_end]
  if args.video:
    times = np.arange(time_start+10.0, time_end, 10.0)

  rcParams["font.family"] = "serif"
  rcParams["xtick.labelsize"] = 6
  rcParams["ytick.labelsize"] = 6
  rcParams["axes.labelsize"] = 6
  rcParams["axes.titlesize"] = 6

  for i in range(len(times)):
    print i, "/", len(times)
    plt.clf()
    plot(data, times[i], args)
    plt.suptitle("{:.2f}".format(times[i]))
    plt.gcf().set_size_inches((12, 9))
    plt.savefig("frame"+str(i).zfill(6)+".jpg", dpi=100)
<commit_msg>added loading models from file<commit_after>#!/usr/bin/env python

import cPickle
import sys
from matplotlib import pyplot as plt
from matplotlib import rcParams
import matplotlib.gridspec as gridspec
import numpy as np
import math
import argparse
import copy
import random
import wifi_im

def plot(data, time, args):
  xyr = [point for point in data["xyr"] if point[0] < time]

  states = [state for state in data["state_log"] if state[0] < time]
  fit_model = [state for state in states if state[1] == "fit_model"]

  if args.model:
    xx = np.linspace(-10, 0, 20)
    yy = np.linspace(-15, 20, 70)
    rssi_plot = []
    try:
      rssi_plot = [model[1] for model in data["model_log"] if model[0] < time][-1]
    except:
      xyr_s = []
      if len(fit_model) > 0:
        xyr_s = copy.deepcopy([point for point in xyr if point[0] <= fit_model[-1][0]])
        random.shuffle(xyr_s)
      rssi_plot = np.zeros((len(xx), len(yy)))
      if len(xyr_s) > 0:
        model = wifi_im.ScaledModel(wifi_im.FNN([10, 5]))
        model.fit([point[1:3] for point in xyr_s], [point[3] for point in xyr_s])
        for i in range(len(xx)):
          for j in range(len(yy)):
            rssi_plot[i, j] = model.predict([[xx[i], yy[j]]])

  gs = gridspec.GridSpec(2, 3)
  plt.subplot(gs[:, 0])
  plt.title("hexbin")
  plt.hexbin([x for t, x, y, r in xyr], [y for t, x, y, r in xyr], [r for t, x, y, r in xyr], gridsize=20, cmap=plt.get_cmap("gnuplot2"), vmin=-80, vmax=-20, extent=(-10, 0, -15, 20))
  plt.plot([x for t, x, y, r in xyr], [y for t, x, y, r in xyr], color="#39FF14")
  plt.gca().set_xlim((-10, 0))
  plt.gca().set_xlabel("x [m]")
  plt.gca().set_ylim((-15, 20))
  plt.gca().set_ylabel("y [m]")
  cbar = plt.colorbar()
  cbar.set_label("mean rssi [dB]")
  if args.model:
    plt.subplot(gs[:, 1])
    plt.title("model")
    plt.pcolormesh(xx, yy, rssi_plot.T, cmap=plt.get_cmap("gnuplot2"), vmin=-80, vmax=-20)
    plt.plot([x for t, x, y, r in xyr], [y for t, x, y, r in xyr], color="#39FF14")
    plt.gca().set_xlim((-10, 0))
    plt.gca().set_xlabel("x [m]")
    plt.gca().set_ylim((-15, 20))
    plt.gca().set_ylabel("y [m]")
    cbar = plt.colorbar()
    cbar.set_label("mean rssi [dB]")
  if args.model:
    plt.subplot(gs[0, 2:])
  else:
    plt.subplot(gs[0, 1:])
  plt.title("log by time")
  plt.plot([rssi[0]-data["xyr"][0][0] for rssi in xyr], [rssi[3] for rssi in xyr], "b.")
  plt.gca().set_xlabel("time [s]")
  plt.gca().set_ylabel("rssi [dB]")
  plt.gca().set_ylim((-80, -20))
  if args.model:
    plt.subplot(gs[1, 2:])
  else:
    plt.subplot(gs[1, 1:])
  plt.title("log by distance")
  trajectory = [xyr[i][1:3] for i in range(len(xyr))]
  lengths = [math.hypot(trajectory[i][0] - trajectory[i-1][0], trajectory[i][1]-trajectory[i][1]) for i in range(1, len(trajectory))]
  dist = [0.0]
  for length in lengths:
    dist.append(dist[-1]+length)
  plt.plot(dist, [rssi[3] for rssi in xyr], "b.")
  plt.gca().set_xlabel("distance travelled [m]")
  plt.gca().set_ylabel("rssi [dB]")
  plt.gca().set_ylim((-80, -20))

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="analyze pickle files generated by wifi_im")
  parser.add_argument("-v", "--video", help="generate a video", action="store_true")
  parser.add_argument("-m", "--model", help="also show the model", action="store_true")
  parser.add_argument("filename", help="pickle file")

  args = parser.parse_args()

  data = cPickle.load(open(args.filename, "rb"))

  time_start = data["xyr"][0][0]
  time_end = data["xyr"][-1][0]
  times = [time_end]
  if args.video:
    times = np.arange(time_start+10.0, time_end, 10.0)

  rcParams["font.family"] = "serif"
  rcParams["xtick.labelsize"] = 6
  rcParams["ytick.labelsize"] = 6
  rcParams["axes.labelsize"] = 6
  rcParams["axes.titlesize"] = 6

  for i in range(len(times)):
    print i, "/", len(times)
    plt.clf()
    plot(data, times[i], args)
    plt.suptitle("{:.2f}".format(times[i]))
    plt.gcf().set_size_inches((12, 9))
    plt.savefig("frame"+str(i).zfill(6)+".jpg", dpi=100)
<|endoftext|>


<commit_before>#!/usr/bin/env python
from optparse import OptionParser
import os
import subprocess
import sys
import tempfile

import h5py
import numpy as np
import pysam
import tensorflow as tf

import basenji

'''
basenji_sed.py

Compute SNP expression difference scores for variants in a VCF file.

Note:
 -I'm having trouble verifying that I'm not double counting the scenario where
  two transcripts map to the same TSS. But I should verify that.
'''

################################################################################
# main
################################################################################
def main():
    usage = 'usage: %prog [options] <params_file> <model_file> <genes_hdf5_file> <vcf_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='batch_size', default=None, type='int', help='Batch size [Default: %default]')
    parser.add_option('-c', dest='csv', default=False, action='store_true', help='Print table as CSV [Default: %default]')
    parser.add_option('-i', dest='index_snp', default=False, action='store_true', help='SNPs are labeled with their index SNP as column 6 [Default: %default]')
    parser.add_option('-o', dest='out_dir', default='sed', help='Output directory for tables and plots [Default: %default]')
    parser.add_option('-s', dest='score', default=False, action='store_true', help='SNPs are labeled with scores as column 7 [Default: %default]')
    parser.add_option('-t', dest='targets_file', default=None, help='File specifying target indexes and labels in table format')
    (options,args) = parser.parse_args()

    if len(args) != 4:
        parser.error('Must provide parameters and model files, genes HDF5 file, and QTL VCF file')
    else:
        params_file = args[0]
        model_file = args[1]
        genes_hdf5_file = args[2]
        vcf_file = args[3]

    if not os.path.isdir(options.out_dir):
        os.mkdir(options.out_dir)

    #################################################################
    # reads in genes HDF5

    genes_hdf5_in = h5py.File(genes_hdf5_file)

    seq_chrom = [chrom.decode('UTF-8') for chrom in genes_hdf5_in['seq_chrom']]
    seq_start = list(genes_hdf5_in['seq_start'])
    seq_end = list(genes_hdf5_in['seq_end'])
    seq_coords = list(zip(seq_chrom,seq_start,seq_end))

    seqs_1hot = genes_hdf5_in['seqs_1hot']

    transcripts = [tx.decode('UTF-8') for tx in genes_hdf5_in['transcripts']]
    transcript_index = list(genes_hdf5_in['transcript_index'])
    transcript_pos = list(genes_hdf5_in['transcript_pos'])

    transcript_map = {}
    for ti in range(len(transcripts)):
        transcript_map[transcripts[ti]] = (transcript_index[ti], transcript_pos[ti])

    #################################################################
    # prep SNPs

    # load SNPs
    snps = basenji.vcf.vcf_snps(vcf_file, options.index_snp, options.score, False)

    # intersect w/ segments
    snps_segs = intersect_snps(vcf_file, seq_coords)

    #################################################################
    # construct SNP sequences

    snp_seqs_1hot = []

    for snp_i in range(len(snps)):
        for seg_i in snps_segs[snp_i]:
            seq_chrom, seq_start, seq_end = seq_coords[seg_i]

            # determine the SNP's position in the segment
            snp_seq_pos = snps[snp_i].pos - seq_start

            # write reference allele
            snp_seqs_1hot.append(seqs_1hot[snp_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].ref_allele)

            # write alternative allele
            snp_seqs_1hot.append(seqs_1hot[snp_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].alt_alleles[0])

    snp_seqs_1hot = np.array(snp_seqs_1hot)

    #################################################################
    # setup model

    job = basenji.dna_io.read_job_params(params_file)

    job['batch_length'] = snp_seqs_1hot.shape[1]
    job['seq_depth'] = snp_seqs_1hot.shape[2]

    if 'num_targets' not in job:
        print("Must specify number of targets (num_targets) in the parameters file. I know, it's annoying. Sorry.", file=sys.stderr)
        exit(1)

    # build model
    dr = basenji.rnn.RNN()
    dr.build(job)

    if options.batch_size is not None:
        dr.batch_size = options.batch_size

    #################################################################
    # predict

    # initialize batcher
    batcher = basenji.batcher.Batcher(snp_seqs_1hot, batch_size=dr.batch_size)

    # initialie saver
    saver = tf.train.Saver()

    with tf.Session() as sess:
        # load variables into session
        saver.restore(sess, model_file)

        # predict
        seq_preds = dr.predict(sess, batcher)

    #################################################################
    # collect and pred SEDs

    if options.targets_file is None:
        target_labels = ['t%d' % ti for ti in range(seq_preds.shape[1])]
    else:
        target_labels = [line.split()[0] for line in open(options.targets_file)]

    header_cols = ('rsid', 'index', 'score', 'ref', 'alt', 'gene', 'target', 'ref_pred', 'alt pred', 'sed')
    if options.csv:
        sed_out = open('%s/sed_table.csv' % options.out_dir, 'w')
        print(','.join(header_cols), file=sed_out)
    else:
        sed_out = open('%s/sed_table.txt' % options.out_dir, 'w')
        print(' '.join(header_cols), file=sed_out)

    pi = 0
    for snp_i in range(len(snps)):
        snp = snps[snp_i]

        # set index SNP
        snp_is = '%-13s' % '.'
        if options.index_snp:
            snp_is = '%-13s' % snp.index_snp

        # set score
        snp_score = '%5s' % '.'
        if options.score:
            snp_score = '%5.3f' % snp.score

        for seg_i in snps_segs[snp_i]:
            # get reference prediction (LxT)
            ref_preds = seq_preds[pi]
            pi += 1

            # get alternate prediction (LxT)
            alt_preds = seq_preds[pi]
            pi += 1

            # find genes
            for transcript in transcript_map:
                tx_index, tx_pos = transcript_map[transcript]
                if tx_index == seg_i:
                    for ti in range(seq_preds.shape[2]):
                        snp_gene_sed = (alt_preds[tx_pos,ti] - ref_preds[tx_pos,ti])

                        cols = (snp.rsid, snp_is, snp_score, basenji.vcf.cap_allele(snp.ref_allele), basenji.vcf.cap_allele(snp.alt_alleles[0]), transcript, target_labels[ti], ref_preds[tx_pos,ti], alt_preds[tx_pos,ti], snp_gene_sed)
                        if options.csv:
                            print(','.join([str(c) for c in cols]), file=sed_out)
                        else:
                            print('%-13s %s %5s %6s %6s %12s %12s %6.4f %6.4f %7.4f' % cols, file=sed_out)

    sed_out.close()

    #################################################################
    # clean up

    genes_hdf5_in.close()

def intersect_snps(vcf_file, seq_coords):
    # print segments to BED
    # hash segments to indexes
    seg_temp = tempfile.NamedTemporaryFile()
    seg_bed_file = seg_temp.name
    seg_bed_out = open(seg_bed_file, 'w')
    segment_indexes = {}
    for si in range(len(seq_coords)):
        segment_indexes[seq_coords[si]] = si
        print('%s\t%d\t%d' % seq_coords[si], file=seg_bed_out)
    seg_bed_out.close()

    # hash SNPs to indexes
    snp_indexes = {}
    si = 0

    vcf_in = open(vcf_file)
    line = vcf_in.readline()
    while line[0] == '#':
        line = vcf_in.readline()
    while line:
        a = line.split()
        snp_id = a[2]
        if snp_id in snp_indexes:
            print('Duplicate SNP id %s will break the script' % snp_id, file=sys.stderr)
            exit(1)
        snp_indexes[snp_id] = si
        si += 1
        line = vcf_in.readline()
    vcf_in.close()

    # initialize list of lists
    snp_segs = []
    for i in range(len(snp_indexes)):
        snp_segs.append([])

    # intersect
    p = subprocess.Popen('bedtools intersect -wo -a %s -b %s' % (vcf_file, seg_bed_file), shell=True, stdout=subprocess.PIPE)
    for line in p.stdout:
        line = line.decode('UTF-8')
        a = line.split()
        snp_id = a[2]
        seg_chrom = a[-4]
        seg_start = int(a[-3])
        seg_end = int(a[-2])
        seg_key = (seg_chrom,seg_start,seg_end)

        snp_segs[snp_indexes[snp_id]].append(segment_indexes[seg_key])

    p.communicate()

    return snp_segs

################################################################################
# __main__
################################################################################
if __name__ == '__main__':
    main()
<commit_msg>substantial enhancements- streaming predictions, bigwigs<commit_after>#!/usr/bin/env python
from optparse import OptionParser
import os
import subprocess
import sys
import tempfile

import h5py
import numpy as np
import pyBigWig
import tensorflow as tf

import basenji

from basenji_test import bigwig_open

'''
basenji_sed.py

Compute SNP expression difference scores for variants in a VCF file.

Note:
 -I'm having trouble verifying that I'm not double counting the scenario where
  two transcripts map to the same TSS. But I should verify that.
'''

################################################################################
# main
################################################################################
def main():
    usage = 'usage: %prog [options] <params_file> <model_file> <genes_hdf5_file> <vcf_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='batch_size', default=None, type='int', help='Batch size [Default: %default]')
    parser.add_option('-c', dest='csv', default=False, action='store_true', help='Print table as CSV [Default: %default]')
    parser.add_option('-g', dest='genome_file', default='%s/assembly/human.hg19.genome'%os.environ['HG19'], help='Chromosome lengths file [Default: %default]')
    parser.add_option('-i', dest='index_snp', default=False, action='store_true', help='SNPs are labeled with their index SNP as column 6 [Default: %default]')
    parser.add_option('-o', dest='out_dir', default='sed', help='Output directory for tables and plots [Default: %default]')
    parser.add_option('-s', dest='score', default=False, action='store_true', help='SNPs are labeled with scores as column 7 [Default: %default]')
    parser.add_option('-t', dest='target_wigs_file', default=None, help='Store target values, extracted from this list of WIG files')
    parser.add_option('--ti', dest='track_indexes', help='Comma-separated list of target indexes to output BigWig tracks')
    (options,args) = parser.parse_args()

    if len(args) != 4:
        parser.error('Must provide parameters and model files, genes HDF5 file, and QTL VCF file')
    else:
        params_file = args[0]
        model_file = args[1]
        genes_hdf5_file = args[2]
        vcf_file = args[3]

    if not os.path.isdir(options.out_dir):
        os.mkdir(options.out_dir)

    if options.track_indexes is None:
        options.track_indexes = []
    else:
        options.track_indexes = [int(ti) for ti in options.track_indexes.split(',')]
        if not os.path.isdir('%s/tracks' % options.out_dir):
            os.mkdir('%s/tracks' % options.out_dir)

    #################################################################
    # reads in genes HDF5

    genes_hdf5_in = h5py.File(genes_hdf5_file)

    seq_chrom = [chrom.decode('UTF-8') for chrom in genes_hdf5_in['seq_chrom']]
    seq_start = list(genes_hdf5_in['seq_start'])
    seq_end = list(genes_hdf5_in['seq_end'])
    seq_coords = list(zip(seq_chrom,seq_start,seq_end))

    seqs_1hot = genes_hdf5_in['seqs_1hot']
    print(seqs_1hot.shape)

    transcripts = [tx.decode('UTF-8') for tx in genes_hdf5_in['transcripts']]
    transcript_index = list(genes_hdf5_in['transcript_index'])
    transcript_pos = list(genes_hdf5_in['transcript_pos'])

    transcript_map = {}
    for ti in range(len(transcripts)):
        transcript_map[transcripts[ti]] = (transcript_index[ti], transcript_pos[ti])

    if 'transcript_targets' in genes_hdf5_in:
        transcript_targets = genes_hdf5_in['transcript_targets']
        target_labels = [tl.decode('UTF-8') for tl in genes_hdf5_in['target_labels']]
    else:
        transcript_targets = None
        target_labels = None

    #################################################################
    # prep SNPs

    # load SNPs
    snps = basenji.vcf.vcf_snps(vcf_file, options.index_snp, options.score, False)

    # intersect w/ segments
    snps_segs = intersect_snps(vcf_file, seq_coords)

    #################################################################
    # construct SNP sequences

    snp_seqs_1hot = []

    for snp_i in range(len(snps)):
        for seg_i in snps_segs[snp_i]:
            seq_chrom, seq_start, seq_end = seq_coords[seg_i]

            # determine the SNP's position in the segment
            snp_seq_pos = snps[snp_i].pos-1 - seq_start

            # write reference allele
            snp_seqs_1hot.append(seqs_1hot[seg_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].ref_allele)

            # write alternative allele
            snp_seqs_1hot.append(seqs_1hot[seg_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].alt_alleles[0])

    snp_seqs_1hot = np.array(snp_seqs_1hot)

    #################################################################
    # setup model

    job = basenji.dna_io.read_job_params(params_file)

    job['batch_length'] = snp_seqs_1hot.shape[1]
    job['seq_depth'] = snp_seqs_1hot.shape[2]
    job['target_pool'] = int(np.array(genes_hdf5_in['pool_width']))

    if transcript_targets is not None:
        job['num_targets'] = transcript_targets.shape[1]

    if 'num_targets' not in job:
        print("Must specify number of targets (num_targets) in the parameters file. I know, it's annoying. Sorry.", file=sys.stderr)
        exit(1)

    # build model
    model = basenji.rnn.RNN()
    model.build(job)

    if options.batch_size is not None:
        model.batch_size = options.batch_size

    # label targets
    if target_labels is None:
        if options.targets_file is None:
            target_labels = ['t%d' % ti for ti in range(job['num_targets'])]
        else:
            target_labels = [line.split()[0] for line in open(options.targets_file)]

    #################################################################
    # predict

    # initialize saver
    saver = tf.train.Saver()

    with tf.Session() as sess:
        # load variables into session
        saver.restore(sess, model_file)

        # initialize prediction stream
        seq_preds = Pred_Stream(sess, model, snp_seqs_1hot, 128)

        # determine prediction buffer
        pred_buffer = model.batch_buffer // model.target_pool

        #################################################################
        # collect and print SEDs

        header_cols = ('rsid', 'index', 'score', 'ref', 'alt', 'gene', 'target', 'ref_pred', 'alt pred', 'ser', 'sed')
        if options.csv:
            sed_out = open('%s/sed_table.csv' % options.out_dir, 'w')
            print(','.join(header_cols), file=sed_out)
        else:
            sed_out = open('%s/sed_table.txt' % options.out_dir, 'w')
            print(' '.join(header_cols), file=sed_out)

        pi = 0
        for snp_i in range(len(snps)):
            snp = snps[snp_i]

            # set index SNP
            snp_is = '%-13s' % '.'
            if options.index_snp:
                snp_is = '%-13s' % snp.index_snp

            # set score
            snp_score = '%5s' % '.'
            if options.score:
                snp_score = '%5.3f' % snp.score

            for seg_i in snps_segs[snp_i]:
                # get reference prediction (LxT)
                ref_preds = seq_preds.get(pi)
                pi += 1

                # get alternate prediction (LxT)
                alt_preds = seq_preds.get(pi)
                pi += 1

                # find genes
                for transcript in transcript_map:
                    tx_index, tx_pos = transcript_map[transcript]

                    if tx_index == seg_i:
                        # compute distance between SNP and gene
                        tx_gpos = seq_coords[1] + (tx_pos + 0.5)*model.target_pool
                        snp_dist = abs(tx_gpos - snp.pos)

                        # compute transcript pos in predictions
                        tx_pos_buf = tx_pos - pred_buffer

                        for ti in range(job['num_targets']):
                            # compute SED score
                            snp_gene_sed = (alt_preds[tx_pos_buf,ti] - ref_preds[tx_pos_buf,ti])
                            snp_gene_ser = np.divide(alt_preds[tx_pos_buf,ti]+1, ref_preds[tx_pos_buf,ti]+1)

                            # print to table
                            cols = (snp.rsid, snp_is, snp_score, basenji.vcf.cap_allele(snp.ref_allele), basenji.vcf.cap_allele(snp.alt_alleles[0]), transcript, target_labels[ti], ref_preds[tx_pos_buf,ti], alt_preds[tx_pos_buf,ti], snp_gene_ser, snp_gene_sed)
                            if options.csv:
                                print(','.join([str(c) for c in cols]), file=sed_out)
                            else:
                                print('%-13s %s %5s %6s %6s %12s %12s %6.4f %6.4f %7.4f' % cols, file=sed_out)

                # print tracks
                for ti in options.track_indexes:
                    ref_bw_file = '%s/tracks/%s_%s_t%d_ref.bw' % (options.out_dir, snp.rsid, seg_i, ti)
                    alt_bw_file = '%s/tracks/%s_%s_t%d_alt.bw' % (options.out_dir, snp.rsid, seg_i, ti)
                    ref_bw_open = bigwig_open(ref_bw_file, options.genome_file)
                    alt_bw_open = bigwig_open(alt_bw_file, options.genome_file)

                    seq_chrom, seq_start, seq_end = seq_coords[seg_i]
                    bw_chroms = [seq_chrom]*ref_preds.shape[0]
                    bw_starts = [int(seq_start + model.batch_buffer + bi*model.target_pool) for bi in range(ref_preds.shape[0])]
                    bw_ends = [int(bws + model.target_pool) for bws in bw_starts]

                    ref_values = [float(p) for p in ref_preds[:,ti]]
                    # print(ref_bw_open.chroms())
                    # print('bw_chroms', len(bw_chroms), type(bw_chroms[0]), bw_chroms[:10])
                    # print('bw_starts', len(bw_starts), type(bw_starts[0]), bw_starts[:10])
                    # print('bw_ends', len(bw_ends), type(bw_ends[0]), bw_ends[:10])
                    # print('ref_values', len(ref_values), type(ref_values[0]), ref_values[:10])
                    ref_bw_open.addEntries(bw_chroms, bw_starts, ends=bw_ends, values=ref_values)

                    alt_values = [float(p) for p in alt_preds[:,ti]]
                    alt_bw_open.addEntries(bw_chroms, bw_starts, ends=bw_ends, values=alt_values)

                    ref_bw_open.close()
                    alt_bw_open.close()

    sed_out.close()

    #################################################################
    # clean up

    genes_hdf5_in.close()

def intersect_snps(vcf_file, seq_coords):
    # print segments to BED
    # hash segments to indexes
    seg_temp = tempfile.NamedTemporaryFile()
    seg_bed_file = seg_temp.name
    seg_bed_out = open(seg_bed_file, 'w')
    segment_indexes = {}
    for si in range(len(seq_coords)):
        segment_indexes[seq_coords[si]] = si
        print('%s\t%d\t%d' % seq_coords[si], file=seg_bed_out)
    seg_bed_out.close()

    # hash SNPs to indexes
    snp_indexes = {}
    si = 0

    vcf_in = open(vcf_file)
    line = vcf_in.readline()
    while line[0] == '#':
        line = vcf_in.readline()
    while line:
        a = line.split()
        snp_id = a[2]
        if snp_id in snp_indexes:
            print('Duplicate SNP id %s will break the script' % snp_id, file=sys.stderr)
            exit(1)
        snp_indexes[snp_id] = si
        si += 1
        line = vcf_in.readline()
    vcf_in.close()

    # initialize list of lists
    snp_segs = []
    for i in range(len(snp_indexes)):
        snp_segs.append([])

    # intersect
    p = subprocess.Popen('bedtools intersect -wo -a %s -b %s' % (vcf_file, seg_bed_file), shell=True, stdout=subprocess.PIPE)
    for line in p.stdout:
        line = line.decode('UTF-8')
        a = line.split()
        snp_id = a[2]
        seg_chrom = a[-4]
        seg_start = int(a[-3])
        seg_end = int(a[-2])
        seg_key = (seg_chrom,seg_start,seg_end)

        snp_segs[snp_indexes[snp_id]].append(segment_indexes[seg_key])

    p.communicate()

    return snp_segs

class Pred_Stream:
    ''' Interface to acquire predictions via a buffered stream mechanism
         rather than getting them all at once and using excessive memory. '''

    def __init__(self, sess, model, seqs_1hot, stream_length):
        self.sess = sess
        self.model = model

        self.seqs_1hot = seqs_1hot

        self.stream_length = stream_length
        self.stream_start = 0
        self.stream_end = 0

        if self.stream_length % self.model.batch_size != 0:
            print('Make the stream length a multiple of the batch size', file=sys.stderr)
            exit(1)

    def get(self, i):
        # acquire predictions, if needed
        if i >= self.stream_end:
            self.stream_start = self.stream_end
            self.stream_end = min(self.stream_start + self.stream_length, self.seqs_1hot.shape[0])

            # subset sequences
            stream_seqs_1hot = self.seqs_1hot[self.stream_start:self.stream_end]

            # initialize batcher
            batcher = basenji.batcher.Batcher(stream_seqs_1hot, batch_size=self.model.batch_size)

            # predict
            self.stream_preds = self.model.predict(self.sess, batcher)
            # print('Acquired predictions ', self.stream_preds.shape, flush=True)

        return self.stream_preds[i - self.stream_start]

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


<commit_before>from datetime import datetime
from slugify import slugify
from sqlalchemy import Column, String, DateTime, ForeignKey, Integer, Text
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.orm import backref, relationship
from sqlalchemy.schema import UniqueConstraint
from uuid import uuid4

from changes.config import db
from changes.db.types.guid import GUID

class Project(db.Model):
    __tablename__ = 'project'

    id = Column(GUID, primary_key=True, default=uuid4)
    slug = Column(String(64), unique=True, nullable=False)
    repository_id = Column(GUID, ForeignKey('repository.id', ondelete="RESTRICT"), nullable=False)
    name = Column(String(64))
    date_created = Column(DateTime, default=datetime.utcnow)
    avg_build_time = Column(Integer)

    repository = relationship('Repository')
    plans = association_proxy('project_plans', 'plan')

    def __init__(self, **kwargs):
        super(Project, self).__init__(**kwargs)
        if not self.id:
            self.id = uuid4()
        if not self.slug:
            self.slug = slugify(self.name)

    def get_entity(self, provider):
        if not hasattr(self, '_entities') or provider not in self._entities:
            raise ValueError('Entity not attached')
        return self._entities[provider]

    def attach_entity(self, entity):
        assert entity.internal_id == self.id

        if not hasattr(self, '_entities'):
            self._entities = {}
        self._entities[entity.provider] = entity

class ProjectOption(db.Model):
    __tablename__ = 'projectoption'
    __table_args__ = (
        UniqueConstraint('project_id', 'name', name='unq_projectoption_name'),
    )

    id = Column(GUID, primary_key=True, default=uuid4)
    project_id = Column(GUID, ForeignKey('project.id', ondelete="CASCADE"), nullable=False)
    name = Column(String(64), nullable=False)
    value = Column(Text, nullable=False)
    date_created = Column(DateTime, default=datetime.utcnow, nullable=False)

    project = relationship('Project')

    def __init__(self, **kwargs):
        super(ProjectOption, self).__init__(**kwargs)
        if self.id is None:
            self.id = uuid4()
        if self.date_created is None:
            self.date_created = datetime.utcnow()

class ProjectPlan(db.Model):
    __tablename__ = 'project_plan'

    project_id = Column(GUID, ForeignKey('project.id', ondelete="CASCADE"),
                        nullable=False, primary_key=True)
    plan_id = Column(GUID, ForeignKey('plan.id', ondelete="CASCADE"),
                     nullable=False, primary_key=True)
    avg_build_time = Column(Integer)

    project = relationship('Project', backref=backref(
        "project_plans", cascade="all, delete-orphan"))
    plan = relationship('Plan', backref=backref(
        "plan_projects", cascade="all, delete-orphan"))

    def __init__(self, project=None, plan=None, **kwargs):
        kwargs.setdefault('project', project)
        kwargs.setdefault('plan', plan)
        super(ProjectPlan, self).__init__(**kwargs)
<commit_msg>Remove obsolete entity helpers<commit_after>from datetime import datetime
from slugify import slugify
from sqlalchemy import Column, String, DateTime, ForeignKey, Integer, Text
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.orm import backref, relationship
from sqlalchemy.schema import UniqueConstraint
from uuid import uuid4

from changes.config import db
from changes.db.types.guid import GUID

class Project(db.Model):
    __tablename__ = 'project'

    id = Column(GUID, primary_key=True, default=uuid4)
    slug = Column(String(64), unique=True, nullable=False)
    repository_id = Column(GUID, ForeignKey('repository.id', ondelete="RESTRICT"), nullable=False)
    name = Column(String(64))
    date_created = Column(DateTime, default=datetime.utcnow)
    avg_build_time = Column(Integer)

    repository = relationship('Repository')
    plans = association_proxy('project_plans', 'plan')

    def __init__(self, **kwargs):
        super(Project, self).__init__(**kwargs)
        if not self.id:
            self.id = uuid4()
        if not self.slug:
            self.slug = slugify(self.name)

class ProjectOption(db.Model):
    __tablename__ = 'projectoption'
    __table_args__ = (
        UniqueConstraint('project_id', 'name', name='unq_projectoption_name'),
    )

    id = Column(GUID, primary_key=True, default=uuid4)
    project_id = Column(GUID, ForeignKey('project.id', ondelete="CASCADE"), nullable=False)
    name = Column(String(64), nullable=False)
    value = Column(Text, nullable=False)
    date_created = Column(DateTime, default=datetime.utcnow, nullable=False)

    project = relationship('Project')

    def __init__(self, **kwargs):
        super(ProjectOption, self).__init__(**kwargs)
        if self.id is None:
            self.id = uuid4()
        if self.date_created is None:
            self.date_created = datetime.utcnow()

class ProjectPlan(db.Model):
    __tablename__ = 'project_plan'

    project_id = Column(GUID, ForeignKey('project.id', ondelete="CASCADE"),
                        nullable=False, primary_key=True)
    plan_id = Column(GUID, ForeignKey('plan.id', ondelete="CASCADE"),
                     nullable=False, primary_key=True)
    avg_build_time = Column(Integer)

    project = relationship('Project', backref=backref(
        "project_plans", cascade="all, delete-orphan"))
    plan = relationship('Plan', backref=backref(
        "plan_projects", cascade="all, delete-orphan"))

    def __init__(self, project=None, plan=None, **kwargs):
        kwargs.setdefault('project', project)
        kwargs.setdefault('plan', plan)
        super(ProjectPlan, self).__init__(**kwargs)
<|endoftext|>


<commit_before>import sys
from pkg_resources import iter_entry_points
from boardserver import server

board_plugins = dict(
    (ep.name, ep.load())
    for ep in iter_entry_points('jrb_board.games')
)

args = sys.argv[1:]
addr, port = None, None

board = board_plugins[args[0]]

if len(args) > 1:
    addr = args[1]
if len(args) > 2:
    port = int(args[2])

api = server.Server(board(), addr, port)
api.run()
<commit_msg>Added a hash-bang line to board-serve.py to make it a standalone script.<commit_after>#!/usr/bin/env python
import sys
from pkg_resources import iter_entry_points
from boardserver import server

board_plugins = dict(
    (ep.name, ep.load())
    for ep in iter_entry_points('jrb_board.games')
)

args = sys.argv[1:]
addr, port = None, None

board = board_plugins[args[0]]

if len(args) > 1:
    addr = args[1]
if len(args) > 2:
    port = int(args[2])

api = server.Server(board(), addr, port)
api.run()
<|endoftext|>


<commit_before>'''
This script creates maps of model outputs at roughly 10km resolution (0.1x0.1 degrees), where each output pixel
represents the total value in the pixel (not the density) (hence, the aggregated results).
This is currently only set up for gross emissions from biomass+soil and net flux from biomass+soil.
It iterates through all the model outputs that are supplied.
First, it rewindows the model output, pixel area tile, and tcd tile into 400x400 (0.1x0.1 degree) windows, instead of the native
40000x1 pixel windows.
Then it calculates the per pixel value for each model output pixel and sums those values within each 0.1x0.1 degree
aggregated pixel.
It converts cumulative carbon gain to CO2 gain per year, converts cumulative CO2 flux to CO2 flux per year, and
converts cumulative gross CO2 emissions to gross CO2 emissions per year.
The user has to supply a tcd threshold for which forest pixels to include in the results.
'''

import multiprocessing
import aggregate_results_to_10_km
import subprocess
from functools import partial
import argparse
import os
import glob
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

def main():

    # Sole argument for the script: the tree cover density threshold (pixels below this will not be aggregated)
    parser = argparse.ArgumentParser(description='Create maps of model output at approximately 10 km resolution')
    parser.add_argument('--tcd-threshold', '-tcd', required=True,
                        help='Tree cover density threshold above which pixels will be included in the aggregation.')
    args = parser.parse_args()

    thresh = args.tcd_threshold
    thresh = int(thresh)

    if thresh < 0 or thresh > 99:
        raise Exception('Invalid tcd. Please provide an integer between 0 and 99.')

    # tile_id_list = ['00N_100E', '00N_110E', '00N_120E'] # test tiles
    # tile_id_list = ['00N_110E'] # test tiles
    # tile_id_list = ['30N_110W']  # test tiles
    # tile_id_list = ['00N_110E', '80N_020E', '30N_080W', '00N_020E'] # test tiles: no mangrove or planted forest, mangrove only, planted forest only, mangrove and planted forest
    # print tile_id_list
    # print "There are {} tiles to process".format(str(len(tile_id_list)))
    #
    # # For copying individual tiles to spot machine for testing
    # for tile_id in tile_id_list:
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.gross_emis_all_gases_all_drivers_biomass_soil_dir, tile_id, cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.annual_gain_AGB_BGB_all_types_dir, tile_id, cn.pattern_annual_gain_AGB_BGB_all_types), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.cumul_gain_AGCO2_BGCO2_all_types_dir, tile_id, cn.pattern_cumul_gain_AGCO2_BGCO2_all_types), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.net_flux_dir, tile_id, cn.pattern_net_flux), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.pixel_area_dir, cn.pattern_pixel_area, tile_id), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.tcd_dir, cn.pattern_tcd, tile_id), '.')

    # # Pixel area tiles-- necessary for calculating sum of pixels for any set of tiles
    # uu.s3_folder_download(cn.pixel_area_dir, '.')
    # # tree cover density tiles-- necessary for filtering sums by tcd
    # uu.s3_folder_download(cn.tcd_dir, '.')

    input_dict = {
             cn.gross_emis_all_gases_all_drivers_biomass_soil_dir: cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil
             # cn.gross_emis_co2_only_all_drivers_biomass_soil_dir: cn.pattern_gross_emis_co2_only_all_drivers_biomass_soil,
             # cn.gross_emis_non_co2_all_drivers_biomass_soil_dir: cn.pattern_gross_emis_non_co2_all_drivers_biomass_soil,
             # cn.cumul_gain_AGCO2_BGCO2_all_types_dir: cn.pattern_cumul_gain_AGCO2_BGCO2_all_types,
             # cn.annual_gain_AGB_BGB_all_types_dir: cn.pattern_annual_gain_AGB_BGB_all_types,
             # cn.net_flux_dir: cn.pattern_net_flux
             }

    print "Model outputs to process are:", input_dict

    for dir, pattern in input_dict.items():

        tile_id_list = uu.tile_list(dir)

        # uu.s3_folder_download(dir, '.')

        print tile_id_list
        print "There are {} tiles to process".format(str(len(tile_id_list)))
        print "Processing:", dir, "; ", pattern

        # Lists the tiles of the particular type that is being iterates through.
        # Excludes all intermediate files
        tile_list = uu.tile_list_spot_machine(".", "{}.tif".format(pattern))
        # from https://stackoverflow.com/questions/12666897/removing-an-item-from-list-matching-a-substring
        tile_list = [i for i in tile_list if not ('hanson_2013' in i)]
        tile_list = [i for i in tile_list if not ('rewindow' in i)]
        tile_list = [i for i in tile_list if not ('10km' in i)]
        print "Tiles to process:", tile_list

        # Converts the 10x10 degree Hansen tiles that are in windows of 40000x1 pixels to windows of 400x400 pixels,
        # which is the resolution of the output tiles. This will allow the 30x30 m pixels in each window to be summed.
        # For multiprocessor use. count/2 used about 400 GB of memory on an r4.16xlarge machine, so that was okay.
        count = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(count/2)
        pool.map(aggregate_results_to_10_km.rewindow, tile_list)
        # Added these in response to error12: Cannot allocate memory error.
        # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
        # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
        pool.close()
        pool.join()

        # Converts the existing (per ha) values to per pixel values (e.g., emissions/ha to emissions/pixel)
        # and sums those values in each 400x400 pixel window.
        # The sum for each 400x400 pixel window is stored in a 2D array, which is then converted back into a raster at
        # 0.1x0.1 degree resolution (approximately 10m in the tropics).
        # Each pixel in that raster is the sum of the 30m pixels converted to value/pixel (instead of value/ha).
        # The 0.1x0.1 degree tile is output.
        # For multiprocessor use. This used about 450 GB of memory with count/2, it's okay on an r4.16xlarge
        count = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(count/2)
        pool.map(partial(aggregate_results_to_10_km.aggregate, thresh=thresh), tile_list)
        # Added these in response to error12: Cannot allocate memory error.
        # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
        # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
        pool.close()
        pool.join()

        # Makes a vrt of all the output 10x10 tiles (10 km resolution)
        out_vrt = "{}_10km.vrt".format(pattern)
        os.system('gdalbuildvrt -tr 0.1 0.1 {0} *{1}_10km*.tif'.format(out_vrt, pattern))

        # Renames outputs
        out_pattern = uu.name_aggregated_output(pattern, thresh)

        # Produces a single raster of all the 10x10 tiles (10 km resolution)
        cmd = ['gdalwarp', '-t_srs', "EPSG:4326", '-overwrite', '-dstnodata', '0', '-co', 'COMPRESS=LZW',
               '-tr', '0.1', '0.1',
               out_vrt, '{}.tif'.format(out_pattern)]
        subprocess.check_call(cmd)

        print "Tiles processed. Uploading to s3 now..."

        # Cleans up the folder before starting on the next raster type
        vrtList = glob.glob('*vrt')
        for vrt in vrtList:
            os.remove(vrt)

        for tile_id in tile_id_list:
            os.remove('{0}_{1}.tif'.format(tile_id, pattern))
            os.remove('{0}_{1}_rewindow.tif'.format(tile_id, pattern))
            os.remove('{0}_{1}_10km.tif'.format(tile_id, pattern))

        # Uploads all output tiles to s3
        uu.upload_final_set(cn.output_aggreg_dir, out_pattern)

if __name__ == '__main__':
    main()<commit_msg>Trouble creating tile list of aggregation script. Trying to debug.<commit_after>'''
This script creates maps of model outputs at roughly 10km resolution (0.1x0.1 degrees), where each output pixel
represents the total value in the pixel (not the density) (hence, the aggregated results).
This is currently only set up for gross emissions from biomass+soil and net flux from biomass+soil.
It iterates through all the model outputs that are supplied.
First, it rewindows the model output, pixel area tile, and tcd tile into 400x400 (0.1x0.1 degree) windows, instead of the native
40000x1 pixel windows.
Then it calculates the per pixel value for each model output pixel and sums those values within each 0.1x0.1 degree
aggregated pixel.
It converts cumulative carbon gain to CO2 gain per year, converts cumulative CO2 flux to CO2 flux per year, and
converts cumulative gross CO2 emissions to gross CO2 emissions per year.
The user has to supply a tcd threshold for which forest pixels to include in the results.
'''

import multiprocessing
import aggregate_results_to_10_km
import subprocess
from functools import partial
import argparse
import os
import glob
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

def main():

    # Sole argument for the script: the tree cover density threshold (pixels below this will not be aggregated)
    parser = argparse.ArgumentParser(description='Create maps of model output at approximately 10 km resolution')
    parser.add_argument('--tcd-threshold', '-tcd', required=True,
                        help='Tree cover density threshold above which pixels will be included in the aggregation.')
    args = parser.parse_args()

    thresh = args.tcd_threshold
    thresh = int(thresh)

    if thresh < 0 or thresh > 99:
        raise Exception('Invalid tcd. Please provide an integer between 0 and 99.')

    # tile_id_list = ['00N_100E', '00N_110E', '00N_120E'] # test tiles
    # tile_id_list = ['00N_110E'] # test tiles
    # tile_id_list = ['30N_110W']  # test tiles
    # tile_id_list = ['00N_110E', '80N_020E', '30N_080W', '00N_020E'] # test tiles: no mangrove or planted forest, mangrove only, planted forest only, mangrove and planted forest
    # print tile_id_list
    # print "There are {} tiles to process".format(str(len(tile_id_list)))
    #
    # # For copying individual tiles to spot machine for testing
    # for tile_id in tile_id_list:
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.gross_emis_all_gases_all_drivers_biomass_soil_dir, tile_id, cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.annual_gain_AGB_BGB_all_types_dir, tile_id, cn.pattern_annual_gain_AGB_BGB_all_types), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.cumul_gain_AGCO2_BGCO2_all_types_dir, tile_id, cn.pattern_cumul_gain_AGCO2_BGCO2_all_types), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.net_flux_dir, tile_id, cn.pattern_net_flux), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.pixel_area_dir, cn.pattern_pixel_area, tile_id), '.')
    #     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.tcd_dir, cn.pattern_tcd, tile_id), '.')

    # # Pixel area tiles-- necessary for calculating sum of pixels for any set of tiles
    # uu.s3_folder_download(cn.pixel_area_dir, '.')
    # # tree cover density tiles-- necessary for filtering sums by tcd
    # uu.s3_folder_download(cn.tcd_dir, '.')

    input_dict = {
             cn.gross_emis_all_gases_all_drivers_biomass_soil_dir: cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil
             # cn.gross_emis_co2_only_all_drivers_biomass_soil_dir: cn.pattern_gross_emis_co2_only_all_drivers_biomass_soil,
             # cn.gross_emis_non_co2_all_drivers_biomass_soil_dir: cn.pattern_gross_emis_non_co2_all_drivers_biomass_soil,
             # cn.cumul_gain_AGCO2_BGCO2_all_types_dir: cn.pattern_cumul_gain_AGCO2_BGCO2_all_types,
             # cn.annual_gain_AGB_BGB_all_types_dir: cn.pattern_annual_gain_AGB_BGB_all_types,
             # cn.net_flux_dir: cn.pattern_net_flux
             }

    print "Model outputs to process are:", input_dict

    for dir, pattern in input_dict.items():

        tile_id_list = uu.tile_list(dir)

        # uu.s3_folder_download(dir, '.')

        print tile_id_list
        print "There are {} tiles to process".format(str(len(tile_id_list)))
        print "Processing:", dir, "; ", pattern

        # Lists the tiles of the particular type that is being iterates through.
        # Excludes all intermediate files
        tile_list = uu.tile_list_spot_machine(".", "{}.tif".format(pattern))
        # from https://stackoverflow.com/questions/12666897/removing-an-item-from-list-matching-a-substring
        tile_list = [i for i in tile_list if not ('hanson_2013' in i)]
        tile_list = [i for i in tile_list if not ('rewindow' in i)]
        tile_list = [i for i in tile_list if not ('10km' in i)]

        print pattern

        print "Tiles to process:", tile_list

        # Converts the 10x10 degree Hansen tiles that are in windows of 40000x1 pixels to windows of 400x400 pixels,
        # which is the resolution of the output tiles. This will allow the 30x30 m pixels in each window to be summed.
        # For multiprocessor use. count/2 used about 400 GB of memory on an r4.16xlarge machine, so that was okay.
        count = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(count/2)
        pool.map(aggregate_results_to_10_km.rewindow, tile_list)
        # Added these in response to error12: Cannot allocate memory error.
        # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
        # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
        pool.close()
        pool.join()

        # Converts the existing (per ha) values to per pixel values (e.g., emissions/ha to emissions/pixel)
        # and sums those values in each 400x400 pixel window.
        # The sum for each 400x400 pixel window is stored in a 2D array, which is then converted back into a raster at
        # 0.1x0.1 degree resolution (approximately 10m in the tropics).
        # Each pixel in that raster is the sum of the 30m pixels converted to value/pixel (instead of value/ha).
        # The 0.1x0.1 degree tile is output.
        # For multiprocessor use. This used about 450 GB of memory with count/2, it's okay on an r4.16xlarge
        count = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(count/2)
        pool.map(partial(aggregate_results_to_10_km.aggregate, thresh=thresh), tile_list)
        # Added these in response to error12: Cannot allocate memory error.
        # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
        # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
        pool.close()
        pool.join()

        # Makes a vrt of all the output 10x10 tiles (10 km resolution)
        out_vrt = "{}_10km.vrt".format(pattern)
        os.system('gdalbuildvrt -tr 0.1 0.1 {0} *{1}_10km*.tif'.format(out_vrt, pattern))

        # Renames outputs
        out_pattern = uu.name_aggregated_output(pattern, thresh)

        # Produces a single raster of all the 10x10 tiles (10 km resolution)
        cmd = ['gdalwarp', '-t_srs', "EPSG:4326", '-overwrite', '-dstnodata', '0', '-co', 'COMPRESS=LZW',
               '-tr', '0.1', '0.1',
               out_vrt, '{}.tif'.format(out_pattern)]
        subprocess.check_call(cmd)

        print "Tiles processed. Uploading to s3 now..."

        # Cleans up the folder before starting on the next raster type
        vrtList = glob.glob('*vrt')
        for vrt in vrtList:
            os.remove(vrt)

        for tile_id in tile_id_list:
            os.remove('{0}_{1}.tif'.format(tile_id, pattern))
            os.remove('{0}_{1}_rewindow.tif'.format(tile_id, pattern))
            os.remove('{0}_{1}_10km.tif'.format(tile_id, pattern))

        # Uploads all output tiles to s3
        uu.upload_final_set(cn.output_aggreg_dir, out_pattern)

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


<commit_before>import json
from os import getenv
from subprocess import check_output, CalledProcessError

JSON_DATA_FILE = getenv('AWS_DB_JSON_DATA_FILE', 'bedrock_db_info.json')
DB_FILE = 'bedrock.db'
CACHE = {}

def get_db_checksum(filename=None):
    filename = filename or DB_FILE
    cache_key = 'db_sum_%s' % filename
    db_sum = CACHE.get(cache_key)
    if not db_sum:
        db_sum = check_output('sha256sum {}'.format(filename), shell=True).split()[0]
        CACHE[cache_key] = db_sum

    return db_sum

def get_git_sha():
    git_sha = CACHE.get('git_sha')
    if not git_sha:
        git_sha = getenv('GIT_SHA')
        if not git_sha:
            try:
                git_sha = check_output('git rev-parse HEAD', shell=True).strip()
            except CalledProcessError:
                git_sha = 'testing'

        CACHE['git_sha'] = git_sha

    return git_sha

def get_prev_db_data():
    try:
        with open(JSON_DATA_FILE, 'r') as fp:
            return json.load(fp)
    except IOError:
        # file does not exist
        return None

def set_db_data(db_data):
    with open(JSON_DATA_FILE, 'w') as fp:
        json.dump(db_data, fp)
<commit_msg>Use python's sha256 instead of requiring a shell program<commit_after>import json
from hashlib import sha256
from os import getenv
from subprocess import check_output, CalledProcessError

JSON_DATA_FILE = getenv('AWS_DB_JSON_DATA_FILE', 'bedrock_db_info.json')
DB_FILE = 'bedrock.db'
CACHE = {}
BLOCKSIZE = 65536

def _sha256_sum(filename):
    hasher = sha256()
    with open(filename, 'rb') as fh:
        for chunk in iter(lambda: fh.read(BLOCKSIZE), b""):
            hasher.update(chunk)

    return hasher.hexdigest()

def get_db_checksum(filename=None):
    filename = filename or DB_FILE
    cache_key = 'db_sum_%s' % filename
    db_sum = CACHE.get(cache_key)
    if not db_sum:
        db_sum = _sha256_sum(filename)
        CACHE[cache_key] = db_sum

    return db_sum

def get_git_sha():
    git_sha = CACHE.get('git_sha')
    if not git_sha:
        git_sha = getenv('GIT_SHA')
        if not git_sha:
            try:
                git_sha = check_output('git rev-parse HEAD', shell=True).strip()
            except CalledProcessError:
                git_sha = 'testing'

        CACHE['git_sha'] = git_sha

    return git_sha

def get_prev_db_data():
    try:
        with open(JSON_DATA_FILE, 'r') as fp:
            return json.load(fp)
    except IOError:
        # file does not exist
        return None

def set_db_data(db_data):
    with open(JSON_DATA_FILE, 'w') as fp:
        json.dump(db_data, fp)
<|endoftext|>


<commit_before>import re
from operator import attrgetter

from six import u
from funcy import imap, flatten

from circuits.protocols.irc import response
from circuits.protocols.irc.replies import RPL_TOPICWHO
from circuits.protocols.irc.replies import RPL_NOTOPIC, RPL_TOPIC, ERR_NOSUCHCHANNEL
from circuits.protocols.irc.replies import ERR_TOOMANYCHANNELS, ERR_USERNOTINCHANNEL
from circuits.protocols.irc.replies import JOIN, TOPIC, KICK, PART, RPL_LIST, RPL_LISTEND
from circuits.protocols.irc.replies import RPL_NAMEREPLY, RPL_ENDOFNAMES, ERR_CHANOPRIVSNEEDED

from .. import models
from ..plugin import BasePlugin
from ..commands import BaseCommands

VALID_CHANNEL_REGEX = re.compile(r"^[&#+!][^\x00\x07\x0a\x0d ,:]*$")

class Commands(BaseCommands):

    def _join(self, sock, source, name):
        user = models.User.objects.filter(sock=sock).first()

        if name and name[0] not in self.parent.chantypes:
            return ERR_NOSUCHCHANNEL(name)

        if VALID_CHANNEL_REGEX.match(name) is None:
            return ERR_NOSUCHCHANNEL(name)

        if len(name) > self.parent.channellen:
            return ERR_NOSUCHCHANNEL(name)

        replies = [JOIN(name, prefix=user.prefix)]

        channel = models.Channel.objects.filter(name=name).first()
        if channel is None:
            channel = models.Channel(name=name)
            channel.save()

        if user in channel.users:
            return

        type = name[0]
        nchannels = len([x for x in user.channels if channel.type == type])
        chanlimit = self.parent.chanlimit.get(type, None)
        if chanlimit and nchannels >= chanlimit:
            return ERR_TOOMANYCHANNELS(name)

        self.notify(
            channel.users[:],
            JOIN(name, prefix=user.prefix)
        )

        user.channels.append(channel)
        user.save()

        if not channel.users:
            channel.operators.append(user)
            channel.save()

        channel.users.append(user)
        channel.save()

        self.fire(response.create("topic", sock, source, channel.name, joined=True), "server")
        self.fire(response.create("names", sock, source, channel.name), "server")

        return replies

    def join(self, sock, source, names, keys=None):
        return filter(
            None,
            flatten(self._join(sock, source, name) for name in names.split(u(",")) if name)
        )

    def part(self, sock, source, name, reason=u("Leaving")):
        user = models.User.objects.filter(sock=sock).first()

        channel = models.Channel.objects.filter(name=name).first()

        if channel is None:
            return

        if user not in channel.users:
            return

        self.notify(channel.users[:], PART(name, reason, prefix=user.prefix))

        user.channels.remove(channel)
        user.save()

        channel.users.remove(user)
        if user in channel.operators:
            channel.operators.remove(user)
        if user in channel.voiced:
            channel.voiced.remove(user)
        channel.save()

        if not channel.users:
            channel.delete()

    def names(self, sock, source, name):
        channel = models.Channel.objects.filter(name=name).first()

        if channel is None:
            return ERR_NOSUCHCHANNEL(name)

        return [
            RPL_NAMEREPLY(channel.name, channel.userprefixes),
            RPL_ENDOFNAMES(name),
        ]

    def topic(self, sock, source, name, topic=None, **kwargs):
        user = models.User.objects.filter(sock=sock).first()

        channel = models.Channel.objects.filter(name=name).first()
        if channel is None:
            return ERR_NOSUCHCHANNEL(name)

        _topic, _topic_setter, _topic_timestamp = channel.topic

        if topic is None and not _topic:
            if kwargs.get("joined", False):
                return
            return RPL_NOTOPIC(channel.name)

        if topic is None:
            return (
                RPL_TOPIC(channel.name, _topic),
                RPL_TOPICWHO(channel.name, _topic_setter, _topic_timestamp)
            )

        if not user.oper and u("t") in channel.modes and user not in channel.operators:
            return ERR_CHANOPRIVSNEEDED(channel.name)

        channel.topic = (topic, user.prefix)

        self.notify(channel.users[:], TOPIC(channel.name, topic, prefix=user.prefix))

    def list(self, sock, source):
        replies = []

        for channel in models.Channel.objects.all():
            nvisible = len([x for x in channel.users if x.visible])
            replies.append(RPL_LIST(channel.name, nvisible, channel.topic))

        replies.append(RPL_LISTEND())

        return replies

    def kick(self, sock, source, name, nick, reason=None):
        user = models.User.objects.filter(sock=sock).first()

        channel = models.Channel.objects.filter(name=name).first()
        if channel is None:
            return ERR_NOSUCHCHANNEL(name)

        if not user.oper and user not in channel.operators:
            return ERR_CHANOPRIVSNEEDED(channel.name)

        if nick not in imap(attrgetter("nick"), channel.users):
            return ERR_USERNOTINCHANNEL(nick, channel.name)

        nick = models.User.objects.filter(nick=nick).first()

        self.notify(
            channel.users[:],
            KICK(channel.name, nick.nick, reason or nick.nick, prefix=user.prefix)
        )

        nick.channels.remove(channel)
        nick.save()

        channel.users.remove(nick)
        if user in channel.operators:
            channel.operators.remove(user)
        if user in channel.voiced:
            channel.voiced.remove(user)
        channel.save()

        if not channel.users:
            channel.delete()

class Channel(BasePlugin):

    def init(self, *args, **kwargs):
        super(Channel, self).init(*args, **kwargs)

        self.channellen = 50
        self.topiclen = 300
        self.chantypes = u("#&")

        self.chanlimit = {
            u("#"): 120,
            u("&"): 10,
        }

        self.features = (
            u("PREFIX=(ov)@+"),
            u("CHANTYPES={0}").format(self.chantypes),
            u("TOPICLEN={0}").format(self.topiclen),
            u("CHANNELLEN={0}").format(self.channellen),
            u("CHANLIMIT={0}").format(
                u(",").join(u("{0}:{1}").format(k, v) for k, v in self.chanlimit.items())
            ),
        )

        Commands(*args, **kwargs).register(self)

    def supports(self):
        return self.features
<commit_msg>Repsect the TOPICLEN and truncate<commit_after>import re
from operator import attrgetter

from six import u
from funcy import imap, flatten

from circuits.protocols.irc import response
from circuits.protocols.irc.replies import RPL_TOPICWHO
from circuits.protocols.irc.replies import RPL_NOTOPIC, RPL_TOPIC, ERR_NOSUCHCHANNEL
from circuits.protocols.irc.replies import ERR_TOOMANYCHANNELS, ERR_USERNOTINCHANNEL
from circuits.protocols.irc.replies import JOIN, TOPIC, KICK, PART, RPL_LIST, RPL_LISTEND
from circuits.protocols.irc.replies import RPL_NAMEREPLY, RPL_ENDOFNAMES, ERR_CHANOPRIVSNEEDED

from .. import models
from ..plugin import BasePlugin
from ..commands import BaseCommands

VALID_CHANNEL_REGEX = re.compile(r"^[&#+!][^\x00\x07\x0a\x0d ,:]*$")

class Commands(BaseCommands):

    def _join(self, sock, source, name):
        user = models.User.objects.filter(sock=sock).first()

        if name and name[0] not in self.parent.chantypes:
            return ERR_NOSUCHCHANNEL(name)

        if VALID_CHANNEL_REGEX.match(name) is None:
            return ERR_NOSUCHCHANNEL(name)

        if len(name) > self.parent.channellen:
            return ERR_NOSUCHCHANNEL(name)

        replies = [JOIN(name, prefix=user.prefix)]

        channel = models.Channel.objects.filter(name=name).first()
        if channel is None:
            channel = models.Channel(name=name)
            channel.save()

        if user in channel.users:
            return

        type = name[0]
        nchannels = len([x for x in user.channels if channel.type == type])
        chanlimit = self.parent.chanlimit.get(type, None)
        if chanlimit and nchannels >= chanlimit:
            return ERR_TOOMANYCHANNELS(name)

        self.notify(
            channel.users[:],
            JOIN(name, prefix=user.prefix)
        )

        user.channels.append(channel)
        user.save()

        if not channel.users:
            channel.operators.append(user)
            channel.save()

        channel.users.append(user)
        channel.save()

        self.fire(response.create("topic", sock, source, channel.name, joined=True), "server")
        self.fire(response.create("names", sock, source, channel.name), "server")

        return replies

    def join(self, sock, source, names, keys=None):
        return filter(
            None,
            flatten(self._join(sock, source, name) for name in names.split(u(",")) if name)
        )

    def part(self, sock, source, name, reason=u("Leaving")):
        user = models.User.objects.filter(sock=sock).first()

        channel = models.Channel.objects.filter(name=name).first()

        if channel is None:
            return

        if user not in channel.users:
            return

        self.notify(channel.users[:], PART(name, reason, prefix=user.prefix))

        user.channels.remove(channel)
        user.save()

        channel.users.remove(user)
        if user in channel.operators:
            channel.operators.remove(user)
        if user in channel.voiced:
            channel.voiced.remove(user)
        channel.save()

        if not channel.users:
            channel.delete()

    def names(self, sock, source, name):
        channel = models.Channel.objects.filter(name=name).first()

        if channel is None:
            return ERR_NOSUCHCHANNEL(name)

        return [
            RPL_NAMEREPLY(channel.name, channel.userprefixes),
            RPL_ENDOFNAMES(name),
        ]

    def topic(self, sock, source, name, topic=None, **kwargs):
        user = models.User.objects.filter(sock=sock).first()

        channel = models.Channel.objects.filter(name=name).first()
        if channel is None:
            return ERR_NOSUCHCHANNEL(name)

        _topic, _topic_setter, _topic_timestamp = channel.topic

        if topic is None and not _topic:
            if kwargs.get("joined", False):
                return
            return RPL_NOTOPIC(channel.name)

        if topic is None:
            return (
                RPL_TOPIC(channel.name, _topic),
                RPL_TOPICWHO(channel.name, _topic_setter, _topic_timestamp)
            )

        if not user.oper and u("t") in channel.modes and user not in channel.operators:
            return ERR_CHANOPRIVSNEEDED(channel.name)

        topic = topic[:self.parent.topiclen]

        channel.topic = (topic, user.prefix)

        self.notify(channel.users[:], TOPIC(channel.name, topic, prefix=user.prefix))

    def list(self, sock, source):
        replies = []

        for channel in models.Channel.objects.all():
            nvisible = len([x for x in channel.users if x.visible])
            replies.append(RPL_LIST(channel.name, nvisible, channel.topic))

        replies.append(RPL_LISTEND())

        return replies

    def kick(self, sock, source, name, nick, reason=None):
        user = models.User.objects.filter(sock=sock).first()

        channel = models.Channel.objects.filter(name=name).first()
        if channel is None:
            return ERR_NOSUCHCHANNEL(name)

        if not user.oper and user not in channel.operators:
            return ERR_CHANOPRIVSNEEDED(channel.name)

        if nick not in imap(attrgetter("nick"), channel.users):
            return ERR_USERNOTINCHANNEL(nick, channel.name)

        nick = models.User.objects.filter(nick=nick).first()

        self.notify(
            channel.users[:],
            KICK(channel.name, nick.nick, reason or nick.nick, prefix=user.prefix)
        )

        nick.channels.remove(channel)
        nick.save()

        channel.users.remove(nick)
        if user in channel.operators:
            channel.operators.remove(user)
        if user in channel.voiced:
            channel.voiced.remove(user)
        channel.save()

        if not channel.users:
            channel.delete()

class Channel(BasePlugin):

    def init(self, *args, **kwargs):
        super(Channel, self).init(*args, **kwargs)

        self.channellen = 50
        self.topiclen = 300
        self.chantypes = u("#&")

        self.chanlimit = {
            u("#"): 120,
            u("&"): 10,
        }

        self.features = (
            u("PREFIX=(ov)@+"),
            u("CHANTYPES={0}").format(self.chantypes),
            u("TOPICLEN={0}").format(self.topiclen),
            u("CHANNELLEN={0}").format(self.channellen),
            u("CHANLIMIT={0}").format(
                u(",").join(u("{0}:{1}").format(k, v) for k, v in self.chanlimit.items())
            ),
        )

        Commands(*args, **kwargs).register(self)

    def supports(self):
        return self.features
<|endoftext|>