<commit_before># Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Routine for decoding the CIFAR-10 binary file format."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

# Process images of this size. Note that this differs from the original CIFAR
# image size of 32 x 32. If one alters this number, then the entire model
# architecture will change and any model would need to be retrained.
IMAGE_SIZE = 24

# Global constants describing the CIFAR-10 data set.
NUM_CLASSES = 10
NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 50000
NUM_EXAMPLES_PER_EPOCH_FOR_EVAL = 10000

def read_cifar10(filename_queue):
  """Reads and parses examples from CIFAR10 data files.

  Recommendation: if you want N-way read parallelism, call this function
  N times.  This will give you N independent Readers reading different
  files & positions within those files, which will give better mixing of
  examples.

  Args:
    filename_queue: A queue of strings with the filenames to read from.

  Returns:
    An object representing a single example, with the following fields:
      height: number of rows in the result (32)
      width: number of columns in the result (32)
      depth: number of color channels in the result (3)
      key: a scalar string Tensor describing the filename & record number
        for this example.
      label: an int32 Tensor with the label in the range 0..9.
      uint8image: a [height, width, depth] uint8 Tensor with the image data
  """

  class CIFAR10Record(object):
    pass
  result = CIFAR10Record()

  # Dimensions of the images in the CIFAR-10 dataset.
  # See http://www.cs.toronto.edu/~kriz/cifar.html for a description of the
  # input format.
  label_bytes = 1  # 2 for CIFAR-100
  result.height = 32
  result.width = 32
  result.depth = 3
  image_bytes = result.height * result.width * result.depth
  # Every record consists of a label followed by the image, with a
  # fixed number of bytes for each.
  record_bytes = label_bytes + image_bytes

  # Read a record, getting filenames from the filename_queue.  No
  # header or footer in the CIFAR-10 format, so we leave header_bytes
  # and footer_bytes at their default of 0.
  reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
  result.key, value = reader.read(filename_queue)

  # Convert from a string to a vector of uint8 that is record_bytes long.
  record_bytes = tf.decode_raw(value, tf.uint8)

  # The first bytes represent the label, which we convert from uint8->int32.
  result.label = tf.cast(
      tf.strided_slice(record_bytes, [0], [label_bytes]), tf.int32)

  # The remaining bytes after the label represent the image, which we reshape
  # from [depth * height * width] to [depth, height, width].
  depth_major = tf.reshape(
      tf.strided_slice(record_bytes, [label_bytes],
                       [label_bytes + image_bytes]),
      [result.depth, result.height, result.width])
  # Convert from [depth, height, width] to [height, width, depth].
  result.uint8image = tf.transpose(depth_major, [1, 2, 0])

  return result

def _generate_image_and_label_batch(image, label, min_queue_examples,
                                    batch_size, shuffle):
  """Construct a queued batch of images and labels.

  Args:
    image: 3-D Tensor of [height, width, 3] of type.float32.
    label: 1-D Tensor of type.int32
    min_queue_examples: int32, minimum number of samples to retain
      in the queue that provides of batches of examples.
    batch_size: Number of images per batch.
    shuffle: boolean indicating whether to use a shuffling queue.

  Returns:
    images: Images. 4D tensor of [batch_size, height, width, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.
  """
  # Create a queue that shuffles the examples, and then
  # read 'batch_size' images + labels from the example queue.
  num_preprocess_threads = 16
  if shuffle:
    images, label_batch = tf.train.shuffle_batch(
        [image, label],
        batch_size=batch_size,
        num_threads=num_preprocess_threads,
        capacity=min_queue_examples + 3 * batch_size,
        min_after_dequeue=min_queue_examples)
  else:
    images, label_batch = tf.train.batch(
        [image, label],
        batch_size=batch_size,
        num_threads=num_preprocess_threads,
        capacity=min_queue_examples + 3 * batch_size)

  # Display the training images in the visualizer.
  tf.contrib.deprecated.image_summary('images', images)

  return images, tf.reshape(label_batch, [batch_size])

def distorted_inputs(data_dir, batch_size):
  """Construct distorted input for CIFAR training using the Reader ops.

  Args:
    data_dir: Path to the CIFAR-10 data directory.
    batch_size: Number of images per batch.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.
  """
  filenames = [os.path.join(data_dir, 'data_batch_%d.bin' % i)
               for i in xrange(1, 6)]
  for f in filenames:
    if not tf.gfile.Exists(f):
      raise ValueError('Failed to find file: ' + f)

  # Create a queue that produces the filenames to read.
  filename_queue = tf.train.string_input_producer(filenames)

  # Read examples from files in the filename queue.
  read_input = read_cifar10(filename_queue)
  reshaped_image = tf.cast(read_input.uint8image, tf.float32)

  height = IMAGE_SIZE
  width = IMAGE_SIZE

  # Image processing for training the network. Note the many random
  # distortions applied to the image.

  # Randomly crop a [height, width] section of the image.
  distorted_image = tf.random_crop(reshaped_image, [height, width, 3])

  # Randomly flip the image horizontally.
  distorted_image = tf.image.random_flip_left_right(distorted_image)

  # Because these operations are not commutative, consider randomizing
  # the order their operation.
  distorted_image = tf.image.random_brightness(distorted_image,
                                               max_delta=63)
  distorted_image = tf.image.random_contrast(distorted_image,
                                             lower=0.2, upper=1.8)

  # Subtract off the mean and divide by the variance of the pixels.
  float_image = tf.image.per_image_standardization(distorted_image)

  # Set the shapes of tensors.
  float_image.set_shape([height, width, 3])
  read_input.label.set_shape([1])

  # Ensure that the random shuffling has good mixing properties.
  min_fraction_of_examples_in_queue = 0.4
  min_queue_examples = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN *
                           min_fraction_of_examples_in_queue)
  print ('Filling queue with %d CIFAR images before starting to train. '
         'This will take a few minutes.' % min_queue_examples)

  # Generate a batch of images and labels by building up a queue of examples.
  return _generate_image_and_label_batch(float_image, read_input.label,
                                         min_queue_examples, batch_size,
                                         shuffle=True)

def inputs(eval_data, data_dir, batch_size):
  """Construct input for CIFAR evaluation using the Reader ops.

  Args:
    eval_data: bool, indicating if one should use the train or eval data set.
    data_dir: Path to the CIFAR-10 data directory.
    batch_size: Number of images per batch.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.
  """
  if not eval_data:
    filenames = [os.path.join(data_dir, 'data_batch_%d.bin' % i)
                 for i in xrange(1, 6)]
    num_examples_per_epoch = NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN
  else:
    filenames = [os.path.join(data_dir, 'test_batch.bin')]
    num_examples_per_epoch = NUM_EXAMPLES_PER_EPOCH_FOR_EVAL

  for f in filenames:
    if not tf.gfile.Exists(f):
      raise ValueError('Failed to find file: ' + f)

  # Create a queue that produces the filenames to read.
  filename_queue = tf.train.string_input_producer(filenames)

  # Read examples from files in the filename queue.
  read_input = read_cifar10(filename_queue)
  reshaped_image = tf.cast(read_input.uint8image, tf.float32)

  height = IMAGE_SIZE
  width = IMAGE_SIZE

  # Image processing for evaluation.
  # Crop the central [height, width] of the image.
  resized_image = tf.image.resize_image_with_crop_or_pad(reshaped_image,
                                                         width, height)

  # Subtract off the mean and divide by the variance of the pixels.
  float_image = tf.image.per_image_standardization(resized_image)

  # Ensure that the random shuffling has good mixing properties.
  min_fraction_of_examples_in_queue = 0.4
  min_queue_examples = int(num_examples_per_epoch *
                           min_fraction_of_examples_in_queue)

  # Generate a batch of images and labels by building up a queue of examples.
  return _generate_image_and_label_batch(float_image, read_input.label,
                                         min_queue_examples, batch_size,
                                         shuffle=False)
<commit_msg>Update cifar10_input.py<commit_after># Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Routine for decoding the CIFAR-10 binary file format."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

# Process images of this size. Note that this differs from the original CIFAR
# image size of 32 x 32. If one alters this number, then the entire model
# architecture will change and any model would need to be retrained.
IMAGE_SIZE = 24

# Global constants describing the CIFAR-10 data set.
NUM_CLASSES = 10
NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 50000
NUM_EXAMPLES_PER_EPOCH_FOR_EVAL = 10000

