<commit_before># TODO: clean up
# TODO: update the readme
# TODO: better CLI descriptions/syntax
# TODO: restoring a model without recreating it (use constants / op names in the code?)
# TODO: move all the training parameters inside the training parser
# TODO: switch to https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn instead of buckets
# TODO: cannot process after a few iterations

import sys
import argparse
import logging

import tensorflow as tf

from .model.model import Model
from .defaults import Config
from .util import dataset
from .util.data_gen import DataGen
from .util.export import Exporter

tf.logging.set_verbosity(tf.logging.ERROR)

def process_args(args, defaults):

    parser = argparse.ArgumentParser()
    parser.prog = 'aocr'
    subparsers = parser.add_subparsers()

    # Global arguments
    parser_base = argparse.ArgumentParser(add_help=False)
    parser_base.add_argument('--log-path', dest="log_path",
                             metavar=defaults.LOG_PATH,
                             type=str, default=defaults.LOG_PATH,
                             help=('log file path (default: %s)'
                                   % (defaults.LOG_PATH)))

    # Dataset generation
    parser_dataset = subparsers.add_parser('dataset', parents=[parser_base],
                                           help='create a dataset in the TFRecords format')
    parser_dataset.set_defaults(phase='dataset')
    parser_dataset.add_argument('annotations_path', metavar='annotations',
                                type=str,
                                help=('path to the annotation file'))
    parser_dataset.add_argument('output_path', nargs='?', metavar='output',
                                type=str, default=defaults.NEW_DATASET_PATH,
                                help=('output path (default: %s)'
                                      % defaults.NEW_DATASET_PATH))
    parser_dataset.add_argument('--log-step', dest='log_step',
                                type=int, default=defaults.LOG_STEP,
                                metavar=defaults.LOG_STEP,
                                help=('print log messages every N steps (default: %s)'
                                      % defaults.LOG_STEP))
    parser_dataset.add_argument('--no-force-uppercase', dest='force_uppercase',
                                action='store_false', default=defaults.FORCE_UPPERCASE,
                                help='do not force uppercase on label values')
    parser_dataset.add_argument('--save-filename', dest='save_filename',
                                action='store_true', default=defaults.SAVE_FILENAME,
                                help='save filename as a field in the dataset')

    # Shared model arguments
    parser_model = argparse.ArgumentParser(add_help=False)
    parser_model.set_defaults(visualize=defaults.VISUALIZE)
    parser_model.set_defaults(load_model=defaults.LOAD_MODEL)
    parser_model.add_argument('--max-width', dest="max_width",
                              metavar=defaults.MAX_WIDTH,
                              type=int, default=defaults.MAX_WIDTH,
                              help=('max image width (default: %s)'
                                    % (defaults.MAX_WIDTH)))
    parser_model.add_argument('--max-height', dest="max_height",
                              metavar=defaults.MAX_HEIGHT,
                              type=int, default=defaults.MAX_HEIGHT,
                              help=('max image height (default: %s)'
                                    % (defaults.MAX_HEIGHT)))
    parser_model.add_argument('--max-prediction', dest="max_prediction",
                              metavar=defaults.MAX_PREDICTION,
                              type=int, default=defaults.MAX_PREDICTION,
                              help=('max length of predicted strings (default: %s)'
                                    % (defaults.MAX_PREDICTION)))
    parser_model.add_argument('--full-ascii', dest='full_ascii', action='store_true',
                              help=('use lowercase in addition to uppercase'))
    parser_model.set_defaults(full_ascii=defaults.FULL_ASCII)
    parser_model.add_argument('--color', dest="channels", action='store_const', const=3,
                              default=defaults.CHANNELS,
                              help=('do not convert source images to grayscale'))
    parser_model.add_argument('--no-distance', dest="use_distance", action="store_false",
                              default=defaults.USE_DISTANCE,
                              help=('require full match when calculating accuracy'))
    parser_model.add_argument('--gpu-id', dest="gpu_id", metavar=defaults.GPU_ID,
                              type=int, default=defaults.GPU_ID,
                              help='specify a GPU ID')
    parser_model.add_argument('--use-gru', dest='use_gru', action='store_true',
                              help='use GRU instead of LSTM')
    parser_model.add_argument('--attn-num-layers', dest="attn_num_layers",
                              type=int, default=defaults.ATTN_NUM_LAYERS,
                              metavar=defaults.ATTN_NUM_LAYERS,
                              help=('hidden layers in attention decoder cell (default: %s)'
                                    % (defaults.ATTN_NUM_LAYERS)))
    parser_model.add_argument('--attn-num-hidden', dest="attn_num_hidden",
                              type=int, default=defaults.ATTN_NUM_HIDDEN,
                              metavar=defaults.ATTN_NUM_HIDDEN,
                              help=('hidden units in attention decoder cell (default: %s)'
                                    % (defaults.ATTN_NUM_HIDDEN)))
    parser_model.add_argument('--initial-learning-rate', dest="initial_learning_rate",
                              type=float, default=defaults.INITIAL_LEARNING_RATE,
                              metavar=defaults.INITIAL_LEARNING_RATE,
                              help=('initial learning rate (default: %s)'
                                    % (defaults.INITIAL_LEARNING_RATE)))
    parser_model.add_argument('--model-dir', '--job-dir', dest="model_dir",
                              type=str, default=defaults.MODEL_DIR,
                              metavar=defaults.MODEL_DIR,
                              help=('directory for the model '
                                    '(default: %s)' % (defaults.MODEL_DIR)))
    parser_model.add_argument('--target-embedding-size', dest="target_embedding_size",
                              type=int, default=defaults.TARGET_EMBEDDING_SIZE,
                              metavar=defaults.TARGET_EMBEDDING_SIZE,
                              help=('embedding dimension for each target (default: %s)'
                                    % (defaults.TARGET_EMBEDDING_SIZE)))
    parser_model.add_argument('--output-dir', dest="output_dir",
                              type=str, default=defaults.OUTPUT_DIR,
                              metavar=defaults.OUTPUT_DIR,
                              help=('output directory (default: %s)'
                                    % (defaults.OUTPUT_DIR)))
    parser_model.add_argument('--max-gradient-norm', dest="max_gradient_norm",
                              type=int, default=defaults.MAX_GRADIENT_NORM,
                              metavar=defaults.MAX_GRADIENT_NORM,
                              help=('clip gradients to this norm (default: %s)'
                                    % (defaults.MAX_GRADIENT_NORM)))
    parser_model.add_argument('--no-gradient-clipping', dest='clip_gradients', action='store_false',
                              help=('do not perform gradient clipping'))
    parser_model.set_defaults(clip_gradients=defaults.CLIP_GRADIENTS)

    # Training
    parser_train = subparsers.add_parser('train', parents=[parser_base, parser_model],
                                         help='Train the model and save checkpoints.')
    parser_train.set_defaults(phase='train')
    parser_train.add_argument('dataset_path', metavar='dataset',
                              type=str, default=defaults.DATA_PATH,
                              help=('training dataset in the TFRecords format'
                                    ' (default: %s)'
                                    % (defaults.DATA_PATH)))
    parser_train.add_argument('--steps-per-checkpoint', dest="steps_per_checkpoint",
                              type=int, default=defaults.STEPS_PER_CHECKPOINT,
                              metavar=defaults.STEPS_PER_CHECKPOINT,
                              help=('steps between saving the model'
                                    ' (default: %s)'
                                    % (defaults.STEPS_PER_CHECKPOINT)))
    parser_train.add_argument('--batch-size', dest="batch_size",
                              type=int, default=defaults.BATCH_SIZE,
                              metavar=defaults.BATCH_SIZE,
                              help=('batch size (default: %s)'
                                    % (defaults.BATCH_SIZE)))
    parser_train.add_argument('--num-epoch', dest="num_epoch",
                              type=int, default=defaults.NUM_EPOCH,
                              metavar=defaults.NUM_EPOCH,
                              help=('number of training epochs (default: %s)'
                                    % (defaults.NUM_EPOCH)))
    parser_train.add_argument('--no-resume', dest='load_model', action='store_false',
                              help=('create a new model even if checkpoints already exist'))

    # Testing
    parser_test = subparsers.add_parser('test', parents=[parser_base, parser_model],
                                        help='Test the saved model.')
    parser_test.set_defaults(phase='test', steps_per_checkpoint=0, batch_size=1,
                             max_width=defaults.MAX_WIDTH, max_height=defaults.MAX_HEIGHT,
                             max_prediction=defaults.MAX_PREDICTION, full_ascii=defaults.FULL_ASCII)
    parser_test.add_argument('dataset_path', metavar='dataset',
                             type=str, default=defaults.DATA_PATH,
                             help=('Testing dataset in the TFRecords format'
                                   ', default=%s'
                                   % (defaults.DATA_PATH)))
    parser_test.add_argument('--visualize', dest='visualize', action='store_true',
                             help=('visualize attentions'))

    # Exporting
    parser_export = subparsers.add_parser('export', parents=[parser_base, parser_model],
                                          help='Export the model with weights for production use.')
    parser_export.set_defaults(phase='export', steps_per_checkpoint=0, batch_size=1)
    parser_export.add_argument('export_path', nargs='?', metavar='dir',
                               type=str, default=defaults.EXPORT_PATH,
                               help=('Directory to save the exported model to,'
                                     'default=%s'
                                     % (defaults.EXPORT_PATH)))
    parser_export.add_argument('--format', dest="format",
                               type=str, default=defaults.EXPORT_FORMAT,
                               choices=['frozengraph', 'savedmodel'],
                               help=('export format'
                                     ' (default: %s)'
                                     % (defaults.EXPORT_FORMAT)))

    # Predicting
    parser_predict = subparsers.add_parser('predict', parents=[parser_base, parser_model],
                                           help='Predict text from files (feed through stdin).')
    parser_predict.set_defaults(phase='predict', steps_per_checkpoint=0, batch_size=1)

    parameters = parser.parse_args(args)
    return parameters

