<commit_before>import inspect
import six
from collections import OrderedDict

from graphql_relay import from_global_id

from google.appengine.ext import ndb

from graphene import ObjectType, Field, String
from graphene.types.objecttype import ObjectTypeMeta, merge, yank_fields_from_attrs
from graphene.utils.is_base_type import is_base_type
from graphene.types.options import Options

from graphene_gae.ndb.converter import convert_ndb_property

__author__ = 'ekampf'

class NdbObjectTypeMeta(ObjectTypeMeta):
    REGISTRY = {}  # Maps between ndb.Model to its GraphQL type

    def __new__(mcs, name, bases, attrs):
        if not is_base_type(bases, NdbObjectTypeMeta):
            return type.__new__(mcs, name, bases, attrs)

        options = Options(
            attrs.pop('Meta', None),
            name=name,
            description=attrs.pop('__doc__', None),
            model=None,
            local_fields=None,
            only_fields=(),
            exclude_fields=(),
            interfaces=(),
            registry=None
        )

        if not options.model:
            raise Exception('NdbObjectType %s must have a model in the Meta class attr' % name)

        if not inspect.isclass(options.model) or not issubclass(options.model, ndb.Model):
            raise Exception('Provided model in %s is not an NDB model' % name)

        new_cls = ObjectTypeMeta.__new__(mcs, name, bases, dict(attrs, _meta=options))
        mcs.register(new_cls)

        ndb_fields = mcs.fields_for_ndb_model(options)
        options.ndb_fields = yank_fields_from_attrs(
            ndb_fields,
            _as=Field,
        )
        options.fields = merge(
            options.interface_fields,
            options.ndb_fields,
            options.base_fields,
            options.local_fields
        )

        return new_cls

    @classmethod
    def register(mcs, object_type_meta):
        mcs.REGISTRY[object_type_meta._meta.model.__name__] = object_type_meta

    @staticmethod
    def fields_for_ndb_model(options):
        ndb_model = options.model
        only_fields = options.only_fields
        already_created_fields = {name for name, _ in options.local_fields.iteritems()}

        ndb_fields = OrderedDict()
        for prop_name, prop in ndb_model._properties.iteritems():
            name = prop._code_name

            is_not_in_only = only_fields and name not in only_fields
            is_excluded = name in options.exclude_fields or name in already_created_fields
            if is_not_in_only or is_excluded:
                continue

            results = convert_ndb_property(prop)
            if not results:
                continue

            if not isinstance(results, list):
                results = [results]

            for r in results:
                ndb_fields[r.name] = r.field

        ndb_fields['ndb_id'] = Field(String, resolver=lambda entity, *_: str(entity.key.id()))

        return ndb_fields

class NdbObjectType(six.with_metaclass(NdbObjectTypeMeta, ObjectType)):
    @classmethod
    def is_type_of(cls, root, context, info):
        if isinstance(root, cls):
            return True

        if not cls.is_valid_ndb_model(type(root)):
            raise Exception(('Received incompatible instance "{}".').format(root))

        return type(root) == cls._meta.model

    @classmethod
    def get_node(cls, global_id, *_):
        try:
            _type, url_safe_key = from_global_id(global_id)
            assert _type == cls.__name__
        except ValueError:
            return None

        model = cls._meta.model
        try:
            key = ndb.Key(urlsafe=url_safe_key)
        except:
            return None

        assert key.kind() == model.__name__
        return key.get()

    @classmethod
    def resolve_id(cls, entity, args, context, info):
        from graphql_relay import to_global_id
        return to_global_id(cls.__name__, entity.key.urlsafe())

    @staticmethod
    def is_valid_ndb_model(model):
        return inspect.isclass(model) and issubclass(model, ndb.Model)

<commit_msg>resolve_id should return local id not global_id<commit_after>import inspect
import six
from collections import OrderedDict

from graphql_relay import from_global_id

from google.appengine.ext import ndb

from graphene import ObjectType, Field, String
from graphene.types.objecttype import ObjectTypeMeta, merge, yank_fields_from_attrs
from graphene.utils.is_base_type import is_base_type
from graphene.types.options import Options

from graphene_gae.ndb.converter import convert_ndb_property

__author__ = 'ekampf'

class NdbObjectTypeMeta(ObjectTypeMeta):
    REGISTRY = {}  # Maps between ndb.Model to its GraphQL type

    def __new__(mcs, name, bases, attrs):
        if not is_base_type(bases, NdbObjectTypeMeta):
            return type.__new__(mcs, name, bases, attrs)

        options = Options(
            attrs.pop('Meta', None),
            name=name,
            description=attrs.pop('__doc__', None),
            model=None,
            local_fields=None,
            only_fields=(),
            exclude_fields=(),
            interfaces=(),
            registry=None
        )

        if not options.model:
            raise Exception('NdbObjectType %s must have a model in the Meta class attr' % name)

        if not inspect.isclass(options.model) or not issubclass(options.model, ndb.Model):
            raise Exception('Provided model in %s is not an NDB model' % name)

        new_cls = ObjectTypeMeta.__new__(mcs, name, bases, dict(attrs, _meta=options))
        mcs.register(new_cls)

        ndb_fields = mcs.fields_for_ndb_model(options)
        options.ndb_fields = yank_fields_from_attrs(
            ndb_fields,
            _as=Field,
        )
        options.fields = merge(
            options.interface_fields,
            options.ndb_fields,
            options.base_fields,
            options.local_fields
        )

        return new_cls

    @classmethod
    def register(mcs, object_type_meta):
        mcs.REGISTRY[object_type_meta._meta.model.__name__] = object_type_meta

    @staticmethod
    def fields_for_ndb_model(options):
        ndb_model = options.model
        only_fields = options.only_fields
        already_created_fields = {name for name, _ in options.local_fields.iteritems()}

        ndb_fields = OrderedDict()
        for prop_name, prop in ndb_model._properties.iteritems():
            name = prop._code_name

            is_not_in_only = only_fields and name not in only_fields
            is_excluded = name in options.exclude_fields or name in already_created_fields
            if is_not_in_only or is_excluded:
                continue

            results = convert_ndb_property(prop)
            if not results:
                continue

            if not isinstance(results, list):
                results = [results]

            for r in results:
                ndb_fields[r.name] = r.field

        ndb_fields['ndb_id'] = Field(String, resolver=lambda entity, *_: str(entity.key.id()))

        return ndb_fields

class NdbObjectType(six.with_metaclass(NdbObjectTypeMeta, ObjectType)):
    @classmethod
    def is_type_of(cls, root, context, info):
        if isinstance(root, cls):
            return True

        if not cls.is_valid_ndb_model(type(root)):
            raise Exception(('Received incompatible instance "{}".').format(root))

        return type(root) == cls._meta.model

    @classmethod
    def get_node(cls, global_id, *_):
        try:
            _type, url_safe_key = from_global_id(global_id)
            assert _type == cls.__name__
        except ValueError:
            return None

        model = cls._meta.model
        try:
            key = ndb.Key(urlsafe=url_safe_key)
        except:
            return None

        assert key.kind() == model.__name__
        return key.get()

    @classmethod
    def resolve_id(cls, entity, args, context, info):
        return entity.key.urlsafe()

    @staticmethod
    def is_valid_ndb_model(model):
        return inspect.isclass(model) and issubclass(model, ndb.Model)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

ALLELE_RANGE        = (0, 3)
FITNESS_DOMINANT    = (0.8, 1.0)
FITNESS_RECESSIVE   = (0.7, 0.9)
MAX_POPULATION      = 100
MAXNUMLOCI          = 10
PYRIMS              = ["C", "T"]
PURINES             = ["A", "G"]
NUCLEOTIDES         = ["A", "C", "G", "T"]
AMINO_ACIDS         = ["A", "C", "D", "E", "F", "G", "H", "I", "K", "L", "M", "N", "P", "Q", "R", "S", "T", "V", "W", "Y"]
GENETIC_CODE        = [["GCA", "GCC", "GCG", "GCT"], ["TGC","TGT"], ["GAC", "GAT"], ["GAA", "GAG"], ["TTC", "TTT"], ["GGA", "GGC", "GGG", "GGT"], ["CAC", "CAT"], ["ATA", "ATC", "ATT"], ["AAA", "AAG"], ["CTA", "CTC", "CTG", "CTT", "TTA", "TTG"], ["ATG"], ["AAC", "AAT"], ["CCA", "CCC", "CCG", "CCT"], ["CAA", "CAG"], ["AGA", "AGG", "CGA", "CGC", "CGG", "CGT"] , ["AGC", "AGT", "TCA", "TCC", "TCG", "TCT"], ["ACA", "ACC", "ACG", "ACT"], ["GTA", "GTC", "GTG", "GTT"], ["TGG"], ["TAC", "TAT"]]
CODON_DICT          = {"AAA":"K", "AAC":"N", "AAG":"K", "AAT":"N", "ACA":"T", "ACC":"T", "ACG":"T", "ACT":"T", "AGA":"R", "AGC":"S", "AGG":"R", "AGT":"S", "ATA":"I", "ATC":"I", "ATG":"M", "ATT":"I", "CAA":"Q", "CAC":"H", "CAG":"Q", "CAT":"H", "CCA":"P", "CCC":"P", "CCG":"P", "CCT":"P", "CGA":"R", "CGC":"R", "CGG":"R", "CGT":"R", "CTA":"L", "CTC":"L", "CTG":"L", "CTT":"L", "GAA":"E", "GAC":"D", "GAG":"E", "GAT":"D", "GCA":"A", "GCC":"A", "GCG":"A", "GCT":"A", "GGA":"G", "GGC":"G", "GGG":"G", "GGT":"G", "GTA":"V", "GTC":"V", "GTG":"V", "GTT":"V", "TAC":"Y", "TAT":"Y", "TCA":"S", "TCC":"S", "TCG":"S", "TCT":"S", "TGC":"C", "TGG":"W", "TGT":"C", "TTA":"L", "TTC":"F", "TTG":"L", "TTT":"F"}
CODONS              = ["AAA", "AAC", "AAG", "AAT", "ACA", "ACC", "ACG", "ACT", "AGA", "AGC", "AGG", "AGT", "ATA", "ATC", "ATG", "ATT", "CAA", "CAC", "CAG", "CAT", "CCA", "CCC", "CCG", "CCT", "CGA", "CGC", "CGG", "CGT", "CTA", "CTC", "CTG", "CTT", "GAA", "GAC", "GAG", "GAT", "GCA", "GCC", "GCG", "GCT", "GGA", "GGC", "GGG", "GGT", "GTA", "GTC", "GTG", "GTT", "TAC", "TAT", "TCA", "TCC", "TCG", "TCT", "TGC", "TGG", "TGT", "TTA", "TTC", "TTG", "TTT"]
STOP_CODONS         = ["TAA", "TAG", "TGA"]
DNA_UPPER           = ''.join([x for x in NUCLEOTIDES])
DNA_LOWER           = ''.join([x for x in NUCLEOTIDES]).lower()
<commit_msg>tweaks fitness range<commit_after># -*- coding: utf-8 -*-