def read_cifar10(filename_queue):
  """Reads and parses examples from CIFAR10 data files.

  Recommendation: if you want N-way read parallelism, call this function
  N times.  This will give you N independent Readers reading different
  files & positions within those files, which will give better mixing of
  examples.

  Args:
    filename_queue: A queue of strings with the filenames to read from.

  Returns:
    An object representing a single example, with the following fields:
      height: number of rows in the result (32)
      width: number of columns in the result (32)
      depth: number of color channels in the result (3)
      key: a scalar string Tensor describing the filename & record number
        for this example.
      label: an int32 Tensor with the label in the range 0..9.
      uint8image: a [height, width, depth] uint8 Tensor with the image data
  """

  class CIFAR10Record(object):
    pass
  result = CIFAR10Record()

  # Dimensions of the images in the CIFAR-10 dataset.
  # See http://www.cs.toronto.edu/~kriz/cifar.html for a description of the
  # input format.
  label_bytes = 1  # 2 for CIFAR-100
  result.height = 32
  result.width = 32
  result.depth = 3
  image_bytes = result.height * result.width * result.depth
  # Every record consists of a label followed by the image, with a
  # fixed number of bytes for each.
  record_bytes = label_bytes + image_bytes

  # Read a record, getting filenames from the filename_queue.  No
  # header or footer in the CIFAR-10 format, so we leave header_bytes
  # and footer_bytes at their default of 0.
  reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
  result.key, value = reader.read(filename_queue)

  # Convert from a string to a vector of uint8 that is record_bytes long.
  record_bytes = tf.decode_raw(value, tf.uint8)

  # The first bytes represent the label, which we convert from uint8->int32.
  result.label = tf.cast(
      tf.strided_slice(record_bytes, [0], [label_bytes]), tf.int32)

  # The remaining bytes after the label represent the image, which we reshape
  # from [depth * height * width] to [depth, height, width].
  depth_major = tf.reshape(
      tf.strided_slice(record_bytes, [label_bytes],
                       [label_bytes + image_bytes]),
      [result.depth, result.height, result.width])
  # Convert from [depth, height, width] to [height, width, depth].
  result.uint8image = tf.transpose(depth_major, [1, 2, 0])

  return result

def _generate_image_and_label_batch(image, label, min_queue_examples,
                                    batch_size, shuffle):
  """Construct a queued batch of images and labels.

  Args:
    image: 3-D Tensor of [height, width, 3] of type.float32.
    label: 1-D Tensor of type.int32
    min_queue_examples: int32, minimum number of samples to retain
      in the queue that provides of batches of examples.
    batch_size: Number of images per batch.
    shuffle: boolean indicating whether to use a shuffling queue.

  Returns:
    images: Images. 4D tensor of [batch_size, height, width, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.
  """
  # Create a queue that shuffles the examples, and then
  # read 'batch_size' images + labels from the example queue.
  num_preprocess_threads = 16
  if shuffle:
    images, label_batch = tf.train.shuffle_batch(
        [image, label],
        batch_size=batch_size,
        num_threads=num_preprocess_threads,
        capacity=min_queue_examples + 3 * batch_size,
        min_after_dequeue=min_queue_examples)
  else:
    images, label_batch = tf.train.batch(
        [image, label],
        batch_size=batch_size,
        num_threads=num_preprocess_threads,
        capacity=min_queue_examples + 3 * batch_size)

  # Display the training images in the visualizer.
  tf.image_summary('images', images)

  return images, tf.reshape(label_batch, [batch_size])

def distorted_inputs(data_dir, batch_size):
  """Construct distorted input for CIFAR training using the Reader ops.

  Args:
    data_dir: Path to the CIFAR-10 data directory.
    batch_size: Number of images per batch.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.
  """
  filenames = [os.path.join(data_dir, 'data_batch_%d.bin' % i)
               for i in xrange(1, 6)]
  for f in filenames:
    if not tf.gfile.Exists(f):
      raise ValueError('Failed to find file: ' + f)

  # Create a queue that produces the filenames to read.
  filename_queue = tf.train.string_input_producer(filenames)

  # Read examples from files in the filename queue.
  read_input = read_cifar10(filename_queue)
  reshaped_image = tf.cast(read_input.uint8image, tf.float32)

  height = IMAGE_SIZE
  width = IMAGE_SIZE

  # Image processing for training the network. Note the many random
  # distortions applied to the image.

  # Randomly crop a [height, width] section of the image.
  distorted_image = tf.random_crop(reshaped_image, [height, width, 3])

  # Randomly flip the image horizontally.
  distorted_image = tf.image.random_flip_left_right(distorted_image)

  # Because these operations are not commutative, consider randomizing
  # the order their operation.
  distorted_image = tf.image.random_brightness(distorted_image,
                                               max_delta=63)
  distorted_image = tf.image.random_contrast(distorted_image,
                                             lower=0.2, upper=1.8)

  # Subtract off the mean and divide by the variance of the pixels.
  float_image = tf.image.per_image_standardization(distorted_image)

  # Set the shapes of tensors.
  float_image.set_shape([height, width, 3])
  read_input.label.set_shape([1])

  # Ensure that the random shuffling has good mixing properties.
  min_fraction_of_examples_in_queue = 0.4
  min_queue_examples = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN *
                           min_fraction_of_examples_in_queue)
  print ('Filling queue with %d CIFAR images before starting to train. '
         'This will take a few minutes.' % min_queue_examples)

  # Generate a batch of images and labels by building up a queue of examples.
  return _generate_image_and_label_batch(float_image, read_input.label,
                                         min_queue_examples, batch_size,
                                         shuffle=True)

def inputs(eval_data, data_dir, batch_size):
  """Construct input for CIFAR evaluation using the Reader ops.

  Args:
    eval_data: bool, indicating if one should use the train or eval data set.
    data_dir: Path to the CIFAR-10 data directory.
    batch_size: Number of images per batch.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.
  """
  if not eval_data:
    filenames = [os.path.join(data_dir, 'data_batch_%d.bin' % i)
                 for i in xrange(1, 6)]
    num_examples_per_epoch = NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN
  else:
    filenames = [os.path.join(data_dir, 'test_batch.bin')]
    num_examples_per_epoch = NUM_EXAMPLES_PER_EPOCH_FOR_EVAL

  for f in filenames:
    if not tf.gfile.Exists(f):
      raise ValueError('Failed to find file: ' + f)

  # Create a queue that produces the filenames to read.
  filename_queue = tf.train.string_input_producer(filenames)

  # Read examples from files in the filename queue.
  read_input = read_cifar10(filename_queue)
  reshaped_image = tf.cast(read_input.uint8image, tf.float32)

  height = IMAGE_SIZE
  width = IMAGE_SIZE

  # Image processing for evaluation.
  # Crop the central [height, width] of the image.
  resized_image = tf.image.resize_image_with_crop_or_pad(reshaped_image,
                                                         width, height)

  # Subtract off the mean and divide by the variance of the pixels.
  float_image = tf.image.per_image_standardization(resized_image)

  # Ensure that the random shuffling has good mixing properties.
  min_fraction_of_examples_in_queue = 0.4
  min_queue_examples = int(num_examples_per_epoch *
                           min_fraction_of_examples_in_queue)

  # Generate a batch of images and labels by building up a queue of examples.
  return _generate_image_and_label_batch(float_image, read_input.label,
                                         min_queue_examples, batch_size,
                                         shuffle=False)
<|endoftext|>


<commit_before>from .imopen import imopen
import numpy as np

def imread(uri, *args, index=None, plugin=None, **kwargs):
    """ Read an ndimage from a URI.

    Opens the given URI and reads an ndimage from it. The exact behavior
    depends on both, the file type and plugin used to open the file. To learn
    about the exact behavior, check the documentation of the relevant plugin.
    Typically, imread attempts to read all data stored in the URI.

    Parameters
    ----------
    uri : {str, pathlib.Path, bytes, file}
        The resource to load the image from, e.g. a filename, pathlib.Path,
        http address or file object, see the docs for more info.
    *args :
        Additional positional arguments will be passed to the plugin instance
        upon object creation.
    index : {int, None}
        If the URI contains multiple ndimages, select the index-th ndimage
        from among them and return it. The exact behavior is plugin dependent.
    plugin : {str, None}
        The plugin to be used. If None, performs a search for a matching plugin.
    **kwargs :
        Additional keyword arguments will be passed to the plugin instance.

    Returns
    -------
    image : ndimage
    """

    with imopen()(uri, *args, plugin=plugin, **kwargs) as img_file:
        return np.asarray(img_file.read(index=index))

def imiter(uri, *args, plugin=None, **kwargs):
    """ Read a sequence of ndimages from a URI.

    Returns an iterable that yields ndimages from the given URI. The exact
    behavior depends on both, the file type and plugin used to open the file. To
    learn about the exact behavior, check the documentation of the relevant
    plugin.

    Parameters
    ----------
    uri : {str, pathlib.Path, bytes, file}
        The resource to load the image from, e.g. a filename, pathlib.Path,
        http address or file object, see the docs for more info.
    *args :
        Additional positional arguments will be passed to the plugin instance
        upon object creation.
    plugin : {str, None}
        The plugin to be used. If None, performs a search for a matching plugin.
    **kwargs :
        Additional keyword arguments will be passed to the plugin instance.

    Returns
    -------
    image : ndimage
    """

    with imopen()(uri, *args, plugin=plugin, **kwargs) as img_file:
        for image in img_file.iter():
            yield np.asarray(image)