def main(args=None):

    if args is None:
        args = sys.argv[1:]

    parameters = process_args(args, Config)
    logging.basicConfig(
        level=logging.DEBUG,
        format='%(asctime)-15s %(name)-5s %(levelname)-8s %(message)s',
        filename=parameters.log_path)
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)-15s %(name)-5s %(levelname)-8s %(message)s')
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)

    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:

        if parameters.phase == 'dataset':
            dataset.generate(
                parameters.annotations_path,
                parameters.output_path,
                parameters.log_step,
                parameters.force_uppercase,
                parameters.save_filename
            )
            return

        if parameters.full_ascii:
            DataGen.set_full_ascii_charmap()

        model = Model(
            phase=parameters.phase,
            visualize=parameters.visualize,
            output_dir=parameters.output_dir,
            batch_size=parameters.batch_size,
            initial_learning_rate=parameters.initial_learning_rate,
            steps_per_checkpoint=parameters.steps_per_checkpoint,
            model_dir=parameters.model_dir,
            target_embedding_size=parameters.target_embedding_size,
            attn_num_hidden=parameters.attn_num_hidden,
            attn_num_layers=parameters.attn_num_layers,
            clip_gradients=parameters.clip_gradients,
            max_gradient_norm=parameters.max_gradient_norm,
            session=sess,
            load_model=parameters.load_model,
            gpu_id=parameters.gpu_id,
            use_gru=parameters.use_gru,
            use_distance=parameters.use_distance,
            max_image_width=parameters.max_width,
            max_image_height=parameters.max_height,
            max_prediction_length=parameters.max_prediction,
            channels=parameters.channels,
        )

        if parameters.phase == 'train':
            model.train(
                data_path=parameters.dataset_path,
                num_epoch=parameters.num_epoch
            )
        elif parameters.phase == 'test':
            model.test(
                data_path=parameters.dataset_path
            )
        elif parameters.phase == 'predict':
            for line in sys.stdin:
                filename = line.rstrip()
                try:
                    with open(filename, 'rb') as img_file:
                        img_file_data = img_file.read()
                except IOError:
                    print('result: err opening file', filename)
                    continue
                text, probability = model.predict(img_file_data)
                print('result: ok', '{:.2f}'.format(probability), text)
        elif parameters.phase == 'export':
            exporter = Exporter(model)
            exporter.save(parameters.export_path, parameters.format)
            return
        else:
            raise NotImplementedError

if __name__ == "__main__":
    main()
<commit_msg>Remove completed TODOs<commit_after># TODO: clean up
# TODO: update the readme
# TODO: restoring a model without recreating it (use constants / op names in the code?)
# TODO: move all the training parameters inside the training parser
# TODO: switch to https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn instead of buckets

import sys
import argparse
import logging

import tensorflow as tf

from .model.model import Model
from .defaults import Config
from .util import dataset
from .util.data_gen import DataGen
from .util.export import Exporter

tf.logging.set_verbosity(tf.logging.ERROR)

def process_args(args, defaults):

    parser = argparse.ArgumentParser()
    parser.prog = 'aocr'
    subparsers = parser.add_subparsers()

    # Global arguments
    parser_base = argparse.ArgumentParser(add_help=False)
    parser_base.add_argument('--log-path', dest="log_path",
                             metavar=defaults.LOG_PATH,
                             type=str, default=defaults.LOG_PATH,
                             help=('log file path (default: %s)'
                                   % (defaults.LOG_PATH)))

    # Dataset generation
    parser_dataset = subparsers.add_parser('dataset', parents=[parser_base],
                                           help='create a dataset in the TFRecords format')
    parser_dataset.set_defaults(phase='dataset')
    parser_dataset.add_argument('annotations_path', metavar='annotations',
                                type=str,
                                help=('path to the annotation file'))
    parser_dataset.add_argument('output_path', nargs='?', metavar='output',
                                type=str, default=defaults.NEW_DATASET_PATH,
                                help=('output path (default: %s)'
                                      % defaults.NEW_DATASET_PATH))
    parser_dataset.add_argument('--log-step', dest='log_step',
                                type=int, default=defaults.LOG_STEP,
                                metavar=defaults.LOG_STEP,
                                help=('print log messages every N steps (default: %s)'
                                      % defaults.LOG_STEP))
    parser_dataset.add_argument('--no-force-uppercase', dest='force_uppercase',
                                action='store_false', default=defaults.FORCE_UPPERCASE,
                                help='do not force uppercase on label values')
    parser_dataset.add_argument('--save-filename', dest='save_filename',
                                action='store_true', default=defaults.SAVE_FILENAME,
                                help='save filename as a field in the dataset')

    # Shared model arguments
    parser_model = argparse.ArgumentParser(add_help=False)
    parser_model.set_defaults(visualize=defaults.VISUALIZE)
    parser_model.set_defaults(load_model=defaults.LOAD_MODEL)
    parser_model.add_argument('--max-width', dest="max_width",
                              metavar=defaults.MAX_WIDTH,
                              type=int, default=defaults.MAX_WIDTH,
                              help=('max image width (default: %s)'
                                    % (defaults.MAX_WIDTH)))
    parser_model.add_argument('--max-height', dest="max_height",
                              metavar=defaults.MAX_HEIGHT,
                              type=int, default=defaults.MAX_HEIGHT,
                              help=('max image height (default: %s)'
                                    % (defaults.MAX_HEIGHT)))
    parser_model.add_argument('--max-prediction', dest="max_prediction",
                              metavar=defaults.MAX_PREDICTION,
                              type=int, default=defaults.MAX_PREDICTION,
                              help=('max length of predicted strings (default: %s)'
                                    % (defaults.MAX_PREDICTION)))
    parser_model.add_argument('--full-ascii', dest='full_ascii', action='store_true',
                              help=('use lowercase in addition to uppercase'))
    parser_model.set_defaults(full_ascii=defaults.FULL_ASCII)
    parser_model.add_argument('--color', dest="channels", action='store_const', const=3,
                              default=defaults.CHANNELS,
                              help=('do not convert source images to grayscale'))
    parser_model.add_argument('--no-distance', dest="use_distance", action="store_false",
                              default=defaults.USE_DISTANCE,
                              help=('require full match when calculating accuracy'))
    parser_model.add_argument('--gpu-id', dest="gpu_id", metavar=defaults.GPU_ID,
                              type=int, default=defaults.GPU_ID,
                              help='specify a GPU ID')
    parser_model.add_argument('--use-gru', dest='use_gru', action='store_true',
                              help='use GRU instead of LSTM')
    parser_model.add_argument('--attn-num-layers', dest="attn_num_layers",
                              type=int, default=defaults.ATTN_NUM_LAYERS,
                              metavar=defaults.ATTN_NUM_LAYERS,
                              help=('hidden layers in attention decoder cell (default: %s)'
                                    % (defaults.ATTN_NUM_LAYERS)))
    parser_model.add_argument('--attn-num-hidden', dest="attn_num_hidden",
                              type=int, default=defaults.ATTN_NUM_HIDDEN,
                              metavar=defaults.ATTN_NUM_HIDDEN,
                              help=('hidden units in attention decoder cell (default: %s)'
                                    % (defaults.ATTN_NUM_HIDDEN)))
    parser_model.add_argument('--initial-learning-rate', dest="initial_learning_rate",
                              type=float, default=defaults.INITIAL_LEARNING_RATE,
                              metavar=defaults.INITIAL_LEARNING_RATE,
                              help=('initial learning rate (default: %s)'
                                    % (defaults.INITIAL_LEARNING_RATE)))
    parser_model.add_argument('--model-dir', '--job-dir', dest="model_dir",
                              type=str, default=defaults.MODEL_DIR,
                              metavar=defaults.MODEL_DIR,
                              help=('directory for the model '
                                    '(default: %s)' % (defaults.MODEL_DIR)))
    parser_model.add_argument('--target-embedding-size', dest="target_embedding_size",
                              type=int, default=defaults.TARGET_EMBEDDING_SIZE,
                              metavar=defaults.TARGET_EMBEDDING_SIZE,
                              help=('embedding dimension for each target (default: %s)'
                                    % (defaults.TARGET_EMBEDDING_SIZE)))
    parser_model.add_argument('--output-dir', dest="output_dir",
                              type=str, default=defaults.OUTPUT_DIR,
                              metavar=defaults.OUTPUT_DIR,
                              help=('output directory (default: %s)'
                                    % (defaults.OUTPUT_DIR)))
    parser_model.add_argument('--max-gradient-norm', dest="max_gradient_norm",
                              type=int, default=defaults.MAX_GRADIENT_NORM,
                              metavar=defaults.MAX_GRADIENT_NORM,
                              help=('clip gradients to this norm (default: %s)'
                                    % (defaults.MAX_GRADIENT_NORM)))
    parser_model.add_argument('--no-gradient-clipping', dest='clip_gradients', action='store_false',
                              help=('do not perform gradient clipping'))
    parser_model.set_defaults(clip_gradients=defaults.CLIP_GRADIENTS)

    # Training
    parser_train = subparsers.add_parser('train', parents=[parser_base, parser_model],
                                         help='Train the model and save checkpoints.')
    parser_train.set_defaults(phase='train')
    parser_train.add_argument('dataset_path', metavar='dataset',
                              type=str, default=defaults.DATA_PATH,
                              help=('training dataset in the TFRecords format'
                                    ' (default: %s)'
                                    % (defaults.DATA_PATH)))
    parser_train.add_argument('--steps-per-checkpoint', dest="steps_per_checkpoint",
                              type=int, default=defaults.STEPS_PER_CHECKPOINT,
                              metavar=defaults.STEPS_PER_CHECKPOINT,
                              help=('steps between saving the model'
                                    ' (default: %s)'
                                    % (defaults.STEPS_PER_CHECKPOINT)))
    parser_train.add_argument('--batch-size', dest="batch_size",
                              type=int, default=defaults.BATCH_SIZE,
                              metavar=defaults.BATCH_SIZE,
                              help=('batch size (default: %s)'
                                    % (defaults.BATCH_SIZE)))
    parser_train.add_argument('--num-epoch', dest="num_epoch",
                              type=int, default=defaults.NUM_EPOCH,
                              metavar=defaults.NUM_EPOCH,
                              help=('number of training epochs (default: %s)'
                                    % (defaults.NUM_EPOCH)))
    parser_train.add_argument('--no-resume', dest='load_model', action='store_false',
                              help=('create a new model even if checkpoints already exist'))

    # Testing
    parser_test = subparsers.add_parser('test', parents=[parser_base, parser_model],
                                        help='Test the saved model.')
    parser_test.set_defaults(phase='test', steps_per_checkpoint=0, batch_size=1,
                             max_width=defaults.MAX_WIDTH, max_height=defaults.MAX_HEIGHT,
                             max_prediction=defaults.MAX_PREDICTION, full_ascii=defaults.FULL_ASCII)
    parser_test.add_argument('dataset_path', metavar='dataset',
                             type=str, default=defaults.DATA_PATH,
                             help=('Testing dataset in the TFRecords format'
                                   ', default=%s'
                                   % (defaults.DATA_PATH)))
    parser_test.add_argument('--visualize', dest='visualize', action='store_true',
                             help=('visualize attentions'))

    # Exporting
    parser_export = subparsers.add_parser('export', parents=[parser_base, parser_model],
                                          help='Export the model with weights for production use.')
    parser_export.set_defaults(phase='export', steps_per_checkpoint=0, batch_size=1)
    parser_export.add_argument('export_path', nargs='?', metavar='dir',
                               type=str, default=defaults.EXPORT_PATH,
                               help=('Directory to save the exported model to,'
                                     'default=%s'
                                     % (defaults.EXPORT_PATH)))
    parser_export.add_argument('--format', dest="format",
                               type=str, default=defaults.EXPORT_FORMAT,
                               choices=['frozengraph', 'savedmodel'],
                               help=('export format'
                                     ' (default: %s)'
                                     % (defaults.EXPORT_FORMAT)))

    # Predicting
    parser_predict = subparsers.add_parser('predict', parents=[parser_base, parser_model],
                                           help='Predict text from files (feed through stdin).')
    parser_predict.set_defaults(phase='predict', steps_per_checkpoint=0, batch_size=1)

    parameters = parser.parse_args(args)
    return parameters