ALLELE_RANGE        = (0, 3)
FITNESS_DOMINANT    = (0.9, 1.0)
FITNESS_RECESSIVE   = (0.7, 0.8)
MAX_POPULATION      = 100
MAXNUMLOCI          = 10
PYRIMS              = ["C", "T"]
PURINES             = ["A", "G"]
NUCLEOTIDES         = ["A", "C", "G", "T"]
AMINO_ACIDS         = ["A", "C", "D", "E", "F", "G", "H", "I", "K", "L", "M", "N", "P", "Q", "R", "S", "T", "V", "W", "Y"]
GENETIC_CODE        = [["GCA", "GCC", "GCG", "GCT"], ["TGC","TGT"], ["GAC", "GAT"], ["GAA", "GAG"], ["TTC", "TTT"], ["GGA", "GGC", "GGG", "GGT"], ["CAC", "CAT"], ["ATA", "ATC", "ATT"], ["AAA", "AAG"], ["CTA", "CTC", "CTG", "CTT", "TTA", "TTG"], ["ATG"], ["AAC", "AAT"], ["CCA", "CCC", "CCG", "CCT"], ["CAA", "CAG"], ["AGA", "AGG", "CGA", "CGC", "CGG", "CGT"] , ["AGC", "AGT", "TCA", "TCC", "TCG", "TCT"], ["ACA", "ACC", "ACG", "ACT"], ["GTA", "GTC", "GTG", "GTT"], ["TGG"], ["TAC", "TAT"]]
CODON_DICT          = {"AAA":"K", "AAC":"N", "AAG":"K", "AAT":"N", "ACA":"T", "ACC":"T", "ACG":"T", "ACT":"T", "AGA":"R", "AGC":"S", "AGG":"R", "AGT":"S", "ATA":"I", "ATC":"I", "ATG":"M", "ATT":"I", "CAA":"Q", "CAC":"H", "CAG":"Q", "CAT":"H", "CCA":"P", "CCC":"P", "CCG":"P", "CCT":"P", "CGA":"R", "CGC":"R", "CGG":"R", "CGT":"R", "CTA":"L", "CTC":"L", "CTG":"L", "CTT":"L", "GAA":"E", "GAC":"D", "GAG":"E", "GAT":"D", "GCA":"A", "GCC":"A", "GCG":"A", "GCT":"A", "GGA":"G", "GGC":"G", "GGG":"G", "GGT":"G", "GTA":"V", "GTC":"V", "GTG":"V", "GTT":"V", "TAC":"Y", "TAT":"Y", "TCA":"S", "TCC":"S", "TCG":"S", "TCT":"S", "TGC":"C", "TGG":"W", "TGT":"C", "TTA":"L", "TTC":"F", "TTG":"L", "TTT":"F"}
CODONS              = ["AAA", "AAC", "AAG", "AAT", "ACA", "ACC", "ACG", "ACT", "AGA", "AGC", "AGG", "AGT", "ATA", "ATC", "ATG", "ATT", "CAA", "CAC", "CAG", "CAT", "CCA", "CCC", "CCG", "CCT", "CGA", "CGC", "CGG", "CGT", "CTA", "CTC", "CTG", "CTT", "GAA", "GAC", "GAG", "GAT", "GCA", "GCC", "GCG", "GCT", "GGA", "GGC", "GGG", "GGT", "GTA", "GTC", "GTG", "GTT", "TAC", "TAT", "TCA", "TCC", "TCG", "TCT", "TGC", "TGG", "TGT", "TTA", "TTC", "TTG", "TTT"]
STOP_CODONS         = ["TAA", "TAG", "TGA"]
DNA_UPPER           = ''.join([x for x in NUCLEOTIDES])
DNA_LOWER           = ''.join([x for x in NUCLEOTIDES]).lower()
<|endoftext|>


<commit_before><commit_msg>fixed syntax error in constants.py<commit_after><|endoftext|>


<commit_before><commit_msg>Changed debug mode to be the default<commit_after>import os

env = os.environ

DEPLOYED = False

if not DEPLOYED: DEBUG = True

<|endoftext|>


<commit_before><commit_msg>First implementation of class Dimensions, passes the tests<commit_after>

class Dimensions(object):

    def __init__(self, **kwargs):

        object.__setattr__(self, "_Dimensions__dimensions", [])

        for key, value in kwargs.items():
            self.__dimensions.append(key)
            setattr(self, key, value)

    def __setattr__(self, name, value):
        if name not in self.__dimensions:
            raise AttributeError("Cannot set attribute %s" %name)
        elif isinstance(value, int):
            value = float(value)
        elif isinstance(value, (float, type(None))):
            pass
        else:
            raise TypeError("Cannot set dimension to value of type %s" %type(value))

        object.__setattr__(self, name, value)

    def __getattribute__(self, name):
        value = object.__getattribute__(self, name)

        if name.startswith("_Dimensions"):
            return value
        elif name in self.__dimensions and value is None:
            raise ValueError("Dimension '%s' is not set" %name)
        else:
            return value
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Thu Aug 24 15:40:25 2017

@author: sakurai
"""

import copy
import os

import matplotlib.pyplot as plt
import numpy as np
import six
import time

import chainer
from chainer import cuda
from chainer import optimizers
from chainer import Variable
from tqdm import tqdm

from . import utils
from ..datasets import data_provider
from ..models.modified_googlenet import ModifiedGoogLeNet

def train(main_script_path, func_train_one_batch, param_dict,
          save_distance_matrix=False,):
    script_filename = os.path.splitext(os.path.basename(main_script_path))[0]

    chainer.config.train = False
    device = 0
    xp = chainer.cuda.cupy
    config_parser = six.moves.configparser.ConfigParser()
    config_parser.read('config')
    log_dir_path = os.path.expanduser(config_parser.get('logs', 'dir_path'))

    p = utils.Logger(log_dir_path, **param_dict)  # hyperparameters

    ##########################################################
    # load database
    ##########################################################
    streams = data_provider.get_streams(p.batch_size, dataset=p.dataset,
                                        method=p.method, crop_size=p.crop_size)
    stream_train, stream_train_eval, stream_test = streams
    iter_train = stream_train.get_epoch_iterator()

    ##########################################################
    # construct the model
    ##########################################################
    model = ModifiedGoogLeNet(p.out_dim, p.normalize_output)
    if device >= 0:
        model.to_gpu()
    model.cleargrads()
    xp = model.xp
    if p.optimizer == 'Adam':
        optimizer = optimizers.Adam(p.learning_rate)
    elif p.optimizer == 'RMSProp':
        optimizer = optimizers.RMSprop(p.learning_rate)
    else:
        raise ValueError
    optimizer.setup(model)
    optimizer.add_hook(chainer.optimizer.WeightDecay(p.l2_weight_decay))

    print(p)
    stop = False
    logger = utils.Logger(log_dir_path)
    logger.soft_test_best = [0]
    time_origin = time.time()
    try:
        for epoch in range(p.num_epochs):
            time_begin = time.time()
            epoch_losses = []

            for i in tqdm(range(p.num_batches_per_epoch),
                          desc='# {}'.format(epoch)):
                with chainer.using_config('train', True):
                    loss = func_train_one_batch(model, p, next(iter_train))
                    loss.backward()
                optimizer.update()
                model.cleargrads()
                epoch_losses.append(loss.data)
                del loss

            loss_average = cuda.to_cpu(xp.array(
                xp.hstack(epoch_losses).mean()))

            # average accuracy and distance matrix for training data
            D, soft, hard, retrieval = evaluate(
                model, stream_train_eval.get_epoch_iterator(), p.distance_type,
                return_distance_matrix=save_distance_matrix)

            # average accuracy and distance matrix for testing data
            D_test, soft_test, hard_test, retrieval_test = evaluate(
                model, stream_test.get_epoch_iterator(), p.distance_type,
                return_distance_matrix=save_distance_matrix)

            time_end = time.time()
            epoch_time = time_end - time_begin
            total_time = time_end - time_origin

            logger.epoch = epoch
            logger.total_time = total_time
            logger.loss_log.append(loss_average)
            logger.train_log.append([soft[0], hard[0], retrieval[0]])
            logger.test_log.append(
                [soft_test[0], hard_test[0], retrieval_test[0]])

            # retain the model if it scored the best test acc. ever
            if soft_test[0] > logger.soft_test_best[0]:
                logger.model_best = copy.deepcopy(model)
                logger.optimizer_best = copy.deepcopy(optimizer)
                logger.epoch_best = epoch
                logger.D_best = D
                logger.D_test_best = D_test
                logger.soft_best = soft
                logger.soft_test_best = soft_test
                logger.hard_best = hard
                logger.hard_test_best = hard_test
                logger.retrieval_best = retrieval
                logger.retrieval_test_best = retrieval_test

            print("#", epoch)
            print("time: {} ({})".format(epoch_time, total_time))
            print("[train] loss:", loss_average)
            print("[train] soft:", soft)
            print("[train] hard:", hard)
            print("[train] retr:", retrieval)
            print("[test]  soft:", soft_test)
            print("[test]  hard:", hard_test)
            print("[test]  retr:", retrieval_test)
            print("[best]  soft: {} (at # {})".format(logger.soft_test_best,
                                                      logger.epoch_best))
            print(p)
            # print norms of the weights
            params = xp.hstack([xp.linalg.norm(param.data)
                                for param in model.params()]).tolist()
            print("|W|", map(lambda param: float('%0.2f' % param), params))
            print()

            # Draw plots
            if save_distance_matrix:
                plt.figure(figsize=(8, 4))
                plt.subplot(1, 2, 1)
                mat = plt.matshow(D, fignum=0, cmap=plt.cm.gray)
                plt.colorbar(mat, fraction=0.045)
                plt.subplot(1, 2, 2)
                mat = plt.matshow(D_test, fignum=0, cmap=plt.cm.gray)
                plt.colorbar(mat, fraction=0.045)
                plt.tight_layout()

            plt.figure(figsize=(8, 4))
            plt.subplot(1, 2, 1)
            plt.plot(logger.loss_log, label="tr-loss")
            plt.grid()
            plt.legend(loc='best')
            plt.subplot(1, 2, 2)
            plt.plot(logger.train_log)
            plt.plot(logger.test_log)
            plt.grid()
            plt.legend(["tr-soft", "tr-hard", "tr-retr",
                        "te-soft", "te-hard", "te-retr"],
                       bbox_to_anchor=(1.4, 1))
            plt.ylim([0.0, 1.0])
            plt.xlim([0, p.num_epochs])
            plt.tight_layout()
            plt.show()
            plt.draw()

            del D
            del D_test

    except KeyboardInterrupt:
        stop = True

    dir_name = "-".join([p.dataset, script_filename,
                         time.strftime("%Y%m%d%H%M%S"),
                         str(logger.soft_test_best[0])])

    logger.save(dir_name)
    p.save(dir_name)

    print("total epochs: {} ({} [s])".format(logger.epoch, logger.total_time))
    print("best test score (at # {})".format(logger.epoch_best))
    print("[test]  soft:", logger.soft_test_best)
    print("[test]  hard:", logger.hard_test_best)
    print("[test]  retr:", logger.retrieval_test_best)
    print(str(p).replace(', ', '\n'))
    print()

    return stop

def iterate_forward(model, epoch_iterator, normalize=False):
    xp = model.xp
    y_batches = []
    c_batches = []
    for batch in tqdm(copy.copy(epoch_iterator)):
        x_batch_data, c_batch_data = batch
        x_batch = Variable(xp.asarray(x_batch_data))
        y_batch = model(x_batch)
        if normalize:
            y_batch_data = y_batch.data / xp.linalg.norm(
                y_batch.data, axis=1, keepdims=True)
        else:
            y_batch_data = y_batch.data
        y_batches.append(y_batch_data)
        y_batch = None
        c_batches.append(c_batch_data)
    y_data = cuda.to_cpu(xp.concatenate(y_batches))
    c_data = np.concatenate(c_batches)
    return y_data, c_data

# memory friendly average accuracy for test data
def evaluate(model, epoch_iterator, distance='euclidean', normalize=False,
             batch_size=10, return_distance_matrix=False):
    if distance not in ('cosine', 'euclidean'):
        raise ValueError("distance must be 'euclidean' or 'cosine'.")

    with chainer.no_backprop_mode():
        with chainer.using_config('train', False):
            y_data, c_data = iterate_forward(
                    model, epoch_iterator, normalize=normalize)

    add_epsilon = True
    xp = cuda.get_array_module(y_data)
    num_examples = len(y_data)

    D_batches = []
    softs = []
    hards = []
    retrievals = []
    yy = xp.sum(y_data ** 2.0, axis=1)

    if distance == 'cosine':
        y_data = y_data / yy[:, None]  # L2 normalization

    for start in range(0, num_examples, batch_size):
        end = start + batch_size
        if end > num_examples:
            end = num_examples
        y_batch = y_data[start:end]
        yy_batch = yy[start:end]
        c_batch = c_data[start:end]

        D_batch = yy + yy_batch[:, None] - 2.0 * xp.dot(y_batch, y_data.T)
        xp.maximum(D_batch, 0, out=D_batch)
        if add_epsilon:
            D_batch += 1e-40
        # ensure the diagonal components are zero
        xp.fill_diagonal(D_batch[:, start:end], 0)

        soft, hard, retr = compute_soft_hard_retrieval(
            D_batch, c_data, c_batch)

        softs.append(len(y_batch) * soft)
        hards.append(len(y_batch) * hard)
        retrievals.append(len(y_batch) * retr)
        if return_distance_matrix:
            D_batches.append(D_batch)

    avg_softs = xp.sum(softs, axis=0) / num_examples
    avg_hards = xp.sum(hards, axis=0) / num_examples
    avg_retrievals = xp.sum(retrievals, axis=0) / num_examples

    if return_distance_matrix:
        D = cuda.to_cpu(xp.vstack(D_batches))
    else:
        D = None
    return D, avg_softs, avg_hards, avg_retrievals

def compute_soft_hard_retrieval(distance_matrix, labels, label_batch=None):
    softs = []
    hards = []
    retrievals = []

    if label_batch is None:
        label_batch = labels
    distance_matrix = cuda.to_cpu(distance_matrix)
    labels = cuda.to_cpu(labels)
    label_batch = cuda.to_cpu(label_batch)

    K = 11  # "K" for top-K
    for d_i, label_i in zip(distance_matrix, label_batch):
        top_k_indexes = np.argpartition(d_i, K)[:K]
        sorted_top_k_indexes = top_k_indexes[np.argsort(d_i[top_k_indexes])]
        ranked_labels = labels[sorted_top_k_indexes]
        # 0th entry is excluded since it is always 0
        ranked_hits = ranked_labels[1:] == label_i

        # soft top-k, k = 1, 2, 5, 10
        soft = [np.any(ranked_hits[:k]) for k in [1, 2, 5, 10]]
        softs.append(soft)
        # hard top-k, k = 2, 3, 4
        hard = [np.all(ranked_hits[:k]) for k in [2, 3, 4]]
        hards.append(hard)
        # retrieval top-k, k = 2, 3, 4
        retrieval = [np.mean(ranked_hits[:k]) for k in [2, 3, 4]]
        retrievals.append(retrieval)

    average_soft = np.array(softs).mean(axis=0)
    average_hard = np.array(hards).mean(axis=0)
    average_retrieval = np.array(retrievals).mean(axis=0)
    return average_soft, average_hard, average_retrieval
<commit_msg>Accept any optimizer<commit_after># -*- coding: utf-8 -*-
"""
Created on Thu Aug 24 15:40:25 2017