<commit_msg>proper plumbing<commit_after>from .imopen import imopen
import numpy as np

def imread(uri, *args, index=None, plugin=None, legacy_api=True,
           extra_kwargs_plugin=None, **extra_kwargs_read):
    """ Read an ndimage from a URI.

    Opens the given URI and reads an ndimage from it. The exact behavior
    depends on both, the file type and plugin used to open the file. To learn
    about the exact behavior, check the documentation of the relevant plugin.
    Typically, imread attempts to read all data stored in the URI.

    Parameters
    ----------
    uri : {str, pathlib.Path, bytes, file}
        The resource to load the image from, e.g. a filename, pathlib.Path,
        http address or file object, see the docs for more info.
    *args :
        Additional positional arguments will be passed to the plugin instance
        upon object creation.
    index : {int, None}
        If the URI contains multiple ndimages, select the index-th ndimage
        from among them and return it. The exact behavior is plugin dependent.
    plugin : {str, None}
        The plugin to be used. If None, performs a search for a matching plugin.
    legacy_api : {bool}
        The API to be used. If true use the v2.9 API, otherwise use the new
        v3.0 API.
    extra_kwargs_plugin : {dict, None}
        Additional keyword arguments to pass to the plugin's constructor.
    **extra_kwargs_read :
        Additional keyword arguments will be passed to the plugin's read call.

    Returns
    -------
    image : ndimage
    """

    plugin_kwargs = {
        "legacy_api": legacy_api,
        "plugin": plugin
    }

    if extra_kwargs_plugin is not None:
        plugin_kwargs.update(extra_kwargs_plugin)

    with imopen()(uri, *args, **plugin_kwargs) as img_file:
        return np.asarray(img_file.read(index=index, **extra_kwargs_read))

def imiter(uri, *args, plugin=None, legacy_api=True,
           extra_kwargs_plugin=None, **extra_kwargs_iter):
    """ Read a sequence of ndimages from a URI.

    Returns an iterable that yields ndimages from the given URI. The exact
    behavior depends on both, the file type and plugin used to open the file.
    To learn about the exact behavior, check the documentation of the relevant
    plugin.

    Parameters
    ----------
    uri : {str, pathlib.Path, bytes, file}
        The resource to load the image from, e.g. a filename, pathlib.Path,
        http address or file object, see the docs for more info.
    *args :
        Additional positional arguments will be passed to the plugin instance
        upon object creation.
    plugin : {str, None}
        The plugin to be used. If None, performs a search for a matching
        plugin.
    legacy_api : {bool}
        The API to be used. If true use the v2.9 API, otherwise use the new
        v3.0 API.
    extra_kwargs_plugin : {dict, None}
        Additional keyword arguments to pass to the plugin's constructor.
    **extra_kwargs_iter :
        Additional keyword arguments will be passed to the plugin's ``iter``
        call.

    Returns
    -------
    image : ndimage
    """

    plugin_kwargs = {
        "legacy_api": legacy_api,
        "plugin": plugin
    }

    if extra_kwargs_plugin is not None:
        plugin_kwargs.update(extra_kwargs_plugin)

    with imopen()(uri, *args, **plugin_kwargs) as img_file:
        for image in img_file.iter(**extra_kwargs_iter):
            # Note: casting to ndarray here to ensure compatibility
            # with the v2.9 API
            yield np.asarray(image)
<|endoftext|>


<commit_before>#-*- python -*-
import logging

# Non-stdlib imports
import pkg_resources
from tg import expose, validate, redirect, response
from pylons import g, c, request

from pymongo.bson import ObjectId

# Pyforge-specific imports
from allura.app import Application, ConfigOption, SitemapEntry, DefaultAdminController
from allura.lib import helpers as h
from allura.lib.security import require, has_artifact_access
from allura.model import ProjectRole
from allura.controllers import BaseController

# Local imports
from forgelink import version

log = logging.getLogger(__name__)

class ForgeLinkApp(Application):
    '''This is the Link app for PyForge'''
    __version__ = version.__version__
    permissions = [ 'configure', 'read' ]
    config_options = Application.config_options + [
        ConfigOption('url', str, None)
    ]
    searchable=True
    tool_label='Downloads'
    default_mount_label='Downloads'
    default_mount_point='downloads'
    ordinal=8

    def __init__(self, project, config):
        Application.__init__(self, project, config)
        self.root = RootController()
        self.admin = LinkAdminController(self)

    @property
    @h.exceptionless([], log)
    def sitemap(self):
        menu_id = self.config.options.mount_label.title()
        return [SitemapEntry(menu_id, '.')[self.sidebar_menu()] ]

    def sidebar_menu(self):
        return []

    def admin_menu(self):
        return super(ForgeLinkApp, self).admin_menu()

    @property
    def templates(self):
         return pkg_resources.resource_filename('forgelink', 'templates')

    def install(self, project):
        'Set up any default permissions and roles here'
        self.config.options['project_name'] = project.name
        super(ForgeLinkApp, self).install(project)
        # Setup permissions
        role_anon = ProjectRole.query.get(name='*anonymous')._id
        self.config.acl.update(
            configure=c.project.acl['tool'],
            read=[role_anon])

    def uninstall(self, project):
        "Remove all the tool's artifacts from the database"
        super(ForgeLinkApp, self).uninstall(project)

class RootController(BaseController):

    def _check_security(self):
        require(has_artifact_access('read'))

    @expose('forgelink.templates.index')
    def index(self, **kw):
        url = c.app.config.options.get('url')
        if url:
            redirect(url)
        return dict()

class LinkAdminController(DefaultAdminController):

    @expose()
    def index(self, **kw):
        redirect('options')