def main(args=None):

    if args is None:
        args = sys.argv[1:]

    parameters = process_args(args, Config)
    logging.basicConfig(
        level=logging.DEBUG,
        format='%(asctime)-15s %(name)-5s %(levelname)-8s %(message)s',
        filename=parameters.log_path)
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)-15s %(name)-5s %(levelname)-8s %(message)s')
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)

    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:

        if parameters.phase == 'dataset':
            dataset.generate(
                parameters.annotations_path,
                parameters.output_path,
                parameters.log_step,
                parameters.force_uppercase,
                parameters.save_filename
            )
            return

        if parameters.full_ascii:
            DataGen.set_full_ascii_charmap()

        model = Model(
            phase=parameters.phase,
            visualize=parameters.visualize,
            output_dir=parameters.output_dir,
            batch_size=parameters.batch_size,
            initial_learning_rate=parameters.initial_learning_rate,
            steps_per_checkpoint=parameters.steps_per_checkpoint,
            model_dir=parameters.model_dir,
            target_embedding_size=parameters.target_embedding_size,
            attn_num_hidden=parameters.attn_num_hidden,
            attn_num_layers=parameters.attn_num_layers,
            clip_gradients=parameters.clip_gradients,
            max_gradient_norm=parameters.max_gradient_norm,
            session=sess,
            load_model=parameters.load_model,
            gpu_id=parameters.gpu_id,
            use_gru=parameters.use_gru,
            use_distance=parameters.use_distance,
            max_image_width=parameters.max_width,
            max_image_height=parameters.max_height,
            max_prediction_length=parameters.max_prediction,
            channels=parameters.channels,
        )

        if parameters.phase == 'train':
            model.train(
                data_path=parameters.dataset_path,
                num_epoch=parameters.num_epoch
            )
        elif parameters.phase == 'test':
            model.test(
                data_path=parameters.dataset_path
            )
        elif parameters.phase == 'predict':
            for line in sys.stdin:
                filename = line.rstrip()
                try:
                    with open(filename, 'rb') as img_file:
                        img_file_data = img_file.read()
                except IOError:
                    print('result: err opening file', filename)
                    continue
                text, probability = model.predict(img_file_data)
                print('result: ok', '{:.2f}'.format(probability), text)
        elif parameters.phase == 'export':
            exporter = Exporter(model)
            exporter.save(parameters.export_path, parameters.format)
            return
        else:
            raise NotImplementedError

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


<commit_before># Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

from hamcrest import (assert_that, has_entries, has_item, only_contains,
                      is_not, empty, greater_than, has_length)

from vapor.helpers import contrail_status
from vapor.helpers import dpdk
from vapor.helpers import nodes_steps
from vapor.helpers import vrouter_steps
from vapor import settings

def test_contrail_vrouter_dpdk(os_faults_steps):
    """Validate existence and status for contrail-vrouter-dpdk service."""
    statuses = contrail_status.get_services_statuses(os_faults_steps)
    for fqdn in settings.CONTRAIL_ROLES_DISTRIBUTION[
            settings.ROLE_CONTRAIL_COMPUTE]:
        compute_services = statuses[fqdn]
        assert_that(
            compute_services,
            has_item(
                has_entries(name='contrail-vrouter-dpdk', status='active')))

def test_bound_network_interfaces(os_faults_steps, computes):
    """Verify if DPDK vRouter binds network interfaces."""
    devices = dpdk.get_devices(os_faults_steps, computes)
    assert_that(devices.values(),
                only_contains(
                    has_entries('Network devices using DPDK-compatible driver',
                                is_not(empty()))))

def test_huge_pages_usage(os_faults_steps, computes):
    """Verify if vRrouter uses Huge Pages."""
    result = os_faults_steps.execute_cmd(computes, 'grep Huge /proc/meminfo')
    for node_result in result:
        node = nodes_steps.get_node_by_result(node_result, os_faults_steps)
        data = {}
        for line in node_result.payload['stdout_lines']:
            key, value = line.split(':', 1)
            data[key] = value.strip()
        assert_that(data,
                    has_entries(
                        HugePages_Total=is_not('0'), AnonHugePages='0 kB'),
                    node.fqdn)

def test_contrail_vrouter_dpdk_cpu_usage(os_faults_steps, computes):
    """Verify if vRouter uses CPU."""
    cmd = "ps -c contrail-vrouter-dpdk -o %cpu="
    result = os_faults_steps.execute_cmd(computes, cmd)
    for node_result in result:
        usage = node_result.payload['stdout']
        usage = float(usage.strip())
        assert_that(usage, greater_than(50))

def test_vrouter_create_interface(request, os_faults_steps, computes):
    """Verify if vRouter creates interface after creation of a virtual machine.

    Steps:
        #. Remember count of virtual interfaces on vRouters
        #. Create server
        #. Check that count of virtual interfaces on server's compute vRouter
            is greater than it was.
    """
    before_ifaces = vrouter_steps.get_interface_table(os_faults_steps,
                                                      computes)
    server = request.getfixturevalue('server')
    compute_fqdn = getattr(server, settings.SERVER_ATTR_HYPERVISOR_HOSTNAME)
    after_ifaces = vrouter_steps.get_interface_table(os_faults_steps, computes)
    assert_that(after_ifaces[compute_fqdn],
                has_length(greater_than(len(before_ifaces[compute_fqdn]))))
<commit_msg>Fix test_contrail_vrouter_dpdk test<commit_after># Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

from hamcrest import (assert_that, has_entries, has_item, only_contains,
                      is_not, empty, greater_than, has_length)

from vapor.helpers import contrail_status
from vapor.helpers import dpdk
from vapor.helpers import nodes_steps
from vapor.helpers import vrouter_steps
from vapor import settings

def test_contrail_vrouter_dpdk(os_faults_steps):
    """Validate existence and status for contrail-vrouter-dpdk service."""
    statuses = contrail_status.get_services_statuses(os_faults_steps)
    for fqdn in settings.CONTRAIL_ROLES_DISTRIBUTION[
            settings.ROLE_CONTRAIL_COMPUTE]:
        compute_services = statuses[fqdn]
        assert_that(
            compute_services,
            has_item(
                has_entries(service='contrail-vrouter-dpdk', status='active')))

def test_bound_network_interfaces(os_faults_steps, computes):
    """Verify if DPDK vRouter binds network interfaces."""
    devices = dpdk.get_devices(os_faults_steps, computes)
    assert_that(devices.values(),
                only_contains(
                    has_entries('Network devices using DPDK-compatible driver',
                                is_not(empty()))))

def test_huge_pages_usage(os_faults_steps, computes):
    """Verify if vRrouter uses Huge Pages."""
    result = os_faults_steps.execute_cmd(computes, 'grep Huge /proc/meminfo')
    for node_result in result:
        node = nodes_steps.get_node_by_result(node_result, os_faults_steps)
        data = {}
        for line in node_result.payload['stdout_lines']:
            key, value = line.split(':', 1)
            data[key] = value.strip()
        assert_that(data,
                    has_entries(
                        HugePages_Total=is_not('0'), AnonHugePages='0 kB'),
                    node.fqdn)

def test_contrail_vrouter_dpdk_cpu_usage(os_faults_steps, computes):
    """Verify if vRouter uses CPU."""
    cmd = "ps -c contrail-vrouter-dpdk -o %cpu="
    result = os_faults_steps.execute_cmd(computes, cmd)
    for node_result in result:
        usage = node_result.payload['stdout']
        usage = float(usage.strip())
        assert_that(usage, greater_than(50))

def test_vrouter_create_interface(request, os_faults_steps, computes):
    """Verify if vRouter creates interface after creation of a virtual machine.

    Steps:
        #. Remember count of virtual interfaces on vRouters
        #. Create server
        #. Check that count of virtual interfaces on server's compute vRouter
            is greater than it was.
    """
    before_ifaces = vrouter_steps.get_interface_table(os_faults_steps,
                                                      computes)
    server = request.getfixturevalue('server')
    compute_fqdn = getattr(server, settings.SERVER_ATTR_HYPERVISOR_HOSTNAME)
    after_ifaces = vrouter_steps.get_interface_table(os_faults_steps, computes)
    assert_that(after_ifaces[compute_fqdn],
                has_length(greater_than(len(before_ifaces[compute_fqdn]))))
<|endoftext|>


<commit_before>#!/usr/bin/python

## Binary Analysis Tool
## Copyright 2010-2011 Armijn Hemel for LOCO (LOOHUIS CONSULTING)
## Licensed under Apache 2.0, see LICENSE file for details

import sys, os, string, re
from optparse import OptionParser
import lucene

## helper method for processing stuff
## output is a tuple (file/dirname, config) for later processing
def matchconfig(filename, dirname, config, kerneldirlen):
	if filename.endswith(".o"):
		try:
			os.stat("%s/%s" % (dirname, filename[:-2] + ".c"))
			return ("%s/%s" % (dirname[kerneldirlen:], filename[:-2] + ".c"), config)
		except:
			return None
	else:
		if "arch" in filename:
			if dirname.split("/")[-2:] == filename.split("/")[:2]:
				try:
					newpath = dirname.split("/") + filename.split("/")[2:]
					os.stat(reduce(lambda x, y: x + "/" + y, newpath))
					return ("%s/%s" % (dirname[kerneldirlen:], filename), config)
				except:
					return None
			else:
				return None
		else:
			try:
				os.stat("%s/%s" % (dirname, filename))
				return ("%s/%s" % (dirname[kerneldirlen:], filename), config)
			except:
				return None