@author: sakurai
"""

import copy
import os

import matplotlib.pyplot as plt
import numpy as np
import six
import time

import chainer
from chainer import cuda
from chainer import Variable
from tqdm import tqdm

from . import utils
from ..datasets import data_provider
from ..models.modified_googlenet import ModifiedGoogLeNet

def train(main_script_path, func_train_one_batch, param_dict,
          save_distance_matrix=False,):
    script_filename = os.path.splitext(os.path.basename(main_script_path))[0]

    chainer.config.train = False
    device = 0
    xp = chainer.cuda.cupy
    config_parser = six.moves.configparser.ConfigParser()
    config_parser.read('config')
    log_dir_path = os.path.expanduser(config_parser.get('logs', 'dir_path'))

    p = utils.Logger(log_dir_path, **param_dict)  # hyperparameters

    ##########################################################
    # load database
    ##########################################################
    streams = data_provider.get_streams(p.batch_size, dataset=p.dataset,
                                        method=p.method, crop_size=p.crop_size)
    stream_train, stream_train_eval, stream_test = streams
    iter_train = stream_train.get_epoch_iterator()

    ##########################################################
    # construct the model
    ##########################################################
    model = ModifiedGoogLeNet(p.out_dim, p.normalize_output)
    if device >= 0:
        model.to_gpu()
    model.cleargrads()
    xp = model.xp
    optimizer_class = getattr(chainer.optimizers, p.optimizer)
    optimizer = optimizer_class(p.learning_rate)
    optimizer.setup(model)
    optimizer.add_hook(chainer.optimizer.WeightDecay(p.l2_weight_decay))

    print(p)
    stop = False
    logger = utils.Logger(log_dir_path)
    logger.soft_test_best = [0]
    time_origin = time.time()
    try:
        for epoch in range(p.num_epochs):
            time_begin = time.time()
            epoch_losses = []

            for i in tqdm(range(p.num_batches_per_epoch),
                          desc='# {}'.format(epoch)):
                with chainer.using_config('train', True):
                    loss = func_train_one_batch(model, p, next(iter_train))
                    loss.backward()
                optimizer.update()
                model.cleargrads()
                epoch_losses.append(loss.data)
                del loss

            loss_average = cuda.to_cpu(xp.array(
                xp.hstack(epoch_losses).mean()))

            # average accuracy and distance matrix for training data
            D, soft, hard, retrieval = evaluate(
                model, stream_train_eval.get_epoch_iterator(), p.distance_type,
                return_distance_matrix=save_distance_matrix)

            # average accuracy and distance matrix for testing data
            D_test, soft_test, hard_test, retrieval_test = evaluate(
                model, stream_test.get_epoch_iterator(), p.distance_type,
                return_distance_matrix=save_distance_matrix)

            time_end = time.time()
            epoch_time = time_end - time_begin
            total_time = time_end - time_origin

            logger.epoch = epoch
            logger.total_time = total_time
            logger.loss_log.append(loss_average)
            logger.train_log.append([soft[0], hard[0], retrieval[0]])
            logger.test_log.append(
                [soft_test[0], hard_test[0], retrieval_test[0]])

            # retain the model if it scored the best test acc. ever
            if soft_test[0] > logger.soft_test_best[0]:
                logger.model_best = copy.deepcopy(model)
                logger.optimizer_best = copy.deepcopy(optimizer)
                logger.epoch_best = epoch
                logger.D_best = D
                logger.D_test_best = D_test
                logger.soft_best = soft
                logger.soft_test_best = soft_test
                logger.hard_best = hard
                logger.hard_test_best = hard_test
                logger.retrieval_best = retrieval
                logger.retrieval_test_best = retrieval_test

            print("#", epoch)
            print("time: {} ({})".format(epoch_time, total_time))
            print("[train] loss:", loss_average)
            print("[train] soft:", soft)
            print("[train] hard:", hard)
            print("[train] retr:", retrieval)
            print("[test]  soft:", soft_test)
            print("[test]  hard:", hard_test)
            print("[test]  retr:", retrieval_test)
            print("[best]  soft: {} (at # {})".format(logger.soft_test_best,
                                                      logger.epoch_best))
            print(p)
            # print norms of the weights
            params = xp.hstack([xp.linalg.norm(param.data)
                                for param in model.params()]).tolist()
            print("|W|", map(lambda param: float('%0.2f' % param), params))
            print()

            # Draw plots
            if save_distance_matrix:
                plt.figure(figsize=(8, 4))
                plt.subplot(1, 2, 1)
                mat = plt.matshow(D, fignum=0, cmap=plt.cm.gray)
                plt.colorbar(mat, fraction=0.045)
                plt.subplot(1, 2, 2)
                mat = plt.matshow(D_test, fignum=0, cmap=plt.cm.gray)
                plt.colorbar(mat, fraction=0.045)
                plt.tight_layout()

            plt.figure(figsize=(8, 4))
            plt.subplot(1, 2, 1)
            plt.plot(logger.loss_log, label="tr-loss")
            plt.grid()
            plt.legend(loc='best')
            plt.subplot(1, 2, 2)
            plt.plot(logger.train_log)
            plt.plot(logger.test_log)
            plt.grid()
            plt.legend(["tr-soft", "tr-hard", "tr-retr",
                        "te-soft", "te-hard", "te-retr"],
                       bbox_to_anchor=(1.4, 1))
            plt.ylim([0.0, 1.0])
            plt.xlim([0, p.num_epochs])
            plt.tight_layout()
            plt.show()
            plt.draw()

            del D
            del D_test

    except KeyboardInterrupt:
        stop = True

    dir_name = "-".join([p.dataset, script_filename,
                         time.strftime("%Y%m%d%H%M%S"),
                         str(logger.soft_test_best[0])])

    logger.save(dir_name)
    p.save(dir_name)

    print("total epochs: {} ({} [s])".format(logger.epoch, logger.total_time))
    print("best test score (at # {})".format(logger.epoch_best))
    print("[test]  soft:", logger.soft_test_best)
    print("[test]  hard:", logger.hard_test_best)
    print("[test]  retr:", logger.retrieval_test_best)
    print(str(p).replace(', ', '\n'))
    print()

    return stop

def iterate_forward(model, epoch_iterator, normalize=False):
    xp = model.xp
    y_batches = []
    c_batches = []
    for batch in tqdm(copy.copy(epoch_iterator)):
        x_batch_data, c_batch_data = batch
        x_batch = Variable(xp.asarray(x_batch_data))
        y_batch = model(x_batch)
        if normalize:
            y_batch_data = y_batch.data / xp.linalg.norm(
                y_batch.data, axis=1, keepdims=True)
        else:
            y_batch_data = y_batch.data
        y_batches.append(y_batch_data)
        y_batch = None
        c_batches.append(c_batch_data)
    y_data = cuda.to_cpu(xp.concatenate(y_batches))
    c_data = np.concatenate(c_batches)
    return y_data, c_data

# memory friendly average accuracy for test data
def evaluate(model, epoch_iterator, distance='euclidean', normalize=False,
             batch_size=10, return_distance_matrix=False):
    if distance not in ('cosine', 'euclidean'):
        raise ValueError("distance must be 'euclidean' or 'cosine'.")

    with chainer.no_backprop_mode():
        with chainer.using_config('train', False):
            y_data, c_data = iterate_forward(
                    model, epoch_iterator, normalize=normalize)

    add_epsilon = True
    xp = cuda.get_array_module(y_data)
    num_examples = len(y_data)

    D_batches = []
    softs = []
    hards = []
    retrievals = []
    yy = xp.sum(y_data ** 2.0, axis=1)

    if distance == 'cosine':
        y_data = y_data / yy[:, None]  # L2 normalization

    for start in range(0, num_examples, batch_size):
        end = start + batch_size
        if end > num_examples:
            end = num_examples
        y_batch = y_data[start:end]
        yy_batch = yy[start:end]
        c_batch = c_data[start:end]

        D_batch = yy + yy_batch[:, None] - 2.0 * xp.dot(y_batch, y_data.T)
        xp.maximum(D_batch, 0, out=D_batch)
        if add_epsilon:
            D_batch += 1e-40
        # ensure the diagonal components are zero
        xp.fill_diagonal(D_batch[:, start:end], 0)

        soft, hard, retr = compute_soft_hard_retrieval(
            D_batch, c_data, c_batch)

        softs.append(len(y_batch) * soft)
        hards.append(len(y_batch) * hard)
        retrievals.append(len(y_batch) * retr)
        if return_distance_matrix:
            D_batches.append(D_batch)

    avg_softs = xp.sum(softs, axis=0) / num_examples
    avg_hards = xp.sum(hards, axis=0) / num_examples
    avg_retrievals = xp.sum(retrievals, axis=0) / num_examples

    if return_distance_matrix:
        D = cuda.to_cpu(xp.vstack(D_batches))
    else:
        D = None
    return D, avg_softs, avg_hards, avg_retrievals

def compute_soft_hard_retrieval(distance_matrix, labels, label_batch=None):
    softs = []
    hards = []
    retrievals = []

    if label_batch is None:
        label_batch = labels
    distance_matrix = cuda.to_cpu(distance_matrix)
    labels = cuda.to_cpu(labels)
    label_batch = cuda.to_cpu(label_batch)

    K = 11  # "K" for top-K
    for d_i, label_i in zip(distance_matrix, label_batch):
        top_k_indexes = np.argpartition(d_i, K)[:K]
        sorted_top_k_indexes = top_k_indexes[np.argsort(d_i[top_k_indexes])]
        ranked_labels = labels[sorted_top_k_indexes]
        # 0th entry is excluded since it is always 0
        ranked_hits = ranked_labels[1:] == label_i

        # soft top-k, k = 1, 2, 5, 10
        soft = [np.any(ranked_hits[:k]) for k in [1, 2, 5, 10]]
        softs.append(soft)
        # hard top-k, k = 2, 3, 4
        hard = [np.all(ranked_hits[:k]) for k in [2, 3, 4]]
        hards.append(hard)
        # retrieval top-k, k = 2, 3, 4
        retrieval = [np.mean(ranked_hits[:k]) for k in [2, 3, 4]]
        retrievals.append(retrieval)

    average_soft = np.array(softs).mean(axis=0)
    average_hard = np.array(hards).mean(axis=0)
    average_retrieval = np.array(retrievals).mean(axis=0)
    return average_soft, average_hard, average_retrieval
<|endoftext|>


<commit_before>import os
import time
import ping
import threading
import ConfigParser
import paho.mqtt.client as mqtt

class Pinger(threading.Thread):

    def _readConfig(self):
        update = False

        if not os.path.isdir(self._homeDir):
            print "Creating homeDir"
            os.makedirs(self._homeDir)

        if os.path.isfile(self._configFileName):
            self._config.read(self._configFileName)
        else:
            print "Config file not found"
            update = True

        if not self._config.has_section('MQTT'):
            print "Adding MQTT part"
            update = True
            self._config.add_section("MQTT")

        if not self._config.has_option("MQTT", "ServerAddress"):
            print "No Server Address"
            update = True
            self._config.set("MQTT", "ServerAddress", "<ServerAddress>")

        if not self._config.has_option("MQTT", "ServerPort"):
            print "No Server Port"
            update = True
            self._config.set("MQTT", "ServerPort", "1883")

        if not self._config.has_section('PINGER'):
            print "Adding Pinger part"
            update = True
            self._config.add_section("PINGER")

        if not self._config.has_option("PINGER", "AnsiIP"):
            print "No AnsiAddress Address"
            update = True
            self._config.set("PINGER", "AnsiIP", "192.168.1.10")

        if not self._config.has_option("PINGER", "TiffyIP"):
            print "No Tiffy IP Address"
            update = True
            self._config.set("PINGER", "TiffyIP", "192.168.1.20")

        if update:
            with open(self._configFileName, 'w') as f:
                self._config.write(f)

    def __init__(self):
        threading.Thread.__init__(self)
        self.setDaemon(True)

        self._homeDir        = os.path.expanduser("~/.sensomatic")
        self._configFileName = self._homeDir + '/config.ini'
        self._config         = ConfigParser.ConfigParser()
        self._readConfig()

        self._mqclient = mqtt.Client("Pinger", clean_session=True)

    def run(self):
        self._mqclient.connect(self._config.get("MQTT","ServerAddress"), self._config.get("MQTT","ServerPort"), 60)
        self._mqclient.loop_start()
        while True:

            a = ping.quiet_ping('192.168.1.103', timeout=1, count=1)
            if a[0] == 0:
                print "Ansi ist da"
                self._mqclient.publish("ansi/wlanPresents","True")

            t = ping.quiet_ping('192.168.1.112', timeout=1, count=1)
            if t[0] == 0:
                print "Tiffy ist da"
                self._mqclient.publish("tiffy/wlanPresents","True")

            time.sleep(30)

if __name__ == "__main__":

    p = Pinger()
    p.start()

    time.sleep(100)
<commit_msg>Config file<commit_after>import os
import time
import ping
import threading
import ConfigParser
import paho.mqtt.client as mqtt

class Pinger(threading.Thread):

    def _readConfig(self):
        update = False

        if not os.path.isdir(self._homeDir):
            print "Creating homeDir"
            os.makedirs(self._homeDir)

        if os.path.isfile(self._configFileName):
            self._config.read(self._configFileName)
        else:
            print "Config file not found"
            update = True

        if not self._config.has_section('MQTT'):
            print "Adding MQTT part"
            update = True
            self._config.add_section("MQTT")

        if not self._config.has_option("MQTT", "ServerAddress"):
            print "No Server Address"
            update = True
            self._config.set("MQTT", "ServerAddress", "<ServerAddress>")

        if not self._config.has_option("MQTT", "ServerPort"):
            print "No Server Port"
            update = True
            self._config.set("MQTT", "ServerPort", "1883")

        if not self._config.has_section('PINGER'):
            print "Adding Pinger part"
            update = True
            self._config.add_section("PINGER")

        if not self._config.has_option("PINGER", "AnsiIP"):
            print "No AnsiAddress Address"
            update = True
            self._config.set("PINGER", "AnsiIP", "192.168.1.10")

        if not self._config.has_option("PINGER", "TiffyIP"):
            print "No Tiffy IP Address"
            update = True
            self._config.set("PINGER", "TiffyIP", "192.168.1.20")

        if update:
            with open(self._configFileName, 'w') as f:
                self._config.write(f)

    def __init__(self):
        threading.Thread.__init__(self)
        self.setDaemon(True)

        self._homeDir        = os.path.expanduser("~/.sensomatic")
        self._configFileName = self._homeDir + '/config.ini'
        self._config         = ConfigParser.ConfigParser()
        self._readConfig()

        self._mqclient = mqtt.Client("Pinger", clean_session=True)

    def run(self):
        self._mqclient.connect(self._config.get("MQTT","ServerAddress"), self._config.get("MQTT","ServerPort"), 60)
        self._mqclient.loop_start()
        while True:

            a = ping.quiet_ping(self._config.get("PINGER","AnsiIP"), timeout=1, count=1)
            if a[0] == 0:
                print "Ansi ist da"
                self._mqclient.publish("ansi/wlanPresents","True")

            t = ping.quiet_ping(self._config.get("PINGER","TiffyIP"), timeout=1, count=1)
            if t[0] == 0:
                print "Tiffy ist da"
                self._mqclient.publish("tiffy/wlanPresents","True")

            time.sleep(30)

if __name__ == "__main__":

    p = Pinger()
    p.start()

    time.sleep(100)
<|endoftext|>


<commit_before>import os
import time
import ping
import threading
import ConfigParser
import paho.mqtt.client as mqtt

class Pinger(threading.Thread):

    def _readConfig(self):
        update = False

        if not os.path.isdir(self._homeDir):
            print "Creating homeDir"
            os.makedirs(self._homeDir)

        if os.path.isfile(self._configFileName):
            self._config.read(self._configFileName)
        else:
            print "Config file not found"
            update = True

        if not self._config.has_section('MQTT'):
            print "Adding MQTT part"
            update = True
            self._config.add_section("MQTT")

        if not self._config.has_option("MQTT", "ServerAddress"):
            print "No Server Address"
            update = True
            self._config.set("MQTT", "ServerAddress", "<ServerAddress>")

        if not self._config.has_option("MQTT", "ServerPort"):
            print "No Server Port"
            update = True
            self._config.set("MQTT", "ServerPort", "1883")

        if not self._config.has_section('PINGER'):
            print "Adding Pinger part"
            update = True
            self._config.add_section("PINGER")

        if not self._config.has_option("PINGER", "AnsiIP"):
            print "No AnsiAddress Address"
            update = True
            self._config.set("PINGER", "AnsiIP", "192.168.1.10")

        if not self._config.has_option("PINGER", "TiffyIP"):
            print "No Tiffy IP Address"
            update = True
            self._config.set("PINGER", "TiffyIP", "192.168.1.20")

        if update:
            with open(self._configFileName, 'w') as f:
                self._config.write(f)

    def __init__(self):
        threading.Thread.__init__(self)
        self.setDaemon(True)

        self._homeDir        = os.path.expanduser("~/.sensomatic")
        self._configFileName = self._homeDir + '/config.ini'
        self._config         = ConfigParser.ConfigParser()
        self._readConfig()

        self._mqclient = mqtt.Client("Pinger", clean_session=True)

    def run(self):
        self._mqclient.connect(self._config.get("MQTT","ServerAddress"), self._config.get("MQTT","ServerPort"), 60)
        self._mqclient.loop_start()
        while True:

            a = ping.quiet_ping(self._config.get("PINGER","AnsiIP"), timeout=1, count=1)
            if a[0] == 0:
                print "Ansi ist da"
                self._mqclient.publish("ansi/wlanPresents","True")

            t = ping.quiet_ping(self._config.get("PINGER","TiffyIP"), timeout=1, count=1)
            if t[0] == 0:
                print "Tiffy ist da"
                self._mqclient.publish("tiffy/wlanPresents","True")

            time.sleep(30)

if __name__ == "__main__":

    p = Pinger()
    p.start()

    time.sleep(100)
<commit_msg>Replace True with 1<commit_after>import os
import time
import ping
import threading
import ConfigParser
import paho.mqtt.client as mqtt

class Pinger(threading.Thread):

    def _readConfig(self):
        update = False

        if not os.path.isdir(self._homeDir):
            print "Creating homeDir"
            os.makedirs(self._homeDir)

        if os.path.isfile(self._configFileName):
            self._config.read(self._configFileName)
        else:
            print "Config file not found"
            update = True

        if not self._config.has_section('MQTT'):
            print "Adding MQTT part"
            update = True
            self._config.add_section("MQTT")

        if not self._config.has_option("MQTT", "ServerAddress"):
            print "No Server Address"
            update = True
            self._config.set("MQTT", "ServerAddress", "<ServerAddress>")

        if not self._config.has_option("MQTT", "ServerPort"):
            print "No Server Port"
            update = True
            self._config.set("MQTT", "ServerPort", "1883")

        if not self._config.has_section('PINGER'):
            print "Adding Pinger part"
            update = True
            self._config.add_section("PINGER")

        if not self._config.has_option("PINGER", "AnsiIP"):
            print "No AnsiAddress Address"
            update = True
            self._config.set("PINGER", "AnsiIP", "192.168.1.10")

        if not self._config.has_option("PINGER", "TiffyIP"):
            print "No Tiffy IP Address"
            update = True
            self._config.set("PINGER", "TiffyIP", "192.168.1.20")

        if update:
            with open(self._configFileName, 'w') as f:
                self._config.write(f)

    def __init__(self):
        threading.Thread.__init__(self)
        self.setDaemon(True)

        self._homeDir        = os.path.expanduser("~/.sensomatic")
        self._configFileName = self._homeDir + '/config.ini'
        self._config         = ConfigParser.ConfigParser()
        self._readConfig()

        self._mqclient = mqtt.Client("Pinger", clean_session=True)

    def run(self):
        self._mqclient.connect(self._config.get("MQTT", "ServerAddress"), self._config.get("MQTT", "ServerPort"), 60)
        self._mqclient.loop_start()
        while True:

            a = ping.quiet_ping(self._config.get("PINGER", "AnsiIP"), timeout=1, count=1)
            if a[0] == 0:
                print "Ansi ist da"
                self._mqclient.publish("ansi/wlanPresents", "1")

            t = ping.quiet_ping(self._config.get("PINGER", "TiffyIP"), timeout=1, count=1)
            if t[0] == 0:
                print "Tiffy ist da"
                self._mqclient.publish("tiffy/wlanPresents", "1")

            time.sleep(30)

if __name__ == "__main__":

    p = Pinger()
    p.start()

    time.sleep(100)
<|endoftext|>


<commit_before>#!/usr/bin/python -O
import threading
import logging
import sqlite3
from datetime import datetime, timedelta, date

g_db = {}
g_params = {}

# This is a way to get the column names after grabbing everything
# I guess it's also good practice
SCHEMA = {
  'intents': [
     ('id', 'INTEGER PRIMARY KEY'),
     ('key', 'TEXT UNIQUE'),
     ('start', 'INTEGER'),
     ('end', 'INTEGER'),
     ('read_count', 'INTEGER DEFAULT 0'),
     ('created_at', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
     ('accessed_at', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
   ],
  'kv': [
     ('id', 'INTEGER PRIMARY KEY'),
     ('key', 'TEXT UNIQUE'),
     ('value', 'TEXT'),
     ('created_at', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
   ],
  'streams': [
     ('id', 'INTEGER PRIMARY KEY'),
     ('name', 'TEXT UNIQUE'),
     ('start_unix', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
     ('end_unix', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
     ('start_minute', 'REAL DEFAULT 0'),
     ('end_minute', 'REAL DEFAULT 0'),
     ('week_number', 'INTEGER DEFAULT 0'),
     ('size', 'INTEGER DEFAULT 0'),
   ]
}

# Ok so if column order or type changes, this isn't found ... nor
# are we doing formal migrations where you can roll back or whatever
# because way too fancy ...
def upgrade():
  my_set = __builtins__['set']
  db = connect()

  for table, schema in SCHEMA.items():
    existing_schema = db['c'].execute('pragma table_info(%s)' % table).fetchall()
    existing_column_names = [str(row[1]) for row in existing_schema]

    our_column_names = [row[0] for row in schema]

    # print table, existing_column_names, our_column_names

    to_add = my_set(our_column_names).difference(my_set(existing_column_names))

    # These are the things we should add ... this can be an empty set, that's fine.
    for key in to_add:
      #
      # sqlite doesn't support adding things into positional places (add column after X)
      # they just get tacked on at the end ... which is fine - you'd have to rebuild
      # everything to achieve positional columns - that's not worth it - we just always
      # tack on at the end as a policy in our schema and we'll be fine.
      #
      # However, given all of that, we still need the schema
      #
      our_schema = schema[our_column_names.index(key)][1]
      # print 'alter table %s add column %s %s' % (table, key, our_schema)
      db['c'].execute('alter table %s add column %s %s' % (table, key, our_schema))
      db['conn'].commit()

    to_remove = my_set(existing_column_names).difference(my_set(our_column_names))

    if len(to_remove) > 0:
      our_schema = ','.join(["%s %s" % (key, klass) for key, klass in schema])
      our_columns = ','.join(our_column_names)

      drop_column_sql = """
      CREATE TEMPORARY TABLE my_backup(%s);
      INSERT INTO my_backup SELECT %s FROM %s;
      DROP TABLE %s;
      CREATE TABLE %s(%s);
      INSERT INTO %s SELECT %s FROM my_backup;
      DROP TABLE my_backup;
      """ % (our_schema, our_columns, table, table,    table, our_schema, table, our_columns)

      for sql_line in drop_column_sql.strip().split('\n'):
        db['c'].execute(sql_line)

      db['conn'].commit()

def map(row_list, table):
  """
  Using the schema of a table, map the row_list to a list of dicts.
  """
  mapped = []
  my_schema = schema(table)

  if not row_list:
    return row_list

  if type(row_list[0]) is str:
    row_list = [row_list]

  for row in row_list:
    mapped_row = {}
    for ix in range(len(my_schema)):
      mapped_row[my_schema[ix]] = row[ix]

    mapped.append(mapped_row)

  return mapped

def all(table, field_list='*', sort_by='id'):
  """ Returns all entries from the sqlite3 database for a given table. """
  db = connect()

  column_count = 1
  if type(field_list) is not str:
    column_count = len(field_list)
    field_list = ','.join(field_list)

  query = db['c'].execute('select %s from %s order by %s asc' % (field_list, table, sort_by))
  if column_count is 1 and field_list != '*':
    return [record[0] for record in query.fetchall()]

  else:
    return [record for record in query.fetchall()]

def schema(table):
  """ Returns the schema for a given table. """
  if table not in SCHEMA:
    return None

  return [ key for key, value in SCHEMA[table] ]

def connect():
  """
  A "singleton pattern" or some other fancy $10-world style of maintaining
  the database connection throughout the execution of the script.

  Returns the database instance.
  """
  global g_db

  #
  # We need to have one instance per thread, as this is what
  # sqlite's driver dictates ... so we do this based on thread id.
  #
  # We don't have to worry about the different memory sharing models here.
  # Really, just think about it ... it's totally irrelevant.
  #
  thread_id = threading.current_thread().ident
  if thread_id not in g_db:
    g_db[thread_id] = {}

  instance = g_db[thread_id]

  if 'conn' not in instance:
    conn = sqlite3.connect('config.db')
    instance['conn'] = conn
    instance['c'] = conn.cursor()

    for table, schema in SCHEMA.items():
      dfn = ','.join(["%s %s" % (key, klass) for key, klass in schema])
      instance['c'].execute("CREATE TABLE IF NOT EXISTS %s(%s)" % (table, dfn))

    instance['conn'].commit()

  return instance

def shutdown():
  """ Closes the individual database connections in each thread. """
  global g_db

  for instance in g_db.items():
    if 'conn' in instance:
      instance['conn'].close()

def incr(key, value=1):
  """
  Increments some key in the database by some value.  It is used
  to maintain statistical counters.
  """
  db = connect()

  try:
    db['c'].execute('insert into kv(value, key) values(?, ?)', (value, key, ))

  except Exception as exc:
    db['c'].execute('update kv set value = value + ? where key = ?', (value, key, ))

  db['conn'].commit()

def set(key, value):
  """
  Sets (or replaces) a given key to a specific value.

  Returns the value that was sent.
  """
  global g_params

  db = connect()

  # From http://stackoverflow.com/questions/418898/sqlite-upsert-not-insert-or-replace
  res = db['c'].execute('''
    INSERT OR REPLACE INTO kv (key, value, created_at)
      VALUES (
        COALESCE((SELECT key FROM kv WHERE key = ?), ?),
        ?,
        current_timestamp
    )''', (key, key, value, ))

  db['conn'].commit()

  g_params[key] = value

  return value

def get(key, expiry=0, use_cache=True):
  """
  Retrieves a value from the database, tentative on the expiry.
  If the cache is set to true then it retrieves it from in-memory if available, otherwise
  it goes out to the db. Other than directly hitting up the g_params parameter which is
  used internally, there is no way to invalidate the cache.
  """
  global g_params

  if use_cache and key in g_params:
    return g_params[key]

  db = connect()

  if expiry > 0:
    # If we let things expire, we first sweep for it
    db['c'].execute('delete from kv where key = ? and created_at < (current_timestamp - ?)', (key, expiry, ))
    db['conn'].commit()

  res = db['c'].execute('select value, created_at from kv where key = ?', (key, )).fetchone()

  if res:
    g_params[key] = res[0]
    return res[0]

  return None

def unregister_stream(name, do_all=False):
  """ Deletes a stream by name, contingent on it existing only once """

  db = connect()
  res = db['c'].execute('select id from streams where name = ?', (name, )).fetchall()

  if res and (len(res) == 1 or do_all):
    logging.info("Removing our reference of %s" % name)
    res = db['c'].execute('delete from streams where id = %d' % res[0][0])
    db['conn'].commit()

  return True

def register_stream(info):
  """
  Registers a stream as existing to be found later when trying to stitch and slice files together.
  This is all that ought to be needed to know if the streams should attempt to be stitched. The input
  info is grabbed from audio.stream_info(filename)
  """

  name = info['name']
  start_unix = info['start_date']
  end_unix = info['start_date'] + timedelta(seconds=info['duration_sec'])
  start_minute = float(info['start_minute'])
  end_minute = float(info['end_minute'])
  week_number = info['week_number']
  size = info['size']

  db = connect()
  res = db['c'].execute('select id from streams where name = ?', (name, )).fetchone()

  # If something exists then we remove it and reinsert ... this is not as effecient
  # as an upsert with coalesce, but this is done on rare occasions in heavy workloads ...
  # so we don't really need to be entirely efficient about it.
  if res:
    unregister_stream(name, do_all=True)

  db['c'].execute("""insert into streams
  (name, start_unix, end_unix, start_minute, end_minute, week_number, size) values
  (   ?,          ?,        ?,            ?,          ?,           ?,    ?) """,
  (name, start_unix, end_unix, start_minute, end_minute, week_number, size))

  db['conn'].commit()
  return db['c'].lastrowid

  return res[0]

def register_intent(minute_list, duration):
  """
  Tells the server to record on a specific minute for a specific duration when
  not in full mode.  Otherwise, this is just here for statistical purposes.
  """
  db = connect()

  for minute in minute_list:
    key = str(minute) + ':' + str(duration)
    res = db['c'].execute('select id from intents where key = ?', (key, )).fetchone()

    if res == None:
      db['c'].execute('insert into intents(key, start, end) values(?, ?, ?)', (key, minute, minute + duration, ))

    else:
      db['c'].execute('update intents set read_count = read_count + 1, accessed_at = (current_timestamp) where id = ?', (res[0], ))

    db['conn'].commit()
    return db['c'].lastrowid

  return None

<commit_msg>Permitting other dbs to be connected<commit_after>#!/usr/bin/python -O
import threading
import logging
import sqlite3
import os
from datetime import datetime, timedelta, date

g_db = {}
g_params = {}

# This is a way to get the column names after grabbing everything
# I guess it's also good practice
SCHEMA = {
  'intents': [
     ('id', 'INTEGER PRIMARY KEY'),
     ('key', 'TEXT UNIQUE'),
     ('start', 'INTEGER'),
     ('end', 'INTEGER'),
     ('read_count', 'INTEGER DEFAULT 0'),
     ('created_at', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
     ('accessed_at', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
   ],
  'kv': [
     ('id', 'INTEGER PRIMARY KEY'),
     ('key', 'TEXT UNIQUE'),
     ('value', 'TEXT'),
     ('created_at', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
   ],
  'streams': [
     ('id', 'INTEGER PRIMARY KEY'),
     ('name', 'TEXT UNIQUE'),
     ('start_unix', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
     ('end_unix', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP'),
     ('start_minute', 'REAL DEFAULT 0'),
     ('end_minute', 'REAL DEFAULT 0'),
     ('week_number', 'INTEGER DEFAULT 0'),
     ('size', 'INTEGER DEFAULT 0'),
   ]
}

# Ok so if column order or type changes, this isn't found ... nor
# are we doing formal migrations where you can roll back or whatever
# because way too fancy ...
def upgrade():
  my_set = __builtins__['set']
  db = connect()

  for table, schema in SCHEMA.items():
    existing_schema = db['c'].execute('pragma table_info(%s)' % table).fetchall()
    existing_column_names = [str(row[1]) for row in existing_schema]

    our_column_names = [row[0] for row in schema]

    # print table, existing_column_names, our_column_names

    to_add = my_set(our_column_names).difference(my_set(existing_column_names))

    # These are the things we should add ... this can be an empty set, that's fine.
    for key in to_add:
      #
      # sqlite doesn't support adding things into positional places (add column after X)
      # they just get tacked on at the end ... which is fine - you'd have to rebuild
      # everything to achieve positional columns - that's not worth it - we just always
      # tack on at the end as a policy in our schema and we'll be fine.
      #
      # However, given all of that, we still need the schema
      #
      our_schema = schema[our_column_names.index(key)][1]
      # print 'alter table %s add column %s %s' % (table, key, our_schema)
      db['c'].execute('alter table %s add column %s %s' % (table, key, our_schema))
      db['conn'].commit()

    to_remove = my_set(existing_column_names).difference(my_set(our_column_names))

    if len(to_remove) > 0:
      our_schema = ','.join(["%s %s" % (key, klass) for key, klass in schema])
      our_columns = ','.join(our_column_names)

      drop_column_sql = """
      CREATE TEMPORARY TABLE my_backup(%s);
      INSERT INTO my_backup SELECT %s FROM %s;
      DROP TABLE %s;
      CREATE TABLE %s(%s);
      INSERT INTO %s SELECT %s FROM my_backup;
      DROP TABLE my_backup;
      """ % (our_schema, our_columns, table, table,    table, our_schema, table, our_columns)

      for sql_line in drop_column_sql.strip().split('\n'):
        db['c'].execute(sql_line)

      db['conn'].commit()

def map(row_list, table):
  """
  Using the schema of a table, map the row_list to a list of dicts.
  """
  mapped = []
  my_schema = schema(table)

  if not row_list:
    return row_list

  if type(row_list[0]) is str:
    row_list = [row_list]

  for row in row_list:
    mapped_row = {}
    for ix in range(len(my_schema)):
      mapped_row[my_schema[ix]] = row[ix]

    mapped.append(mapped_row)

  return mapped

def all(table, field_list='*', sort_by='id'):
  """ Returns all entries from the sqlite3 database for a given table. """
  db = connect()

  column_count = 1
  if type(field_list) is not str:
    column_count = len(field_list)
    field_list = ','.join(field_list)

  query = db['c'].execute('select %s from %s order by %s asc' % (field_list, table, sort_by))
  if column_count is 1 and field_list != '*':
    return [record[0] for record in query.fetchall()]

  else:
    return [record for record in query.fetchall()]

def schema(table):
  """ Returns the schema for a given table. """
  if table not in SCHEMA:
    return None

  return [ key for key, value in SCHEMA[table] ]

def connect(db_file='config.db'):
  """
  A "singleton pattern" or some other fancy $10-world style of maintaining
  the database connection throughout the execution of the script.

  Returns the database instance.
  """
  global g_db

  #
  # We need to have one instance per thread, as this is what
  # sqlite's driver dictates ... so we do this based on thread id.
  #
  # We don't have to worry about the different memory sharing models here.
  # Really, just think about it ... it's totally irrelevant.
  #
  thread_id = threading.current_thread().ident
  if thread_id not in g_db:
    g_db[thread_id] = {}

  instance = g_db[thread_id]

  if 'conn' not in instance:

    if not os.path.exists(db_file):
      print "Info: Creating db file %s" % db_file

    conn = sqlite3.connect(db_file)
    instance['conn'] = conn
    instance['c'] = conn.cursor()

    for table, schema in SCHEMA.items():
      dfn = ','.join(["%s %s" % (key, klass) for key, klass in schema])
      instance['c'].execute("CREATE TABLE IF NOT EXISTS %s(%s)" % (table, dfn))

    instance['conn'].commit()

  return instance

def shutdown():
  """ Closes the individual database connections in each thread. """
  global g_db

  for instance in g_db.items():
    if 'conn' in instance:
      instance['conn'].close()

def incr(key, value=1):
  """
  Increments some key in the database by some value.  It is used
  to maintain statistical counters.
  """
  db = connect()

  try:
    db['c'].execute('insert into kv(value, key) values(?, ?)', (value, key, ))

  except Exception as exc:
    db['c'].execute('update kv set value = value + ? where key = ?', (value, key, ))

  db['conn'].commit()

def set(key, value):
  """
  Sets (or replaces) a given key to a specific value.

  Returns the value that was sent.
  """
  global g_params

  db = connect()

  # From http://stackoverflow.com/questions/418898/sqlite-upsert-not-insert-or-replace
  res = db['c'].execute('''
    INSERT OR REPLACE INTO kv (key, value, created_at)
      VALUES (
        COALESCE((SELECT key FROM kv WHERE key = ?), ?),
        ?,
        current_timestamp
    )''', (key, key, value, ))

  db['conn'].commit()

  g_params[key] = value

  return value

def get(key, expiry=0, use_cache=True):
  """
  Retrieves a value from the database, tentative on the expiry.
  If the cache is set to true then it retrieves it from in-memory if available, otherwise
  it goes out to the db. Other than directly hitting up the g_params parameter which is
  used internally, there is no way to invalidate the cache.
  """
  global g_params

  if use_cache and key in g_params:
    return g_params[key]

  db = connect()

  if expiry > 0:
    # If we let things expire, we first sweep for it
    db['c'].execute('delete from kv where key = ? and created_at < (current_timestamp - ?)', (key, expiry, ))
    db['conn'].commit()

  res = db['c'].execute('select value, created_at from kv where key = ?', (key, )).fetchone()

  if res:
    g_params[key] = res[0]
    return res[0]

  return None

def unregister_stream(name, do_all=False):
  """ Deletes a stream by name, contingent on it existing only once """

  db = connect()
  res = db['c'].execute('select id from streams where name = ?', (name, )).fetchall()

  if res and (len(res) == 1 or do_all):
    logging.info("Removing our reference of %s" % name)
    res = db['c'].execute('delete from streams where id = %d' % res[0][0])
    db['conn'].commit()

  return True

def register_stream(info):
  """
  Registers a stream as existing to be found later when trying to stitch and slice files together.
  This is all that ought to be needed to know if the streams should attempt to be stitched. The input
  info is grabbed from audio.stream_info(filename)
  """

  name = info['name']
  start_unix = info['start_date']
  end_unix = info['start_date'] + timedelta(seconds=info['duration_sec'])
  start_minute = float(info['start_minute'])
  end_minute = float(info['end_minute'])
  week_number = info['week_number']
  size = info['size']

  db = connect()
  res = db['c'].execute('select id from streams where name = ?', (name, )).fetchone()

  # If something exists then we remove it and reinsert ... this is not as effecient
  # as an upsert with coalesce, but this is done on rare occasions in heavy workloads ...
  # so we don't really need to be entirely efficient about it.
  if res:
    unregister_stream(name, do_all=True)

  db['c'].execute("""insert into streams
  (name, start_unix, end_unix, start_minute, end_minute, week_number, size) values
  (   ?,          ?,        ?,            ?,          ?,           ?,    ?) """,
  (name, start_unix, end_unix, start_minute, end_minute, week_number, size))

  db['conn'].commit()
  return db['c'].lastrowid

  return res[0]

def register_intent(minute_list, duration):
  """
  Tells the server to record on a specific minute for a specific duration when
  not in full mode.  Otherwise, this is just here for statistical purposes.
  """
  db = connect()

  for minute in minute_list:
    key = str(minute) + ':' + str(duration)
    res = db['c'].execute('select id from intents where key = ?', (key, )).fetchone()

    if res == None:
      db['c'].execute('insert into intents(key, start, end) values(?, ?, ?)', (key, minute, minute + duration, ))

    else:
      db['c'].execute('update intents set read_count = read_count + 1, accessed_at = (current_timestamp) where id = ?', (res[0], ))

    db['conn'].commit()
    return db['c'].lastrowid

  return None

<|endoftext|>


<commit_before>import base64
import bz2
import plistlib
import pytz
import random
import string
from datetime import datetime
from xml.parsers.expat import ExpatError

from dateutil.parser import *

from django.contrib.auth.models import User
from django.db import models

from server import text_utils

OS_CHOICES = (
    ('Darwin', 'macOS'),
    ('Windows', 'Windows'),
    ('Linux', 'Linux'),
    ('ChromeOS', 'Chrome OS'),
)

REPORT_CHOICES = (
    ('base64', 'base64'),
    ('base64bz2', 'base64bz2'),
    ('bz2', 'bz2'),
)

class ProfileLevel():
    stats_only = 'SO'
    read_only = 'RO'
    read_write = 'RW'
    global_admin = 'GA'

def GenerateKey():
    key = ''.join(random.choice(string.ascii_lowercase + string.digits) for x in range(128))
    try:
        MachineGroup.objects.get(key=key)
        return GenerateKey()
    except MachineGroup.DoesNotExist:
        return key

def GenerateAPIKey():
    key = ''.join(random.choice(string.ascii_lowercase + string.digits) for x in range(24))
    try:
        ApiKey.objects.get(public_key=key)
        return GenerateAPIKey()
    except ApiKey.DoesNotExist:
        return key

class UserProfile(models.Model):
    user = models.OneToOneField(User, unique=True)

    LEVEL_CHOICES = (
        ('SO', 'Stats Only'),
        ('RO', 'Read Only'),
        ('RW', 'Read Write'),
        ('GA', 'Global Admin'),
    )
    level = models.CharField(max_length=2, choices=LEVEL_CHOICES, default='RO')

    def __unicode__(self):
        return self.user.username

User.userprofile = property(lambda u: UserProfile.objects.get_or_create(user=u)[0])

class BusinessUnit(models.Model):
    name = models.CharField(max_length=100)
    users = models.ManyToManyField(User, blank=True)

    def __unicode__(self):
        return self.name

    @classmethod
    def display_name(cls):
        return text_utils.class_to_title(cls.__name__)

    class Meta:
        ordering = ['name']

class MachineGroup(models.Model):
    business_unit = models.ForeignKey(BusinessUnit)
    name = models.CharField(max_length=100)
    key = models.CharField(db_index=True, max_length=255, unique=True,
                           blank=True, null=True, editable=False)

    def save(self, **kwargs):
        if not self.id:
            self.key = GenerateKey()
        super(MachineGroup, self).save()

    def __unicode__(self):
        return self.name

    @classmethod
    def display_name(cls):
        return text_utils.class_to_title(cls.__name__)

    class Meta:
        ordering = ['name']

class DeployedManager(models.Manager):
    def get_queryset(self):
        return super(DeployedManager, self).get_queryset().filter(deployed=True)

class Machine(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine_group = models.ForeignKey(MachineGroup)
    serial = models.CharField(db_index=True, max_length=100, unique=True)
    hostname = models.CharField(max_length=256, null=True, blank=True)
    operating_system = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    memory = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    memory_kb = models.IntegerField(db_index=True, default=0)
    munki_version = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    manifest = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    hd_space = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    hd_total = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    hd_percent = models.CharField(max_length=256, null=True, blank=True)
    console_user = models.CharField(max_length=256, null=True, blank=True)
    machine_model = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    machine_model_friendly = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    cpu_type = models.CharField(max_length=256, null=True, blank=True)
    cpu_speed = models.CharField(max_length=256, null=True, blank=True)
    os_family = models.CharField(db_index=True, max_length=256,
                                 choices=OS_CHOICES, verbose_name="OS Family", default="Darwin")
    last_checkin = models.DateTimeField(db_index=True, blank=True, null=True)
    first_checkin = models.DateTimeField(db_index=True, blank=True, null=True, auto_now_add=True)
    report = models.TextField(editable=True, null=True)
    errors = models.IntegerField(default=0)
    warnings = models.IntegerField(default=0)
    activity = models.BooleanField(editable=True, default=False)
    puppet_version = models.CharField(db_index=True, null=True, blank=True, max_length=256)
    sal_version = models.CharField(db_index=True, null=True, blank=True, max_length=255)
    last_puppet_run = models.DateTimeField(db_index=True, blank=True, null=True)
    puppet_errors = models.IntegerField(db_index=True, default=0)
    deployed = models.BooleanField(default=True)
    broken_client = models.BooleanField(default=False)

    objects = models.Manager()  # The default manager.
    deployed_objects = DeployedManager()

    def get_fields(self):
        return [(field.name, field.value_to_string(self)) for field in Machine._meta.fields]

    def get_report(self):
        try:
            return plistlib.loads(self.report)
        except TypeError:
            return {}

    def __unicode__(self):
        if self.hostname:
            return self.hostname
        else:
            return self.serial

    @classmethod
    def display_name(cls):
        return text_utils.class_to_title(cls.__name__)

    class Meta:
        ordering = ['hostname']

GROUP_NAMES = {
    'all': None,
    'machine_group': MachineGroup,
    'business_unit': BusinessUnit,
    'machine': Machine}

class UpdateHistory(models.Model):
    id = models.BigAutoField(primary_key=True)
    UPDATE_TYPE = (
        ('third_party', '3rd Party'),
        ('apple', 'Apple'),
    )
    machine = models.ForeignKey(Machine)
    update_type = models.CharField(max_length=254, choices=UPDATE_TYPE, verbose_name="Update Type")
    name = models.CharField(max_length=255, db_index=True)
    version = models.CharField(max_length=254, db_index=True)

    def __unicode__(self):
        return "%s: %s %s" % (self.machine, self.name, self.version)

    class Meta:
        ordering = ['name']
        unique_together = (("machine", "name", "version", "update_type"),)

class UpdateHistoryItem(models.Model):
    id = models.BigAutoField(primary_key=True)
    UPDATE_STATUS = (
        ('pending', 'Pending'),
        ('error', 'Error'),
        ('install', 'Install'),
        ('removal', 'Removal')
    )
    update_history = models.ForeignKey(UpdateHistory)
    recorded = models.DateTimeField()
    uuid = models.CharField(null=True, blank=True, max_length=100)
    status = models.CharField(max_length=254, choices=UPDATE_STATUS, verbose_name="Status")
    extra = models.TextField(blank=True, null=True)

    def __unicode__(self):
        return "%s: %s %s %s %s" % (
            self.update_history.machine,
            self.update_history.name,
            self.update_history.version,
            self.status, self.recorded
        )

    class Meta:
        ordering = ['-recorded']
        unique_together = (("update_history", "recorded", "status"),)

class Fact(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='facts')
    fact_name = models.TextField()
    fact_data = models.TextField()

    def __unicode__(self):
        return '%s: %s' % (self.fact_name, self.fact_data)

    class Meta:
        ordering = ['fact_name']

class HistoricalFact(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='historical_facts')
    fact_name = models.CharField(max_length=255)
    fact_data = models.TextField()
    fact_recorded = models.DateTimeField()

    def __unicode__(self):
        return self.fact_name

    class Meta:
        ordering = ['fact_name', 'fact_recorded']

class Condition(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='conditions')
    condition_name = models.CharField(max_length=255)
    condition_data = models.TextField()

    def __unicode__(self):
        return self.condition_name

    class Meta:
        ordering = ['condition_name']

class PluginScriptSubmission(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine)
    plugin = models.CharField(max_length=255)
    historical = models.BooleanField(default=False)
    recorded = models.DateTimeField(auto_now_add=True)

    def __unicode__(self):
        return '%s: %s' % (self.machine, self.plugin)

    class Meta:
        ordering = ['recorded', 'plugin']

class PluginScriptRow(models.Model):
    id = models.BigAutoField(primary_key=True)
    submission = models.ForeignKey(PluginScriptSubmission)
    pluginscript_name = models.TextField()
    pluginscript_data = models.TextField(blank=True, null=True)
    pluginscript_data_string = models.TextField(blank=True, null=True)
    pluginscript_data_int = models.IntegerField(default=0)
    pluginscript_data_date = models.DateTimeField(blank=True, null=True)
    submission_and_script_name = models.TextField()

    def save(self):
        try:
            self.pluginscript_data_int = int(self.pluginscript_data)
        except Exception:
            self.pluginscript_data_int = 0

        try:
            self.pluginscript_data_string = str(self.pluginscript_data)
        except Exception:
            self.pluginscript_data_string = ""

        try:
            date_data = parse(self.pluginscript_data)
            if not date_data.tzinfo:
                date_data = date_data.replace(tzinfo=pytz.UTC)
            self.pluginscript_data_date = date_data
        except Exception:
            # Try converting it to an int if we're here
            try:
                if int(self.pluginscript_data) != 0:

                    try:
                        self.pluginscript_data_date = datetime.fromtimestamp(
                            int(self.pluginscript_data), tz=pytz.UTC)
                    except Exception:
                        self.pluginscript_data_date = None
                else:
                    self.pluginscript_data_date = None
            except Exception:
                self.pluginscript_data_date = None

        super(PluginScriptRow, self).save()

    def __unicode__(self):
        return '%s: %s' % (self.pluginscript_name, self.pluginscript_data)

    class Meta:
        ordering = ['pluginscript_name']

class PendingUpdate(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='pending_updates')
    update = models.CharField(db_index=True, max_length=255, null=True, blank=True)
    update_version = models.CharField(max_length=255, null=True, blank=True)
    display_name = models.CharField(max_length=255, null=True, blank=True)

    def __unicode__(self):
        return self.update

    class Meta:
        ordering = ['display_name']

class InstalledUpdate(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='installed_updates')
    update = models.CharField(db_index=True, max_length=255, null=True, blank=True)
    update_version = models.CharField(max_length=255, null=True, blank=True)
    display_name = models.CharField(max_length=255, null=True, blank=True)
    installed = models.BooleanField()

    def __unicode__(self):
        return self.update

    class Meta:
        ordering = ['display_name']
        unique_together = ("machine", "update", "update_version")

class PendingAppleUpdate(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='pending_apple_updates')
    update = models.CharField(db_index=True, max_length=255, null=True, blank=True)
    update_version = models.CharField(max_length=256, null=True, blank=True)
    display_name = models.CharField(max_length=256, null=True, blank=True)

    def __unicode__(self):
        return str(self.update) or ''

    class Meta:
        ordering = ['display_name']

class Plugin(models.Model):
    name = models.CharField(max_length=255, unique=True)
    order = models.IntegerField()

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['order']

class MachineDetailPlugin(models.Model):
    name = models.CharField(max_length=255, unique=True)
    order = models.IntegerField()

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['order']

class Report(models.Model):
    name = models.CharField(max_length=255, unique=True)

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['name']

class SalSetting(models.Model):
    name = models.CharField(max_length=255, unique=True)
    value = models.TextField()

    def __unicode__(self):
        return self.name

class ApiKey(models.Model):
    public_key = models.CharField(max_length=255)
    private_key = models.CharField(max_length=255)
    name = models.CharField(max_length=255)
    has_been_seen = models.BooleanField(default=False)
    read_write = models.BooleanField(default=False)

    def save(self, *args, **kwargs):
        if not self.id:
            self.public_key = GenerateAPIKey()
            self.private_key = ''.join(random.choice(
                string.ascii_lowercase + string.digits) for x in range(64))
        super(ApiKey, self).save(*args, **kwargs)

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['name']
        unique_together = ("public_key", "private_key")

class FriendlyNameCache(models.Model):
    serial_stub = models.CharField(max_length=5)
    friendly_name = models.CharField(max_length=255)
<commit_msg>Put more specific exception handling into plugin processing.<commit_after>import plistlib
import random
import string
from datetime import datetime

import pytz
from dateutil.parser import parse

from django.contrib.auth.models import User
from django.db import models

from server import text_utils

OS_CHOICES = (
    ('Darwin', 'macOS'),
    ('Windows', 'Windows'),
    ('Linux', 'Linux'),
    ('ChromeOS', 'Chrome OS'),
)

REPORT_CHOICES = (
    ('base64', 'base64'),
    ('base64bz2', 'base64bz2'),
    ('bz2', 'bz2'),
)

class ProfileLevel():
    stats_only = 'SO'
    read_only = 'RO'
    read_write = 'RW'
    global_admin = 'GA'

def GenerateKey():
    key = ''.join(random.choice(string.ascii_lowercase + string.digits) for x in range(128))
    try:
        MachineGroup.objects.get(key=key)
        return GenerateKey()
    except MachineGroup.DoesNotExist:
        return key

def GenerateAPIKey():
    key = ''.join(random.choice(string.ascii_lowercase + string.digits) for x in range(24))
    try:
        ApiKey.objects.get(public_key=key)
        return GenerateAPIKey()
    except ApiKey.DoesNotExist:
        return key

class UserProfile(models.Model):
    user = models.OneToOneField(User, unique=True)

    LEVEL_CHOICES = (
        ('SO', 'Stats Only'),
        ('RO', 'Read Only'),
        ('RW', 'Read Write'),
        ('GA', 'Global Admin'),
    )
    level = models.CharField(max_length=2, choices=LEVEL_CHOICES, default='RO')

    def __unicode__(self):
        return self.user.username

User.userprofile = property(lambda u: UserProfile.objects.get_or_create(user=u)[0])

class BusinessUnit(models.Model):
    name = models.CharField(max_length=100)
    users = models.ManyToManyField(User, blank=True)

    def __unicode__(self):
        return self.name

    @classmethod
    def display_name(cls):
        return text_utils.class_to_title(cls.__name__)

    class Meta:
        ordering = ['name']

class MachineGroup(models.Model):
    business_unit = models.ForeignKey(BusinessUnit)
    name = models.CharField(max_length=100)
    key = models.CharField(db_index=True, max_length=255, unique=True,
                           blank=True, null=True, editable=False)

    def save(self, **kwargs):
        if not self.id:
            self.key = GenerateKey()
        super(MachineGroup, self).save()

    def __unicode__(self):
        return self.name

    @classmethod
    def display_name(cls):
        return text_utils.class_to_title(cls.__name__)

    class Meta:
        ordering = ['name']

class DeployedManager(models.Manager):
    def get_queryset(self):
        return super(DeployedManager, self).get_queryset().filter(deployed=True)

class Machine(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine_group = models.ForeignKey(MachineGroup)
    serial = models.CharField(db_index=True, max_length=100, unique=True)
    hostname = models.CharField(max_length=256, null=True, blank=True)
    operating_system = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    memory = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    memory_kb = models.IntegerField(db_index=True, default=0)
    munki_version = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    manifest = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    hd_space = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    hd_total = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    hd_percent = models.CharField(max_length=256, null=True, blank=True)
    console_user = models.CharField(max_length=256, null=True, blank=True)
    machine_model = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    machine_model_friendly = models.CharField(db_index=True, max_length=256, null=True, blank=True)
    cpu_type = models.CharField(max_length=256, null=True, blank=True)
    cpu_speed = models.CharField(max_length=256, null=True, blank=True)
    os_family = models.CharField(db_index=True, max_length=256,
                                 choices=OS_CHOICES, verbose_name="OS Family", default="Darwin")
    last_checkin = models.DateTimeField(db_index=True, blank=True, null=True)
    first_checkin = models.DateTimeField(db_index=True, blank=True, null=True, auto_now_add=True)
    report = models.TextField(editable=True, null=True)
    errors = models.IntegerField(default=0)
    warnings = models.IntegerField(default=0)
    activity = models.BooleanField(editable=True, default=False)
    puppet_version = models.CharField(db_index=True, null=True, blank=True, max_length=256)
    sal_version = models.CharField(db_index=True, null=True, blank=True, max_length=255)
    last_puppet_run = models.DateTimeField(db_index=True, blank=True, null=True)
    puppet_errors = models.IntegerField(db_index=True, default=0)
    deployed = models.BooleanField(default=True)
    broken_client = models.BooleanField(default=False)

    objects = models.Manager()  # The default manager.
    deployed_objects = DeployedManager()

    def get_fields(self):
        return [(field.name, field.value_to_string(self)) for field in Machine._meta.fields]

    def get_report(self):
        try:
            return plistlib.loads(self.report)
        except TypeError:
            return {}

    def __unicode__(self):
        if self.hostname:
            return self.hostname
        else:
            return self.serial

    @classmethod
    def display_name(cls):
        return text_utils.class_to_title(cls.__name__)

    class Meta:
        ordering = ['hostname']

GROUP_NAMES = {
    'all': None,
    'machine_group': MachineGroup,
    'business_unit': BusinessUnit,
    'machine': Machine}

class UpdateHistory(models.Model):
    id = models.BigAutoField(primary_key=True)
    UPDATE_TYPE = (
        ('third_party', '3rd Party'),
        ('apple', 'Apple'),
    )
    machine = models.ForeignKey(Machine)
    update_type = models.CharField(max_length=254, choices=UPDATE_TYPE, verbose_name="Update Type")
    name = models.CharField(max_length=255, db_index=True)
    version = models.CharField(max_length=254, db_index=True)

    def __unicode__(self):
        return "%s: %s %s" % (self.machine, self.name, self.version)

    class Meta:
        ordering = ['name']
        unique_together = (("machine", "name", "version", "update_type"),)

class UpdateHistoryItem(models.Model):
    id = models.BigAutoField(primary_key=True)
    UPDATE_STATUS = (
        ('pending', 'Pending'),
        ('error', 'Error'),
        ('install', 'Install'),
        ('removal', 'Removal')
    )
    update_history = models.ForeignKey(UpdateHistory)
    recorded = models.DateTimeField()
    uuid = models.CharField(null=True, blank=True, max_length=100)
    status = models.CharField(max_length=254, choices=UPDATE_STATUS, verbose_name="Status")
    extra = models.TextField(blank=True, null=True)

    def __unicode__(self):
        return "%s: %s %s %s %s" % (
            self.update_history.machine,
            self.update_history.name,
            self.update_history.version,
            self.status, self.recorded
        )

    class Meta:
        ordering = ['-recorded']
        unique_together = (("update_history", "recorded", "status"),)

class Fact(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='facts')
    fact_name = models.TextField()
    fact_data = models.TextField()

    def __unicode__(self):
        return '%s: %s' % (self.fact_name, self.fact_data)

    class Meta:
        ordering = ['fact_name']

class HistoricalFact(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='historical_facts')
    fact_name = models.CharField(max_length=255)
    fact_data = models.TextField()
    fact_recorded = models.DateTimeField()

    def __unicode__(self):
        return self.fact_name

    class Meta:
        ordering = ['fact_name', 'fact_recorded']

class Condition(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='conditions')
    condition_name = models.CharField(max_length=255)
    condition_data = models.TextField()

    def __unicode__(self):
        return self.condition_name

    class Meta:
        ordering = ['condition_name']

class PluginScriptSubmission(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine)
    plugin = models.CharField(max_length=255)
    historical = models.BooleanField(default=False)
    recorded = models.DateTimeField(auto_now_add=True)

    def __unicode__(self):
        return '%s: %s' % (self.machine, self.plugin)

    class Meta:
        ordering = ['recorded', 'plugin']

class PluginScriptRow(models.Model):
    id = models.BigAutoField(primary_key=True)
    submission = models.ForeignKey(PluginScriptSubmission)
    pluginscript_name = models.TextField()
    pluginscript_data = models.TextField(blank=True, null=True)
    pluginscript_data_string = models.TextField(blank=True, null=True)
    pluginscript_data_int = models.IntegerField(default=0)
    pluginscript_data_date = models.DateTimeField(blank=True, null=True)
    submission_and_script_name = models.TextField()

    def save(self):
        try:
            self.pluginscript_data_int = int(self.pluginscript_data)
        except (ValueError, TypeError):
            self.pluginscript_data_int = 0

        self.pluginscript_data_string = str(self.pluginscript_data)

        try:
            date_data = parse(self.pluginscript_data)
            if not date_data.tzinfo:
                date_data = date_data.replace(tzinfo=pytz.UTC)
            self.pluginscript_data_date = date_data
        except ValueError:
            # Try converting it to an int if we're here
            try:
                if int(self.pluginscript_data) != 0:

                    try:
                        self.pluginscript_data_date = datetime.fromtimestamp(
                            int(self.pluginscript_data), tz=pytz.UTC)
                    except (ValueError, TypeError):
                        self.pluginscript_data_date = None
                else:
                    self.pluginscript_data_date = None
            except (ValueError, TypeError):
                self.pluginscript_data_date = None

        super(PluginScriptRow, self).save()

    def __unicode__(self):
        return '%s: %s' % (self.pluginscript_name, self.pluginscript_data)

    class Meta:
        ordering = ['pluginscript_name']

class PendingUpdate(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='pending_updates')
    update = models.CharField(db_index=True, max_length=255, null=True, blank=True)
    update_version = models.CharField(max_length=255, null=True, blank=True)
    display_name = models.CharField(max_length=255, null=True, blank=True)

    def __unicode__(self):
        return self.update

    class Meta:
        ordering = ['display_name']

class InstalledUpdate(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='installed_updates')
    update = models.CharField(db_index=True, max_length=255, null=True, blank=True)
    update_version = models.CharField(max_length=255, null=True, blank=True)
    display_name = models.CharField(max_length=255, null=True, blank=True)
    installed = models.BooleanField()

    def __unicode__(self):
        return self.update

    class Meta:
        ordering = ['display_name']
        unique_together = ("machine", "update", "update_version")

class PendingAppleUpdate(models.Model):
    id = models.BigAutoField(primary_key=True)
    machine = models.ForeignKey(Machine, related_name='pending_apple_updates')
    update = models.CharField(db_index=True, max_length=255, null=True, blank=True)
    update_version = models.CharField(max_length=256, null=True, blank=True)
    display_name = models.CharField(max_length=256, null=True, blank=True)

    def __unicode__(self):
        return str(self.update) or ''

    class Meta:
        ordering = ['display_name']

class Plugin(models.Model):
    name = models.CharField(max_length=255, unique=True)
    order = models.IntegerField()

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['order']

class MachineDetailPlugin(models.Model):
    name = models.CharField(max_length=255, unique=True)
    order = models.IntegerField()

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['order']

class Report(models.Model):
    name = models.CharField(max_length=255, unique=True)

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['name']

class SalSetting(models.Model):
    name = models.CharField(max_length=255, unique=True)
    value = models.TextField()

    def __unicode__(self):
        return self.name

class ApiKey(models.Model):
    public_key = models.CharField(max_length=255)
    private_key = models.CharField(max_length=255)
    name = models.CharField(max_length=255)
    has_been_seen = models.BooleanField(default=False)
    read_write = models.BooleanField(default=False)

    def save(self, *args, **kwargs):
        if not self.id:
            self.public_key = GenerateAPIKey()
            self.private_key = ''.join(random.choice(
                string.ascii_lowercase + string.digits) for x in range(64))
        super(ApiKey, self).save(*args, **kwargs)

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ['name']
        unique_together = ("public_key", "private_key")

class FriendlyNameCache(models.Model):
    serial_stub = models.CharField(max_length=5)
    friendly_name = models.CharField(max_length=255)
<|endoftext|>