<commit_msg>[#1027] Refer to (already converted) Jinja template.<commit_after>#-*- python -*-
import logging

# Non-stdlib imports
import pkg_resources
from tg import expose, validate, redirect, response
from pylons import g, c, request

from pymongo.bson import ObjectId

# Pyforge-specific imports
from allura.app import Application, ConfigOption, SitemapEntry, DefaultAdminController
from allura.lib import helpers as h
from allura.lib.security import require, has_artifact_access
from allura.model import ProjectRole
from allura.controllers import BaseController

# Local imports
from forgelink import version

log = logging.getLogger(__name__)

class ForgeLinkApp(Application):
    '''This is the Link app for PyForge'''
    __version__ = version.__version__
    permissions = [ 'configure', 'read' ]
    config_options = Application.config_options + [
        ConfigOption('url', str, None)
    ]
    searchable=True
    tool_label='Downloads'
    default_mount_label='Downloads'
    default_mount_point='downloads'
    ordinal=8

    def __init__(self, project, config):
        Application.__init__(self, project, config)
        self.root = RootController()
        self.admin = LinkAdminController(self)

    @property
    @h.exceptionless([], log)
    def sitemap(self):
        menu_id = self.config.options.mount_label.title()
        return [SitemapEntry(menu_id, '.')[self.sidebar_menu()] ]

    def sidebar_menu(self):
        return []

    def admin_menu(self):
        return super(ForgeLinkApp, self).admin_menu()

    @property
    def templates(self):
         return pkg_resources.resource_filename('forgelink', 'templates')

    def install(self, project):
        'Set up any default permissions and roles here'
        self.config.options['project_name'] = project.name
        super(ForgeLinkApp, self).install(project)
        # Setup permissions
        role_anon = ProjectRole.query.get(name='*anonymous')._id
        self.config.acl.update(
            configure=c.project.acl['tool'],
            read=[role_anon])

    def uninstall(self, project):
        "Remove all the tool's artifacts from the database"
        super(ForgeLinkApp, self).uninstall(project)

class RootController(BaseController):

    def _check_security(self):
        require(has_artifact_access('read'))

    @expose('jinja:link/index.html')
    def index(self, **kw):
        url = c.app.config.options.get('url')
        if url:
            redirect(url)
        return dict()

class LinkAdminController(DefaultAdminController):

    @expose()
    def index(self, **kw):
        redirect('options')
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
from scipy.interpolate import UnivariateSpline
from scipy.optimize import leastsq

import abel

# see https://github.com/PyAbel/PyAbel/issues/186 for discussion

def circularize_image(IM, method="argmax", center=None, radial_range=None,
                      zoom=1, smooth=1.0e-7, nslices=32, inverse=False,
                      return_correction=False):
    """
    Remove radial distortion from a velocity-map-image through radial scaling
    of the Newton-rings to enforce circularity.

    The image is divided into angular slices, adjacent radial intensity
    profile slices compared to give a radial scaling factor that best
    aligns common structure.

    The resultant radial scaling factor vs angle is interpolated to apply the
    correction to the whole image grid. The image is remapped to the
    correct cartesian grid.

    Parameters
    ----------
    IM : numpy 2D array

    method: str
        method to align slice profiles,
        "argmax" - compare intensity-profile.argmax() of each radial slice.
        "lsq" - determine radial scaling factor by minimizing the
                intensity-prfoile with an adjacent slice intensity-profile.

    center: str, or None
        pre-center image using PyAbel centering methods,
        "com", "convolution", "gaussian", "image_center", "slice"

    radial_range: tuple, or None
        limit slice comparison to the radial range (rmin, rmax).

    zoom: float
        nmimage.zoom image before analysis. This may help improve the low
        radius slice profile.

    smooth: float
        smoothing for spline interpolation of the determined radial scaling
        factor vs angle

    nslices: int
        divide the image into nslices. Ideally, this should be a divisor of
        the image column width.

    inverse: bool
        inverse Abel transform the *polar* image, to remove the background
        intensity. This helps follow a given peak intensityr, for all angles,
        when the anisotropy parameter is large.

    return_correction: bool
        additional outputs, see below

    Returns
    -------
    IMcirc : numpy 2D array, same size as input
        Circularized imput image

    (if return_correction is True)
    slice_angles: numpy 1D array
        Mid-point angle (radians) of image slice

    radial_corrections: numpy 1D array
        radial correction scale factors, at each angle

    spline_function: numpy func()
        spline function used to evaluate the radial correction at each angle

    """

    if zoom > 1:
        # zoom may improve the low radii angular resolution
        IM = ndimage.zoom(IM, zoom)

    if center is not None:
        # convenience function for the case image is not centered
        IM = abel.tools.center.center_image(IM, center=center)

    # map image into polar coordinates - much easier to slice
    # cartesian (Y, X) -> polar (Radius, Theta)
    polarIM, radial_coord, angle_coord =\
                           abel.tools.polar.reproject_image_into_polar(IM)

    if inverse:
        # pseudo inverse Abel transform of the polar image, removes background
        # to enhance transition peaks
        polarIM = abel.dasch.two_point_transform(polarIM.T).T

    # limit radial range of polar image, if selected
    radial = radial_coord[:, 0]
    if radial_range is not None:
        subr = np.logical_and(radial > radial_range[0]*zoom,
                              radial < radial_range[1]*zoom)
        polarIM = polarIM[subr]
        radial = radial[subr]

    # split image into n-slices
    slices = np.array_split(polarIM, nslices, axis=1)
    slice_angles = np.array(np.hsplit(angle_coord[0], nslices)).mean(axis=1)

    # evaluate radial scaling factor for each slice
    radcorr = correction(slice_angles, slices, radial, method=method)

    # spline radial scaling vs angle
    radcorrspl = UnivariateSpline(slice_angles, radcorr, s=smooth, ext=3)

    # apply the correction
    IMcirc = circularize(IM, radcorrspl)

    if zoom > 1:
        # return to original image size
        IMcirc = ndimage.zoom(IMcirc, 1/zoom)

    if return_correction:
        return IMcirc, slice_angles, radcorr, radcorrspl
    else:
        return IMcirc

def circularize(IM, radcorrspl):
    """
    Remap image from its distorted grid to the true cartesian grid.

    Parameters
    ----------
    IM : numpy 2D array
        original image

    radcorrspl: numpy func(theta)
        spline function to evaluate radial correction at a given angle

    """
    # cartesian coordinate system
    Y, X = np.indices(IM.shape)

    row, col = IM.shape
    origin = (col//2 + col % 2, row//2 + row % 2)  # % handles odd size image

    # coordinates relative to center
    X -= origin[0]
    Y -= origin[1]
    theta = np.arctan2(X, Y) - np.pi/2 # referenced to vertical direction

    # radial correction
    # O2- this works better for O2-
    # Xactual = X*radcorrspl(theta)
    # Yactual = Y*radcorrspl(theta)
    # this corrects the sample image ??
    Xactual = X/radcorrspl(theta)
    Yactual = Y/radcorrspl(theta)

    # @DanHickstein magic
    # https://github.com/PyAbel/PyAbel/issues/186#issuecomment-275471271
    IMcirc = ndimage.interpolation.map_coordinates(IM,
                                  (Yactual+origin[1], Xactual+origin[0]))

    return IMcirc

def residual(param, radial, profile, previous):
    """ compare radially scaled intensity profile with adjacent 'previous'
        profile, to determine the radial scale factor param[0], and
        amplitude param[1].

    """
    newradial = radial*param[0]
    spline_prof = UnivariateSpline(newradial, profile, s=0, ext=3)
    newprof = spline_prof(radial)*param[1]

    # residual cf adjacent slice profile
    return newprof - previous

def correction(slice_angles, slices, radial, method):
    """ evaluate a radial scaling factor to align intensity profile features,
        between adjacent angular slices.

    """

    if method == "argmax":
        # follow position of intensity maximum
        pkpos = []

        for ang, aslice in zip(slice_angles, slices):
            profile = aslice.sum(axis=1)  # intensity vs radius for slice
            pkpos.append(profile.argmax())  # store index of peak position

        # radial scaling factor relative to peak max in first angular slice
        sf = radial[pkpos[0]]/radial[pkpos]

    elif method == "lsq":
        # least-squares radially scale intensity profile to match previous slice

        fitpar = np.array([1.0, 1.0])  # radial scale factor, amplitude
        sf = []
        sf.append(1)  # first slice nothing to compare with
        previous = slices[0].sum(axis=1)

        for ang, aslice in zip(slice_angles[1:], slices[1:]):
            profile = aslice.sum(axis=1)  # intensity vs radius for slice

            result = leastsq(residual, fitpar, args=(radial, profile, previous))

            sf.append(result[0][0])  # radial scale factor direct from lsq

            previous += residual(result[0], radial, profile, previous)
            fitpar = result[0]

        #sf[0] = sf[-1]

    else:
        raise ValueError("method variable must be one of 'argmax' or 'lsq',"
                         " not '{}'".format(method))

    return sf
<commit_msg>origin[1] - Y to give the correct cartesian grid, relative to indices<commit_after># -*- coding: utf-8 -*-
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
from scipy.interpolate import UnivariateSpline
from scipy.optimize import leastsq

import abel

# see https://github.com/PyAbel/PyAbel/issues/186 for discussion

def circularize_image(IM, method="argmax", center=None, radial_range=None,
                      zoom=1, smooth=1.0e-7, nslices=32, inverse=False,
                      return_correction=False):
    """
    Remove radial distortion from a velocity-map-image through radial scaling
    of the Newton-rings to enforce circularity.

    The image is divided into angular slices, adjacent radial intensity
    profile slices compared to give a radial scaling factor that best
    aligns common structure.

    The resultant radial scaling factor vs angle is interpolated to apply the
    correction to the whole image grid. The image is remapped to the
    correct cartesian grid.

    Parameters
    ----------
    IM : numpy 2D array

    method: str
        method to align slice profiles,
        "argmax" - compare intensity-profile.argmax() of each radial slice.
        "lsq" - determine radial scaling factor by minimizing the
                intensity-prfoile with an adjacent slice intensity-profile.

    center: str, or None
        pre-center image using PyAbel centering methods,
        "com", "convolution", "gaussian", "image_center", "slice"

    radial_range: tuple, or None
        limit slice comparison to the radial range (rmin, rmax).

    zoom: float
        nmimage.zoom image before analysis. This may help improve the low
        radius slice profile.

    smooth: float
        smoothing for spline interpolation of the determined radial scaling
        factor vs angle

    nslices: int
        divide the image into nslices. Ideally, this should be a divisor of
        the image column width.

    inverse: bool
        inverse Abel transform the *polar* image, to remove the background
        intensity. This helps follow a given peak intensityr, for all angles,
        when the anisotropy parameter is large.

    return_correction: bool
        additional outputs, see below

    Returns
    -------
    IMcirc : numpy 2D array, same size as input
        Circularized imput image

    (if return_correction is True)
    slice_angles: numpy 1D array
        Mid-point angle (radians) of image slice

    radial_corrections: numpy 1D array
        radial correction scale factors, at each angle

    spline_function: numpy func()
        spline function used to evaluate the radial correction at each angle

    """

    if zoom > 1:
        # zoom may improve the low radii angular resolution
        IM = ndimage.zoom(IM, zoom)

    if center is not None:
        # convenience function for the case image is not centered
        IM = abel.tools.center.center_image(IM, center=center)

    # map image into polar coordinates - much easier to slice
    # cartesian (Y, X) -> polar (Radius, Theta)
    polarIM, radial_coord, angle_coord =\
                           abel.tools.polar.reproject_image_into_polar(IM)

    if inverse:
        # pseudo inverse Abel transform of the polar image, removes background
        # to enhance transition peaks
        polarIM = abel.dasch.two_point_transform(polarIM.T).T

    # limit radial range of polar image, if selected
    radial = radial_coord[:, 0]
    if radial_range is not None:
        subr = np.logical_and(radial > radial_range[0]*zoom,
                              radial < radial_range[1]*zoom)
        polarIM = polarIM[subr]
        radial = radial[subr]

    # split image into n-slices
    slices = np.array_split(polarIM, nslices, axis=1)
    slice_angles = np.array(np.hsplit(angle_coord[0], nslices)).mean(axis=1)

    # evaluate radial scaling factor for each slice
    radcorr = correction(slice_angles, slices, radial, method=method)

    # spline radial scaling vs angle
    radcorrspl = UnivariateSpline(slice_angles, radcorr, s=smooth, ext=3)

    # apply the correction
    IMcirc = circularize(IM, radcorrspl)

    if zoom > 1:
        # return to original image size
        IMcirc = ndimage.zoom(IMcirc, 1/zoom)

    if return_correction:
        return IMcirc, slice_angles, radcorr, radcorrspl
    else:
        return IMcirc

def circularize(IM, radcorrspl):
    """
    Remap image from its distorted grid to the true cartesian grid.

    Parameters
    ----------
    IM : numpy 2D array
        original image

    radcorrspl: numpy func(theta)
        spline function to evaluate radial correction at a given angle

    """
    # cartesian coordinate system
    Y, X = np.indices(IM.shape)

    row, col = IM.shape
    origin = (col//2, row//2)  # odd image

    # coordinates relative to center
    X -= origin[0]
    Y = origin[1] - Y   # negative values below the axis
    theta = np.arctan2(X, Y)  # referenced to vertical direction

    # radial correction
    Xactual = X/radcorrspl(theta)
    Yactual = Y/radcorrspl(theta)

    # @DanHickstein magic
    # https://github.com/PyAbel/PyAbel/issues/186#issuecomment-275471271
    IMcirc = ndimage.interpolation.map_coordinates(IM,
                                  (Yactual+origin[1], Xactual+origin[0]))

    return IMcirc

def residual(param, radial, profile, previous):
    """ compare radially scaled intensity profile with adjacent 'previous'
        profile, to determine the radial scale factor param[0], and
        amplitude param[1].

    """
    newradial = radial*param[0]
    spline_prof = UnivariateSpline(newradial, profile, s=0, ext=3)
    newprof = spline_prof(radial)*param[1]

    # residual cf adjacent slice profile
    return newprof - previous

def correction(slice_angles, slices, radial, method):
    """ evaluate a radial scaling factor to align intensity profile features,
        between adjacent angular slices.

    """

    if method == "argmax":
        # follow position of intensity maximum
        pkpos = []

        for ang, aslice in zip(slice_angles, slices):
            profile = aslice.sum(axis=1)  # intensity vs radius for slice
            pkpos.append(profile.argmax())  # store index of peak position

        # radial scaling factor relative to peak max in first angular slice
        sf = radial[pkpos[0]]/radial[pkpos]

    elif method == "lsq":
        # least-squares radially scale intensity profile to match previous slice

        fitpar = np.array([1.0, 1.0])  # radial scale factor, amplitude
        sf = []
        sf.append(1)  # first slice nothing to compare with
        previous = slices[0].sum(axis=1)

        for ang, aslice in zip(slice_angles[1:], slices[1:]):
            profile = aslice.sum(axis=1)  # intensity vs radius for slice

            result = leastsq(residual, fitpar, args=(radial, profile, previous))

            sf.append(result[0][0])  # radial scale factor direct from lsq

            previous += residual(result[0], radial, profile, previous)
            fitpar = result[0]

    else:
        raise ValueError("method variable must be one of 'argmax' or 'lsq',"
                         " not '{}'".format(method))

    return sf
<|endoftext|>


<commit_before>from idiokit.xmlcore import Element
from abusehelper.core import serialize

class RuleError(Exception):
    pass

class _Rule(object):
    @classmethod
    def serialize_register(cls):
        name = "rule-" + cls.__name__.lower()
        serialize.register(cls.dump_rule, cls.load_rule, cls, name)

    def __init__(self, *args, **keys):
        self.arguments = tuple(args), frozenset(keys.items())

    def __eq__(self, other):
        if other.__class__ is not self.__class__:
            return NotImplemented
        return self.arguments == other.arguments

    def __ne__(self, other):
        result = self.__eq__(other)
        if result is NotImplemented:
            return result
        return not result

    def __hash__(self):
        return hash(self.__class__) ^ hash(self.arguments)

class NOT(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(dump(rule.child))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 1:
            raise RuleError(element)
        return cls(load(children[0]))

    def __init__(self, child):
        _Rule.__init__(self, child)
        self.child = child

    def __call__(self, *args, **keys):
        return not self.child(*args, **keys)
NOT.serialize_register()

class OR(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if child(*args, **keys):
                return True
        return False
OR.serialize_register()

class AND(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if not child(*args, **keys):
                return False
        return True
AND.serialize_register()

class CONTAINS(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(serialize.dump_list(dump, "keys", rule.keys))
        element.add(serialize.dump_dict(dump, "key-values", rule.key_values))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 2:
            raise RuleError(element)
        keys = set(serialize.load_list(load, children[0]))
        key_values = serialize.load_dict(load, children[1])
        return cls(*keys, **key_values)

    def __init__(self, *keys, **key_values):
        _Rule.__init__(self, *keys, **key_values)
        self.keys = keys
        self.key_values = key_values

    def __call__(self, event):
        for key in self.keys:
            if not event.contains_key(key):
                return False
        for key, value in self.key_values.items():
            if not event.contains_key_value(key, value):
                return False
        return True
CONTAINS.serialize_register()

import struct
from socket import inet_aton, error

class NETBLOCK(_Rule):
    unpack_ip = struct.Struct("!I").unpack

    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name, ip=rule.ip, bits=rule.bits)
        if rule.keys is not None:
            element.add(serialize.dump_list(dump, "keys", rule.keys))
        return element

    @classmethod
    def load_rule(cls, load, element):
        ip = element.get_attr("ip", None)
        bits = element.get_attr("bits", None)
        if None in (ip, bits):
            raise RuleError(element)
        try:
            bits = int(bits)
        except ValueError:
            raise RuleError(element)

        keys = None
        for child in element.children():
            keys = serialize.load_list(load, child)
        return cls(ip, bits, keys)

    def ip_to_num(self, ip):
        try:
            packed = inet_aton(ip)
        except error:
            return None
        return self.unpack_ip(packed)[0]

    def __init__(self, ip, bits, keys=None):
        if keys is not None:
            keys = frozenset(keys)

        _Rule.__init__(self, ip, bits, keys)

        self.ip = ip
        self.bits = bits
        self.keys = keys

        self.mask = ((1<<32)-1) ^ ((1<<(32-bits))-1)
        self.ip_num = self.ip_to_num(ip) & self.mask

    def __call__(self, event):
        if self.keys is None:
            keys = event.attrs.keys()
        else:
            keys = self.keys

        for key in keys:
            values = event.attrs.get(key, ())
            for value in values:
                ip_num = self.ip_to_num(value)
                if ip_num is None:
                    continue
                if ip_num & self.mask == self.ip_num:
                    return True
        return False
NETBLOCK.serialize_register()

import unittest

class MockEvent(object):
    def __init__(self, **keys):
        self.attrs = dict(keys)

class NetblockTests(unittest.TestCase):
    def test_match_arbitrary_key(self):
        rule = NETBLOCK("1.2.3.4", 24)
        assert rule(MockEvent(somekey=["1.2.3.255"]))

    def test_non_match_arbitrary_key(self):
        rule = NETBLOCK("1.2.3.4", 24)
        assert not rule(MockEvent(somekey=["1.2.4.255"]))

    def test_match_given_key(self):
        rule = NETBLOCK("1.2.3.4", 24, keys=["ip"])
        assert rule(MockEvent(somekey=["4.5.6.255"], ip=["1.2.3.255"]))

    def test_nonmatch_given_key(self):
        rule = NETBLOCK("1.2.3.4", 24, keys=["ip"])
        assert not rule(MockEvent(somekey=["4.5.6.255"], ip=["1.2.4.255"]))

    def test_non_match_non_ip_data(self):
        rule = NETBLOCK("1.2.3.4", 24)
        assert not rule(MockEvent(somekey=["this is just some data"]))

if __name__ == "__main__":
    unittest.main()
<commit_msg>Added IPv6 support for abusehelper.core.rules.NETBLOCK.<commit_after>from idiokit.xmlcore import Element
from abusehelper.core import serialize

class RuleError(Exception):
    pass

class _Rule(object):
    @classmethod
    def serialize_register(cls):
        name = "rule-" + cls.__name__.lower()
        serialize.register(cls.dump_rule, cls.load_rule, cls, name)

    def __init__(self, *args, **keys):
        self.arguments = tuple(args), frozenset(keys.items())

    def __eq__(self, other):
        if other.__class__ is not self.__class__:
            return NotImplemented
        return self.arguments == other.arguments

    def __ne__(self, other):
        result = self.__eq__(other)
        if result is NotImplemented:
            return result
        return not result

    def __hash__(self):
        return hash(self.__class__) ^ hash(self.arguments)

class NOT(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(dump(rule.child))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 1:
            raise RuleError(element)
        return cls(load(children[0]))

    def __init__(self, child):
        _Rule.__init__(self, child)
        self.child = child

    def __call__(self, *args, **keys):
        return not self.child(*args, **keys)
NOT.serialize_register()

class OR(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if child(*args, **keys):
                return True
        return False
OR.serialize_register()

class AND(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if not child(*args, **keys):
                return False
        return True
AND.serialize_register()

class CONTAINS(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(serialize.dump_list(dump, "keys", rule.keys))
        element.add(serialize.dump_dict(dump, "key-values", rule.key_values))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 2:
            raise RuleError(element)
        keys = set(serialize.load_list(load, children[0]))
        key_values = serialize.load_dict(load, children[1])
        return cls(*keys, **key_values)

    def __init__(self, *keys, **key_values):
        _Rule.__init__(self, *keys, **key_values)
        self.keys = keys
        self.key_values = key_values

    def __call__(self, event):
        for key in self.keys:
            if not event.contains_key(key):
                return False
        for key, value in self.key_values.items():
            if not event.contains_key_value(key, value):
                return False
        return True
CONTAINS.serialize_register()

import struct
from socket import inet_pton, AF_INET, AF_INET6, error

class NETBLOCKError(Exception):
    pass

class NETBLOCK(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name, ip=rule.ip, bits=rule.bits)
        if rule.keys is not None:
            element.add(serialize.dump_list(dump, "keys", rule.keys))
        return element

    @classmethod
    def load_rule(cls, load, element):
        ip = element.get_attr("ip", None)
        bits = element.get_attr("bits", None)
        if None in (ip, bits):
            raise RuleError(element)
        try:
            bits = int(bits)
        except ValueError:
            raise RuleError(element)

        keys = None
        for child in element.children():
            keys = serialize.load_list(load, child)
        return cls(ip, bits, keys)

    _unpack_ipv4 = struct.Struct("!I").unpack
    _unpack_ipv6 = struct.Struct("!QQ").unpack
    def parse_ip(self, ip):
        try:
            packed = inet_pton(AF_INET, ip)
            return 4, self._unpack_ipv4(packed)[0]
        except error:
            try:
                packed = inet_pton(AF_INET6, ip)
                hi, lo = self._unpack_ipv6(packed)
                return 6, ((hi << 64) | lo)
            except error:
                return None

    def __init__(self, ip, bits, keys=["ip"]):
        if keys is not None:
            keys = frozenset(keys)

        _Rule.__init__(self, ip, bits, keys)

        self.ip = ip
        self.bits = bits
        self.keys = keys

        parsed = self.parse_ip(ip)
        if parsed is None:
            raise NETBLOCKError("could not parse IP %r" % ip)
        self.version, self.ip_num = parsed
        assert self.version in (4, 6)

        if self.version == 4:
            self.mask = ((1<<32)-1) ^ ((1<<(32-bits))-1)
        elif self.version == 6:
            self.mask = ((1<<128)-1) ^ ((1<<(128-bits))-1)
        self.ip_num &= self.mask

    def __call__(self, event):
        if self.keys is None:
            keys = event.attrs.keys()
        else:
            keys = self.keys

        for key in keys:
            values = event.attrs.get(key, ())
            for value in values:
                parsed = self.parse_ip(value)
                if parsed is None:
                    continue
                version, ip_num = parsed
                if version != self.version:
                    continue
                if ip_num & self.mask == self.ip_num:
                    return True
        return False
NETBLOCK.serialize_register()

import unittest

class MockEvent(object):
    def __init__(self, **keys):
        self.attrs = dict(keys)

class NetblockTests(unittest.TestCase):
    def test_match_ipv6(self):
        rule = NETBLOCK("2001:0db8:ac10:fe01::", 32)
        assert rule(MockEvent(ip=["2001:0db8:aaaa:bbbb:cccc::"]))

    def test_non_match_ipv6(self):
        rule = NETBLOCK("2001:0db8:ac10:fe01::", 32)
        assert not rule(MockEvent(ip=["::1"]))

    def test_match_arbitrary_key(self):
        rule = NETBLOCK("1.2.3.4", 24, keys=None)
        assert rule(MockEvent(somekey=["1.2.3.255"]))

    def test_non_match_arbitrary_key(self):
        rule = NETBLOCK("1.2.3.4", 24, keys=None)
        assert not rule(MockEvent(somekey=["1.2.4.255"]))

    def test_match_ip_key(self):
        rule = NETBLOCK("1.2.3.4", 24)
        assert rule(MockEvent(somekey=["4.5.6.255"], ip=["1.2.3.255"]))

    def test_nonmatch_ip_key(self):
        rule = NETBLOCK("1.2.3.4", 24, keys=["ip"])
        assert not rule(MockEvent(somekey=["4.5.6.255"], ip=["1.2.4.255"]))

    def test_non_match_non_ip_data(self):
        rule = NETBLOCK("1.2.3.4", 24)
        assert not rule(MockEvent(ip=["this is just some data"]))

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


<commit_before># Copyright 2016 Google LLC All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Pools managing shared Session objects."""

import datetime

from six.moves import queue

from google.cloud.exceptions import NotFound
from google.cloud.spanner_v1._helpers import _metadata_with_prefix

_NOW = datetime.datetime.utcnow  # unit tests may replace

class AbstractSessionPool(object):
    """Specifies required API for concrete session pool implementations.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    _database = None

    def __init__(self, labels=None):
        if labels is None:
            labels = {}
        self._labels = labels

    @property
    def labels(self):
        """User-assigned labels for sesions created by the pool.

        :rtype: dict (str -> str)
        :returns: labels assigned by the user
        """
        return self._labels

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.

        Concrete implementations of this method may pre-fill the pool
        using the database.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def get(self):
        """Check a session out from the pool.

        Concrete implementations of this method are allowed to raise an
        error to signal that the pool is exhausted, or to block until a
        session is available.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def put(self, session):
        """Return a session to the pool.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        Concrete implementations of this method are allowed to raise an
        error to signal that the pool is full, or to block until it is
        not full.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def clear(self):
        """Delete all sessions in the pool.

        Concrete implementations of this method are allowed to raise an
        error to signal that the pool is full, or to block until it is
        not full.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def _new_session(self):
        """Helper for concrete methods creating session instances.

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: new session instance.
        """
        if self.labels:
            return self._database.session(labels=self.labels)
        return self._database.session()

    def session(self, **kwargs):
        """Check out a session from the pool.

        :param kwargs: (optional) keyword arguments, passed through to
                       the returned checkout.

        :rtype: :class:`~google.cloud.spanner_v1.session.SessionCheckout`
        :returns: a checkout instance, to be used as a context manager for
                  accessing the session and returning it to the pool.
        """
        return SessionCheckout(self, **kwargs)

class FixedSizePool(AbstractSessionPool):
    """Concrete session pool implementation:

    - Pre-allocates / creates a fixed number of sessions.

    - "Pings" existing sessions via :meth:`session.exists` before returning
      them, and replaces expired sessions.

    - Blocks, with a timeout, when :meth:`get` is called on an empty pool.
      Raises after timing out.

    - Raises when :meth:`put` is called on a full pool.  That error is
      never expected in normal practice, as users should be calling
      :meth:`get` followed by :meth:`put` whenever in need of a session.

    :type size: int
    :param size: fixed pool size

    :type default_timeout: int
    :param default_timeout: default timeout, in seconds, to wait for
                                 a returned session.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    DEFAULT_SIZE = 10
    DEFAULT_TIMEOUT = 10

    def __init__(self, size=DEFAULT_SIZE, default_timeout=DEFAULT_TIMEOUT, labels=None):
        super(FixedSizePool, self).__init__(labels=labels)
        self.size = size
        self.default_timeout = default_timeout
        self._sessions = queue.LifoQueue(size)

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        self._database = database
        api = database.spanner_api
        metadata = _metadata_with_prefix(database.name)

        while not self._sessions.full():
            resp = api.batch_create_sessions(
                database.name,
                self.size - self._sessions.qsize(),
                timeout=self.default_timeout,
                metadata=metadata,
            )
            for session_pb in resp.session:
                session = self._new_session()
                session._session_id = session_pb.name.split("/")[-1]
                self._sessions.put(session)

    def get(self, timeout=None):  # pylint: disable=arguments-differ
        """Check a session out from the pool.

        :type timeout: int
        :param timeout: seconds to block waiting for an available session

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: an existing session from the pool, or a newly-created
                  session.
        :raises: :exc:`six.moves.queue.Empty` if the queue is empty.
        """
        if timeout is None:
            timeout = self.default_timeout

        session = self._sessions.get(block=True, timeout=timeout)

        if not session.exists():
            session = self._database.session()
            session.create()

        return session

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, raises.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        :raises: :exc:`six.moves.queue.Full` if the queue is full.
        """
        self._sessions.put_nowait(session)

    def clear(self):
        """Delete all sessions in the pool."""

        while True:
            try:
                session = self._sessions.get(block=False)
            except queue.Empty:
                break
            else:
                session.delete()

class BurstyPool(AbstractSessionPool):
    """Concrete session pool implementation:

    - "Pings" existing sessions via :meth:`session.exists` before returning
      them.

    - Creates a new session, rather than blocking, when :meth:`get` is called
      on an empty pool.

    - Discards the returned session, rather than blocking, when :meth:`put`
      is called on a full pool.

    :type target_size: int
    :param target_size: max pool size

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    def __init__(self, target_size=10, labels=None):
        super(BurstyPool, self).__init__(labels=labels)
        self.target_size = target_size
        self._database = None
        self._sessions = queue.LifoQueue(target_size)

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        self._database = database

    def get(self):
        """Check a session out from the pool.

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: an existing session from the pool, or a newly-created
                  session.
        """
        try:
            session = self._sessions.get_nowait()
        except queue.Empty:
            session = self._new_session()
            session.create()
        else:
            if not session.exists():
                session = self._new_session()
                session.create()
        return session

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, the returned session is
        discarded.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.
        """
        try:
            self._sessions.put_nowait(session)
        except queue.Full:
            try:
                session.delete()
            except NotFound:
                pass

    def clear(self):
        """Delete all sessions in the pool."""

        while True:
            try:
                session = self._sessions.get(block=False)
            except queue.Empty:
                break
            else:
                session.delete()

class PingingPool(AbstractSessionPool):
    """Concrete session pool implementation:

    - Pre-allocates / creates a fixed number of sessions.

    - Sessions are used in "round-robin" order (LRU first).

    - "Pings" existing sessions in the background after a specified interval
      via an API call (``session.exists()``).

    - Blocks, with a timeout, when :meth:`get` is called on an empty pool.
      Raises after timing out.

    - Raises when :meth:`put` is called on a full pool.  That error is
      never expected in normal practice, as users should be calling
      :meth:`get` followed by :meth:`put` whenever in need of a session.

    The application is responsible for calling :meth:`ping` at appropriate
    times, e.g. from a background thread.

    :type size: int
    :param size: fixed pool size

    :type default_timeout: int
    :param default_timeout: default timeout, in seconds, to wait for
                            a returned session.

    :type ping_interval: int
    :param ping_interval: interval at which to ping sessions.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
        super(PingingPool, self).__init__(labels=labels)
        self.size = size
        self.default_timeout = default_timeout
        self._delta = datetime.timedelta(seconds=ping_interval)
        self._sessions = queue.PriorityQueue(size)

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        self._database = database
        api = database.spanner_api
        metadata = _metadata_with_prefix(database.name)
        created_session_count = 0

        while created_session_count < self.size:
            resp = api.batch_create_sessions(
                database.name,
                self.size - created_session_count,
                timeout=self.default_timeout,
                metadata=metadata,
            )
            for session_pb in resp.session:
                session = self._new_session()
                session._session_id = session_pb.name.split("/")[-1]
                self.put(session)
            created_session_count += len(resp.session)

    def get(self, timeout=None):  # pylint: disable=arguments-differ
        """Check a session out from the pool.

        :type timeout: int
        :param timeout: seconds to block waiting for an available session

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: an existing session from the pool, or a newly-created
                  session.
        :raises: :exc:`six.moves.queue.Empty` if the queue is empty.
        """
        if timeout is None:
            timeout = self.default_timeout

        ping_after, session = self._sessions.get(block=True, timeout=timeout)

        if _NOW() > ping_after:
            if not session.exists():
                session = self._new_session()
                session.create()

        return session

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, raises.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        :raises: :exc:`six.moves.queue.Full` if the queue is full.
        """
        self._sessions.put_nowait((_NOW() + self._delta, session))

    def clear(self):
        """Delete all sessions in the pool."""
        while True:
            try:
                _, session = self._sessions.get(block=False)
            except queue.Empty:
                break
            else:
                session.delete()

    def ping(self):
        """Refresh maybe-expired sessions in the pool.

        This method is designed to be called from a background thread,
        or during the "idle" phase of an event loop.
        """
        while True:
            try:
                ping_after, session = self._sessions.get(block=False)
            except queue.Empty:  # all sessions in use
                break
            if ping_after > _NOW():  # oldest session is fresh
                # Re-add to queue with existing expiration
                self._sessions.put((ping_after, session))
                break
            if not session.exists():  # stale
                session = self._new_session()
                session.create()
            # Re-add to queue with new expiration
            self.put(session)

class TransactionPingingPool(PingingPool):
    """Concrete session pool implementation:

    In addition to the features of :class:`PingingPool`, this class
    creates and begins a transaction for each of its sessions at startup.

    When a session is returned to the pool, if its transaction has been
    committed or rolled back, the pool creates a new transaction for the
    session and pushes the transaction onto a separate queue of "transactions
    to begin."  The application is responsible for flushing this queue
    as appropriate via the pool's :meth:`begin_pending_transactions` method.

    :type size: int
    :param size: fixed pool size

    :type default_timeout: int
    :param default_timeout: default timeout, in seconds, to wait for
                            a returned session.

    :type ping_interval: int
    :param ping_interval: interval at which to ping sessions.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
        self._pending_sessions = queue.Queue()

        super(TransactionPingingPool, self).__init__(
            size, default_timeout, ping_interval, labels=labels
        )

        self.begin_pending_transactions()

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        super(TransactionPingingPool, self).bind(database)
        self.begin_pending_transactions()

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, raises.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        :raises: :exc:`six.moves.queue.Full` if the queue is full.
        """
        if self._sessions.full():
            raise queue.Full

        txn = session._transaction
        if txn is None or txn.committed or txn.rolled_back:
            session.transaction()
            self._pending_sessions.put(session)
        else:
            super(TransactionPingingPool, self).put(session)

    def begin_pending_transactions(self):
        """Begin all transactions for sessions added to the pool."""
        while not self._pending_sessions.empty():
            session = self._pending_sessions.get()
            session._transaction.begin()
            super(TransactionPingingPool, self).put(session)

class SessionCheckout(object):
    """Context manager: hold session checked out from a pool.

    :type pool: concrete subclass of
        :class:`~google.cloud.spanner_v1.session.AbstractSessionPool`
    :param pool: Pool from which to check out a session.

    :param kwargs: extra keyword arguments to be passed to :meth:`pool.get`.
    """

    _session = None  # Not checked out until '__enter__'.

    def __init__(self, pool, **kwargs):
        self._pool = pool
        self._kwargs = kwargs.copy()

    def __enter__(self):
        self._session = self._pool.get(**self._kwargs)
        return self._session

    def __exit__(self, *ignored):
        self._pool.put(self._session)
<commit_msg>fix: remove erroneous timeouts for batch_create_session calls (#18)<commit_after># Copyright 2016 Google LLC All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Pools managing shared Session objects."""

import datetime

from six.moves import queue

from google.cloud.exceptions import NotFound
from google.cloud.spanner_v1._helpers import _metadata_with_prefix

_NOW = datetime.datetime.utcnow  # unit tests may replace

class AbstractSessionPool(object):
    """Specifies required API for concrete session pool implementations.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    _database = None

    def __init__(self, labels=None):
        if labels is None:
            labels = {}
        self._labels = labels

    @property
    def labels(self):
        """User-assigned labels for sesions created by the pool.

        :rtype: dict (str -> str)
        :returns: labels assigned by the user
        """
        return self._labels

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.

        Concrete implementations of this method may pre-fill the pool
        using the database.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def get(self):
        """Check a session out from the pool.

        Concrete implementations of this method are allowed to raise an
        error to signal that the pool is exhausted, or to block until a
        session is available.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def put(self, session):
        """Return a session to the pool.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        Concrete implementations of this method are allowed to raise an
        error to signal that the pool is full, or to block until it is
        not full.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def clear(self):
        """Delete all sessions in the pool.

        Concrete implementations of this method are allowed to raise an
        error to signal that the pool is full, or to block until it is
        not full.

        :raises NotImplementedError: abstract method
        """
        raise NotImplementedError()

    def _new_session(self):
        """Helper for concrete methods creating session instances.

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: new session instance.
        """
        if self.labels:
            return self._database.session(labels=self.labels)
        return self._database.session()

    def session(self, **kwargs):
        """Check out a session from the pool.

        :param kwargs: (optional) keyword arguments, passed through to
                       the returned checkout.

        :rtype: :class:`~google.cloud.spanner_v1.session.SessionCheckout`
        :returns: a checkout instance, to be used as a context manager for
                  accessing the session and returning it to the pool.
        """
        return SessionCheckout(self, **kwargs)

class FixedSizePool(AbstractSessionPool):
    """Concrete session pool implementation:

    - Pre-allocates / creates a fixed number of sessions.

    - "Pings" existing sessions via :meth:`session.exists` before returning
      them, and replaces expired sessions.

    - Blocks, with a timeout, when :meth:`get` is called on an empty pool.
      Raises after timing out.

    - Raises when :meth:`put` is called on a full pool.  That error is
      never expected in normal practice, as users should be calling
      :meth:`get` followed by :meth:`put` whenever in need of a session.

    :type size: int
    :param size: fixed pool size

    :type default_timeout: int
    :param default_timeout: default timeout, in seconds, to wait for
                                 a returned session.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    DEFAULT_SIZE = 10
    DEFAULT_TIMEOUT = 10

    def __init__(self, size=DEFAULT_SIZE, default_timeout=DEFAULT_TIMEOUT, labels=None):
        super(FixedSizePool, self).__init__(labels=labels)
        self.size = size
        self.default_timeout = default_timeout
        self._sessions = queue.LifoQueue(size)

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        self._database = database
        api = database.spanner_api
        metadata = _metadata_with_prefix(database.name)

        while not self._sessions.full():
            resp = api.batch_create_sessions(
                database.name, self.size - self._sessions.qsize(), metadata=metadata
            )
            for session_pb in resp.session:
                session = self._new_session()
                session._session_id = session_pb.name.split("/")[-1]
                self._sessions.put(session)

    def get(self, timeout=None):  # pylint: disable=arguments-differ
        """Check a session out from the pool.

        :type timeout: int
        :param timeout: seconds to block waiting for an available session

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: an existing session from the pool, or a newly-created
                  session.
        :raises: :exc:`six.moves.queue.Empty` if the queue is empty.
        """
        if timeout is None:
            timeout = self.default_timeout

        session = self._sessions.get(block=True, timeout=timeout)

        if not session.exists():
            session = self._database.session()
            session.create()

        return session

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, raises.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        :raises: :exc:`six.moves.queue.Full` if the queue is full.
        """
        self._sessions.put_nowait(session)

    def clear(self):
        """Delete all sessions in the pool."""

        while True:
            try:
                session = self._sessions.get(block=False)
            except queue.Empty:
                break
            else:
                session.delete()

class BurstyPool(AbstractSessionPool):
    """Concrete session pool implementation:

    - "Pings" existing sessions via :meth:`session.exists` before returning
      them.

    - Creates a new session, rather than blocking, when :meth:`get` is called
      on an empty pool.

    - Discards the returned session, rather than blocking, when :meth:`put`
      is called on a full pool.

    :type target_size: int
    :param target_size: max pool size

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    def __init__(self, target_size=10, labels=None):
        super(BurstyPool, self).__init__(labels=labels)
        self.target_size = target_size
        self._database = None
        self._sessions = queue.LifoQueue(target_size)

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        self._database = database

    def get(self):
        """Check a session out from the pool.

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: an existing session from the pool, or a newly-created
                  session.
        """
        try:
            session = self._sessions.get_nowait()
        except queue.Empty:
            session = self._new_session()
            session.create()
        else:
            if not session.exists():
                session = self._new_session()
                session.create()
        return session

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, the returned session is
        discarded.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.
        """
        try:
            self._sessions.put_nowait(session)
        except queue.Full:
            try:
                session.delete()
            except NotFound:
                pass

    def clear(self):
        """Delete all sessions in the pool."""

        while True:
            try:
                session = self._sessions.get(block=False)
            except queue.Empty:
                break
            else:
                session.delete()

class PingingPool(AbstractSessionPool):
    """Concrete session pool implementation:

    - Pre-allocates / creates a fixed number of sessions.

    - Sessions are used in "round-robin" order (LRU first).

    - "Pings" existing sessions in the background after a specified interval
      via an API call (``session.exists()``).

    - Blocks, with a timeout, when :meth:`get` is called on an empty pool.
      Raises after timing out.

    - Raises when :meth:`put` is called on a full pool.  That error is
      never expected in normal practice, as users should be calling
      :meth:`get` followed by :meth:`put` whenever in need of a session.

    The application is responsible for calling :meth:`ping` at appropriate
    times, e.g. from a background thread.

    :type size: int
    :param size: fixed pool size

    :type default_timeout: int
    :param default_timeout: default timeout, in seconds, to wait for
                            a returned session.

    :type ping_interval: int
    :param ping_interval: interval at which to ping sessions.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
        super(PingingPool, self).__init__(labels=labels)
        self.size = size
        self.default_timeout = default_timeout
        self._delta = datetime.timedelta(seconds=ping_interval)
        self._sessions = queue.PriorityQueue(size)

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        self._database = database
        api = database.spanner_api
        metadata = _metadata_with_prefix(database.name)
        created_session_count = 0

        while created_session_count < self.size:
            resp = api.batch_create_sessions(
                database.name, self.size - created_session_count, metadata=metadata
            )
            for session_pb in resp.session:
                session = self._new_session()
                session._session_id = session_pb.name.split("/")[-1]
                self.put(session)
            created_session_count += len(resp.session)

    def get(self, timeout=None):  # pylint: disable=arguments-differ
        """Check a session out from the pool.

        :type timeout: int
        :param timeout: seconds to block waiting for an available session

        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
        :returns: an existing session from the pool, or a newly-created
                  session.
        :raises: :exc:`six.moves.queue.Empty` if the queue is empty.
        """
        if timeout is None:
            timeout = self.default_timeout

        ping_after, session = self._sessions.get(block=True, timeout=timeout)

        if _NOW() > ping_after:
            if not session.exists():
                session = self._new_session()
                session.create()

        return session

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, raises.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        :raises: :exc:`six.moves.queue.Full` if the queue is full.
        """
        self._sessions.put_nowait((_NOW() + self._delta, session))

    def clear(self):
        """Delete all sessions in the pool."""
        while True:
            try:
                _, session = self._sessions.get(block=False)
            except queue.Empty:
                break
            else:
                session.delete()

    def ping(self):
        """Refresh maybe-expired sessions in the pool.

        This method is designed to be called from a background thread,
        or during the "idle" phase of an event loop.
        """
        while True:
            try:
                ping_after, session = self._sessions.get(block=False)
            except queue.Empty:  # all sessions in use
                break
            if ping_after > _NOW():  # oldest session is fresh
                # Re-add to queue with existing expiration
                self._sessions.put((ping_after, session))
                break
            if not session.exists():  # stale
                session = self._new_session()
                session.create()
            # Re-add to queue with new expiration
            self.put(session)

class TransactionPingingPool(PingingPool):
    """Concrete session pool implementation:

    In addition to the features of :class:`PingingPool`, this class
    creates and begins a transaction for each of its sessions at startup.

    When a session is returned to the pool, if its transaction has been
    committed or rolled back, the pool creates a new transaction for the
    session and pushes the transaction onto a separate queue of "transactions
    to begin."  The application is responsible for flushing this queue
    as appropriate via the pool's :meth:`begin_pending_transactions` method.

    :type size: int
    :param size: fixed pool size

    :type default_timeout: int
    :param default_timeout: default timeout, in seconds, to wait for
                            a returned session.

    :type ping_interval: int
    :param ping_interval: interval at which to ping sessions.

    :type labels: dict (str -> str) or None
    :param labels: (Optional) user-assigned labels for sessions created
                    by the pool.
    """

    def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
        self._pending_sessions = queue.Queue()

        super(TransactionPingingPool, self).__init__(
            size, default_timeout, ping_interval, labels=labels
        )

        self.begin_pending_transactions()

    def bind(self, database):
        """Associate the pool with a database.

        :type database: :class:`~google.cloud.spanner_v1.database.Database`
        :param database: database used by the pool:  used to create sessions
                         when needed.
        """
        super(TransactionPingingPool, self).bind(database)
        self.begin_pending_transactions()

    def put(self, session):
        """Return a session to the pool.

        Never blocks:  if the pool is full, raises.

        :type session: :class:`~google.cloud.spanner_v1.session.Session`
        :param session: the session being returned.

        :raises: :exc:`six.moves.queue.Full` if the queue is full.
        """
        if self._sessions.full():
            raise queue.Full

        txn = session._transaction
        if txn is None or txn.committed or txn.rolled_back:
            session.transaction()
            self._pending_sessions.put(session)
        else:
            super(TransactionPingingPool, self).put(session)

    def begin_pending_transactions(self):
        """Begin all transactions for sessions added to the pool."""
        while not self._pending_sessions.empty():
            session = self._pending_sessions.get()
            session._transaction.begin()
            super(TransactionPingingPool, self).put(session)

class SessionCheckout(object):
    """Context manager: hold session checked out from a pool.

    :type pool: concrete subclass of
        :class:`~google.cloud.spanner_v1.session.AbstractSessionPool`
    :param pool: Pool from which to check out a session.

    :param kwargs: extra keyword arguments to be passed to :meth:`pool.get`.
    """

    _session = None  # Not checked out until '__enter__'.

    def __init__(self, pool, **kwargs):
        self._pool = pool
        self._kwargs = kwargs.copy()

    def __enter__(self):
        self._session = self._pool.get(**self._kwargs)
        return self._session

    def __exit__(self, *ignored):
        self._pool.put(self._session)
<|endoftext|>