def extractkernelstrings(kerneldir):
	kerneldirlen = len(kerneldir)+1
	osgen = os.walk(kerneldir)
	searchresults = []

	try:
		while True:
                	i = osgen.next()
			## some top level dirs are not interesting
			if 'Documentation' in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('Documentation')
			if "scripts" in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('scripts')
			if "usr" in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('usr')
			if "samples" in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('samples')
                	for p in i[2]:
				## we only want Makefiles
				if p != 'Makefile':
					continue
				## not interested in the top level Makefile
				if i[0][kerneldirlen:] == "":
					continue
				source = open("%s/%s" % (i[0], p)).readlines()

				continued = False
				inif = False
				iniflevel = 0
				currentconfig = ""
				for line in source:
					if line.strip().startswith('#'):
						continue
					if line.strip().startswith('echo'):
						continue
					if line.strip().startswith('@'):
						continue
					if line.strip() == "":
						continue
					# if statements can be nested, so keep track of levels
					if line.strip() == "endif":
						inif = False
						iniflevel = iniflevel -1
						continue
					# if statements can be nested, so keep track of levels
					if re.match("ifn?\w+", line.strip()):
						inif = True
						iniflevel = iniflevel +1
					if not continued and line.strip().endswith("\\") and "=" in line.strip():
						## weed out more stuff
						## we are interested in three cases:
						## +=
						## :=
						## =  but only if there are object files or dirs defined in the right hand part
						continued = True
						currentconfig = ""
					elif continued and currentconfig != "":
						if line.strip().endswith("\\"):
							continued = True
							files = line.strip()[:-1].split()
						else:
							continued = False
							files = line.strip().split()
						for f in files:
							match = matchconfig(f, i[0], currentconfig, kerneldirlen)
							if match != None:
								searchresults.append(match)
					else:
						continued = False
						currentconfig = ""

					res = re.match("([\w\.]+)\-\$\(CONFIG_(\w+)\)\s*[:+]=\s*([\w\-\.\s/]*)", line.strip())
					if res != None:
						## current issues: ARCH (SH, Xtensa, h8300) is giving some issues
						if "flags" in res.groups()[0]:
							continue
						if "FLAGS" in res.groups()[0]:
							continue
						if "zimage" in res.groups()[0]:
							continue
						if res.groups()[0] == "defaultimage":
							continue
						if res.groups()[0] == "cacheflag":
							continue
						if res.groups()[0] == "cpuincdir":
							continue
						if res.groups()[0] == "cpuclass":
							continue
						if res.groups()[0] == "cpu":
							continue
						if res.groups()[0] == "machine":
							continue
						if res.groups()[0] == "model":
							continue
						if res.groups()[0] == "load":
							continue
						if res.groups()[0] == "dataoffset":
							continue
						if res.groups()[0] == "CPP_MODE":
							continue
						if res.groups()[0] == "LINK":
							continue
						config = "CONFIG_" + res.groups()[1]
						files = res.groups()[2].split()
						currentconfig = config
						for f in files:
							match = matchconfig(f, i[0], currentconfig, kerneldirlen)
							if match != None:
								searchresults.append(match)
				continue
	except StopIteration:
		return searchresults

def storematch(results, lucenewriter):
	## store two things:
	## 1. if we have a path/subdir, we store subdir + configuration
	##    making it searchable by subdir
	## 2. if we have an objectfile, we store name of source(!) file+ configuration
	##    making it searchable by source file
	## these can and will overlap
	for res in results:
		pathstring = res[0]
		configstring = res[1]
		doc = lucene.Document()
		doc.add(lucene.Field("name", pathstring,
			lucene.Field.Store.YES,
			lucene.Field.Index.NOT_ANALYZED))
		doc.add(lucene.Field("configstring", configstring,
			lucene.Field.Store.YES,
			lucene.Field.Index.NOT_ANALYZED))
		lucenewriter.addDocument(doc)

def main(argv):
        parser = OptionParser()
        parser.add_option("-d", "--directory", dest="kd", help="path to Linux kernel directory", metavar="DIR")
        parser.add_option("-i", "--index", dest="id", help="path to Lucene index directory", metavar="DIR")
        (options, args) = parser.parse_args()
        if options.kd == None:
                parser.error("Path to Linux kernel directory needed")
        if options.id == None:
                parser.error("Path to Lucene index directory needed")
        #try:
        	## open the Linux kernel directory and do some sanity checks
                #kernel_path = open(options.kd, 'rb')
        #except:
                #print "No valid Linux kernel directory"
                #sys.exit(1)
	# strip trailing slash, will not work this way if there are tons of slashes
	if options.kd.endswith('/'):
		kerneldir = options.kd[:-1]
	else:
		kerneldir = options.kd
	lucene.initVM()

	storeDir = options.id
        store = lucene.SimpleFSDirectory(lucene.File(storeDir))
	analyzer = lucene.StandardAnalyzer(lucene.Version.LUCENE_CURRENT)
        writer = lucene.IndexWriter(store, analyzer, True,
                                    lucene.IndexWriter.MaxFieldLength.LIMITED)
        writer.setMaxFieldLength(1048576)

	results = extractkernelstrings(kerneldir)
	storematch(results, writer)

        writer.optimize()
        writer.close()

if __name__ == "__main__":
        main(sys.argv)
<commit_msg>remove Lucene dependency, add documentation strings<commit_after>#!/usr/bin/python

## Binary Analysis Tool
## Copyright 2010-2011 Armijn Hemel for LOCO (LOOHUIS CONSULTING)
## Licensed under Apache 2.0, see LICENSE file for details

import sys, os, string, re
from optparse import OptionParser
import sqlite3

'''
This tool extracts configurations from Makefiles in kernels and tries to
determine which files are included by a certain configuration directive.
This information is useful to try and determine a reverse mapping from a
binary kernel image to a configuration.
'''

## helper method for processing stuff
## output is a tuple (file/dirname, config) for later processing
def matchconfig(filename, dirname, config, kerneldirlen):
	if filename.endswith(".o"):
		try:
			os.stat("%s/%s" % (dirname, filename[:-2] + ".c"))
			return ("%s/%s" % (dirname[kerneldirlen:], filename[:-2] + ".c"), config)
		except:
			return None
	else:
		if "arch" in filename:
			if dirname.split("/")[-2:] == filename.split("/")[:2]:
				try:
					newpath = dirname.split("/") + filename.split("/")[2:]
					os.stat(reduce(lambda x, y: x + "/" + y, newpath))
					return ("%s/%s" % (dirname[kerneldirlen:], filename), config)
				except:
					return None
			else:
				return None
		else:
			try:
				os.stat("%s/%s" % (dirname, filename))
				return ("%s/%s" % (dirname[kerneldirlen:], filename), config)
			except:
				return None

def extractkernelstrings(kerneldir):
	kerneldirlen = len(kerneldir)+1
	osgen = os.walk(kerneldir)
	searchresults = []

	try:
		while True:
                	i = osgen.next()
			## some top level dirs are not interesting
			if 'Documentation' in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('Documentation')
			if "scripts" in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('scripts')
			if "usr" in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('usr')
			if "samples" in i[1] and i[0][kerneldirlen:] == "":
				i[1].remove('samples')
                	for p in i[2]:
				## we only want Makefiles
				if p != 'Makefile':
					continue
				## not interested in the top level Makefile
				if i[0][kerneldirlen:] == "":
					continue
				source = open("%s/%s" % (i[0], p)).readlines()

				continued = False
				inif = False
				iniflevel = 0
				currentconfig = ""
				for line in source:
					if line.strip().startswith('#'):
						continue
					if line.strip().startswith('echo'):
						continue
					if line.strip().startswith('@'):
						continue
					if line.strip() == "":
						continue
					# if statements can be nested, so keep track of levels
					if line.strip() == "endif":
						inif = False
						iniflevel = iniflevel -1
						continue
					# if statements can be nested, so keep track of levels
					if re.match("ifn?\w+", line.strip()):
						inif = True
						iniflevel = iniflevel +1
					if not continued and line.strip().endswith("\\") and "=" in line.strip():
						## weed out more stuff
						## we are interested in three cases:
						## +=
						## :=
						## =  but only if there are object files or dirs defined in the right hand part
						continued = True
						currentconfig = ""
					elif continued and currentconfig != "":
						if line.strip().endswith("\\"):
							continued = True
							files = line.strip()[:-1].split()
						else:
							continued = False
							files = line.strip().split()
						for f in files:
							match = matchconfig(f, i[0], currentconfig, kerneldirlen)
							if match != None:
								searchresults.append(match)
					else:
						continued = False
						currentconfig = ""

					res = re.match("([\w\.]+)\-\$\(CONFIG_(\w+)\)\s*[:+]=\s*([\w\-\.\s/]*)", line.strip())
					if res != None:
						## current issues: ARCH (SH, Xtensa, h8300) is giving some issues
						if "flags" in res.groups()[0]:
							continue
						if "FLAGS" in res.groups()[0]:
							continue
						if "zimage" in res.groups()[0]:
							continue
						if res.groups()[0] == "defaultimage":
							continue
						if res.groups()[0] == "cacheflag":
							continue
						if res.groups()[0] == "cpuincdir":
							continue
						if res.groups()[0] == "cpuclass":
							continue
						if res.groups()[0] == "cpu":
							continue
						if res.groups()[0] == "machine":
							continue
						if res.groups()[0] == "model":
							continue
						if res.groups()[0] == "load":
							continue
						if res.groups()[0] == "dataoffset":
							continue
						if res.groups()[0] == "CPP_MODE":
							continue
						if res.groups()[0] == "LINK":
							continue
						config = "CONFIG_" + res.groups()[1]
						files = res.groups()[2].split()
						currentconfig = config
						for f in files:
							match = matchconfig(f, i[0], currentconfig, kerneldirlen)
							if match != None:
								searchresults.append(match)
				continue
	except StopIteration:
		return searchresults

def storematch(results, dbcursor):
	## store two things:
	## 1. if we have a path/subdir, we store subdir + configuration
	##    making it searchable by subdir
	## 2. if we have an objectfile, we store name of source(!) file+ configuration
	##    making it searchable by source file
	## these can and will overlap
	for res in results:
		pathstring = res[0]
		configstring = res[1]
		dbcursor.execute('''insert into config (configstring, filename) values (?, ?)''', (configstring, pathstring))

