<commit_before>import tensorflow as tf

import sys

sys.path.append('../')
import tfutil as t

tf.set_random_seed(777)

# In image_utils, up/down_sampling
def image_sampling(img, sampling_type='down'):
    shape = img.get_shape()  # [batch, height, width, channels]

    if sampling_type == 'down':
        h = int(shape[1] // 2)
        w = int(shape[2] // 2)
    else:  # 'up'
        h = int(shape[1] * 2)
        w = int(shape[2] * 2)

    return tf.image.resize_images(img, [h, w], tf.image.ResizeMethod.BILINEAR)

class LAPGAN:

    def __init__(self, s, batch_size=128, height=32, width=32, channel=3, n_classes=10,
                 sample_num=10 * 10, sample_size=10,
                 z_dim=128, gf_dim=64, df_dim=64, d_fc_unit=512, g_fc_unit=1024):

        """
        # General Settings
        :param s: TF Session
        :param batch_size: training batch size, default 128
        :param height: input image height, default 32
        :param width: input image width, default 32
        :param channel: input image channel, default 3 (RGB)
        :param n_classes: the number of classes, default 10
        - in case of CIFAR, image size is 32x32x3(HWC), classes are 10.

        # Output Settings
        :param sample_size: sample image size, default 8
        :param sample_num: the number of sample images, default 64

        # Model Settings
        :param z_dim: z noise dimension, default 128
        :param gf_dim: the number of generator filters, default 64
        :param df_dim: the number of discriminator filters, default 64
        :param d_fc_unit: the number of fully connected filters used at Disc, default 512
        :param g_fc_unit: the number of fully connected filters used at Gen, default 1024
        """

        self.s = s
        self.batch_size = batch_size
        self.height = height
        self.width = width
        self.channel = channel
        self.n_classes = n_classes

        self.sample_size = sample_size
        self.sample_num = sample_num

        self.image_shape = [self.height, self.width, self.channel]

        self.z_dim = z_dim

        self.gf_dim = gf_dim
        self.df_dim = df_dim
        self.d_fc_unit = d_fc_unit
        self.g_fc_unit = g_fc_unit

        # Placeholders
        self.y = tf.placeholder(tf.float32, shape=[None, self.n_classes],
                                name='y-classes')  # one_hot
        self.x1_fine = tf.placeholder(tf.float32, shape=[None, self.height, self.width, self.channel],
                                      name='x-images')

        self.x1_scaled = image_sampling(self.x1_fine, 'down')
        self.x1_coarse = image_sampling(self.x1_scaled, 'up')
        self.x1_diff = self.x1_fine - self.x1_coarse

        self.x2_fine = self.x1_scaled  # [16, 16]
        self.x2_scaled = image_sampling(self.x2_fine, 'down')
        self.x2_coarse = image_sampling(self.x2_scaled, 'up')
        self.x2_diff = self.x2_fine - self.x2_coarse

        self.x3_fine = self.x2_scaled  # [8, 8]

        self.z = []
        self.z_noises = [32 * 32, 16 * 16, self.z_dim]
        for i in range(3):
            self.z.append(tf.placeholder(tf.float32,
                                         shape=[None, self.z_noises[i]],
                                         name='z-noise_{0}'.format(i)))

        self.do_rate = tf.placeholder(tf.float32, None, name='do-rate')

        self.g = []       # generators
        self.g_loss = []  # generator losses

        self.d_reals = []       # discriminator_real logits
        self.d_fakes = []       # discriminator_fake logits
        self.d_loss = []        # discriminator_real losses

        # Training Options
        self.d_op = []
        self.g_op = []

        self.beta1 = 0.5
        self.beta2 = 0.9
        self.lr = 8e-4

        self.saver = None
        self.merged = None
        self.writer = None

        self.bulid_lapgan()  # build LAPGAN model

    def discriminator(self, x1, x2, y, scale=32, reuse=None):
        """
        :param x1: image to discriminate
        :param x2: down-up sampling-ed images
        :param y: classes
        :param scale: image size
        :param reuse: variable re-use
        :return: logits
        """

        assert (scale % 8 == 0)  # 32, 16, 8

        with tf.variable_scope('discriminator_{0}'.format(scale), reuse=reuse):
            if scale == 8:
                x1 = tf.reshape(x1, [-1, scale * scale * 3])  # tf.layers.flatten(x1)

                h = tf.concat([x1, y], axis=1)

                h = t.dense(h, self.d_fc_unit, name='disc-fc-1')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-1')

                h = t.dense(h, self.d_fc_unit // 2, name='d-fc-2')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-2')

                h = t.dense(h, 1, name='disc-fc-3')
            else:
                x = x1 + x2

                y = t.dense(y, scale * scale, name='disc-fc-y')
                y = tf.nn.relu(y)

                y = tf.reshape(y, [-1, scale, scale, 1])

                h = tf.concat([x, y], axis=3)  # (-1, scale, scale, channel + 1)

                h = t.conv2d(h, self.df_dim * 1, 5, 1, pad='SAME', name='disc-conv2d-1')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-1')

                h = t.conv2d(h, self.df_dim * 1, 5, 1, pad='SAME', name='disc-conv2d-2')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-2')

                h = tf.layers.flatten(h)

                h = t.dense(h, 1, name='disc-fc-2')

            return h

    def generator(self, x, y, z, scale=32, reuse=None, do_rate=0.5):
        """
        :param x: images to fake
        :param y: classes
        :param z: noise
        :param scale: image size
        :param reuse: variable re-use
        :param do_rate: dropout rate
        :return: logits
        """

        assert (scale % 8 == 0)  # 32, 16, 8

        with tf.variable_scope('generator_{0}'.format(scale), reuse=reuse):
            if scale == 8:
                h = tf.concat([z, y], axis=1)

                h = t.dense(h, self.g_fc_unit, name='gen-fc-1')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, do_rate, name='gen-dropout-1')

                h = t.dense(h, self.g_fc_unit, name='gen-fc-2')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, do_rate, name='gen-dropout-2')

                h = t.dense(h, self.channel * 8 * 8, name='gen-fc-3')

                h = tf.reshape(h, [-1, 8, 8, self.channel])
            else:
                y = t.dense(y, scale * scale, name='gen-fc-y')

                y = tf.reshape(y, [-1, scale, scale, 1])
                z = tf.reshape(z, [-1, scale, scale, 1])

                h = tf.concat([z, y, x], axis=3)  # concat into 5 dims

                h = t.conv2d(h, self.gf_dim * 1, 5, 1, name='gen-deconv2d-1')
                h = tf.nn.relu(h)

                h = t.conv2d(h, self.gf_dim * 1, 5, 1, name='gen-deconv2d-2')
                h = tf.nn.relu(h)

                h = t.conv2d(h, self.channel, 5, 1, name='gen-deconv2d-3')

            h = tf.nn.tanh(h)

            return h

    def bulid_lapgan(self):
        # Generator & Discriminator
        g1 = self.generator(x=self.x1_coarse, y=self.y, z=self.z[0], scale=32, do_rate=self.do_rate)
        d1_fake = self.discriminator(x1=g1, x2=self.x1_coarse, y=self.y, scale=32)
        d1_real = self.discriminator(x1=self.x1_diff, x2=self.x1_coarse, y=self.y, scale=32, reuse=True)

        g2 = self.generator(x=self.x2_coarse, y=self.y, z=self.z[1], scale=16, do_rate=self.do_rate)
        d2_fake = self.discriminator(x1=g2, x2=self.x2_coarse, y=self.y, scale=16)
        d2_real = self.discriminator(x1=self.x2_diff, x2=self.x2_coarse, y=self.y, scale=16, reuse=True)

        g3 = self.generator(x=None, y=self.y, z=self.z[2], scale=8, do_rate=self.do_rate)
        d3_fake = self.discriminator(x1=g3, x2=None, y=self.y, scale=8)
        d3_real = self.discriminator(x1=self.x3_fine, x2=None, y=self.y, scale=8, reuse=True)

        self.g = [g1, g2, g3]
        self.d_reals = [d1_real, d2_real, d3_real]
        self.d_fakes = [d1_fake, d2_fake, d3_fake]

        # Losses
        with tf.variable_scope('loss'):
            for i in range(len(self.g)):
                self.d_loss.append(t.sce_loss(self.d_reals[i], tf.ones_like(self.d_reals[i])) +
                                   t.sce_loss(self.d_fakes[i], tf.zeros_like(self.d_fakes[i])))
                self.g_loss.append(t.sce_loss(self.d_fakes[i], tf.ones_like(self.d_fakes[i])))

        # Summary
        for i in range(len(self.g)):
            tf.summary.scalar('loss/d_loss_{0}'.format(i), self.d_loss[i])
            tf.summary.scalar('loss/g_loss_{0}'.format(i), self.g_loss[i])

        # Optimizer
        t_vars = tf.trainable_variables()
        for idx, i in enumerate([32, 16, 8]):
            self.d_op.append(tf.train.AdamOptimizer(learning_rate=self.lr,
                                                    beta1=self.beta1, beta2=self.beta2).
                             minimize(loss=self.d_loss[idx],
                                      var_list=[v for v in t_vars if v.name.startswith('discriminator_{0}'.format(i))]))
            self.g_op.append(tf.train.AdamOptimizer(learning_rate=self.lr,
                                                    beta1=self.beta1, beta2=self.beta2).
                             minimize(loss=self.g_loss[idx],
                                      var_list=[v for v in t_vars if v.name.startswith('generator_{0}'.format(i))]))

        # Merge summary
        self.merged = tf.summary.merge_all()

        # Model Saver
        self.saver = tf.train.Saver(max_to_keep=1)
        self.writer = tf.summary.FileWriter('./model/', self.s.graph)
<commit_msg>fix: layers' name in generator<commit_after>import tensorflow as tf

import sys

sys.path.append('../')
import tfutil as t

tf.set_random_seed(777)

# In image_utils, up/down_sampling
def image_sampling(img, sampling_type='down'):
    shape = img.get_shape()  # [batch, height, width, channels]

    if sampling_type == 'down':
        h = int(shape[1] // 2)
        w = int(shape[2] // 2)
    else:  # 'up'
        h = int(shape[1] * 2)
        w = int(shape[2] * 2)

    return tf.image.resize_images(img, [h, w], tf.image.ResizeMethod.BILINEAR)

class LAPGAN:

    def __init__(self, s, batch_size=128, height=32, width=32, channel=3, n_classes=10,
                 sample_num=10 * 10, sample_size=10,
                 z_dim=128, gf_dim=64, df_dim=64, d_fc_unit=512, g_fc_unit=1024):

        """
        # General Settings
        :param s: TF Session
        :param batch_size: training batch size, default 128
        :param height: input image height, default 32
        :param width: input image width, default 32
        :param channel: input image channel, default 3 (RGB)
        :param n_classes: the number of classes, default 10
        - in case of CIFAR, image size is 32x32x3(HWC), classes are 10.

        # Output Settings
        :param sample_size: sample image size, default 8
        :param sample_num: the number of sample images, default 64

        # Model Settings
        :param z_dim: z noise dimension, default 128
        :param gf_dim: the number of generator filters, default 64
        :param df_dim: the number of discriminator filters, default 64
        :param d_fc_unit: the number of fully connected filters used at Disc, default 512
        :param g_fc_unit: the number of fully connected filters used at Gen, default 1024
        """

        self.s = s
        self.batch_size = batch_size
        self.height = height
        self.width = width
        self.channel = channel
        self.n_classes = n_classes

        self.sample_size = sample_size
        self.sample_num = sample_num

        self.image_shape = [self.height, self.width, self.channel]

        self.z_dim = z_dim

        self.gf_dim = gf_dim
        self.df_dim = df_dim
        self.d_fc_unit = d_fc_unit
        self.g_fc_unit = g_fc_unit

        # Placeholders
        self.y = tf.placeholder(tf.float32, shape=[None, self.n_classes],
                                name='y-classes')  # one_hot
        self.x1_fine = tf.placeholder(tf.float32, shape=[None, self.height, self.width, self.channel],
                                      name='x-images')

        self.x1_scaled = image_sampling(self.x1_fine, 'down')
        self.x1_coarse = image_sampling(self.x1_scaled, 'up')
        self.x1_diff = self.x1_fine - self.x1_coarse

        self.x2_fine = self.x1_scaled  # [16, 16]
        self.x2_scaled = image_sampling(self.x2_fine, 'down')
        self.x2_coarse = image_sampling(self.x2_scaled, 'up')
        self.x2_diff = self.x2_fine - self.x2_coarse

        self.x3_fine = self.x2_scaled  # [8, 8]

        self.z = []
        self.z_noises = [32 * 32, 16 * 16, self.z_dim]
        for i in range(3):
            self.z.append(tf.placeholder(tf.float32,
                                         shape=[None, self.z_noises[i]],
                                         name='z-noise_{0}'.format(i)))

        self.do_rate = tf.placeholder(tf.float32, None, name='do-rate')

        self.g = []       # generators
        self.g_loss = []  # generator losses

        self.d_reals = []       # discriminator_real logits
        self.d_fakes = []       # discriminator_fake logits
        self.d_loss = []        # discriminator_real losses

        # Training Options
        self.d_op = []
        self.g_op = []

        self.beta1 = 0.5
        self.beta2 = 0.9
        self.lr = 8e-4

        self.saver = None
        self.merged = None
        self.writer = None

        self.bulid_lapgan()  # build LAPGAN model

    def discriminator(self, x1, x2, y, scale=32, reuse=None):
        """
        :param x1: image to discriminate
        :param x2: down-up sampling-ed images
        :param y: classes
        :param scale: image size
        :param reuse: variable re-use
        :return: logits
        """

        assert (scale % 8 == 0)  # 32, 16, 8

        with tf.variable_scope('discriminator_{0}'.format(scale), reuse=reuse):
            if scale == 8:
                x1 = tf.reshape(x1, [-1, scale * scale * 3])  # tf.layers.flatten(x1)

                h = tf.concat([x1, y], axis=1)

                h = t.dense(h, self.d_fc_unit, name='disc-fc-1')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-1')

                h = t.dense(h, self.d_fc_unit // 2, name='d-fc-2')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-2')

                h = t.dense(h, 1, name='disc-fc-3')
            else:
                x = x1 + x2

                y = t.dense(y, scale * scale, name='disc-fc-y')
                y = tf.nn.relu(y)

                y = tf.reshape(y, [-1, scale, scale, 1])

                h = tf.concat([x, y], axis=3)  # (-1, scale, scale, channel + 1)

                h = t.conv2d(h, self.df_dim * 1, 5, 1, pad='SAME', name='disc-conv2d-1')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-1')

                h = t.conv2d(h, self.df_dim * 1, 5, 1, pad='SAME', name='disc-conv2d-2')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, 0.5, name='disc-dropout-2')

                h = tf.layers.flatten(h)

                h = t.dense(h, 1, name='disc-fc-2')

            return h

    def generator(self, x, y, z, scale=32, reuse=None, do_rate=0.5):
        """
        :param x: images to fake
        :param y: classes
        :param z: noise
        :param scale: image size
        :param reuse: variable re-use
        :param do_rate: dropout rate
        :return: logits
        """

        assert (scale % 8 == 0)  # 32, 16, 8

        with tf.variable_scope('generator_{0}'.format(scale), reuse=reuse):
            if scale == 8:
                h = tf.concat([z, y], axis=1)

                h = t.dense(h, self.g_fc_unit, name='gen-fc-1')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, do_rate, name='gen-dropout-1')

                h = t.dense(h, self.g_fc_unit, name='gen-fc-2')
                h = tf.nn.relu(h)
                h = tf.layers.dropout(h, do_rate, name='gen-dropout-2')

                h = t.dense(h, self.channel * 8 * 8, name='gen-fc-3')

                h = tf.reshape(h, [-1, 8, 8, self.channel])
            else:
                y = t.dense(y, scale * scale, name='gen-fc-y')

                y = tf.reshape(y, [-1, scale, scale, 1])
                z = tf.reshape(z, [-1, scale, scale, 1])

                h = tf.concat([z, y, x], axis=3)  # concat into 5 dims

                h = t.conv2d(h, self.gf_dim * 1, 5, 1, name='gen-conv2d-1')
                h = tf.nn.relu(h)

                h = t.conv2d(h, self.gf_dim * 1, 5, 1, name='gen-conv2d-2')
                h = tf.nn.relu(h)

                h = t.conv2d(h, self.channel, 5, 1, name='gen-conv2d-3')

            h = tf.nn.tanh(h)

            return h

    def bulid_lapgan(self):
        # Generator & Discriminator
        g1 = self.generator(x=self.x1_coarse, y=self.y, z=self.z[0], scale=32, do_rate=self.do_rate)
        d1_fake = self.discriminator(x1=g1, x2=self.x1_coarse, y=self.y, scale=32)
        d1_real = self.discriminator(x1=self.x1_diff, x2=self.x1_coarse, y=self.y, scale=32, reuse=True)

        g2 = self.generator(x=self.x2_coarse, y=self.y, z=self.z[1], scale=16, do_rate=self.do_rate)
        d2_fake = self.discriminator(x1=g2, x2=self.x2_coarse, y=self.y, scale=16)
        d2_real = self.discriminator(x1=self.x2_diff, x2=self.x2_coarse, y=self.y, scale=16, reuse=True)

        g3 = self.generator(x=None, y=self.y, z=self.z[2], scale=8, do_rate=self.do_rate)
        d3_fake = self.discriminator(x1=g3, x2=None, y=self.y, scale=8)
        d3_real = self.discriminator(x1=self.x3_fine, x2=None, y=self.y, scale=8, reuse=True)

        self.g = [g1, g2, g3]
        self.d_reals = [d1_real, d2_real, d3_real]
        self.d_fakes = [d1_fake, d2_fake, d3_fake]

        # Losses
        with tf.variable_scope('loss'):
            for i in range(len(self.g)):
                self.d_loss.append(t.sce_loss(self.d_reals[i], tf.ones_like(self.d_reals[i])) +
                                   t.sce_loss(self.d_fakes[i], tf.zeros_like(self.d_fakes[i])))
                self.g_loss.append(t.sce_loss(self.d_fakes[i], tf.ones_like(self.d_fakes[i])))

        # Summary
        for i in range(len(self.g)):
            tf.summary.scalar('loss/d_loss_{0}'.format(i), self.d_loss[i])
            tf.summary.scalar('loss/g_loss_{0}'.format(i), self.g_loss[i])

        # Optimizer
        t_vars = tf.trainable_variables()
        for idx, i in enumerate([32, 16, 8]):
            self.d_op.append(tf.train.AdamOptimizer(learning_rate=self.lr,
                                                    beta1=self.beta1, beta2=self.beta2).
                             minimize(loss=self.d_loss[idx],
                                      var_list=[v for v in t_vars if v.name.startswith('discriminator_{0}'.format(i))]))
            self.g_op.append(tf.train.AdamOptimizer(learning_rate=self.lr,
                                                    beta1=self.beta1, beta2=self.beta2).
                             minimize(loss=self.g_loss[idx],
                                      var_list=[v for v in t_vars if v.name.startswith('generator_{0}'.format(i))]))

        # Merge summary
        self.merged = tf.summary.merge_all()

        # Model Saver
        self.saver = tf.train.Saver(max_to_keep=1)
        self.writer = tf.summary.FileWriter('./model/', self.s.graph)
<|endoftext|>


<commit_before>from pathlib import Path
from io import BytesIO
import json
from contextlib import contextmanager
import tempfile
import shutil
from django.db import models, transaction
from django.contrib.postgres.fields import JSONField
from django.conf import settings

def cache(model, keyfunc):

    def decorator(func):

        if not settings.SNOOP_CACHE:
            return func

        @transaction.atomic
        def wrapper(*args, **kwargs):
            key = keyfunc(*args, **kwargs)

            row, created = model.objects.get_or_create(pk=key)
            if not created:
                return json.loads(row.value)

            value = func(*args, **kwargs)

            row.value = json.dumps(value)
            row.save()

            return value

        wrapper.no_cache = func
        return wrapper

    return decorator

class EmailCache(models.Model):
    id = models.IntegerField(primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

class Collection(models.Model):
    path = models.CharField(max_length=4000)
    slug = models.CharField(max_length=100, unique=True)
    title = models.CharField(max_length=200)
    es_index = models.CharField(max_length=200)
    description = models.TextField(blank=True)

class Document(models.Model):
    collection = models.ForeignKey('Collection')
    container = models.ForeignKey('Document',
                                  related_name='contained_set',
                                  null=True)
    parent = models.ForeignKey('Document',
                               related_name='child_set',
                               null=True)
    path = models.CharField(max_length=4000)
    content_type = models.CharField(max_length=100, blank=True)
    filename = models.CharField(max_length=1000)
    disk_size = models.BigIntegerField()
    md5 = models.CharField(max_length=40, blank=True, db_index=True)
    sha1 = models.CharField(max_length=50, blank=True, db_index=True)
    broken = models.CharField(max_length=100, blank=True)
    rev = models.IntegerField(null=True)
    flags = JSONField(default=dict, blank=True)

    class Meta:
        # TODO: constraint does not apply to container=None rows
        unique_together = ('container', 'path')

    @property
    def absolute_path(self):
        assert self.container is None
        return Path(self.collection.path) / self.path

    def _open_file(self):
        if self.content_type == 'application/x-directory':
            return BytesIO()

        if self.container is None:
            return self.absolute_path.open('rb')

        else:
            from . import emails, archives, pst

            if emails.is_email(self.container):
                return emails.get_email_part(self.container, self.path)

            if archives.is_archive(self.container):
                return archives.open_file(self.container, self.path)

            if pst.is_pst_file(self.container):
                return pst.open_file(self.container, self.path)

        raise RuntimeError

    @contextmanager
    def open(self, filesystem=False):
        """ Open the document as a file. If the document is inside an email or
        archive, it will be copied to a temporary file:

            with doc.open() as f:
                f.read()

        If ``filesystem`` is True, ``f`` will have a ``path`` attribute, which
        is the absolute path of the file on disk.
        """

        with self._open_file() as f:
            if filesystem:
                if self.container:
                    MB = 1024*1024
                    suffix = Path(self.filename).suffix
                    with tempfile.NamedTemporaryFile(suffix=suffix) as tmp:
                        shutil.copyfileobj(f, tmp, length=4*MB)
                        tmp.flush()
                        tmp.path = Path(tmp.name)
                        yield tmp

                else:
                    f.path = self.absolute_path
                    yield f

            else:
                yield f

class Ocr(models.Model):
    tag = models.CharField(max_length=100)
    md5 = models.CharField(max_length=40, db_index=True)
    path = models.CharField(max_length=4000)
    text = models.TextField(blank=True)

    class Meta:
        unique_together = ('tag', 'md5')

    @property
    def absolute_path(self):
        return Path(settings.SNOOP_OCR_ROOT) / self.tag / self.path

class Digest(models.Model):
    id = models.IntegerField(primary_key=True)
    data = models.TextField()
    updated_at = models.DateTimeField(auto_now=True)

class Job(models.Model):
    queue = models.CharField(max_length=100)
    data = JSONField(null=True)
    started = models.BooleanField(default=False)

    class Meta:
        unique_together = ('queue', 'data')
        index_together = ('queue', 'started')

class TikaCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

class TikaLangCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.CharField(max_length=20)
    time = models.DateTimeField(auto_now=True)

class HtmlTextCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)
<commit_msg>pretty-print documents<commit_after>from pathlib import Path
from io import BytesIO
import json
from contextlib import contextmanager
import tempfile
import shutil
from django.db import models, transaction
from django.contrib.postgres.fields import JSONField
from django.conf import settings

def cache(model, keyfunc):

    def decorator(func):

        if not settings.SNOOP_CACHE:
            return func

        @transaction.atomic
        def wrapper(*args, **kwargs):
            key = keyfunc(*args, **kwargs)

            row, created = model.objects.get_or_create(pk=key)
            if not created:
                return json.loads(row.value)

            value = func(*args, **kwargs)

            row.value = json.dumps(value)
            row.save()

            return value

        wrapper.no_cache = func
        return wrapper

    return decorator

class EmailCache(models.Model):
    id = models.IntegerField(primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

class Collection(models.Model):
    path = models.CharField(max_length=4000)
    slug = models.CharField(max_length=100, unique=True)
    title = models.CharField(max_length=200)
    es_index = models.CharField(max_length=200)
    description = models.TextField(blank=True)

class Document(models.Model):
    collection = models.ForeignKey('Collection')
    container = models.ForeignKey('Document',
                                  related_name='contained_set',
                                  null=True)
    parent = models.ForeignKey('Document',
                               related_name='child_set',
                               null=True)
    path = models.CharField(max_length=4000)
    content_type = models.CharField(max_length=100, blank=True)
    filename = models.CharField(max_length=1000)
    disk_size = models.BigIntegerField()
    md5 = models.CharField(max_length=40, blank=True, db_index=True)
    sha1 = models.CharField(max_length=50, blank=True, db_index=True)
    broken = models.CharField(max_length=100, blank=True)
    rev = models.IntegerField(null=True)
    flags = JSONField(default=dict, blank=True)

    class Meta:
        # TODO: constraint does not apply to container=None rows
        unique_together = ('container', 'path')

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

    @property
    def absolute_path(self):
        assert self.container is None
        return Path(self.collection.path) / self.path

    def _open_file(self):
        if self.content_type == 'application/x-directory':
            return BytesIO()

        if self.container is None:
            return self.absolute_path.open('rb')

        else:
            from . import emails, archives, pst

            if emails.is_email(self.container):
                return emails.get_email_part(self.container, self.path)

            if archives.is_archive(self.container):
                return archives.open_file(self.container, self.path)

            if pst.is_pst_file(self.container):
                return pst.open_file(self.container, self.path)

        raise RuntimeError

    @contextmanager
    def open(self, filesystem=False):
        """ Open the document as a file. If the document is inside an email or
        archive, it will be copied to a temporary file:

            with doc.open() as f:
                f.read()

        If ``filesystem`` is True, ``f`` will have a ``path`` attribute, which
        is the absolute path of the file on disk.
        """

        with self._open_file() as f:
            if filesystem:
                if self.container:
                    MB = 1024*1024
                    suffix = Path(self.filename).suffix
                    with tempfile.NamedTemporaryFile(suffix=suffix) as tmp:
                        shutil.copyfileobj(f, tmp, length=4*MB)
                        tmp.flush()
                        tmp.path = Path(tmp.name)
                        yield tmp

                else:
                    f.path = self.absolute_path
                    yield f

            else:
                yield f

class Ocr(models.Model):
    tag = models.CharField(max_length=100)
    md5 = models.CharField(max_length=40, db_index=True)
    path = models.CharField(max_length=4000)
    text = models.TextField(blank=True)

    class Meta:
        unique_together = ('tag', 'md5')

    @property
    def absolute_path(self):
        return Path(settings.SNOOP_OCR_ROOT) / self.tag / self.path

class Digest(models.Model):
    id = models.IntegerField(primary_key=True)
    data = models.TextField()
    updated_at = models.DateTimeField(auto_now=True)

class Job(models.Model):
    queue = models.CharField(max_length=100)
    data = JSONField(null=True)
    started = models.BooleanField(default=False)

    class Meta:
        unique_together = ('queue', 'data')
        index_together = ('queue', 'started')

class TikaCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

class TikaLangCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.CharField(max_length=20)
    time = models.DateTimeField(auto_now=True)

class HtmlTextCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)
<|endoftext|>


<commit_before>VERSION = (0, 0, 1)

from sol.system import *

# Shared twelve-tone object to improve caching
tts = TwelveToneSystem()

from sol.scale import *
from sol import mode
<commit_msg>import sequence into root<commit_after>VERSION = (0, 0, 1)

from sol.system import *

# Shared twelve-tone object to improve caching
tts = TwelveToneSystem()

from sol.scale import *
from sol import mode
from sol import sequence
<|endoftext|>


<commit_before>import numpy as np
from scipy.optimize import basinhopping

from patient_state import PatientState

def solve():
    starting_params = PatientState.schnider_params()

    params_array = convert_params_structure_to_vector(starting_params)

    result = basinhopping(find_lsq_for_all_patients, params_array)
    return result.x

def get_patients():
    return [{
        'age': 60,
        'expected_result': 0,
        'weight': 70,
        'height': 180,
        'sex': 'm'
    }]

def convert_vector_to_params_structure(params_vector):
    # ToDo: need to construct params object in a well-defined way
    pass

def convert_params_structure_to_vector(params):
    # ToDo: order of dictionary items is not well defined
    params_list = [v for k, v in params.items()]
    return np.array(params_list)

def find_lsq_for_all_patients(params_vector):
    params = convert_vector_to_params_structure(params_vector)
    total_lsq = 0
    for patient in get_patients():
        result = solve_for_patient(patient, params)
        residual = result - patient['expected_result']
        total_lsq += residual ** 2

def solve_for_patient(patient, params):
    patient = PatientState(patient['age'], patient['weight'], patient['height'], patient['sex'], params)
    patient.give_drug(92.60001)
    for t in range(130):
        patient.wait_time(1)
    return patient.x1

result = solve()
print result
<commit_msg>Using dummy data, calculate the measurement error per patient<commit_after>import numpy as np
from scipy.optimize import basinhopping

from patient_state import PatientState

def solve():
    starting_params = PatientState.schnider_params()

    params_array = convert_params_structure_to_vector(starting_params)

    result = basinhopping(find_lsq_for_all_patients, params_array)
    return result.x

def get_patients():
    return [{
        'age': 60,
        'expected_result': 0,
        'weight': 70,
        'height': 180,
        'sex': 'm',
        # TODO: Are there patients who receive multiple boluses when their cp is non-zero?
        'propofol_mg': 92.60001,
        'measurements': [
            {
                'time_mins': 2.11,
                'predicted_cp': 3.62
            },
            {
                'time_mins': 4.01,
                'predicted_cp': 1.33
            }
        ]
    }]

def convert_vector_to_params_structure(params_vector):
    # ToDo: need to construct params object in a well-defined way
    pass

def convert_params_structure_to_vector(params):
    # ToDo: order of dictionary items is not well defined
    params_list = [v for k, v in params.items()]
    return np.array(params_list)

def find_lsq_for_all_patients(params_vector):
    params = convert_vector_to_params_structure(params_vector)
    total_lsq = 0
    for patient in get_patients():
        # TODO: Use params rather than schnider_params
        patient_error = solve_for_patient(patient, PatientState.schnider_params())
        total_lsq += patient_error

    return total_lsq

def solve_for_patient(patient, params):
    patient_model = PatientState(patient['age'], patient['weight'], patient['height'], patient['sex'], params)
    patient_model.give_drug(patient['propofol_mg'])

    previous_time_mins = 0

    totalerror = 0
    total_measurements = 0

    for measurement in patient['measurements']:
        for t in range(int((measurement['time_mins'] - previous_time_mins) * 60)):
            patient_model.wait_time(1)

        actual_cp = patient_model.x1
        error = actual_cp - measurement['predicted_cp']

        totalerror += error
        total_measurements += 1

    return totalerror / total_measurements

result = solve()
print result
<|endoftext|>


<commit_before>import argparse
import glob
import json
import logging
import os
import platform
import re
import subprocess
import sys
import zipfile

import github
import requests
import uritemplate

from lbrynet.core import log_support

def main(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--file', help='artifact to publish')
    parser.add_argument('--label', help='text to append to `file`')
    parser.add_argument('--zip', action='store_true')
    parser.add_argument(
        '--force', action='store_true',
        help='ignores whether the repo is currently tagged, publishes a draft release')
    args = parser.parse_args(args)

    gh_token = os.environ['GH_TOKEN']
    auth = github.Github(gh_token)
    current_repo = auth.get_repo(current_repo_name())

    if args.file:
        artifact = args.file
    else:
        artifact = get_artifact()

    current_tag = None
    if not args.force:
        try:
            current_tag = subprocess.check_output(
                ['git', 'describe', '--exact-match', 'HEAD']).strip()
        except subprocess.CalledProcessError:
            log.info('Stopping as we are not currently on a tag')
            return
        if not check_repo_has_tag(current_repo, current_tag):
            log.info('Tag %s is not in repo %s', current_tag, current_repo)
            # TODO: maybe this should be an error
            return

    release = get_release(current_repo, current_tag, args.force)
    asset_to_upload = get_asset(artifact, args.label, args.zip)
    upload_asset(release, asset_to_upload, gh_token)

def get_system():
    system = platform.system()
    if system == 'Darwin':
        return 'macOS'
    else:
        return system

def current_repo_name():
    pattern = 'github.com[:/](.*).git'
    m = re.search(pattern, subprocess.check_output('git remote -v'.split()))
    return m.group(1)

def check_repo_has_tag(repo, target_tag):
    tags = repo.get_tags().get_page(0)
    for tag in tags:
        if tag.name == target_tag:
            return True
    return False

def get_release(current_repo, current_tag=None, draft=False):
    assert current_tag or draft, 'either current_tag or draft must be set'
    need_new_release = True
    if not draft and current_tag:
        try:
            release = current_repo.get_release(current_tag)
            log.info('Using an existing release for %s', current_tag)
        except github.UnknownObjectException:
            need_new_release = True
    if draft or need_new_release:
        log.info('Creating a new release for %s:%s', current_repo, current_tag)
        tag = current_tag or 'draft'
        release_name = current_tag or 'draft'
        msg = 'Release' # TODO: parse changelogs to get a better message
        release = current_repo.create_git_release(tag, release_name, msg, draft)
    return release

def get_artifact():
    system = platform.system()
    if system == 'Darwin':
        return glob.glob('dist/mac/LBRY*.dmg')[0]
    elif system == 'Linux':
        return glob.glob('dist/LBRY*.deb')[0]
    else:
        raise Exception("I don't know about any artifact on {}".format(system))

def get_asset(filename, label=None, use_zip=False):
    if label:
        label = '-{}'.format(label)
    else:
        label = ''
    base, ext = os.path.splitext(filename)
    if use_zip:
        # TODO: probably want to clean this up
        zipfilename = '{}{}.zip'.format(base, label)
        with zipfile.ZipFile(zipfilename, 'w') as myzip:
            myzip.write(filename)
        asset_to_uplaod = zipfilename
    else:
        asset_to_upload = '{}{}{}'.format(base, label, ext)
    return asset_to_upload

def upload_asset(release, asset_to_upload, token):
    basename = os.path.basename(asset_to_upload)

    if is_asset_already_uploaded(release, basename):
        return

    upload_uri = uritemplate.expand(
        release.upload_url, {'name': basename, 'label': get_system()})
    # using requests.post fails miserably with SSL EPIPE errors. I spent
    # half a day trying to debug before deciding to switch to curl.
    cmd = [
        'curl', '-sS', '-X', 'POST', '-u', ':{}'.format(os.environ['GH_TOKEN']),
        '--header', 'Content-Type:application/zip',
        '--data-binary', '@{}'.format(asset_to_upload), upload_uri
    ]
    raw_output = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
    output = json.loads(raw_output)
    if 'errors' in output:
        raise Exception(output)
    else:
        log.info('Successfully uploaded to %s', output['browser_download_url'])

def is_asset_already_uploaded(release, basename):
    for asset in release.raw_data['assets']:
        if asset['name'] == basename:
            log.info('File %s has already been uploaded to %s', basename, release.tag_name)
            return True
    return False

if __name__ == '__main__':
    log = logging.getLogger('release-on-tag')
    log_support.configure_console(level='DEBUG')
    sys.exit(main())
else:
    log = logging.getLogger(__name__)
<commit_msg>better error message<commit_after>import argparse
import glob
import json
import logging
import os
import platform
import re
import subprocess
import sys
import zipfile

import github
import requests
import uritemplate

from lbrynet.core import log_support

def main(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--file', help='artifact to publish')
    parser.add_argument('--label', help='text to append to `file`')
    parser.add_argument('--zip', action='store_true')
    parser.add_argument(
        '--force', action='store_true',
        help='ignores whether the repo is currently tagged, publishes a draft release')
    args = parser.parse_args(args)

    gh_token = os.environ['GH_TOKEN']
    auth = github.Github(gh_token)
    current_repo = auth.get_repo(current_repo_name())

    if args.file:
        artifact = args.file
    else:
        artifact = get_artifact()

    current_tag = None
    if not args.force:
        try:
            current_tag = subprocess.check_output(
                ['git', 'describe', '--exact-match', 'HEAD']).strip()
        except subprocess.CalledProcessError:
            log.info('Stopping as we are not currently on a tag')
            return
        if not check_repo_has_tag(current_repo, current_tag):
            log.info('Tag %s is not in repo %s', current_tag, current_repo)
            # TODO: maybe this should be an error
            return

    release = get_release(current_repo, current_tag, args.force)
    asset_to_upload = get_asset(artifact, args.label, args.zip)
    upload_asset(release, asset_to_upload, gh_token)

def get_system():
    system = platform.system()
    if system == 'Darwin':
        return 'macOS'
    else:
        return system

def current_repo_name():
    pattern = 'github.com[:/](.*)\.git'
    remote = subprocess.check_output('git remote -v'.split())
    m = re.search(pattern, remote)
    if not m:
        raise Exception('Unable to parse repo name from remote: {}'.format(remote))
    return m.group(1)

def check_repo_has_tag(repo, target_tag):
    tags = repo.get_tags().get_page(0)
    for tag in tags:
        if tag.name == target_tag:
            return True
    return False

def get_release(current_repo, current_tag=None, draft=False):
    assert current_tag or draft, 'either current_tag or draft must be set'
    need_new_release = True
    if not draft and current_tag:
        try:
            release = current_repo.get_release(current_tag)
            log.info('Using an existing release for %s', current_tag)
        except github.UnknownObjectException:
            need_new_release = True
    if draft or need_new_release:
        log.info('Creating a new release for %s:%s', current_repo, current_tag)
        tag = current_tag or 'draft'
        release_name = current_tag or 'draft'
        msg = 'Release' # TODO: parse changelogs to get a better message
        release = current_repo.create_git_release(tag, release_name, msg, draft)
    return release

def get_artifact():
    system = platform.system()
    if system == 'Darwin':
        return glob.glob('dist/mac/LBRY*.dmg')[0]
    elif system == 'Linux':
        return glob.glob('dist/LBRY*.deb')[0]
    else:
        raise Exception("I don't know about any artifact on {}".format(system))

def get_asset(filename, label=None, use_zip=False):
    if label:
        label = '-{}'.format(label)
    else:
        label = ''
    base, ext = os.path.splitext(filename)
    if use_zip:
        # TODO: probably want to clean this up
        zipfilename = '{}{}.zip'.format(base, label)
        with zipfile.ZipFile(zipfilename, 'w') as myzip:
            myzip.write(filename)
        asset_to_uplaod = zipfilename
    else:
        asset_to_upload = '{}{}{}'.format(base, label, ext)
    return asset_to_upload

def upload_asset(release, asset_to_upload, token):
    basename = os.path.basename(asset_to_upload)

    if is_asset_already_uploaded(release, basename):
        return

    upload_uri = uritemplate.expand(
        release.upload_url, {'name': basename, 'label': get_system()})
    # using requests.post fails miserably with SSL EPIPE errors. I spent
    # half a day trying to debug before deciding to switch to curl.
    cmd = [
        'curl', '-sS', '-X', 'POST', '-u', ':{}'.format(os.environ['GH_TOKEN']),
        '--header', 'Content-Type:application/zip',
        '--data-binary', '@{}'.format(asset_to_upload), upload_uri
    ]
    raw_output = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
    output = json.loads(raw_output)
    if 'errors' in output:
        raise Exception(output)
    else:
        log.info('Successfully uploaded to %s', output['browser_download_url'])

def is_asset_already_uploaded(release, basename):
    for asset in release.raw_data['assets']:
        if asset['name'] == basename:
            log.info('File %s has already been uploaded to %s', basename, release.tag_name)
            return True
    return False

if __name__ == '__main__':
    log = logging.getLogger('release-on-tag')
    log_support.configure_console(level='DEBUG')
    sys.exit(main())
else:
    log = logging.getLogger(__name__)
<|endoftext|>


<commit_before>import os
from ..enums import GameTag
from .game import *
from .classic import *
from .gvg import *
from .naxxramas import *
from . import cardxml

_PATH = os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir, "data", "TextAsset", "enUS.txt")

tagnames = {
	"Atk": GameTag.ATK,
	"Aura": GameTag.AURA,
	"Charge": GameTag.CHARGE,
	"Chromatic": GameTag.CANT_BE_TARGETED_BY_ABILITIES, # XXX
	"Durability": GameTag.DURABILITY,
	"Deathrattle": GameTag.DEATH_RATTLE,
	"Enrage": GameTag.ENRAGED,
	"Health": GameTag.HEALTH,
	"Recall": GameTag.RECALL,
	"Stealth": GameTag.STEALTH,
	"Taunt": GameTag.TAUNT,
	"Windfury": GameTag.WINDFURY,
	"OneTurnEffect": GameTag.OneTurnEffect,
	"cantAttack": GameTag.CANT_ATTACK, # XXX
}

def _initTags(carddef, cls):
	"""
	Iterate over the class attributes, mapping them to the tags dict
	Note that this only needs to be done once per class, hence why we
	do it here instead of in Card.__new__()
	"""
	for attr, value in carddef.__dict__.items():
		if attr in tagnames:
			cls.tags[tagnames[attr]] = value

def merge(xmlcard, carddef):
	if not carddef:
		cls = type(xmlcard.id, (), {})
	else:
		cls = type(xmlcard.id, (carddef, ), {})
	cls.tags = xmlcard.tags
	if carddef:
		if hasattr(carddef, "Enrage"):
			# Initialize the Enrage virtual card too
			carddef.Enrage.tags = {}
			_initTags(carddef.Enrage, carddef.Enrage)
		_initTags(carddef, cls)
	cls.requirements = xmlcard.requirements
	cls.entourage = xmlcard.entourage
	return cls

# Here we import every card from every set and load the cardxml database.
# For every card, we will "merge" the class with its Python definition if
# it exists.
# This code is only ran once, at initial import.

with open(_PATH, "r") as f:
	db = cardxml.load(_PATH)
	for id, card in db.items():
		carddef = globals().get(id)
		globals()[id] = merge(card, carddef)
<commit_msg>Make card definition IDs available in cls.id<commit_after>import os
from ..enums import GameTag
from .game import *
from .classic import *
from .gvg import *
from .naxxramas import *
from . import cardxml

_PATH = os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir, "data", "TextAsset", "enUS.txt")

tagnames = {
	"Atk": GameTag.ATK,
	"Aura": GameTag.AURA,
	"Charge": GameTag.CHARGE,
	"Chromatic": GameTag.CANT_BE_TARGETED_BY_ABILITIES, # XXX
	"Durability": GameTag.DURABILITY,
	"Deathrattle": GameTag.DEATH_RATTLE,
	"Enrage": GameTag.ENRAGED,
	"Health": GameTag.HEALTH,
	"Recall": GameTag.RECALL,
	"Stealth": GameTag.STEALTH,
	"Taunt": GameTag.TAUNT,
	"Windfury": GameTag.WINDFURY,
	"OneTurnEffect": GameTag.OneTurnEffect,
	"cantAttack": GameTag.CANT_ATTACK, # XXX
}

def _initTags(carddef, cls):
	"""
	Iterate over the class attributes, mapping them to the tags dict
	Note that this only needs to be done once per class, hence why we
	do it here instead of in Card.__new__()
	"""
	for attr, value in carddef.__dict__.items():
		if attr in tagnames:
			cls.tags[tagnames[attr]] = value

def merge(xmlcard, carddef):
	if not carddef:
		cls = type(xmlcard.id, (), {})
	else:
		cls = type(xmlcard.id, (carddef, ), {})
	cls.tags = xmlcard.tags
	if carddef:
		if hasattr(carddef, "Enrage"):
			# Initialize the Enrage virtual card too
			carddef.Enrage.tags = {}
			_initTags(carddef.Enrage, carddef.Enrage)
		_initTags(carddef, cls)
	cls.requirements = xmlcard.requirements
	cls.entourage = xmlcard.entourage
	cls.id = xmlcard.id
	return cls

# Here we import every card from every set and load the cardxml database.
# For every card, we will "merge" the class with its Python definition if
# it exists.
# This code is only ran once, at initial import.

with open(_PATH, "r") as f:
	db = cardxml.load(_PATH)
	for id, card in db.items():
		carddef = globals().get(id)
		globals()[id] = merge(card, carddef)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
gds.burp.models
~~~~~~~~~~~~~~~

This module contains the primary objects that make working with
Burp's IHttpRequestResponse object's more... Pythonic.
'''
from Cookie import SimpleCookie
from cgi import parse_header, parse_qs
from urlparse import urlparse

from .decorators import reify
from .structures import CaseInsensitiveDict

CRLF = '\r\n'
SP = chr(0x20)

class HttpRequest(object):
    '''The :class:`HttpRequest <HttpRequest>` object. Pass Burp's
    IHttpRequestResponse object to the constructor.

    Optional init arguments:
    :param _burp: IBurpExtender implementation
    '''
    def __init__(self, messageInfo=None, _burp=None):
        self._messageInfo = messageInfo
        self._burp = _burp

        self._host = None
        self._port = 80
        self._protocol = 'http'
        self._url = ''

        self.method = None
        self._uri = None
        self.version = None
        self._headers = {}
        self.body = None

        if messageInfo is not None and hasattr(messageInfo, 'request'):
            if messageInfo.getRequest():
                self.method, self._uri, self.version, self._headers, self.body = \
                    _parse_message(messageInfo.getRequest().tostring())

        if hasattr(messageInfo, 'response'):
            self.response = HttpResponse(getattr(messageInfo, 'response', None),
                                         request=self)
        else:
            self.response = HttpResponse(None, request=self)

    def __contains__(self, item):
        return item in self.body if self.body else False

    def __len__(self):
        return int(self.headers.get('content-length', len(self.body or '')))

    def __nonzero__(self):
        return self.raw is not None

    def __repr__(self):
        return '<HttpRequest [%s]>' % (self.url.path,)

    @property
    def host(self):
        '''
        Returns the name of the application host.
        '''
        if self._messageInfo is not None and \
            self._host != self._messageInfo.getHost():
            self._host = self._messageInfo.getHost()

        return self._host

    @host.setter
    def host(self, host):
        '''
        Sets the name of the application host to which the request
        should be sent.

        Note: This method generally can only be used before the
        message has been forwarded to the application, and not in
        read-only contexts.

        :param host: The name of the application host to which the
        request should be sent.
        '''
        if self._messageInfo is not None:
            self._messageInfo.setHost(host)

        return

    @property
    def port(self):
        '''
        Returns the port number used by the application.
        '''
        if self._messageInfo is not None and \
            self._port != self._messageInfo.getPort():
            self._port = self._messageInfo.getPort()

        return self._port

    @port.setter
    def port(self, port):
        '''
        Sets the port number to which the request should be sent.

        Note: This method generally can only be used before the
        message has been forwarded to the application, and not in
        read-only contexts.

        :param port: The port number to which the request should be
        sent.
        '''
        if self._messageInfo is not None:
            self._messageInfo.setPort(port)

        return

    @property
    def protocol(self):
        '''
        Returns the protocol used by the application.
        '''
        if self._messageInfo is not None and \
            self._protocol != self._messageInfo.getProtocol():
            self._protocol = self._messageInfo.getProtocol()

        return self._protocol

    @protocol.setter
    def protocol(self, protocol):
        '''
        Sets the protocol which should be used by the request.

        Note: This method generally can only be used before the
        message has been forwarded to the application, and not in
        read-only contexts.

        :param protocol: The protocol which should be used by the
        request. Valid values are "http" and "https".
        '''
        if self._messageInfo is not None:
            self._messageInfo.setProtocol(protocol)

        return

    @reify
    def url(self):
        '''
        The URL requested in this HTTP request.

        Note: This is a **read-only** attribute.

        :returns: :class:`~urlparse.ParseResult` object.
        '''
        if self._messageInfo is not None:
            self._url = urlparse(self._messageInfo.getUrl().toString())

        return self._url

    @reify
    def cookies(self):
        '''
        The HTTP Cookies sent in this request.

        Note: This is a **read-only** attribute.

        :returns: :class:`~Cookie.SimpleCookie` object.
        '''
        self._cookies = SimpleCookie(self.headers.get('cookie', ''))
        return self._cookies

    @reify
    def headers(self):
        '''
        The HTTP headers sent in this request. Headers are accessible
        by their header names (case insensitive).

        Note: This is a **read-only** attribute.
        '''
        self._headers = CaseInsensitiveDict(self._headers)
        return self._headers

    @reify
    def parameters(self):
        '''
        Parameters parsed into a dictionary based on type (i.e., query,
        body, etc.)

        Note: This is a **read-only** attribute.
        '''
        self._parameters = _parse_parameters(self)
        return self._parameters

    @property
    def content_type(self):
        '''
        Content-Type of the HTTP request.

        Note: This is a **read-only** attribute.
        '''
        return self.headers.get('content-type', '')

    @property
    def is_secure(self):
        '''
        True if the HTTP request was sent over HTTPS.

        Note: This is a **read-only** attribute.
        '''
        return True if self.protocol == 'https' else False

    @property
    def raw(self):
        '''
        Returns the full request contents.
        '''
        if self._messageInfo:
            return self._messageInfo.getRequest().tostring()

        return

    @raw.setter
    def raw(self, message):
        '''
        Sets the request contents which should be sent to the application.

        :param message: The request contents which should be sent to the
        application.
        '''
        if self._messageInfo:
            self._messageInfo.setRequest(message)

        return

    @property
    def comment(self):
        '''
        Returns the user-annotated comment for this item, if applicable.
        '''
        if self._messageInfo:
            return self._messageInfo.getComment()

        return

    @comment.setter
    def comment(self, comment):
        '''
        Sets the user-annotated comment for this item.

        :param comment: The comment to be associated with this item.
        '''
        if self._messageInfo:
            return self._messageInfo.setComment(comment)

        return

    @property
    def highlight(self):
        '''
        Returns the user-annotated color for this item, if applicable.
        '''
        if self._messageInfo:
            return self._messageInfo.getHighlight()

        return

    @highlight.setter
    def highlight(self, color):
        '''
        Sets the user-annotated color for this item.

        :param color: The color to be associated with this item.
        Valid values are: red, orange, yellow, green, cyan, blue, pink,
        magenta, gray.
        '''
        if self._messageInfo:
            self._messageInfo.setHighlight(color)

        return

class HttpResponse(object):
    def __init__(self, message=None, request=None):
        self.request = request

        self.version = None
        self.status_code = None
        self.reason = None
        self.encoding = None
        self._headers = {}
        self.body = None

        if message is not None:
            self.version, self.status_code, self.reason, self._headers, self.body = \
                _parse_message(message.tostring())

    def __contains__(self, item):
        return item in self.body if self.body else False

    def __len__(self):
        return int(self.headers.get('content-length', len(self.body or '')))

    def __nonzero__(self):
        return self.raw is not None

    def __repr__(self):
        return '<HttpResponse [%s]>' % (self.status_code, )

    @reify
    def cookies(self):
        '''
        The HTTP Cookies set in this response.

        Note: This is a **read-only** attribute.

        :returns: :class:`~Cookie.SimpleCookie` object.
        '''
        self._cookies = SimpleCookie(self.headers.get('set-cookie', ''))
        return self._cookies

    @reify
    def headers(self):
        '''
        The HTTP headers received in this response. Headers are accessible
        by their header names (case insensitive).

        Note: This is a **read-only** attribute.
        '''
        self._headers = CaseInsensitiveDict(self._headers)
        return self._headers

    @property
    def content_type(self):
        '''
        Content-Type of the HTTP response.

        Note: This is a **read-only** attribute.
        '''
        return self.headers.get('content-type', '')

    @property
    def raw(self):
        '''
        Returns the full response contents.
        '''
        if self.request._messageInfo:
            return self.request._messageInfo.getResponse().tostring()

        return

    @raw.setter
    def raw(self, message):
        '''
        Sets the response contents which should be processed by the
        invoking Burp tool.

        :param message: The response contents which should be processed
        by the invoking Burp tool.
        '''
        if self.request._messageInfo:
            return self.request._messageInfo.setResponse(message)

        return

def _parse_message(message):
    is_response = False
    pos = idx = 0

    idx = message.find(CRLF, pos)

    if idx != -1:
        start_line = message[pos:idx]

        if start_line.startswith('HTTP/'):
            is_response = True

        _idx = start_line.find(SP)

        if _idx != -1:
            if is_response:
                version = start_line[0:_idx]
            else:
                method = start_line[0:_idx]

            _pos = _idx + 1

        if is_response:
            _idx = start_line.find(SP, _pos)

            status = start_line[_pos:_idx]
            if not status.isdigit():
                raise ValueError('status code %r is not a number' % (status,))

            status = int(status)

            _pos = _idx + 1
            reason = start_line[_pos:]

        else:
            # work out the http version by looking in reverse
            _ridx = start_line.rfind(SP)
            version = start_line[_ridx + 1:]
            if not version.startswith('HTTP/'):
                raise ValueError('Invalid HTTP version: %r' % (version,))

            # request-uri will be everything in-between.
            # some clients might not encode space into a plus or %20
            uri = start_line[_pos:_ridx]
            if not uri or uri.isspace():
                raise ValueError('Invalid URI: %r' % (uri,))

        pos = idx + 2
    else:
        raise ValueError('Could not parse start-line from message')

    headers = CaseInsensitiveDict()

    while (idx != -1):
        idx = message.find(CRLF, pos)

        if idx == pos:
            # we've reached the end of the request headers
            # advance 4 bytes (2 * CRLF)
            pos = idx + 2
            break

        if idx != -1:
            header = message[pos:idx]
            _idx = header.find(':')

            if _idx != -1:
                name = header[:_idx].strip()
                value = header[_idx+1:].strip()

                has_value = headers.get(name)

                if has_value and has_value != value:
                    value = ', '.join([has_value, value])

                headers[name] = value
            else:
                raise ValueError('Error parsing header: %r' % (header,))

            pos = idx + 2
        else:
            # looks like we reached the end of the message before EOL
            break

    if idx < len(message):
        body = message[pos:]
    else:
        raise ValueError('Parsed past message body??')

    if not is_response:
        return method, uri, version, headers, body
    else:
        return version, status, reason, headers, body

def _parse_parameters(request):
    parameters = {}

    if request.url.query:
        parameters['query'] = parse_qs(request.url.query, keep_blank_values=True)

    _type, _options = parse_header(request.headers.get('content-type', ''))

    if _type == 'application/x-www-form-urlencoded':
        parameters['body'] = parse_qs(request.body, keep_blank_values=True)

    elif _type in ('application/json', ):
        try:
            parameters['body'] = json.loads(request.body)
        except (NameError, TypeError):
            pass

    elif _type == 'application/x-amf':
        pass

    elif _type == 'text/x-gwt-rpc':
        pass

    elif _type == 'application/xml':
        pass

    else:
        pass

    return parameters
<commit_msg>parse multipart form data into a cgi.FieldStorage object<commit_after># -*- coding: utf-8 -*-
'''
gds.burp.models
~~~~~~~~~~~~~~~

This module contains the primary objects that make working with
Burp's IHttpRequestResponse object's more... Pythonic.
'''
from Cookie import SimpleCookie
from cStringIO import StringIO
from cgi import FieldStorage, parse_header, parse_qs
from urlparse import urlparse

from .decorators import reify
from .structures import CaseInsensitiveDict

import json

CRLF = '\r\n'
SP = chr(0x20)

class HttpRequest(object):
    '''The :class:`HttpRequest <HttpRequest>` object. Pass Burp's
    IHttpRequestResponse object to the constructor.

    Optional init arguments:
    :param _burp: IBurpExtender implementation
    '''
    def __init__(self, messageInfo=None, _burp=None):
        self._messageInfo = messageInfo
        self._burp = _burp

        self._host = None
        self._port = 80
        self._protocol = 'http'
        self._url = ''

        self.method = None
        self._uri = None
        self.version = None
        self._headers = {}
        self.body = None

        if messageInfo is not None and hasattr(messageInfo, 'request'):
            if messageInfo.getRequest():
                self.method, self._uri, self.version, self._headers, self.body = \
                    _parse_message(messageInfo.getRequest().tostring())

        if hasattr(messageInfo, 'response'):
            self.response = HttpResponse(getattr(messageInfo, 'response', None),
                                         request=self)
        else:
            self.response = HttpResponse(None, request=self)

    def __contains__(self, item):
        return item in self.body if self.body else False

    def __len__(self):
        return int(self.headers.get('content-length', len(self.body or '')))

    def __nonzero__(self):
        return self.raw is not None

    def __repr__(self):
        return '<HttpRequest [%s]>' % (self.url.path,)

    @property
    def host(self):
        '''
        Returns the name of the application host.
        '''
        if self._messageInfo is not None and \
            self._host != self._messageInfo.getHost():
            self._host = self._messageInfo.getHost()

        return self._host

    @host.setter
    def host(self, host):
        '''
        Sets the name of the application host to which the request
        should be sent.

        Note: This method generally can only be used before the
        message has been forwarded to the application, and not in
        read-only contexts.

        :param host: The name of the application host to which the
        request should be sent.
        '''
        if self._messageInfo is not None:
            self._messageInfo.setHost(host)

        return

    @property
    def port(self):
        '''
        Returns the port number used by the application.
        '''
        if self._messageInfo is not None and \
            self._port != self._messageInfo.getPort():
            self._port = self._messageInfo.getPort()

        return self._port

    @port.setter
    def port(self, port):
        '''
        Sets the port number to which the request should be sent.

        Note: This method generally can only be used before the
        message has been forwarded to the application, and not in
        read-only contexts.

        :param port: The port number to which the request should be
        sent.
        '''
        if self._messageInfo is not None:
            self._messageInfo.setPort(port)

        return

    @property
    def protocol(self):
        '''
        Returns the protocol used by the application.
        '''
        if self._messageInfo is not None and \
            self._protocol != self._messageInfo.getProtocol():
            self._protocol = self._messageInfo.getProtocol()

        return self._protocol

    @protocol.setter
    def protocol(self, protocol):
        '''
        Sets the protocol which should be used by the request.

        Note: This method generally can only be used before the
        message has been forwarded to the application, and not in
        read-only contexts.

        :param protocol: The protocol which should be used by the
        request. Valid values are "http" and "https".
        '''
        if self._messageInfo is not None:
            self._messageInfo.setProtocol(protocol)

        return

    @reify
    def url(self):
        '''
        The URL requested in this HTTP request.

        Note: This is a **read-only** attribute.

        :returns: :class:`~urlparse.ParseResult` object.
        '''
        if self._messageInfo is not None:
            self._url = urlparse(self._messageInfo.getUrl().toString())

        return self._url

    @reify
    def cookies(self):
        '''
        The HTTP Cookies sent in this request.

        Note: This is a **read-only** attribute.

        :returns: :class:`~Cookie.SimpleCookie` object.
        '''
        self._cookies = SimpleCookie(self.headers.get('cookie', ''))
        return self._cookies

    @reify
    def headers(self):
        '''
        The HTTP headers sent in this request. Headers are accessible
        by their header names (case insensitive).

        Note: This is a **read-only** attribute.
        '''
        self._headers = CaseInsensitiveDict(self._headers)
        return self._headers

    @reify
    def parameters(self):
        '''
        Parameters parsed into a dictionary based on type (i.e., query,
        body, etc.)

        Note: This is a **read-only** attribute.
        '''
        self._parameters = _parse_parameters(self)
        return self._parameters

    @property
    def content_type(self):
        '''
        Content-Type of the HTTP request.

        Note: This is a **read-only** attribute.
        '''
        return self.headers.get('content-type', '')

    @property
    def is_secure(self):
        '''
        True if the HTTP request was sent over HTTPS.

        Note: This is a **read-only** attribute.
        '''
        return True if self.protocol == 'https' else False

    @property
    def raw(self):
        '''
        Returns the full request contents.
        '''
        if self._messageInfo:
            return self._messageInfo.getRequest().tostring()

        return

    @raw.setter
    def raw(self, message):
        '''
        Sets the request contents which should be sent to the application.

        :param message: The request contents which should be sent to the
        application.
        '''
        if self._messageInfo:
            self._messageInfo.setRequest(message)

        return

    @property
    def comment(self):
        '''
        Returns the user-annotated comment for this item, if applicable.
        '''
        if self._messageInfo:
            return self._messageInfo.getComment()

        return

    @comment.setter
    def comment(self, comment):
        '''
        Sets the user-annotated comment for this item.

        :param comment: The comment to be associated with this item.
        '''
        if self._messageInfo:
            return self._messageInfo.setComment(comment)

        return

    @property
    def highlight(self):
        '''
        Returns the user-annotated color for this item, if applicable.
        '''
        if self._messageInfo:
            return self._messageInfo.getHighlight()

        return

    @highlight.setter
    def highlight(self, color):
        '''
        Sets the user-annotated color for this item.

        :param color: The color to be associated with this item.
        Valid values are: red, orange, yellow, green, cyan, blue, pink,
        magenta, gray.
        '''
        if self._messageInfo:
            self._messageInfo.setHighlight(color)

        return

class HttpResponse(object):
    def __init__(self, message=None, request=None):
        self.request = request

        self.version = None
        self.status_code = None
        self.reason = None
        self.encoding = None
        self._headers = {}
        self.body = None

        if message is not None:
            self.version, self.status_code, self.reason, self._headers, self.body = \
                _parse_message(message.tostring())

    def __contains__(self, item):
        return item in self.body if self.body else False

    def __len__(self):
        return int(self.headers.get('content-length', len(self.body or '')))

    def __nonzero__(self):
        return self.raw is not None

    def __repr__(self):
        return '<HttpResponse [%s]>' % (self.status_code, )

    @reify
    def cookies(self):
        '''
        The HTTP Cookies set in this response.

        Note: This is a **read-only** attribute.

        :returns: :class:`~Cookie.SimpleCookie` object.
        '''
        self._cookies = SimpleCookie(self.headers.get('set-cookie', ''))
        return self._cookies

    @reify
    def headers(self):
        '''
        The HTTP headers received in this response. Headers are accessible
        by their header names (case insensitive).

        Note: This is a **read-only** attribute.
        '''
        self._headers = CaseInsensitiveDict(self._headers)
        return self._headers

    @property
    def content_type(self):
        '''
        Content-Type of the HTTP response.

        Note: This is a **read-only** attribute.
        '''
        return self.headers.get('content-type', '')

    @property
    def raw(self):
        '''
        Returns the full response contents.
        '''
        if self.request._messageInfo:
            return self.request._messageInfo.getResponse().tostring()

        return

    @raw.setter
    def raw(self, message):
        '''
        Sets the response contents which should be processed by the
        invoking Burp tool.

        :param message: The response contents which should be processed
        by the invoking Burp tool.
        '''
        if self.request._messageInfo:
            return self.request._messageInfo.setResponse(message)

        return

def _parse_message(message):
    is_response = False
    pos = idx = 0

    idx = message.find(CRLF, pos)

    if idx != -1:
        start_line = message[pos:idx]

        if start_line.startswith('HTTP/'):
            is_response = True

        _idx = start_line.find(SP)

        if _idx != -1:
            if is_response:
                version = start_line[0:_idx]
            else:
                method = start_line[0:_idx]

            _pos = _idx + 1

        if is_response:
            _idx = start_line.find(SP, _pos)

            status = start_line[_pos:_idx]
            if not status.isdigit():
                raise ValueError('status code %r is not a number' % (status,))

            status = int(status)

            _pos = _idx + 1
            reason = start_line[_pos:]

        else:
            # work out the http version by looking in reverse
            _ridx = start_line.rfind(SP)
            version = start_line[_ridx + 1:]
            if not version.startswith('HTTP/'):
                raise ValueError('Invalid HTTP version: %r' % (version,))

            # request-uri will be everything in-between.
            # some clients might not encode space into a plus or %20
            uri = start_line[_pos:_ridx]
            if not uri or uri.isspace():
                raise ValueError('Invalid URI: %r' % (uri,))

        pos = idx + 2
    else:
        raise ValueError('Could not parse start-line from message')

    headers = CaseInsensitiveDict()

    while (idx != -1):
        idx = message.find(CRLF, pos)

        if idx == pos:
            # we've reached the end of the request headers
            # advance 4 bytes (2 * CRLF)
            pos = idx + 2
            break

        if idx != -1:
            header = message[pos:idx]
            _idx = header.find(':')

            if _idx != -1:
                name = header[:_idx].strip()
                value = header[_idx+1:].strip()

                has_value = headers.get(name)

                if has_value and has_value != value:
                    value = ', '.join([has_value, value])

                headers[name] = value
            else:
                raise ValueError('Error parsing header: %r' % (header,))

            pos = idx + 2
        else:
            # looks like we reached the end of the message before EOL
            break

    if idx < len(message):
        body = message[pos:]
    else:
        raise ValueError('Parsed past message body??')

    if not is_response:
        return method, uri, version, headers, body
    else:
        return version, status, reason, headers, body

def _parse_parameters(request):
    parameters = {}

    if request.url.query:
        parameters['query'] = parse_qs(request.url.query, keep_blank_values=True)

    _type, _options = parse_header(request.headers.get('content-type', ''))

    if _type == 'application/x-www-form-urlencoded':
        parameters['body'] = parse_qs(request.body, keep_blank_values=True)

    elif _type.startswith('multipart/'):
        parameters['body'] = FieldStorage(
            fp = StringIO(request.body),
            headers = request.headers,
            environ = dict(REQUEST_METHOD = request.method),
            keep_blank_values=True)

    elif _type in ('application/json', ):
        try:
            parameters['body'] = json.loads(request.body)
        except (NameError, TypeError):
            pass

    elif _type == 'application/x-amf':
        pass

    elif _type == 'text/x-gwt-rpc':
        pass

    elif _type == 'application/xml':
        pass

    else:
        pass

    return parameters
<|endoftext|>


<commit_before>#! /usr/bin/env python
"""Test program for the fcntl C module.
   Roger E. Masse
"""
import struct
import fcntl
import FCNTL
import os, sys
from test_support import verbose

filename = '/tmp/delete-me'

# the example from the library docs
f = open(filename,'w')
rv = fcntl.fcntl(f.fileno(), FCNTL.F_SETFL, os.O_NONBLOCK)
if verbose:
    print 'Status from fnctl with O_NONBLOCK: ', rv

if sys.platform in ('netbsd1',
                    'freebsd2', 'freebsd3',
                    'bsdos2', 'bsdos3', 'bsdos4'):
    lockdata = struct.pack('lxxxxlxxxxlhh', 0, 0, 0, FCNTL.F_WRLCK, 0)
elif sys.platform in ['aix3', 'aix4']:
    lockdata = struct.pack('hhlllii', FCNTL.F_WRLCK, 0, 0, 0, 0, 0, 0)
else:
    lockdata = struct.pack('hhllhh', FCNTL.F_WRLCK, 0, 0, 0, 0, 0)
if verbose:
    print 'struct.pack: ', `lockdata`

rv = fcntl.fcntl(f.fileno(), FCNTL.F_SETLKW, lockdata)
if verbose:
    print 'String from fcntl with F_SETLKW: ', `rv`

f.close()
os.unlink(filename)
<commit_msg>Daniel Neri: OpenBSD is just as BSD'ish as the other BSD's ;-)<commit_after>#! /usr/bin/env python
"""Test program for the fcntl C module.
   Roger E. Masse
"""
import struct
import fcntl
import FCNTL
import os, sys
from test_support import verbose

filename = '/tmp/delete-me'

# the example from the library docs
f = open(filename,'w')
rv = fcntl.fcntl(f.fileno(), FCNTL.F_SETFL, os.O_NONBLOCK)
if verbose:
    print 'Status from fnctl with O_NONBLOCK: ', rv

if sys.platform in ('netbsd1',
                    'freebsd2', 'freebsd3',
                    'bsdos2', 'bsdos3', 'bsdos4',
                    'openbsd', 'openbsd2'):
    lockdata = struct.pack('lxxxxlxxxxlhh', 0, 0, 0, FCNTL.F_WRLCK, 0)
elif sys.platform in ['aix3', 'aix4']:
    lockdata = struct.pack('hhlllii', FCNTL.F_WRLCK, 0, 0, 0, 0, 0, 0)
else:
    lockdata = struct.pack('hhllhh', FCNTL.F_WRLCK, 0, 0, 0, 0, 0)
if verbose:
    print 'struct.pack: ', `lockdata`

rv = fcntl.fcntl(f.fileno(), FCNTL.F_SETLKW, lockdata)
if verbose:
    print 'String from fcntl with F_SETLKW: ', `rv`

f.close()
os.unlink(filename)
<|endoftext|>


<commit_before># Some simple Queue module tests, plus some failure conditions
# to ensure the Queue locks remain stable.
import Queue
import sys
import threading
import time

from test.test_support import verify, TestFailed, verbose

QUEUE_SIZE = 5

# A thread to run a function that unclogs a blocked Queue.
class _TriggerThread(threading.Thread):
    def __init__(self, fn, args):
        self.fn = fn
        self.args = args
        self.startedEvent = threading.Event()
        threading.Thread.__init__(self)

    def run(self):
        # The sleep isn't necessary, but is intended to give the blocking
        # function in the main thread a chance at actually blocking before
        # we unclog it.  But if the sleep is longer than the timeout-based
        # tests wait in their blocking functions, those tests will fail.
        # So we give them much longer timeout values compared to the
        # sleep here (I aimed at 10 seconds for blocking functions --
        # they should never actually wait that long - they should make
        # progress as soon as we call self.fn()).
        time.sleep(0.1)
        self.startedEvent.set()
        self.fn(*self.args)

# Execute a function that blocks, and in a separate thread, a function that
# triggers the release.  Returns the result of the blocking function.
# Caution:  block_func must guarantee to block until trigger_func is
# called, and trigger_func must guarantee to change queue state so that
# block_func can make enough progress to return.  In particular, a
# block_func that just raises an exception regardless of whether trigger_func
# is called will lead to timing-dependent sporadic failures, and one of
# those went rarely seen but undiagnosed for years.  Now block_func
# must be unexceptional.  If block_func is supposed to raise an exception,
# call _doExceptionalBlockingTest() instead.
def _doBlockingTest(block_func, block_args, trigger_func, trigger_args):
    t = _TriggerThread(trigger_func, trigger_args)
    t.start()
    result = block_func(*block_args)
    # If block_func returned before our thread made the call, we failed!
    if not t.startedEvent.isSet():
        raise TestFailed("blocking function '%r' appeared not to block" %
                         block_func)
    t.join(10) # make sure the thread terminates
    if t.isAlive():
        raise TestFailed("trigger function '%r' appeared to not return" %
                         trigger_func)
    return result

# Call this instead if block_func is supposed to raise an exception.
def _doExceptionalBlockingTest(block_func, block_args, trigger_func,
                               trigger_args, expected_exception_class):
    t = _TriggerThread(trigger_func, trigger_args)
    t.start()
    try:
        try:
            block_func(*block_args)
        except expected_exception_class:
            raise
        else:
            raise TestFailed("expected exception of kind %r" %
                             expected_exception_class)
    finally:
        t.join(10) # make sure the thread terminates
        if t.isAlive():
            raise TestFailed("trigger function '%r' appeared to not return" %
                             trigger_func)
        if not t.startedEvent.isSet():
            raise TestFailed("trigger thread ended but event never set")

# A Queue subclass that can provoke failure at a moment's notice :)
class FailingQueueException(Exception):
    pass

class FailingQueue(Queue.Queue):
    def __init__(self, *args):
        self.fail_next_put = False
        self.fail_next_get = False
        Queue.Queue.__init__(self, *args)
    def _put(self, item):
        if self.fail_next_put:
            self.fail_next_put = False
            raise FailingQueueException, "You Lose"
        return Queue.Queue._put(self, item)
    def _get(self):
        if self.fail_next_get:
            self.fail_next_get = False
            raise FailingQueueException, "You Lose"
        return Queue.Queue._get(self)

def FailingQueueTest(q):
    if not q.empty():
        raise RuntimeError, "Call this function with an empty queue"
    for i in range(QUEUE_SIZE-1):
        q.put(i)
    # Test a failing non-blocking put.
    q.fail_next_put = True
    try:
        q.put("oops", block=0)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    q.fail_next_put = True
    try:
        q.put("oops", timeout=0.1)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    q.put("last")
    verify(q.full(), "Queue should be full")
    # Test a failing blocking put
    q.fail_next_put = True
    try:
        _doBlockingTest(q.put, ("full",), q.get, ())
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    # Check the Queue isn't damaged.
    # put failed, but get succeeded - re-add
    q.put("last")
    # Test a failing timeout put
    q.fail_next_put = True
    try:
        _doExceptionalBlockingTest(q.put, ("full", True, 10), q.get, (),
                                   FailingQueueException)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    # Check the Queue isn't damaged.
    # put failed, but get succeeded - re-add
    q.put("last")
    verify(q.full(), "Queue should be full")
    q.get()
    verify(not q.full(), "Queue should not be full")
    q.put("last")
    verify(q.full(), "Queue should be full")
    # Test a blocking put
    _doBlockingTest( q.put, ("full",), q.get, ())
    # Empty it
    for i in range(QUEUE_SIZE):
        q.get()
    verify(q.empty(), "Queue should be empty")
    q.put("first")
    q.fail_next_get = True
    try:
        q.get()
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    verify(not q.empty(), "Queue should not be empty")
    q.fail_next_get = True
    try:
        q.get(timeout=0.1)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    verify(not q.empty(), "Queue should not be empty")
    q.get()
    verify(q.empty(), "Queue should be empty")
    q.fail_next_get = True
    try:
        _doExceptionalBlockingTest(q.get, (), q.put, ('empty',),
                                   FailingQueueException)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    # put succeeded, but get failed.
    verify(not q.empty(), "Queue should not be empty")
    q.get()
    verify(q.empty(), "Queue should be empty")

def SimpleQueueTest(q):
    if not q.empty():
        raise RuntimeError, "Call this function with an empty queue"
    # I guess we better check things actually queue correctly a little :)
    q.put(111)
    q.put(222)
    verify(q.get() == 111 and q.get() == 222,
           "Didn't seem to queue the correct data!")
    for i in range(QUEUE_SIZE-1):
        q.put(i)
        verify(not q.empty(), "Queue should not be empty")
    verify(not q.full(), "Queue should not be full")
    q.put("last")
    verify(q.full(), "Queue should be full")
    try:
        q.put("full", block=0)
        raise TestFailed("Didn't appear to block with a full queue")
    except Queue.Full:
        pass
    try:
        q.put("full", timeout=0.01)
        raise TestFailed("Didn't appear to time-out with a full queue")
    except Queue.Full:
        pass
    # Test a blocking put
    _doBlockingTest(q.put, ("full",), q.get, ())
    _doBlockingTest(q.put, ("full", True, 10), q.get, ())
    # Empty it
    for i in range(QUEUE_SIZE):
        q.get()
    verify(q.empty(), "Queue should be empty")
    try:
        q.get(block=0)
        raise TestFailed("Didn't appear to block with an empty queue")
    except Queue.Empty:
        pass
    try:
        q.get(timeout=0.01)
        raise TestFailed("Didn't appear to time-out with an empty queue")
    except Queue.Empty:
        pass
    # Test a blocking get
    _doBlockingTest(q.get, (), q.put, ('empty',))
    _doBlockingTest(q.get, (True, 10), q.put, ('empty',))

cum = 0
cumlock = threading.Lock()

def worker(q):
    global cum
    while True:
        x = q.get()
        if x is None:
            q.task_done()
            return
        cumlock.acquire()
        try:
            cum += x
        finally:
            cumlock.release()
        q.task_done()

def QueueJoinTest(q):
    global cum
    cum = 0
    for i in (0,1):
        threading.Thread(target=worker, args=(q,)).start()
    for i in xrange(100):
        q.put(i)
    q.join()
    verify(cum==sum(range(100)), "q.join() did not block until all tasks were done")
    for i in (0,1):
        q.put(None)         # instruct the threads to close
    q.join()                # verify that you can join twice

def QueueTaskDoneTest(q)
    try:
        q.task_done()
    except ValueError:
        pass
    else:
        raise TestFailed("Did not detect task count going negative")

def test():
    q = Queue.Queue()
    QueueTaskDoneTest(q)
    QueueJoinTest(q)
    QueueJoinTest(q)
    QueueTaskDoneTest(q)

    q = Queue.Queue(QUEUE_SIZE)
    # Do it a couple of times on the same queue
    SimpleQueueTest(q)
    SimpleQueueTest(q)
    if verbose:
        print "Simple Queue tests seemed to work"
    q = FailingQueue(QUEUE_SIZE)
    FailingQueueTest(q)
    FailingQueueTest(q)
    if verbose:
        print "Failing Queue tests seemed to work"

test()
<commit_msg>fix typo<commit_after># Some simple Queue module tests, plus some failure conditions
# to ensure the Queue locks remain stable.
import Queue
import sys
import threading
import time

from test.test_support import verify, TestFailed, verbose

QUEUE_SIZE = 5

# A thread to run a function that unclogs a blocked Queue.
class _TriggerThread(threading.Thread):
    def __init__(self, fn, args):
        self.fn = fn
        self.args = args
        self.startedEvent = threading.Event()
        threading.Thread.__init__(self)

    def run(self):
        # The sleep isn't necessary, but is intended to give the blocking
        # function in the main thread a chance at actually blocking before
        # we unclog it.  But if the sleep is longer than the timeout-based
        # tests wait in their blocking functions, those tests will fail.
        # So we give them much longer timeout values compared to the
        # sleep here (I aimed at 10 seconds for blocking functions --
        # they should never actually wait that long - they should make
        # progress as soon as we call self.fn()).
        time.sleep(0.1)
        self.startedEvent.set()
        self.fn(*self.args)

# Execute a function that blocks, and in a separate thread, a function that
# triggers the release.  Returns the result of the blocking function.
# Caution:  block_func must guarantee to block until trigger_func is
# called, and trigger_func must guarantee to change queue state so that
# block_func can make enough progress to return.  In particular, a
# block_func that just raises an exception regardless of whether trigger_func
# is called will lead to timing-dependent sporadic failures, and one of
# those went rarely seen but undiagnosed for years.  Now block_func
# must be unexceptional.  If block_func is supposed to raise an exception,
# call _doExceptionalBlockingTest() instead.
def _doBlockingTest(block_func, block_args, trigger_func, trigger_args):
    t = _TriggerThread(trigger_func, trigger_args)
    t.start()
    result = block_func(*block_args)
    # If block_func returned before our thread made the call, we failed!
    if not t.startedEvent.isSet():
        raise TestFailed("blocking function '%r' appeared not to block" %
                         block_func)
    t.join(10) # make sure the thread terminates
    if t.isAlive():
        raise TestFailed("trigger function '%r' appeared to not return" %
                         trigger_func)
    return result

# Call this instead if block_func is supposed to raise an exception.
def _doExceptionalBlockingTest(block_func, block_args, trigger_func,
                               trigger_args, expected_exception_class):
    t = _TriggerThread(trigger_func, trigger_args)
    t.start()
    try:
        try:
            block_func(*block_args)
        except expected_exception_class:
            raise
        else:
            raise TestFailed("expected exception of kind %r" %
                             expected_exception_class)
    finally:
        t.join(10) # make sure the thread terminates
        if t.isAlive():
            raise TestFailed("trigger function '%r' appeared to not return" %
                             trigger_func)
        if not t.startedEvent.isSet():
            raise TestFailed("trigger thread ended but event never set")

# A Queue subclass that can provoke failure at a moment's notice :)
class FailingQueueException(Exception):
    pass

class FailingQueue(Queue.Queue):
    def __init__(self, *args):
        self.fail_next_put = False
        self.fail_next_get = False
        Queue.Queue.__init__(self, *args)
    def _put(self, item):
        if self.fail_next_put:
            self.fail_next_put = False
            raise FailingQueueException, "You Lose"
        return Queue.Queue._put(self, item)
    def _get(self):
        if self.fail_next_get:
            self.fail_next_get = False
            raise FailingQueueException, "You Lose"
        return Queue.Queue._get(self)

def FailingQueueTest(q):
    if not q.empty():
        raise RuntimeError, "Call this function with an empty queue"
    for i in range(QUEUE_SIZE-1):
        q.put(i)
    # Test a failing non-blocking put.
    q.fail_next_put = True
    try:
        q.put("oops", block=0)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    q.fail_next_put = True
    try:
        q.put("oops", timeout=0.1)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    q.put("last")
    verify(q.full(), "Queue should be full")
    # Test a failing blocking put
    q.fail_next_put = True
    try:
        _doBlockingTest(q.put, ("full",), q.get, ())
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    # Check the Queue isn't damaged.
    # put failed, but get succeeded - re-add
    q.put("last")
    # Test a failing timeout put
    q.fail_next_put = True
    try:
        _doExceptionalBlockingTest(q.put, ("full", True, 10), q.get, (),
                                   FailingQueueException)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    # Check the Queue isn't damaged.
    # put failed, but get succeeded - re-add
    q.put("last")
    verify(q.full(), "Queue should be full")
    q.get()
    verify(not q.full(), "Queue should not be full")
    q.put("last")
    verify(q.full(), "Queue should be full")
    # Test a blocking put
    _doBlockingTest( q.put, ("full",), q.get, ())
    # Empty it
    for i in range(QUEUE_SIZE):
        q.get()
    verify(q.empty(), "Queue should be empty")
    q.put("first")
    q.fail_next_get = True
    try:
        q.get()
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    verify(not q.empty(), "Queue should not be empty")
    q.fail_next_get = True
    try:
        q.get(timeout=0.1)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    verify(not q.empty(), "Queue should not be empty")
    q.get()
    verify(q.empty(), "Queue should be empty")
    q.fail_next_get = True
    try:
        _doExceptionalBlockingTest(q.get, (), q.put, ('empty',),
                                   FailingQueueException)
        raise TestFailed("The queue didn't fail when it should have")
    except FailingQueueException:
        pass
    # put succeeded, but get failed.
    verify(not q.empty(), "Queue should not be empty")
    q.get()
    verify(q.empty(), "Queue should be empty")

def SimpleQueueTest(q):
    if not q.empty():
        raise RuntimeError, "Call this function with an empty queue"
    # I guess we better check things actually queue correctly a little :)
    q.put(111)
    q.put(222)
    verify(q.get() == 111 and q.get() == 222,
           "Didn't seem to queue the correct data!")
    for i in range(QUEUE_SIZE-1):
        q.put(i)
        verify(not q.empty(), "Queue should not be empty")
    verify(not q.full(), "Queue should not be full")
    q.put("last")
    verify(q.full(), "Queue should be full")
    try:
        q.put("full", block=0)
        raise TestFailed("Didn't appear to block with a full queue")
    except Queue.Full:
        pass
    try:
        q.put("full", timeout=0.01)
        raise TestFailed("Didn't appear to time-out with a full queue")
    except Queue.Full:
        pass
    # Test a blocking put
    _doBlockingTest(q.put, ("full",), q.get, ())
    _doBlockingTest(q.put, ("full", True, 10), q.get, ())
    # Empty it
    for i in range(QUEUE_SIZE):
        q.get()
    verify(q.empty(), "Queue should be empty")
    try:
        q.get(block=0)
        raise TestFailed("Didn't appear to block with an empty queue")
    except Queue.Empty:
        pass
    try:
        q.get(timeout=0.01)
        raise TestFailed("Didn't appear to time-out with an empty queue")
    except Queue.Empty:
        pass
    # Test a blocking get
    _doBlockingTest(q.get, (), q.put, ('empty',))
    _doBlockingTest(q.get, (True, 10), q.put, ('empty',))

cum = 0
cumlock = threading.Lock()

def worker(q):
    global cum
    while True:
        x = q.get()
        if x is None:
            q.task_done()
            return
        cumlock.acquire()
        try:
            cum += x
        finally:
            cumlock.release()
        q.task_done()

def QueueJoinTest(q):
    global cum
    cum = 0
    for i in (0,1):
        threading.Thread(target=worker, args=(q,)).start()
    for i in xrange(100):
        q.put(i)
    q.join()
    verify(cum==sum(range(100)), "q.join() did not block until all tasks were done")
    for i in (0,1):
        q.put(None)         # instruct the threads to close
    q.join()                # verify that you can join twice

def QueueTaskDoneTest(q):
    try:
        q.task_done()
    except ValueError:
        pass
    else:
        raise TestFailed("Did not detect task count going negative")

def test():
    q = Queue.Queue()
    QueueTaskDoneTest(q)
    QueueJoinTest(q)
    QueueJoinTest(q)
    QueueTaskDoneTest(q)

    q = Queue.Queue(QUEUE_SIZE)
    # Do it a couple of times on the same queue
    SimpleQueueTest(q)
    SimpleQueueTest(q)
    if verbose:
        print "Simple Queue tests seemed to work"
    q = FailingQueue(QUEUE_SIZE)
    FailingQueueTest(q)
    FailingQueueTest(q)
    if verbose:
        print "Failing Queue tests seemed to work"

test()
<|endoftext|>