def main(argv):
        parser = OptionParser()
        parser.add_option("-d", "--directory", dest="kd", help="path to Linux kernel directory", metavar="DIR")
        parser.add_option("-i", "--index", dest="id", help="path to database", metavar="DIR")
        (options, args) = parser.parse_args()
        if options.kd == None:
                parser.error("Path to Linux kernel directory needed")
        if options.id == None:
                parser.error("Path to database needed")
        #try:
        	## open the Linux kernel directory and do some sanity checks
                #kernel_path = open(options.kd, 'rb')
        #except:
                #print "No valid Linux kernel directory"
                #sys.exit(1)
	# strip trailing slash, will not work this way if there are tons of slashes
	if options.kd.endswith('/'):
		kerneldir = options.kd[:-1]
	else:
		kerneldir = options.kd

        conn = sqlite3.connect(options.id)
        c = conn.cursor()

        try:
                c.execute('''create table config (configstring text, filename text)''')
        except:
                pass

        results = extractkernelstrings(kerneldir)
        storematch(results, c)

        conn.commit()
        c.close()

if __name__ == "__main__":
        main(sys.argv)
<|endoftext|>


<commit_before>def process(api, data):
    """
    Breakthroughs are rewarded to the first team to solve a given problem.

    Data Required: tid, pid
    """

    submission_count = api.problem.count_submissions(pid=data["pid"], correctness=True, eligibility=True)
    if submission_count == 1:
        submissions = api.problem.get_submissions(pid=data["pid"], correctness=True, eligibility=True)
        problem = api.problem.get_problem(pid=data["pid"])
        valid = submissions[0]['tid'] == data['tid']
    else:
        valid = False
    return valid, {
        "name": "Breakthrough!".format(problem["name"]),
        "description": "Your team was the first team to solve {}.".format(problem["name"])
    }
<commit_msg>Fixed breakthrough achievement bug<commit_after>def process(api, data):
    """
    Breakthroughs are rewarded to the first team to solve a given problem.

    Data Required: tid, pid
    """

    submission_count = api.problem.count_submissions(pid=data["pid"], correctness=True, eligibility=True)
    if submission_count == 1:
        submissions = api.problem.get_submissions(pid=data["pid"], correctness=True, eligibility=True)
        problem = api.problem.get_problem(pid=data["pid"])
        valid = submissions[0]['tid'] == data['tid']
        return valid, {
            "name": "Breakthrough!".format(problem["name"]),
            "description": "Your team was the first team to solve {}.".format(problem["name"])
        }
    else:
        return False, {}<|endoftext|>


<commit_before># Copyright (c) 2015 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

from . import STLWriter

from UM.i18n import i18nCatalog
i18n_catalog = i18nCatalog("uranium")

def getMetaData():
    return {
        "type": "mesh_writer",
        "plugin": {
            "name": "STL Writer",
            "author": "Ultimaker",
            "version": "1.0",
            "decription": i18n_catalog.i18nc("STL Writer plugin description", "Provides support for writing STL files.")
        },
        "mesh_reader": {
            "extension": "stl",
            "description": i18n_catalog.i18nc("STL Reader plugin file type", "STL File")
        }
    }

def register(app):
    return { "mesh_writer": STLWriter.STLWriter() }
<commit_msg>Fix STL Writer to properly advertise itself as mesh writer<commit_after># Copyright (c) 2015 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

from . import STLWriter

from UM.i18n import i18nCatalog
i18n_catalog = i18nCatalog("uranium")

def getMetaData():
    return {
        "type": "mesh_writer",
        "plugin": {
            "name": "STL Writer",
            "author": "Ultimaker",
            "version": "1.0",
            "decription": i18n_catalog.i18nc("STL Writer plugin description", "Provides support for writing STL files.")
        },
        "mesh_writer": {
            "extension": "stl",
            "description": i18n_catalog.i18nc("STL Reader plugin file type", "STL File")
        }
    }

def register(app):
    return { "mesh_writer": STLWriter.STLWriter() }
<|endoftext|>


<commit_before>import asyncio
import aioetcd.client

@asyncio.coroutine
def print_leader():
    client = aioetcd.client.Client()
    result = yield from client.leader()
    print(result)

@asyncio.coroutine
def client_set(key, value):
    client = aioetcd.client.Client()
    result = yield from client.set(key, value)
    print(result)

@asyncio.coroutine
def client_get(key):
    client = aioetcd.client.Client()
    result = yield from client.get(key)
    print(result)

@asyncio.coroutine
def print_machines():
    client = aioetcd.client.Client()
    result = yield from client.machines()
    print(result)

@asyncio.coroutine
def watch_eternally():
    client = aioetcd.client.Client(read_timeout=2, allow_reconnect=True)
    result = yield from client.watch('hello')
    print(result)

@asyncio.coroutine
def watch_timeout():
    client = aioetcd.client.Client(read_timeout=2, allow_reconnect=True)
    result = yield from client.watch('hello', timeout=5)
    print(result)

@asyncio.coroutine
def main():
    loop = asyncio.get_event_loop()
    loop.create_task(print_leader())
    loop.create_task(print_machines())
    loop.create_task(client_set('hop', 42))
    loop.create_task(client_get('hop'))
    yield from watch_eternally()
    yield from watch_timeout()

loop = asyncio.get_event_loop()
loop.run_until_complete(main())
loop.close()

<commit_msg>use watch_forever in example client<commit_after>import asyncio
import aioetcd.client

@asyncio.coroutine
def print_leader():
    print('print_leader')
    client = aioetcd.client.Client()
    result = yield from client.leader()
    print(result)

@asyncio.coroutine
def client_set(key, value):
    print('client_set')
    client = aioetcd.client.Client()
    result = yield from client.set(key, value)
    print(result)

@asyncio.coroutine
def client_get(key):
    print('client_get')
    client = aioetcd.client.Client()
    result = yield from client.get(key)
    print(result)

@asyncio.coroutine
def print_machines():
    print('print_machine')
    client = aioetcd.client.Client()
    result = yield from client.machines()
    print(result)

@asyncio.coroutine
def watch_eternally(key):
    print('watch_eternally')
    client = aioetcd.client.Client(read_timeout=2, allow_reconnect=True)
    result = yield from client.watch(key)
    print(result)

@asyncio.coroutine
def watch_timeout(key):
    print('watch_timeout')
    client = aioetcd.client.Client(read_timeout=2, allow_reconnect=True)
    result = yield from client.watch(key, timeout=5)
    print(result)

@asyncio.coroutine
def watch_forever(key):
    print('watch_forever')
    client = aioetcd.client.Client(read_timeout=5, allow_reconnect=True)
    for watcher in client.watch_forever(key):
        resp = yield from watcher
        print(resp.value)
        if resp.value == '42':
            break

@asyncio.coroutine
def main():
    loop = asyncio.get_event_loop()
    print('---------------------')
    yield from print_leader()
    print('---------------------')
    yield from print_machines()
    print('---------------------')
    yield from client_set('hello', 43)
    print('---------------------')
    yield from client_get('hello')
    print('---------------------')
    # run :
    #    etcdctl set hello 0
    # a couple of time and then
    #    etcdctl set hello 42
    # as soon as it becomes boring
    yield from watch_forever('hello')
    print('---------------------')
    # run :
    #    etcdctl set hello 0
    # to unlock
    yield from watch_eternally("hello")
    print('---------------------')
    # run :
    #    etcdctl set hello 0
    # or wait 5 seconds to unlock
    yield from watch_timeout('hello')

loop = asyncio.get_event_loop()
loop.run_until_complete(main())
loop.close()

<|endoftext|>


<commit_before>from apiclient.errors import HttpError
from oauth2client.client import AccessTokenRefreshError
from http_response_builder import HTTP_Response_Builder
from param_definition.parameter import Parameter, Date_Time_Parameter_Type, String_Parameter_Type, Integer_Parameter_Type
from datetime import datetime, timedelta, date
from model.db_session import DB_Session_Factory
from model.interview import Interview
from sqlalchemy import func
from lib.calendar import Google_Calendar, google_ts_to_datetime, LOS_ANGELES_TZ, datetime_to_google_ts
from model.admin_permissions import ROUNDUP_SCHEDULING_PERMISSION_MASK
import sys
import json
import urllib

class Schedule_Roundup_HTTP_Response_Builder(HTTP_Response_Builder):
    required_admin_permissions = ROUNDUP_SCHEDULING_PERMISSION_MASK
    roundup_start_date = Parameter('start_date', required = True, parameter_type = Date_Time_Parameter_Type)
    roundup_length = Parameter('length', required = True, parameter_type = Integer_Parameter_Type)
    candidate_name = Parameter('candidate_name', required = True, parameter_type = String_Parameter_Type)

    def print_body_for_user(self, authenticated_user):
        db_session = DB_Session_Factory.get_db_session()
        panel = []
        for interview in db_session.query(Interview).filter(Interview.candidate_name == self.candidate_name, func.date(Interview.start_time) == func.date(datetime.now())):
            panel.append({
                'email' : interview.interviewer_email,
                'optional' : True,
                'responseStatus' : 'needsAction'
            })
        if len(panel) is 0:
            print json.dumps({'error' : 'Nobody to invite to roundup for ' + self.candidate_name})
        else:
            event = {}
            event['organizer'] = {'self' : True}
            event['location'] = '4440-2-6 Nuthouse'
            event['start'] = datetime_to_google_ts(self.roundup_start_date)
            event['end'] = datetime_to_google_ts(self.roundup_start_date + timedelta(minutes = self.roundup_length))
            event['summary'] = "Roundup for " + self.candidate_name
            event['guestsCanInviteOthers'] = True
            event['attendees'] = panel
            now_date = date.today()
            event['description'] = "https://onsite-inflight.com/#" + str(now_date.year) + "-" + ("%02d"%now_date.month) + "-" + ("%02d"%now_date.day) + "/" + urllib.quote_plus(self.candidate_name)
            calendar = Google_Calendar.get_calendar()
            try:
                calendar.service.events().insert(calendarId = "box.com_k7897kemv3a14qldj8s546jfig@group.calendar.google.com", body = event, sendNotifications=True).execute(calendar.http)
                print json.dumps({'status' : 'success'})
            except HttpError as e:
                error_response = json.dumps({"error" : "Google calendar http error: " + e.content})
                sys.stderr.write(error_response)
                print error_response
            except AccessTokenRefreshError:
                print json.dumps({"error" : "The Google calendar credentials have been revoked or expired"})
<commit_msg>Changed roundup location to Lambda<commit_after>from apiclient.errors import HttpError
from oauth2client.client import AccessTokenRefreshError
from http_response_builder import HTTP_Response_Builder
from param_definition.parameter import Parameter, Date_Time_Parameter_Type, String_Parameter_Type, Integer_Parameter_Type
from datetime import datetime, timedelta, date
from model.db_session import DB_Session_Factory
from model.interview import Interview
from sqlalchemy import func
from lib.calendar import Google_Calendar, google_ts_to_datetime, LOS_ANGELES_TZ, datetime_to_google_ts
from model.admin_permissions import ROUNDUP_SCHEDULING_PERMISSION_MASK
import sys
import json
import urllib

class Schedule_Roundup_HTTP_Response_Builder(HTTP_Response_Builder):
    required_admin_permissions = ROUNDUP_SCHEDULING_PERMISSION_MASK
    roundup_start_date = Parameter('start_date', required = True, parameter_type = Date_Time_Parameter_Type)
    roundup_length = Parameter('length', required = True, parameter_type = Integer_Parameter_Type)
    candidate_name = Parameter('candidate_name', required = True, parameter_type = String_Parameter_Type)

    def print_body_for_user(self, authenticated_user):
        db_session = DB_Session_Factory.get_db_session()
        panel = []
        for interview in db_session.query(Interview).filter(Interview.candidate_name == self.candidate_name, func.date(Interview.start_time) == func.date(datetime.now())):
            panel.append({
                'email' : interview.interviewer_email,
                'optional' : True,
                'responseStatus' : 'needsAction'
            })
        if len(panel) is 0:
            print json.dumps({'error' : 'Nobody to invite to roundup for ' + self.candidate_name})
        else:
            event = {}
            event['organizer'] = {'self' : True}
            event['location'] = 'RWC-A-423-Lambda-5-AV2.1'
            event['start'] = datetime_to_google_ts(self.roundup_start_date)
            event['end'] = datetime_to_google_ts(self.roundup_start_date + timedelta(minutes = self.roundup_length))
            event['summary'] = "Roundup for " + self.candidate_name
            event['guestsCanInviteOthers'] = True
            event['attendees'] = panel
            now_date = date.today()
            event['description'] = "https://onsite-inflight.com/#" + str(now_date.year) + "-" + ("%02d"%now_date.month) + "-" + ("%02d"%now_date.day) + "/" + urllib.quote_plus(self.candidate_name)
            calendar = Google_Calendar.get_calendar()
            try:
                calendar.service.events().insert(calendarId = "box.com_k7897kemv3a14qldj8s546jfig@group.calendar.google.com", body = event, sendNotifications=True).execute(calendar.http)
                print json.dumps({'status' : 'success'})
            except HttpError as e:
                error_response = json.dumps({"error" : "Google calendar http error: " + e.content})
                sys.stderr.write(error_response)
                print error_response
            except AccessTokenRefreshError:
                print json.dumps({"error" : "The Google calendar credentials have been revoked or expired"})
<|endoftext|>


<commit_before>""" Logging configuration.

We configure separate loggers for general server logging and per-user loggers
for mail sync and contact sync.
"""
import os
import sys
import socket
import logging
import traceback

import requests

from colorlog import ColoredFormatter
from gevent import GreenletExit

from inbox.util.file import mkdirp

file_formatter = logging.Formatter(
    "[%(levelname)-.1s %(asctime)s %(module)-8.8s:%(lineno)-4s] %(message)s")

def get_tty_handler():
    tty_formatter = ColoredFormatter(
        # wish we could left-truncate name!
        "%(name)-20.20s %(log_color)s[%(levelname)-.1s %(asctime)s %(module)-8.8s:%(lineno)-4s]%(reset)s %(message)s",
        reset=True,
        log_colors={
            'DEBUG':    'cyan',
            'INFO':     'green',
            'WARNING':  'yellow',
            'ERROR':    'red',
            'CRITICAL': 'red',
        })

    tty_handler = logging.StreamHandler()
    tty_handler.setFormatter(tty_formatter)
    return tty_handler

def get_logger(account_id=None, purpose=None):
    """ Helper for abstracting away our logger names. """
    if account_id is None:
        return logging.getLogger('inbox.server.general')
    else:
        if purpose is None:
            purpose = 'sync'
        return logging.getLogger('inbox.server.{1}.{0}'.format(
            account_id, purpose))

def configure_general_logging():
    """ Configure the general server logger to output to screen if a TTY is
        attached, and server.log always.

        Logs are output to a directory configurable via LOGDIR.
    """
    # import here to avoid import loop from config.py
    from .config import config
    assert 'LOGDIR' in config, "LOGDIR not specified in config file"
    assert 'LOGLEVEL' in config, "LOGLEVEL not specified in config file"
    mkdirp(config['LOGDIR'])

    # configure properties that should cascade
    inbox_root_logger = logging.getLogger('inbox.server')
    inbox_root_logger.setLevel(int(config['LOGLEVEL']))
    # don't pass messages up to the root root logger
    inbox_root_logger.propagate = False

    # log everything to screen
    if sys.stdout.isatty():
        inbox_root_logger.addHandler(get_tty_handler())

    logger = get_logger()

    for handler in logger.handlers:
        logger.removeHandler(handler)

    logfile = os.path.join(config['LOGDIR'], 'server.log')
    file_handler = logging.FileHandler(logfile, encoding='utf-8')
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)

    return logger

def configure_logging(account_id, purpose):
    # avoid import loop from config.py
    from .config import config
    logger = get_logger(account_id, purpose)
    logger.propagate = True

    logdir = os.path.join(config['LOGDIR'], str(account_id))
    mkdirp(logdir)
    logfile = os.path.join(logdir, '{0}.log'.format(purpose))
    handler = logging.FileHandler(logfile, encoding='utf-8')
    handler.setFormatter(file_formatter)
    logger.addHandler(handler)

    return logger

def configure_mailsync_logging(account_id):
    """ We log output for each mail sync instance to a different file than the
        main server log, for ease of debugging. Sync logs still go to screen
        too, for now.
    """
    return configure_logging(account_id, "mailsync")

def configure_contacts_logging(account_id):
    """ We log output for each contacts sync instance to a different file than
        the main server log, for ease of debugging. Contacts sync logs still go
        to screen too, for now.
    """
    return configure_logging(account_id, "contacts")

def email_exception(logger, etype, evalue, tb):
    """ Send stringified exception to configured email address. """

    from inbox.server.config import config

    exc_email_addr = config.get('EXCEPTION_EMAIL_ADDRESS')
    if exc_email_addr is None:
        logger.error("No EXCEPTION_EMAIL_ADDRESS configured!")
    mailgun_api_endpoint = config.get('MAILGUN_API_ENDPOINT')
    if mailgun_api_endpoint is None:
        logger.error("No MAILGUN_API_ENDPOINT configured!")
    mailgun_api_key = config.get('MAILGUN_API_KEY')
    if mailgun_api_key is None:
        logger.error("No MAILGUN_API_KEY configured!")

    r = requests.post(
        mailgun_api_endpoint,
        auth=("api", mailgun_api_key),
        data={"from": "Inbox App Server <{}>".format(exc_email_addr),
              "to": [exc_email_addr],
              "subject": "Uncaught error! {} {}".format(etype, evalue),
              "text": u"""
    Something went wrong on {}. Please investigate. :)

    {}

    """.format(socket.getfqdn(),
               '\t'.join(traceback.format_exception(etype, evalue, tb)))})
    if r.status_code != requests.codes.ok:
        logger.error("Couldn't send exception email: {}".format(r.json()))

class log_uncaught_errors(object):
    """ Helper to log uncaught exceptions raised within the wrapped function.

        Modeled after gevent.util.wrap_errors.

        Parameters
        ----------
        func: function
            The function to wrap.

        logger: logging.Logger, optional
            The logging object to write to.
    """

    def __init__(self, func, logger=None):
        self.func = func
        self.logger = logger or get_logger()

    def _log_failsafe(self, *args, **kwargs):
        # We wrap the logging call in a try/except block so that if it fails
        # for any reason, the *original* error still gets propagated.
        try:
            self.logger.exception(*args, **kwargs)
        except:
            pass

    def __call__(self, *args, **kwargs):
        from inbox.server.config import is_prod, is_staging
        func = self.func
        try:
            return func(*args, **kwargs)
        except Exception, e:
            if not isinstance(e, GreenletExit):
                self._log_failsafe("Uncaught error!")
                exc_type, exc_value, exc_tb = sys.exc_info()
                if is_prod() or is_staging():
                    email_exception(self.logger, exc_type, exc_value, exc_tb)
            raise

    def __str__(self):
        return str(self.func)

    def __repr__(self):
        return repr(self.func)

    def __getattr__(self, item):
        return getattr(self.func, item)
<commit_msg>Fix default logger purpose.<commit_after>""" Logging configuration.

We configure separate loggers for general server logging and per-user loggers
for mail sync and contact sync.
"""
import os
import sys
import socket
import logging
import traceback

import requests

from colorlog import ColoredFormatter
from gevent import GreenletExit

from inbox.util.file import mkdirp

file_formatter = logging.Formatter(
    "[%(levelname)-.1s %(asctime)s %(module)-8.8s:%(lineno)-4s] %(message)s")

def get_tty_handler():
    tty_formatter = ColoredFormatter(
        # wish we could left-truncate name!
        "%(name)-20.20s %(log_color)s[%(levelname)-.1s %(asctime)s %(module)-8.8s:%(lineno)-4s]%(reset)s %(message)s",
        reset=True,
        log_colors={
            'DEBUG':    'cyan',
            'INFO':     'green',
            'WARNING':  'yellow',
            'ERROR':    'red',
            'CRITICAL': 'red',
        })

    tty_handler = logging.StreamHandler()
    tty_handler.setFormatter(tty_formatter)
    return tty_handler

def get_logger(account_id=None, purpose=None):
    """ Helper for abstracting away our logger names. """
    if account_id is None:
        return logging.getLogger('inbox.server.general')
    else:
        if purpose is None:
            purpose = 'mailsync'
        return logging.getLogger('inbox.server.{1}.{0}'.format(
            account_id, purpose))

def configure_general_logging():
    """ Configure the general server logger to output to screen if a TTY is
        attached, and server.log always.

        Logs are output to a directory configurable via LOGDIR.
    """
    # import here to avoid import loop from config.py
    from .config import config
    assert 'LOGDIR' in config, "LOGDIR not specified in config file"
    assert 'LOGLEVEL' in config, "LOGLEVEL not specified in config file"
    mkdirp(config['LOGDIR'])

    # configure properties that should cascade
    inbox_root_logger = logging.getLogger('inbox.server')
    inbox_root_logger.setLevel(int(config['LOGLEVEL']))
    # don't pass messages up to the root root logger
    inbox_root_logger.propagate = False

    # log everything to screen
    if sys.stdout.isatty():
        inbox_root_logger.addHandler(get_tty_handler())

    logger = get_logger()

    for handler in logger.handlers:
        logger.removeHandler(handler)

    logfile = os.path.join(config['LOGDIR'], 'server.log')
    file_handler = logging.FileHandler(logfile, encoding='utf-8')
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)

    return logger

def configure_logging(account_id, purpose):
    # avoid import loop from config.py
    from .config import config
    logger = get_logger(account_id, purpose)
    logger.propagate = True

    logdir = os.path.join(config['LOGDIR'], str(account_id))
    mkdirp(logdir)
    logfile = os.path.join(logdir, '{0}.log'.format(purpose))
    handler = logging.FileHandler(logfile, encoding='utf-8')
    handler.setFormatter(file_formatter)
    logger.addHandler(handler)

    return logger

def configure_mailsync_logging(account_id):
    """ We log output for each mail sync instance to a different file than the
        main server log, for ease of debugging. Sync logs still go to screen
        too, for now.
    """
    return configure_logging(account_id, "mailsync")

def configure_contacts_logging(account_id):
    """ We log output for each contacts sync instance to a different file than
        the main server log, for ease of debugging. Contacts sync logs still go
        to screen too, for now.
    """
    return configure_logging(account_id, "contacts")

def email_exception(logger, etype, evalue, tb):
    """ Send stringified exception to configured email address. """

    from inbox.server.config import config

    exc_email_addr = config.get('EXCEPTION_EMAIL_ADDRESS')
    if exc_email_addr is None:
        logger.error("No EXCEPTION_EMAIL_ADDRESS configured!")
    mailgun_api_endpoint = config.get('MAILGUN_API_ENDPOINT')
    if mailgun_api_endpoint is None:
        logger.error("No MAILGUN_API_ENDPOINT configured!")
    mailgun_api_key = config.get('MAILGUN_API_KEY')
    if mailgun_api_key is None:
        logger.error("No MAILGUN_API_KEY configured!")

    r = requests.post(
        mailgun_api_endpoint,
        auth=("api", mailgun_api_key),
        data={"from": "Inbox App Server <{}>".format(exc_email_addr),
              "to": [exc_email_addr],
              "subject": "Uncaught error! {} {}".format(etype, evalue),
              "text": u"""
    Something went wrong on {}. Please investigate. :)

    {}

    """.format(socket.getfqdn(),
               '\t'.join(traceback.format_exception(etype, evalue, tb)))})
    if r.status_code != requests.codes.ok:
        logger.error("Couldn't send exception email: {}".format(r.json()))

class log_uncaught_errors(object):
    """ Helper to log uncaught exceptions raised within the wrapped function.

        Modeled after gevent.util.wrap_errors.

        Parameters
        ----------
        func: function
            The function to wrap.

        logger: logging.Logger, optional
            The logging object to write to.
    """

    def __init__(self, func, logger=None):
        self.func = func
        self.logger = logger or get_logger()

    def _log_failsafe(self, *args, **kwargs):
        # We wrap the logging call in a try/except block so that if it fails
        # for any reason, the *original* error still gets propagated.
        try:
            self.logger.exception(*args, **kwargs)
        except:
            pass

    def __call__(self, *args, **kwargs):
        from inbox.server.config import is_prod, is_staging
        func = self.func
        try:
            return func(*args, **kwargs)
        except Exception, e:
            if not isinstance(e, GreenletExit):
                self._log_failsafe("Uncaught error!")
                exc_type, exc_value, exc_tb = sys.exc_info()
                if is_prod() or is_staging():
                    email_exception(self.logger, exc_type, exc_value, exc_tb)
            raise

    def __str__(self):
        return str(self.func)

    def __repr__(self):
        return repr(self.func)

    def __getattr__(self, item):
        return getattr(self.func, item)
<|endoftext|>


<commit_before>from unittest import TestCase
from thrift.Thrift import TType
from thriftmodel.protocol import Protocol
from test.helpers import flatten
from thriftmodel.TFlexibleJSONProtocol import ReadValidationException
from thriftmodel.model import (serialize, deserialize,
        ThriftField, ThriftModel, IntField, ListField,
        MapField, StringField, UTF8Field, StructField, serialize, deserialize,
        ValidationException, BoolField, FieldMerger, ThriftFieldMergeException)

class FieldMergerTestCase(TestCase):

    def test_basic_merge(self):
        field1 = IntField(thrift_field_name='a', field_id=2)
        field2 = IntField(thrift_field_name='b', field_id=1)
        fm = FieldMerger(
            [('a', field1)],
            [('b', field2)])
        result = sorted(fm.merge(), key=lambda x:x[0])
        # basic merge scenario: both list have non-conflicting fields
        self.assertEquals(result, [('a', field1), ('b', field2)])

        # test_empty_to_merge_list
        fm = FieldMerger([('a', field1), ('b', field2)], [])
        result = sorted(fm.merge(), key=lambda x:x[0])
        self.assertEquals(result, [('a', field1), ('b', field2)])

        # test empty original list
        fm = FieldMerger([], [('a', field1), ('b', field2)])
        result = sorted(fm.merge(), key=lambda x:x[0])
        self.assertEquals(result, [('a', field1), ('b', field2)])

    def test_field_name_mandatory(self):
        """ the thrift_field_name attribute is mandatory.
        """
        field1 = IntField(field_id=2, thrift_field_name='a')
        field2 = IntField(field_id=1)
        def f():
            fm = FieldMerger(
                [('a', field1)],
                [('b', field2)])
        self.assertRaises(ThriftFieldMergeException, lambda: f())

        field1 = IntField(field_id=2)
        field2 = IntField(field_id=1)
        def f():
            fm = FieldMerger(
                [('a', field1)],
                [('b', field2)])
        self.assertRaises(ThriftFieldMergeException, lambda: f())

    def test_field_id_mandatory(self):
        """ Field id is mandatory in the original field list.
        Optional in to_merge fields."""
        field1 = IntField(field_id=1, thrift_field_name='a')
        field2 = IntField(thrift_field_name='b')
        fm = FieldMerger(
            [('a', field1)],
            [('b', field2)])
        result = sorted(fm.merge(), key=lambda x:x[0])
        self.assertEquals(result, [('a', field1), ('b', field2)])

        field1 = IntField(thrift_field_name='a')
        field2 = IntField(field_id=1, thrift_field_name='b')
        def fun():
            fm = FieldMerger(
                [('b', field1)],
                [('a', field2)])
            fm.merge()
        self.assertRaises(ThriftFieldMergeException, lambda: fun())

class ModelInheritanceTestCase(TestCase):

    def test_inherit_from_generated_python(self):
        from test.compiled.py.ttypes import SomeStruct
        class A(SomeStruct, ThriftModel):
            a = IntField()
        # A will have one more field than SomeStruct
        self.assertTrue(len(A.thrift_spec) == (len(SomeStruct.thrift_spec) + 1))
        # the last field of the thrift_spec will be a
        self.assertEquals(A.thrift_spec[-1][2], "a")

    def test_model_inheritance(self):
        class A(ThriftModel):
            a = IntField()
        class B(A):
            b = UTF8Field()
        class C(B):
            c = BoolField()

        class NoInherit(ThriftModel):
            a = IntField()
            b = UTF8Field()
            c = BoolField()

        self.assertEquals(C.thrift_spec, NoInherit.thrift_spec)

    def atest_field_override(self):
        # TODO: inheritance works based on the thrift_spec of the base class.
        # This means validators are lost and duplicate field names are used.
        # FIX THIS!
        class A(ThriftModel):
            a = IntField()
        class B(A):
            b = UTF8Field()
        class C(B):
            b = BoolField()

        class NoInherit(ThriftModel):
            a = IntField()
            b = BoolField()

        self.assertEquals(C.thrift_spec, NoInherit.thrift_spec)

<commit_msg>removed test_inheritance: struct inherintance is not supported<commit_after><|endoftext|>


<commit_before>from marshmallow_pynamodb import ModelSchema

from pynamodb.attributes import ListAttribute, MapAttribute, NumberAttribute, UnicodeAttribute
from pynamodb.models import Model

from unittest import TestCase

from os import environ

class Location(MapAttribute):
    latitude = NumberAttribute()
    longitude = NumberAttribute()
    name = UnicodeAttribute()

class Person(MapAttribute):
    firstName = UnicodeAttribute()
    lastName = UnicodeAttribute()
    age = NumberAttribute()

class OfficeEmployeeMap(MapAttribute):
    office_employee_id = NumberAttribute()
    person = Person()
    office_location = Location()

class Office(Model):
    class Meta:
        table_name = 'OfficeModel'
        host = "http://localhost:{}".format(environ['DOCKER_PORT'])

    office_id = NumberAttribute(hash_key=True)
    address = Location()
    employees = ListAttribute(of=OfficeEmployeeMap)

class OfficeSchema(ModelSchema):
    class Meta:
        model = Office

class TestValidatedModelSchema(TestCase):
    def setUp(self):
        if not Office.exists():
            Office.create_table(read_capacity_units=1, write_capacity_units=1)

        self.hash_key = 789
        self.attrs = {'address': {'latitude': 6.98454,
                                  'longitude': 172.38832,
                                  'name': 'some_location'},
                      'employees': [{'office_employee_id': 123,
                                     'person': {'firstName': 'John',
                                                'lastName': 'Smith',
                                                'age': 45},
                                     'office_location': {'latitude': -24.0853,
                                                         'longitude': 144.87660,
                                                         'name': 'other_location'}},
                                    {'office_employee_id': 456,
                                     'person': {'firstName': 'Jane',
                                                'lastName': 'Doe',
                                                'age': 33},
                                     'office_location': {'latitude': -20.57989,
                                                         'longitude': 92.30463,
                                                         'name': 'yal'}}]}

    def test_load(self):
        self.attrs['office_id'] = self.hash_key

        data, errors = OfficeSchema().load(self.attrs)

        self.assertTrue(getattr(OfficeSchema, '_declared_fields')['office_id'].required)
        self.assertDictEqual(errors, {})

        self.assertIsInstance(data, Office)
        self.assertEqual(data.attribute_values['office_id'], self.hash_key)
        self.assertIsInstance(data.attribute_values['address'], Location)
        self.assertIsInstance(data.attribute_values['employees'], list)
        self.assertIsInstance(data.attribute_values['employees'][0], OfficeEmployeeMap)
        self.assertIsInstance(data.attribute_values['employees'][0]['person'], Person)
        self.assertIsInstance(data.attribute_values['employees'][0]['office_location'], Location)

    def test_load_fails_dict_level_0(self):
        self.attrs['office_id'] = 'foobar'

        data, errors = OfficeSchema().load(self.attrs)

        self.assertDictEqual(errors, {'office_id': [u'Not a valid number.']})

    def test_load_fails_dict_level_gt_0(self):
        self.attrs['office_id'] = self.hash_key
        self.attrs['employees'][0]['office_location']['latitude'] = 'foobar'

        data, errors = OfficeSchema().load(self.attrs)

        self.assertDictEqual(errors, {'employees': {0: {'office_location': {'latitude': [u'Not a valid number.']}}}})

    def test_dump(self):
        model = Office(hash_key=self.hash_key, **self.attrs)
        self.attrs['office_id'] = self.hash_key

        data, errors = OfficeSchema().dump(model)

        self.assertDictEqual(errors, {})
        self.assertDictEqual(data, self.attrs)

    def test_save_and_get_dict(self):
        self.attrs['office_id'] = self.hash_key

        office = OfficeSchema().load(self.attrs).data
        office.save()

        office_model = Office.get(789)
        office_json = OfficeSchema().dump(office_model).data

        self.assertDictEqual(office_json, self.attrs)

    def test_objs_serialization(self):
        self.maxDiff = None

        self.attrs['office_id'] = self.hash_key

        john = Person(
            firstName='John',
            lastName='Smith',
            age=45,
        )
        jane = Person(
            firstName='Jane',
            lastName='Doe',
            age=33,
        )
        loc = Location(
            latitude=6.98454,
            longitude=172.38832,
            name='some_location'
        )
        loc1 = Location(
            latitude=-24.0853,
            longitude=144.87660,
            name='other_location'
        )
        loc2 = Location(
            latitude=-20.57989,
            longitude=92.30463,
            name='yal'
        )
        emp1 = OfficeEmployeeMap(
            office_employee_id=123,
            person=john,
            office_location=loc1
        )
        emp2 = OfficeEmployeeMap(
            office_employee_id=456,
            person=jane,
            office_location=loc2
        )

        office = Office(
            office_id=self.hash_key,
            address=loc,
            employees=[emp1, emp2]
        )

        office_json = OfficeSchema().dump(office).data
        self.assertDictEqual(office_json, self.attrs)

    def tearDown(self):
        Office(hash_key=self.hash_key).delete()
<commit_msg>Updated integration tests with Set fields<commit_after>from marshmallow_pynamodb import ModelSchema

from pynamodb.attributes import ListAttribute, MapAttribute, NumberAttribute, UnicodeAttribute, UnicodeSetAttribute, NumberSetAttribute
from pynamodb.models import Model

from unittest import TestCase

from os import environ

class Location(MapAttribute):
    latitude = NumberAttribute()
    longitude = NumberAttribute()
    name = UnicodeAttribute()

class Person(MapAttribute):
    firstName = UnicodeAttribute()
    lastName = UnicodeAttribute()
    age = NumberAttribute()

class OfficeEmployeeMap(MapAttribute):
    office_employee_id = NumberAttribute()
    person = Person()
    office_location = Location()

class Office(Model):
    class Meta:
        table_name = 'OfficeModel'
        host = "http://localhost:{}".format(environ['DOCKER_PORT'])

    office_id = NumberAttribute(hash_key=True)
    address = Location()
    employees = ListAttribute(of=OfficeEmployeeMap)
    departments = UnicodeSetAttribute()
    numbers = NumberSetAttribute()

class OfficeSchema(ModelSchema):
    class Meta:
        model = Office

class TestValidatedModelSchema(TestCase):
    def setUp(self):
        if not Office.exists():
            Office.create_table(read_capacity_units=1, write_capacity_units=1)

        self.hash_key = 789
        self.attrs = {'address': {'latitude': 6.98454,
                                  'longitude': 172.38832,
                                  'name': 'some_location'},
                      'employees': [{'office_employee_id': 123,
                                     'person': {'firstName': 'John',
                                                'lastName': 'Smith',
                                                'age': 45},
                                     'office_location': {'latitude': -24.0853,
                                                         'longitude': 144.87660,
                                                         'name': 'other_location'}},
                                    {'office_employee_id': 456,
                                     'person': {'firstName': 'Jane',
                                                'lastName': 'Doe',
                                                'age': 33},
                                     'office_location': {'latitude': -20.57989,
                                                         'longitude': 92.30463,
                                                         'name': 'yal'}}],
                      'departments': ['engineering', 'dev-ops', 'UI/UX', 'sales'],
                      'numbers': [1, 2, 3, 4, 5, 6]}

    def test_load(self):
        self.attrs['office_id'] = self.hash_key

        data, errors = OfficeSchema().load(self.attrs)

        self.assertTrue(getattr(OfficeSchema, '_declared_fields')['office_id'].required)
        self.assertDictEqual(errors, {})

        self.assertIsInstance(data, Office)
        self.assertEqual(data.attribute_values['office_id'], self.hash_key)
        self.assertIsInstance(data.attribute_values['address'], Location)
        self.assertIsInstance(data.attribute_values['departments'], set)
        self.assertIsInstance(data.attribute_values['numbers'], set)
        self.assertIsInstance(data.attribute_values['employees'], list)
        self.assertIsInstance(data.attribute_values['employees'][0], OfficeEmployeeMap)
        self.assertIsInstance(data.attribute_values['employees'][0]['person'], Person)
        self.assertIsInstance(data.attribute_values['employees'][0]['office_location'], Location)

    def test_load_fails_dict_level_0(self):
        self.attrs['office_id'] = 'foobar'

        data, errors = OfficeSchema().load(self.attrs)

        self.assertDictEqual(errors, {'office_id': [u'Not a valid number.']})

    def test_load_fails_dict_level_gt_0(self):
        self.attrs['office_id'] = self.hash_key
        self.attrs['employees'][0]['office_location']['latitude'] = 'foobar'

        data, errors = OfficeSchema().load(self.attrs)

        self.assertDictEqual(errors, {'employees': {0: {'office_location': {'latitude': [u'Not a valid number.']}}}})

    def test_dump(self):
        model = Office(hash_key=self.hash_key, **self.attrs)
        self.attrs['office_id'] = self.hash_key

        data, errors = OfficeSchema().dump(model)

        self.assertDictEqual(errors, {})
        self.assertDictEqual(data, self.attrs)

    def test_save_and_get_dict(self):
        self.maxDiff = None
        self.attrs['office_id'] = self.hash_key

        office = OfficeSchema().load(self.attrs).data
        office.save()

        office_model = Office.get(789)
        office_json = OfficeSchema().dump(office_model).data
        office_json['departments'].sort()
        self.attrs['departments'].sort()

        self.assertDictEqual(office_json, self.attrs)

    def test_objs_serialization(self):
        self.maxDiff = None

        self.attrs['office_id'] = self.hash_key

        john = Person(
            firstName='John',
            lastName='Smith',
            age=45,
        )
        jane = Person(
            firstName='Jane',
            lastName='Doe',
            age=33,
        )
        loc = Location(
            latitude=6.98454,
            longitude=172.38832,
            name='some_location'
        )
        loc1 = Location(
            latitude=-24.0853,
            longitude=144.87660,
            name='other_location'
        )
        loc2 = Location(
            latitude=-20.57989,
            longitude=92.30463,
            name='yal'
        )
        emp1 = OfficeEmployeeMap(
            office_employee_id=123,
            person=john,
            office_location=loc1
        )
        emp2 = OfficeEmployeeMap(
            office_employee_id=456,
            person=jane,
            office_location=loc2
        )

        office = Office(
            office_id=self.hash_key,
            address=loc,
            employees=[emp1, emp2],
            departments=('engineering', 'dev-ops', 'UI/UX', 'sales'),
            numbers=(1, 2, 3, 4, 5, 6)

        )

        office_json = OfficeSchema().dump(office).data
        self.assertDictEqual(office_json, self.attrs)

    def tearDown(self):
        Office(hash_key=self.hash_key).delete()
<|endoftext|>


<commit_before>from plex import Plex
from plex.core.helpers import to_iterable
from plex_metadata.guid import Guid
from plex_metadata.matcher import Default as Matcher
from plex_metadata.metadata import Default as Metadata

class Library(object):
    @classmethod
    def all(cls, types=None, keys=None, titles=None):
        types = to_iterable(types)

        sections = Plex['library'].sections().filter(types, keys, titles)
        result = {}

        for section in sections:
            if section.type not in result:
                result[section.type] = {}

            for item in section.all():
                cls.item_map(result[section.type], item)

        if types and len(types) == 1:
            # Return single type-map if only one was requested
            return result.get(types[0], {})

        return result

    # TODO Library.episodes() `parent`?
    @classmethod
    def episodes(cls, key, show=None):
        result = {}

        container = Plex['library/metadata'].all_leaves(key)

        if not container:
            return None

        for item in container:
            if show is not None:
                item.show = show

            season, episodes = Matcher.process(item)

            if not season or not episodes:
                continue

            for episode in episodes:
                result[season, episode] = item

        return result

    @classmethod
    def item_map(cls, table, item):
        metadata = Metadata.get(item.rating_key)

        # Update with extended information
        item.guid = metadata.guid

        # Parse guid
        guid = Guid.parse(item.guid)
        key = (guid.agent, guid.sid)

        # Map item into table
        if key not in table:
            table[key] = []

        table[key].append(item)
        return True
<commit_msg>Fixed bug with Library.episodes() skipping specials<commit_after>from plex import Plex
from plex.core.helpers import to_iterable
from plex_metadata.guid import Guid
from plex_metadata.matcher import Default as Matcher
from plex_metadata.metadata import Default as Metadata

class Library(object):
    @classmethod
    def all(cls, types=None, keys=None, titles=None):
        types = to_iterable(types)

        sections = Plex['library'].sections().filter(types, keys, titles)
        result = {}

        for section in sections:
            if section.type not in result:
                result[section.type] = {}

            for item in section.all():
                cls.item_map(result[section.type], item)

        if types and len(types) == 1:
            # Return single type-map if only one was requested
            return result.get(types[0], {})

        return result

    @classmethod
    def episodes(cls, key, show=None):
        result = {}

        container = Plex['library/metadata'].all_leaves(key)

        if not container:
            return None

        for item in container:
            if show is not None:
                item.show = show

            season, episodes = Matcher.process(item)

            if season is None or not episodes:
                continue

            for episode in episodes:
                result[season, episode] = item

        return result

    @classmethod
    def item_map(cls, table, item):
        metadata = Metadata.get(item.rating_key)

        # Update with extended information
        item.guid = metadata.guid

        # Parse guid
        guid = Guid.parse(item.guid)
        key = (guid.agent, guid.sid)

        # Map item into table
        if key not in table:
            table[key] = []

        table[key].append(item)
        return True
<|endoftext|>