<commit_before># coding=utf-8
# Copyright 2021 The Tensor2Robot Authors.
#
# 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.

# Lint as python3
"""Default input generators wrapping tfdata, metadata and replay buffer."""

import abc
import json
import os
from typing import Dict, Optional, Text, Union, Sequence

import gin
import six
from tensor2robot.input_generators import abstract_input_generator
from tensor2robot.utils import tensorspec_utils
from tensor2robot.utils import tfdata
import tensorflow.compat.v1 as tf

_TF_CONFIG_ENV = 'TF_CONFIG'
_MULTI_EVAL_NAME = 'multi_eval_name'

def _get_tf_config_env():
  """Returns the TF_CONFIG environment variable as dict."""
  return json.loads(os.environ.get(_TF_CONFIG_ENV, '{}'))

def get_multi_eval_name(tf_config_env=None):
  """Returns the multi eval index set in the TF_CONFIG."""
  tf_config_env = tf_config_env or _get_tf_config_env()
  return tf_config_env.get(_MULTI_EVAL_NAME)

@gin.configurable
class DefaultRecordInputGenerator(
    abstract_input_generator.AbstractInputGenerator):
  """A recordio, sstable or tfrecords based input generator."""

  def __init__(self,
               file_patterns = None,
               dataset_map = None,
               label = '',
               **parent_kwargs):
    """Create an instance.

    Args:
      file_patterns: Comma-separated string of file patterns [recordio, sstable,
        tfrecord] we will load the data from. Only one of `file_patterns` or
        `dataset_map` should be set.
      dataset_map: Dictionary of dataset_key and Comma-separated string of file
        patterns we load data from. Used for loading from multiple datasets.
        Only one of `file_patterns` or `dataset_map` should be set.
      label: Name of the input generator.
      **parent_kwargs: All parent arguments.
    """
    super(DefaultRecordInputGenerator, self).__init__(**parent_kwargs)
    if file_patterns and dataset_map:
      raise ValueError(
          'Only one of `file_patterns` or `dataset_map` should be set.')
    self._file_patterns = file_patterns
    self._dataset_map = dataset_map
    self._label = label

  def create_dataset(self,
                     mode,
                     params=None):
    """Create the actual input_fn.

    This is potentially wrapped in create_dataset_input_fn.

    Args:
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      params: Not used for this implementation but expected by callers. An
        optional dict of hyper parameters that will be passed into input_fn and
        model_fn. Keys are names of parameters, values are basic Python types.
        There are reserved keys for TPUEstimator, including 'batch_size'.

    Returns:
      A valid input_fn for the estimator api.
    """
    input_fn = tfdata.get_input_fn(
        file_patterns=self._file_patterns or self._dataset_map,
        batch_size=self.batch_size,
        feature_spec=self._feature_spec,
        label_spec=self._label_spec,
        mode=mode,
        preprocess_fn=self._preprocess_fn)
    return input_fn(params)

@gin.configurable
class FractionalRecordInputGenerator(DefaultRecordInputGenerator):
  """Fraction of files in dataset (e.g. data ablation experiments)."""

  def __init__(self,
               file_fraction = 1.0,
               **parent_kwargs):
    """Create an instance.

    Args:
      file_fraction: If file_fraction < 1.0, choose first file_fraction percent
        of files, (rounded down to the nearest integer number of files).
      **parent_kwargs: All parent arguments.
    """
    super(FractionalRecordInputGenerator, self).__init__(**parent_kwargs)
    if file_fraction < 1.0:
      data_format, filenames = tfdata.get_data_format_and_filenames(
          self._file_patterns)
      n = int(file_fraction * len(filenames))
      filenames = filenames[:n]
      self._file_patterns = '{}:{}'.format(data_format, ','.join(filenames))

@gin.configurable
class MultiEvalRecordInputGenerator(DefaultRecordInputGenerator):
  """Evaluating on multiple datasets."""

  def __init__(self,
               eval_map = None,
               **parent_kwargs):
    super(MultiEvalRecordInputGenerator, self).__init__(**parent_kwargs)
    # If multi_eval_name is set via TF_CONFIG_ENV variable, override dataset.
    multi_eval_name = get_multi_eval_name()
    if multi_eval_name:
      self._file_patterns = eval_map[multi_eval_name]
    else:
      raise ValueError('multi_eval_name not set in TF_CONFIG env variable')

class GeneratorInputGenerator(
    six.with_metaclass(abc.ABCMeta,
                       abstract_input_generator.AbstractInputGenerator)):
  """Class to use for constructing input generators from Python generator objects."""

  def __init__(self, sequence_length=None, **kwargs):
    self._sequence_length = sequence_length
    super(GeneratorInputGenerator, self).__init__(**kwargs)

  @abc.abstractmethod
  def _generator_fn(self, batch_size):
    """Subclasses override this generator.

    Args:
      batch_size: batch size to generate examples

    Yields:
      training examples
    """

  def create_dataset(self, mode, params=None):
    """Creates a dataset using _generator_fn.

    Args:
      mode: prediction mode
      params: Hyperparameters sent by estimator object.

    Returns:
      A tfdata.Dataset object.
    """
    used_batch_size = tfdata.get_batch_size(params, self._batch_size)

    def shape_transform(x):
      if isinstance(x, tensorspec_utils.ExtendedTensorSpec) and x.is_sequence:
        return [used_batch_size, self._sequence_length] + x.shape.as_list()
      else:
        return [used_batch_size] + x.shape.as_list()

    feature_dtypes = tf.nest.map_structure(lambda x: x.dtype,
                                           self._feature_spec)
    feature_shapes = tf.nest.map_structure(shape_transform, self._feature_spec)
    label_dtypes = tf.nest.map_structure(lambda x: x.dtype, self._label_spec)
    label_shapes = tf.nest.map_structure(shape_transform, self._label_spec)

    dataset = tf.data.Dataset.from_generator(
        lambda: self._generator_fn(used_batch_size),
        (feature_dtypes, label_dtypes), (feature_shapes, label_shapes))
    if self._preprocess_fn is not None:
      dataset = dataset.map(self._preprocess_fn)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset

@gin.configurable
class DefaultRandomInputGenerator(GeneratorInputGenerator):
  """An input generator generating random inputs."""

  def _generator_fn(self, batch_size):
    while True:
      features = tensorspec_utils.make_random_numpy(
          self._feature_spec, batch_size, self._sequence_length)
      labels = tensorspec_utils.make_random_numpy(self._label_spec, batch_size,
                                                  self._sequence_length)
      yield features, labels

@gin.configurable
class DefaultConstantInputGenerator(GeneratorInputGenerator):
  """An input generator generating constant inputs."""

  def __init__(self, constant_value, **kwargs):
    self._constant_value = constant_value
    super(DefaultConstantInputGenerator, self).__init__(**kwargs)

  def _generator_fn(self, batch_size):
    while True:
      features = tensorspec_utils.make_constant_numpy(
          self._feature_spec, self._constant_value, batch_size,
          self._sequence_length)
      labels = tensorspec_utils.make_constant_numpy(
          self._label_spec, self._constant_value, batch_size,
          self._sequence_length)
      yield features, labels

@gin.configurable
class WeightedRecordInputGenerator(DefaultRecordInputGenerator):
  """An input generator sampling multiple datasets with weighting."""

  def __init__(
      self,
      file_patterns,
      num_parallel_calls = 4,
      shuffle_buffer_size = 500,
      prefetch_buffer_size = (tf.data.experimental.AUTOTUNE),
      parallel_shards = 10,
      weights = None,
      seed = None,
      **parent_kwargs):
    """Create an instance.

    Args:
      file_patterns: Comma-separated string of file patterns, where each
        individual file contains data for one task.
      num_parallel_calls: The number elements to process in parallel.
      shuffle_buffer_size: How many examples to shuffle within each file.
      prefetch_buffer_size: How many examples to prefetch.
      parallel_shards: Shards for applying preprocess_fn.
      weights: Weight to sampling each file pattern with. Should be equal in
        length to number of file_patterns.
      seed: Seed for weighted dataset sampling.
      **parent_kwargs: All parent arguments.
    """
    super(WeightedRecordInputGenerator, self).__init__(**parent_kwargs)
    self._file_patterns = file_patterns
    self._num_parallel_calls = num_parallel_calls
    self._shuffle_buffer_size = shuffle_buffer_size
    self._prefetch_buffer_size = prefetch_buffer_size
    self._parallel_shards = parallel_shards
    self._weights = weights
    self._seed = seed

  def create_dataset(self, mode, params, **unused_kwargs):
    """This abstract function is not required for default input generators."""
    batch_size = tfdata.get_batch_size(params, self.batch_size)
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    data_format, filenames_list = tfdata.get_data_format_and_filenames_list(
        self._file_patterns)
    datasets = []
    if self._weights is not None:
      if len(filenames_list) != len(self._weights):
        raise ValueError(
            'Weights need to be same length as number of filenames.')
    for filenames in filenames_list:
      filenames_dataset = tf.data.Dataset.list_files(
          filenames, shuffle=is_training)
      if is_training:
        cycle_length = min(self._parallel_shards, len(filenames))
      else:
        cycle_length = 1
      dataset = filenames_dataset.apply(
          tf.data.experimental.parallel_interleave(
              tfdata.DATA_FORMAT[data_format],
              cycle_length=cycle_length,
              sloppy=is_training))
      if is_training:
        dataset = dataset.shuffle(
            buffer_size=self._shuffle_buffer_size).repeat()
      else:
        dataset = dataset.repeat()
      dataset = dataset.batch(batch_size, drop_remainder=True)
      datasets.append(dataset)
    if self._weights is None:
      weights = [float(1) for _ in range(len(datasets))]
    else:
      weights = [float(w) for w in self._weights]
    dataset = tf.data.experimental.sample_from_datasets(
        datasets=datasets, weights=weights, seed=self._seed)
    # Parse all datasets together.
    dataset = tfdata.serialized_to_parsed(
        dataset,
        self._feature_spec,
        self._label_spec,
        num_parallel_calls=self._num_parallel_calls)
    if self._preprocess_fn is not None:
      dataset = dataset.map(
          self._preprocess_fn, num_parallel_calls=self._parallel_shards)
    if self._prefetch_buffer_size is not None:
      dataset = dataset.prefetch(self._prefetch_buffer_size)
    return dataset
<commit_msg>Fixes two bugs in default_input_generator.WeightedRecordInputGenerator.<commit_after># coding=utf-8
# Copyright 2021 The Tensor2Robot Authors.
#
# 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.

# Lint as python3
"""Default input generators wrapping tfdata, metadata and replay buffer."""

import abc
import json
import os
from typing import Dict, Optional, Text, Union, Sequence

import gin
import numpy as np
import six
from tensor2robot.input_generators import abstract_input_generator
from tensor2robot.utils import tensorspec_utils
from tensor2robot.utils import tfdata
import tensorflow.compat.v1 as tf

_TF_CONFIG_ENV = 'TF_CONFIG'
_MULTI_EVAL_NAME = 'multi_eval_name'

def _get_tf_config_env():
  """Returns the TF_CONFIG environment variable as dict."""
  return json.loads(os.environ.get(_TF_CONFIG_ENV, '{}'))

def get_multi_eval_name(tf_config_env=None):
  """Returns the multi eval index set in the TF_CONFIG."""
  tf_config_env = tf_config_env or _get_tf_config_env()
  return tf_config_env.get(_MULTI_EVAL_NAME)

@gin.configurable
class DefaultRecordInputGenerator(
    abstract_input_generator.AbstractInputGenerator):
  """A recordio, sstable or tfrecords based input generator."""

  def __init__(self,
               file_patterns = None,
               dataset_map = None,
               label = '',
               **parent_kwargs):
    """Create an instance.

    Args:
      file_patterns: Comma-separated string of file patterns [recordio, sstable,
        tfrecord] we will load the data from. Only one of `file_patterns` or
        `dataset_map` should be set.
      dataset_map: Dictionary of dataset_key and Comma-separated string of file
        patterns we load data from. Used for loading from multiple datasets.
        Only one of `file_patterns` or `dataset_map` should be set.
      label: Name of the input generator.
      **parent_kwargs: All parent arguments.
    """
    super(DefaultRecordInputGenerator, self).__init__(**parent_kwargs)
    if file_patterns and dataset_map:
      raise ValueError(
          'Only one of `file_patterns` or `dataset_map` should be set.')
    self._file_patterns = file_patterns
    self._dataset_map = dataset_map
    self._label = label

  def create_dataset(self,
                     mode,
                     params=None):
    """Create the actual input_fn.

    This is potentially wrapped in create_dataset_input_fn.

    Args:
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      params: Not used for this implementation but expected by callers. An
        optional dict of hyper parameters that will be passed into input_fn and
        model_fn. Keys are names of parameters, values are basic Python types.
        There are reserved keys for TPUEstimator, including 'batch_size'.

    Returns:
      A valid input_fn for the estimator api.
    """
    input_fn = tfdata.get_input_fn(
        file_patterns=self._file_patterns or self._dataset_map,
        batch_size=self.batch_size,
        feature_spec=self._feature_spec,
        label_spec=self._label_spec,
        mode=mode,
        preprocess_fn=self._preprocess_fn)
    return input_fn(params)

@gin.configurable
class FractionalRecordInputGenerator(DefaultRecordInputGenerator):
  """Fraction of files in dataset (e.g. data ablation experiments)."""

  def __init__(self,
               file_fraction = 1.0,
               **parent_kwargs):
    """Create an instance.

    Args:
      file_fraction: If file_fraction < 1.0, choose first file_fraction percent
        of files, (rounded down to the nearest integer number of files).
      **parent_kwargs: All parent arguments.
    """
    super(FractionalRecordInputGenerator, self).__init__(**parent_kwargs)
    if file_fraction < 1.0:
      data_format, filenames = tfdata.get_data_format_and_filenames(
          self._file_patterns)
      n = int(file_fraction * len(filenames))
      filenames = filenames[:n]
      self._file_patterns = '{}:{}'.format(data_format, ','.join(filenames))

@gin.configurable
class MultiEvalRecordInputGenerator(DefaultRecordInputGenerator):
  """Evaluating on multiple datasets."""

  def __init__(self,
               eval_map = None,
               **parent_kwargs):
    super(MultiEvalRecordInputGenerator, self).__init__(**parent_kwargs)
    # If multi_eval_name is set via TF_CONFIG_ENV variable, override dataset.
    multi_eval_name = get_multi_eval_name()
    if multi_eval_name:
      self._file_patterns = eval_map[multi_eval_name]
    else:
      raise ValueError('multi_eval_name not set in TF_CONFIG env variable')

class GeneratorInputGenerator(
    six.with_metaclass(abc.ABCMeta,
                       abstract_input_generator.AbstractInputGenerator)):
  """Class to use for constructing input generators from Python generator objects."""

  def __init__(self, sequence_length=None, **kwargs):
    self._sequence_length = sequence_length
    super(GeneratorInputGenerator, self).__init__(**kwargs)

  @abc.abstractmethod
  def _generator_fn(self, batch_size):
    """Subclasses override this generator.

    Args:
      batch_size: batch size to generate examples

    Yields:
      training examples
    """

  def create_dataset(self, mode, params=None):
    """Creates a dataset using _generator_fn.

    Args:
      mode: prediction mode
      params: Hyperparameters sent by estimator object.

    Returns:
      A tfdata.Dataset object.
    """
    used_batch_size = tfdata.get_batch_size(params, self._batch_size)

    def shape_transform(x):
      if isinstance(x, tensorspec_utils.ExtendedTensorSpec) and x.is_sequence:
        return [used_batch_size, self._sequence_length] + x.shape.as_list()
      else:
        return [used_batch_size] + x.shape.as_list()

    feature_dtypes = tf.nest.map_structure(lambda x: x.dtype,
                                           self._feature_spec)
    feature_shapes = tf.nest.map_structure(shape_transform, self._feature_spec)
    label_dtypes = tf.nest.map_structure(lambda x: x.dtype, self._label_spec)
    label_shapes = tf.nest.map_structure(shape_transform, self._label_spec)

    dataset = tf.data.Dataset.from_generator(
        lambda: self._generator_fn(used_batch_size),
        (feature_dtypes, label_dtypes), (feature_shapes, label_shapes))
    if self._preprocess_fn is not None:
      dataset = dataset.map(self._preprocess_fn)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset

@gin.configurable
class DefaultRandomInputGenerator(GeneratorInputGenerator):
  """An input generator generating random inputs."""

  def _generator_fn(self, batch_size):
    while True:
      features = tensorspec_utils.make_random_numpy(
          self._feature_spec, batch_size, self._sequence_length)
      labels = tensorspec_utils.make_random_numpy(self._label_spec, batch_size,
                                                  self._sequence_length)
      yield features, labels

@gin.configurable
class DefaultConstantInputGenerator(GeneratorInputGenerator):
  """An input generator generating constant inputs."""

  def __init__(self, constant_value, **kwargs):
    self._constant_value = constant_value
    super(DefaultConstantInputGenerator, self).__init__(**kwargs)

  def _generator_fn(self, batch_size):
    while True:
      features = tensorspec_utils.make_constant_numpy(
          self._feature_spec, self._constant_value, batch_size,
          self._sequence_length)
      labels = tensorspec_utils.make_constant_numpy(
          self._label_spec, self._constant_value, batch_size,
          self._sequence_length)
      yield features, labels

@gin.configurable
class WeightedRecordInputGenerator(DefaultRecordInputGenerator):
  """An input generator sampling multiple datasets with weighting."""

  def __init__(
      self,
      file_patterns,
      num_parallel_calls = 4,
      shuffle_buffer_size = 500,
      prefetch_buffer_size = (tf.data.experimental.AUTOTUNE),
      parallel_shards = 10,
      weights = None,
      seed = None,
      **parent_kwargs):
    """Create an instance.

    Args:
      file_patterns: Comma-separated string of file patterns, where each
        individual file contains data for one task.
      num_parallel_calls: The number elements to process in parallel.
      shuffle_buffer_size: How many examples to shuffle within each file.
      prefetch_buffer_size: How many examples to prefetch.
      parallel_shards: Shards for applying preprocess_fn.
      weights: Weight to sampling each file pattern with. Should be equal in
        length to number of file_patterns.
      seed: Seed for weighted dataset sampling.
      **parent_kwargs: All parent arguments.
    """
    super(WeightedRecordInputGenerator, self).__init__(**parent_kwargs)
    self._file_patterns = file_patterns
    self._num_parallel_calls = num_parallel_calls
    self._shuffle_buffer_size = shuffle_buffer_size
    self._prefetch_buffer_size = prefetch_buffer_size
    self._parallel_shards = parallel_shards
    self._weights = weights
    self._seed = seed

  def create_dataset(self, mode, params, **unused_kwargs):
    """This abstract function is not required for default input generators."""
    batch_size = tfdata.get_batch_size(params, self.batch_size)
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    data_format, filenames_list = tfdata.get_data_format_and_filenames_list(
        self._file_patterns)
    datasets = []
    if self._weights is not None:
      if len(filenames_list) != len(self._weights):
        raise ValueError(
            'Weights need to be same length as number of filenames.')
    for filenames in filenames_list:
      filenames_dataset = tf.data.Dataset.list_files(
          filenames, shuffle=is_training)
      if is_training:
        cycle_length = min(self._parallel_shards, len(filenames))
      else:
        cycle_length = 1
      dataset = filenames_dataset.apply(
          tf.data.experimental.parallel_interleave(
              tfdata.DATA_FORMAT[data_format],
              cycle_length=cycle_length,
              sloppy=is_training))
      if is_training:
        dataset = dataset.shuffle(
            buffer_size=self._shuffle_buffer_size).repeat()
      else:
        dataset = dataset.repeat()
      datasets.append(dataset)
    if self._weights is None:
      weights = [float(1) for _ in range(len(datasets))]
    else:
      weights = self._weights
    sum_weight = np.sum(weights)
    weights = [float(w) / sum_weight for w in weights]
    dataset = tf.data.experimental.sample_from_datasets(
        datasets=datasets, weights=weights, seed=self._seed)
    dataset = dataset.batch(batch_size, drop_remainder=True)
    # Parse all datasets together.
    dataset = tfdata.serialized_to_parsed(
        dataset,
        self._feature_spec,
        self._label_spec,
        num_parallel_calls=self._num_parallel_calls)
    if self._preprocess_fn is not None:
      dataset = dataset.map(
          self._preprocess_fn, num_parallel_calls=self._parallel_shards)
    if self._prefetch_buffer_size is not None:
      dataset = dataset.prefetch(self._prefetch_buffer_size)
    return dataset
<|endoftext|>


<commit_before>'''The vtk_anim module define a set of classes to generate 3d
   animations with vtk in the form of a series of png images.
'''
import vtk
import os, numpy as np

class vtkAnimationScene:
    def __init__(self, ren, ren_size=(600, 600)):
        self.timer_count = 0
        self.timer_incr = 1
        self.timer_end = -1  # run until 'q' is pressed
        self.save_image = False
        self.prefix = 'prefix'
        self.verbose = False
        self.anims = []
        # Create a window for the renderer
        self.renWin = vtk.vtkRenderWindow()
        self.renWin.AddRenderer(ren)
        self.renWin.SetSize(ren_size)
        # Start the initialization and rendering
        self.iren = vtk.vtkRenderWindowInteractor()
        self.iren.SetRenderWindow(self.renWin)
        self.renWin.Render()
        self.iren.Initialize()

    def add_animation(self, anim):
        anim.scene = self
        anim.save_image = self.save_image
        anim.prefix = self.prefix
        anim.verbose = self.verbose
        self.anims.append(anim)
        self.iren.AddObserver('TimerEvent', anim.execute)

    def write_image(self):
        # capture the display and write a png image
        w2i = vtk.vtkWindowToImageFilter()
        w2i.SetInput(self.iren.GetRenderWindow())
        # the next two liines fix some opacity problems but slow things down...
        # self.renWin.Render()
        # self.iren.Render()
        writer = vtk.vtkPNGWriter()
        writer.SetInputConnection(w2i.GetOutputPort())
        file_name = os.path.join(self.prefix, '%s_%03d.png' % (self.prefix, self.timer_count))
        writer.SetFileName(file_name)
        writer.Write()

    def execute(self, iren, event):
        self.timer_count += self.timer_incr
        if self.verbose:
            print 'animation scene timer_count=', self.timer_count
        if self.timer_end > 0 and self.timer_count > self.timer_end:
            print 'end of animation loop, exiting...'
            self.iren.ExitCallback()
        else:
            self.iren.Render()
        if self.save_image:
            self.write_image()

    def render(self):
        if self.save_image and not os.path.exists(self.prefix):
            os.mkdir(self.prefix)  # create a folder to store the images
        timerId = self.iren.CreateRepeatingTimer(100);  # time in ms
        self.iren.AddObserver('TimerEvent', self.execute)
        self.iren.Start()

'''
Abstract class for all vtk animation stuff.
'''

class vtkAnimation:
    def __init__(self, t):
        self.scene = None
        self.time_anim_starts = t
        self.time_anim_ends = t + 10
        self.verbose = False

    def pre_execute(self):
        if self.verbose:
            print self.__repr__()
        if self.scene.timer_count < self.time_anim_starts or self.scene.timer_count > self.time_anim_ends:
            return 0
        else:
            return 1

    def post_execute(self, iren, event):
        pass

    def __repr__(self):
        out = [self.__class__.__name__,
               ' timer: ' + str(self.scene.timer_count),
               ' anim starts at: ' + str(self.time_anim_starts),
               ' anim ends at: ' + str(self.time_anim_ends)]
        return '\n'.join(out)

class vtkAnimCameraAroundZ(vtkAnimation):
    '''
    Animate the camera around the vertical axis.

    This class can be used to generate a series of images (default 36)
    while the camera rotate around the vertical axis (defined by the
    camera SetViewUp method).
    '''

    def __init__(self, t, cam, turn=360):
        '''Initialize the animation.

        The animation perform a full turn in 36 frames by default.
        '''
        print 'init vtkAnimCameraAroundZ'
        vtkAnimation.__init__(self, t)
        self.turn = turn
        self.time_anim_ends = t + abs(self.turn) / 10
        print 'time_anim_starts', self.time_anim_starts
        print 'time_anim_ends', self.time_anim_ends
        print 'turn', self.turn
        self.camera = cam

    def execute(self, iren, event):
        '''Execute method called to rotate the camera.'''
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        r = self.turn / (t2 - t1)
        if self.scene.verbose: print 'moving azimuth by', r
        self.camera.Azimuth(r)
        vtkAnimation.post_execute(self, iren, event)

class vtkRotateActorAroundZAxis(vtkAnimation):
    def __init__(self, t=0):
        vtkAnimation.__init__(self, t)
        self.time_anim_ends = t + 360
        self.actor = None
        self.actor_position = (0., 0., 0.)

    def execute(self, iren, event):
        '''instruction block exectued when a TimerEvent is captured by the vtkRotateActorAroundZAxis.

        If the time is not in [start, end] it just returns. Otherwise the
        transform matrix corresponding to the rotation around the Z-axis is
        computed and applied to the actor.
        '''
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        # t = (self.scene.timer_count - self.time_anim_starts) * np.pi / 180.
        angle = (self.scene.timer_count - t1) / float(t2 - t1) * 2 * np.pi
        X = ([np.cos(angle), np.sin(angle), 0])
        print self.scene.timer_count, X, angle * 180 / np.pi
        Z = np.array([0, 0, 1])
        Y = np.cross(Z, X)
        m = vtk.vtkMatrix4x4()
        m.Identity()
        # Create the direction cosine matrix
        for i in range(3):
            m.SetElement(i, 0, X[i]);
            m.SetElement(i, 1, Y[i]);
            m.SetElement(i, 2, Z[i]);
        t = vtk.vtkTransform()
        t.Identity()
        t.Concatenate(m)
        t.PostMultiply()
        t.Translate(self.actor_position)
        self.actor.SetUserTransform(t)
        vtkAnimation.post_execute(self, iren, event)

class vtkAnimCameraToZ(vtkAnimation):
    def __init__(self, t, cam):
        vtkAnimation.__init__(self, t)
        self.camera = cam

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        angle = 90 - (t2 - self.scene.timer_count) / float(t2 - t1) * (90 - 15)
        if self.verbose: print self.scene.timer_count, self.camera.GetPosition(), angle
        self.camera.SetPosition(0, -2 * np.cos(angle * np.pi / 180.), 2 * np.sin(angle * np.pi / 180.))
        vtkAnimation.post_execute(self, iren, event)

class vtkZoom(vtkAnimation):
    def __init__(self, t, cam, zoom):
        vtkAnimation.__init__(self, t)
        self.camera = cam
        self.zoom = zoom
        self.timer_end = t + 10

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        z = 1 + (self.zoom - 1) * (self.scene.timer_count - t1) / float(t2 - t1)
        if self.verbose: print 'zooming to', z
        self.camera.Zoom(z)
        vtkAnimation.post_execute(self, iren, event)

class vtkSetVisibility(vtkAnimation):
    def __init__(self, t, actor, visible=1, max_opacity=1, gradually=False):
        vtkAnimation.__init__(self, t)
        self.actor = actor
        self.visible = visible
        self.gradually = gradually
        self.max_opacity = max_opacity

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        if not self.gradually:
            self.actor.SetVisibility(self.visible)
        else:
            t1 = self.time_anim_starts
            t2 = self.time_anim_ends
            if self.scene.timer_count >= t1 and self.scene.timer_count <= t2:  # useless to test this (do == 1 here)
                if self.actor.GetVisibility() == 0:
                    self.actor.SetVisibility(1)  # make the actor visible
                if self.visible:
                    opacity = self.max_opacity * (1 - (t2 - self.scene.timer_count) / float(t2 - t1))
                else:
                    opacity = self.max_opacity * (t2 - self.scene.timer_count) / float(t2 - t1)
                if self.verbose: print 'opacity=', opacity
                # change the opacity for each actor in the assemby
                collection = vtk.vtkPropCollection()
                self.actor.GetActors(collection);
                for i in range(collection.GetNumberOfItems()):
                    collection.GetItemAsObject(i).GetProperty().SetOpacity(opacity)
        vtkAnimation.post_execute(self, iren, event)

class vtkMoveActor(vtkAnimation):
    def __init__(self, t, actor, motion):
        vtkAnimation.__init__(self, t)
        self.actor = actor
        if self.actor.GetUserTransform() == None:
            if self.verbose: print 'setting initial 4x4 matrix'
            t = vtk.vtkTransform()
            t.Identity()
            self.actor.SetUserTransform(t)
        self.motion = np.array(motion)

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        d = self.motion / (t2 - t1)
        if self.verbose: print 'will move actor by', d
        self.actor.GetUserTransform().Translate(d)
        vtkAnimation.post_execute(self, iren, event)

class vtkAnimLine(vtkAnimation):
    def __init__(self, points, t1, t2):
        vtkAnimation.__init__(self, t1)
        self.time_anim_line_end = t2
        self.line_points = points
        self.p0 = np.array(self.line_points.GetPoint(0))
        self.p1 = np.array(self.line_points.GetPoint(1))
        self.grid = None
        self.actor = None
        self.pole = None

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        # if self.scene.timer_count >= t1 and self.scene.timer_count <= t2:
        self.actor.SetVisibility(1)
        point = self.p1 + (t2 - self.scene.timer_count) / float(t2 - t1) * (self.p0 - self.p1)
        self.line_points.SetPoint(1, point)
        if point[2] <= 0 and self.pole != None:
            self.pole.SetVisibility(1)
        self.grid.Modified()
        vtkAnimation.post_execute(self, iren, event)

if __name__ == '__main__':
    cam = vtk.vtkCamera()
    anim = vtkAnimCameraAroundZ(cam, 10)
    anim.verbose = True
<commit_msg>added a new method to render animation at a given time.<commit_after>'''The vtk_anim module define a set of classes to generate 3d
   animations with vtk in the form of a series of png images.
'''
import vtk
import os, numpy as np

class vtkAnimationScene:
    def __init__(self, ren, ren_size=(600, 600)):
        self.timer_count = 0
        self.timer_incr = 1
        self.timer_end = -1  # run until 'q' is pressed
        self.save_image = False
        self.prefix = 'prefix'
        self.verbose = False
        self.anims = []
        # Create a window for the renderer
        self.renWin = vtk.vtkRenderWindow()
        self.renWin.AddRenderer(ren)
        self.renWin.SetSize(ren_size)
        # Start the initialization and rendering
        self.iren = vtk.vtkRenderWindowInteractor()
        self.iren.SetRenderWindow(self.renWin)
        self.renWin.Render()
        self.iren.Initialize()

    def add_animation(self, anim):
        anim.scene = self
        anim.save_image = self.save_image
        anim.prefix = self.prefix
        anim.verbose = self.verbose
        self.anims.append(anim)
        self.iren.AddObserver('TimerEvent', anim.execute)

    def write_image(self):
        # capture the display and write a png image
        w2i = vtk.vtkWindowToImageFilter()
        w2i.SetInput(self.iren.GetRenderWindow())
        # the next two liines fix some opacity problems but slow things down...
        # self.renWin.Render()
        # self.iren.Render()
        writer = vtk.vtkPNGWriter()
        writer.SetInputConnection(w2i.GetOutputPort())
        file_name = os.path.join(self.prefix, '%s_%03d.png' % (self.prefix, self.timer_count))
        writer.SetFileName(file_name)
        writer.Write()

    def execute(self, iren, event):
        self.timer_count += self.timer_incr
        if self.verbose:
            print 'animation scene timer_count=', self.timer_count
        if self.timer_end > 0 and self.timer_count > self.timer_end:
            print 'end of animation loop, exiting...'
            self.iren.ExitCallback()
        else:
            self.iren.Render()
        if self.save_image:
            self.write_image()

    def render(self):
        if self.save_image and not os.path.exists(self.prefix):
            os.mkdir(self.prefix)  # create a folder to store the images
        timerId = self.iren.CreateRepeatingTimer(100)  # time in ms
        self.iren.AddObserver('TimerEvent', self.execute)
        self.iren.Start()

    def render_at(self, time=0.):
        if self.save_image and not os.path.exists(self.prefix):
            os.mkdir(self.prefix)  # create a folder to store the images
        self.timer_count = time
        self.iren.CreateOneShotTimer(100)  # time in ms
        self.iren.AddObserver('TimerEvent', self.execute)
        self.iren.Start()

'''
Abstract class for all vtk animation stuff.
'''

class vtkAnimation:
    def __init__(self, t):
        self.scene = None
        self.time_anim_starts = t
        self.time_anim_ends = t + 10
        self.verbose = False

    def pre_execute(self):
        if self.verbose:
            print self.__repr__()
        if self.scene.timer_count < self.time_anim_starts or self.scene.timer_count > self.time_anim_ends:
            return 0
        else:
            return 1

    def post_execute(self, iren, event):
        pass

    def __repr__(self):
        out = [self.__class__.__name__,
               ' timer: ' + str(self.scene.timer_count),
               ' anim starts at: ' + str(self.time_anim_starts),
               ' anim ends at: ' + str(self.time_anim_ends)]
        return '\n'.join(out)

class vtkAnimCameraAroundZ(vtkAnimation):
    '''
    Animate the camera around the vertical axis.

    This class can be used to generate a series of images (default 36)
    while the camera rotate around the vertical axis (defined by the
    camera SetViewUp method).
    '''

    def __init__(self, t, cam, turn=360):
        '''Initialize the animation.

        The animation perform a full turn in 36 frames by default.
        '''
        print 'init vtkAnimCameraAroundZ'
        vtkAnimation.__init__(self, t)
        self.turn = turn
        self.time_anim_ends = t + abs(self.turn) / 10
        print 'time_anim_starts', self.time_anim_starts
        print 'time_anim_ends', self.time_anim_ends
        print 'turn', self.turn
        self.camera = cam

    def execute(self, iren, event):
        '''Execute method called to rotate the camera.'''
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        r = self.turn / (t2 - t1)
        if self.scene.verbose: print 'moving azimuth by', r
        self.camera.Azimuth(r)
        vtkAnimation.post_execute(self, iren, event)

class vtkRotateActorAroundZAxis(vtkAnimation):
    def __init__(self, t=0):
        vtkAnimation.__init__(self, t)
        self.time_anim_ends = t + 360
        self.actor = None
        self.actor_position = (0., 0., 0.)

    def execute(self, iren, event):
        '''instruction block exectued when a TimerEvent is captured by the vtkRotateActorAroundZAxis.

        If the time is not in [start, end] it just returns. Otherwise the
        transform matrix corresponding to the rotation around the Z-axis is
        computed and applied to the actor.
        '''
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        # t = (self.scene.timer_count - self.time_anim_starts) * np.pi / 180.
        angle = (self.scene.timer_count - t1) / float(t2 - t1) * 2 * np.pi
        X = ([np.cos(angle), np.sin(angle), 0])
        print self.scene.timer_count, X, angle * 180 / np.pi
        Z = np.array([0, 0, 1])
        Y = np.cross(Z, X)
        m = vtk.vtkMatrix4x4()
        m.Identity()
        # Create the direction cosine matrix
        for i in range(3):
            m.SetElement(i, 0, X[i]);
            m.SetElement(i, 1, Y[i]);
            m.SetElement(i, 2, Z[i]);
        t = vtk.vtkTransform()
        t.Identity()
        t.Concatenate(m)
        t.PostMultiply()
        t.Translate(self.actor_position)
        self.actor.SetUserTransform(t)
        vtkAnimation.post_execute(self, iren, event)

class vtkAnimCameraToZ(vtkAnimation):
    def __init__(self, t, cam):
        vtkAnimation.__init__(self, t)
        self.camera = cam

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        angle = 90 - (t2 - self.scene.timer_count) / float(t2 - t1) * (90 - 15)
        if self.verbose: print self.scene.timer_count, self.camera.GetPosition(), angle
        self.camera.SetPosition(0, -2 * np.cos(angle * np.pi / 180.), 2 * np.sin(angle * np.pi / 180.))
        vtkAnimation.post_execute(self, iren, event)

class vtkZoom(vtkAnimation):
    def __init__(self, t, cam, zoom):
        vtkAnimation.__init__(self, t)
        self.camera = cam
        self.zoom = zoom
        self.timer_end = t + 10

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        z = 1 + (self.zoom - 1) * (self.scene.timer_count - t1) / float(t2 - t1)
        if self.verbose: print 'zooming to', z
        self.camera.Zoom(z)
        vtkAnimation.post_execute(self, iren, event)

class vtkSetVisibility(vtkAnimation):
    def __init__(self, t, actor, visible=1, max_opacity=1, gradually=False):
        vtkAnimation.__init__(self, t)
        self.actor = actor
        self.visible = visible
        self.gradually = gradually
        self.max_opacity = max_opacity

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        if not self.gradually:
            self.actor.SetVisibility(self.visible)
        else:
            t1 = self.time_anim_starts
            t2 = self.time_anim_ends
            if self.scene.timer_count >= t1 and self.scene.timer_count <= t2:  # useless to test this (do == 1 here)
                if self.actor.GetVisibility() == 0:
                    self.actor.SetVisibility(1)  # make the actor visible
                if self.visible:
                    opacity = self.max_opacity * (1 - (t2 - self.scene.timer_count) / float(t2 - t1))
                else:
                    opacity = self.max_opacity * (t2 - self.scene.timer_count) / float(t2 - t1)
                if self.verbose: print 'opacity=', opacity
                # change the opacity for each actor in the assemby
                collection = vtk.vtkPropCollection()
                self.actor.GetActors(collection);
                for i in range(collection.GetNumberOfItems()):
                    collection.GetItemAsObject(i).GetProperty().SetOpacity(opacity)
        vtkAnimation.post_execute(self, iren, event)

class vtkMoveActor(vtkAnimation):
    def __init__(self, t, actor, motion):
        vtkAnimation.__init__(self, t)
        self.actor = actor
        if self.actor.GetUserTransform() == None:
            if self.verbose: print 'setting initial 4x4 matrix'
            t = vtk.vtkTransform()
            t.Identity()
            self.actor.SetUserTransform(t)
        self.motion = np.array(motion)

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        d = self.motion / (t2 - t1)
        if self.verbose: print 'will move actor by', d
        self.actor.GetUserTransform().Translate(d)
        vtkAnimation.post_execute(self, iren, event)

class vtkAnimLine(vtkAnimation):
    def __init__(self, points, t1, t2):
        vtkAnimation.__init__(self, t1)
        self.time_anim_line_end = t2
        self.line_points = points
        self.p0 = np.array(self.line_points.GetPoint(0))
        self.p1 = np.array(self.line_points.GetPoint(1))
        self.grid = None
        self.actor = None
        self.pole = None

    def execute(self, iren, event):
        do = vtkAnimation.pre_execute(self)
        if not do: return
        t1 = self.time_anim_starts
        t2 = self.time_anim_ends
        # if self.scene.timer_count >= t1 and self.scene.timer_count <= t2:
        self.actor.SetVisibility(1)
        point = self.p1 + (t2 - self.scene.timer_count) / float(t2 - t1) * (self.p0 - self.p1)
        self.line_points.SetPoint(1, point)
        if point[2] <= 0 and self.pole != None:
            self.pole.SetVisibility(1)
        self.grid.Modified()
        vtkAnimation.post_execute(self, iren, event)

if __name__ == '__main__':
    cam = vtk.vtkCamera()
    anim = vtkAnimCameraAroundZ(cam, 10)
    anim.verbose = True
<|endoftext|>


<commit_before>from math import log10, floor

import datetime
import ago

def fader(val, perc1, perc2, perc3, color1, color2, color3):
    """
    Accepts a decimal (0.1, 0.5, etc) and slots it into one of three categories based
    on the percentage.
    """
    if val > perc3:
        return color3
    if val > perc2:
        return color2
    return color1

def sig_round(x, sig=2):
    if x == 0:
        return "0"
    whole = int(str(x).split(".")[0])
    if whole != x:
        x -= whole
    else:
        whole = 0
    return "{:,f}".format(whole + round(x, sig - int(floor(log10(abs(x)))) - 1)).rstrip('0').rstrip('.')

def duration(seconds):
    # microseconds
    if seconds > 3600:
        return "{}".format(datetime.timedelta(seconds=seconds))
    if seconds > 60:
        return "{:,.2f} mins".format(seconds / 60.0)
    if seconds <= 1.0e-3:
        return "{:,.4f} us".format(seconds * 1000000.0)
    if seconds <= 1.0:
        return "{:,.4f} ms".format(seconds * 1000.0)
    return "{:,.4f} sec".format(seconds)

def humana_date(*args, **kwargs):
    return ago.human(*args, **kwargs)

def hashrate(hashrate, num_fmt="{:,.2f}"):
    if hashrate > 1000000000:
        return "{} GH/s".format(num_fmt.format(hashrate / 1000000000))
    if hashrate > 1000000:
        return "{} MH/s".format(num_fmt.format(hashrate / 1000000))
    if hashrate > 1000:
        return "{} KH/s".format(num_fmt.format(hashrate / 1000))
    return "{} H/s".format(num_fmt.format(hashrate))

def human_date_utc(*args, **kwargs):
    delta = (datetime.datetime.utcnow() - args[0])
    delta = delta - datetime.timedelta(microseconds=delta.microseconds)
    return ago.human(delta, *args[1:], **kwargs)

def currency(value):
    return "{:,.7f}".format(float(value))
<commit_msg>Hack for now to stop 500<commit_after>from math import log10, floor

import datetime
import ago

def fader(val, perc1, perc2, perc3, color1, color2, color3):
    """
    Accepts a decimal (0.1, 0.5, etc) and slots it into one of three categories based
    on the percentage.
    """
    if val > perc3:
        return color3
    if val > perc2:
        return color2
    return color1

def sig_round(x, sig=2):
    try:
        if x == 0:
            return "0"
        whole = int(str(x).split(".")[0])
        if whole != x:
            x -= whole
        else:
            whole = 0
        return "{:,f}".format(whole + round(x, sig - int(floor(log10(abs(x)))) - 1)).rstrip('0').rstrip('.')
    except Exception:
        return "Err. Will fix soon"

def duration(seconds):
    # microseconds
    if seconds > 3600:
        return "{}".format(datetime.timedelta(seconds=seconds))
    if seconds > 60:
        return "{:,.2f} mins".format(seconds / 60.0)
    if seconds <= 1.0e-3:
        return "{:,.4f} us".format(seconds * 1000000.0)
    if seconds <= 1.0:
        return "{:,.4f} ms".format(seconds * 1000.0)
    return "{:,.4f} sec".format(seconds)

def humana_date(*args, **kwargs):
    return ago.human(*args, **kwargs)

def hashrate(hashrate, num_fmt="{:,.2f}"):
    if hashrate > 1000000000:
        return "{} GH/s".format(num_fmt.format(hashrate / 1000000000))
    if hashrate > 1000000:
        return "{} MH/s".format(num_fmt.format(hashrate / 1000000))
    if hashrate > 1000:
        return "{} KH/s".format(num_fmt.format(hashrate / 1000))
    return "{} H/s".format(num_fmt.format(hashrate))

def human_date_utc(*args, **kwargs):
    delta = (datetime.datetime.utcnow() - args[0])
    delta = delta - datetime.timedelta(microseconds=delta.microseconds)
    return ago.human(delta, *args[1:], **kwargs)

def currency(value):
    return "{:,.7f}".format(float(value))
<|endoftext|>


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

from __future__ import division, print_function

import os
import sys
import tempfile

import setuptools
from setuptools.command.build_ext import build_ext as _build_ext

__all__ = ["build_ext"]

def has_flag(compiler, flagname):
    """Return a boolean indicating whether a flag name is supported on
    the specified compiler.
    """
    with tempfile.NamedTemporaryFile('w', suffix='.cpp') as f:
        f.write('int main (int argc, char **argv) { return 0; }')
        try:
            compiler.compile([f.name], extra_postargs=[flagname])
        except setuptools.distutils.errors.CompileError:
            return False
    return True

def cpp_flag(compiler):
    """Return the -std=c++[11/14] compiler flag.

    The c++14 is prefered over c++11 (when it is available).
    """
    if has_flag(compiler, '-std=c++14'):
        return '-std=c++14'
    elif has_flag(compiler, '-std=c++11'):
        return '-std=c++11'
    else:
        raise RuntimeError('Unsupported compiler -- at least C++11 support '
                           'is needed!')

class build_ext(_build_ext):
    """
    A custom extension builder that finds the include directories for Eigen
    before compiling.

    """

    c_opts = {
        'msvc': ['/EHsc'],
        'unix': [],
    }

    if sys.platform == 'darwin':
        c_opts['unix'] += ["-mmacosx-version-min=10.7"]

    def build_extensions(self):
        # Add the pybind11 include directory
        import pybind11
        include_dirs = [
            pybind11.get_include(False),
            pybind11.get_include(True),
        ]
        for ext in self.extensions:
            ext.include_dirs += include_dirs

        # Building on RTDs takes a bit of special care
        if os.environ.get("READTHEDOCS", None) == "True":
            for ext in self.extensions:
                ext.extra_compile_args = ["-std=c++14"]
            _build_ext.build_extensions(self)
            return

        # Set up pybind11
        ct = self.compiler.compiler_type
        opts = self.c_opts.get(ct, [])
        if ct == 'unix':
            opts.append('-DVERSION_INFO="{0:s}"'
                        .format(self.distribution.get_version()))
            opts.append(cpp_flag(self.compiler))
            for flag in ["-stdlib=libc++", "-fvisibility=hidden",
                         "-Wno-unused-function", "-Wno-uninitialized",
                         "-Wno-unused-local-typedefs"]:
                if has_flag(self.compiler, flag):
                    opts.append(flag)
        elif ct == 'msvc':
            opts.append('/DVERSION_INFO=\\"{0:s}\\"'
                        .format(self.distribution.get_version()))

        for ext in self.extensions:
            ext.extra_compile_args = opts

        # Run the standard build procedure.
        _build_ext.build_extensions(self)
<commit_msg>RTDs O0 [ci skip]<commit_after># -*- coding: utf-8 -*-

from __future__ import division, print_function

import os
import sys
import tempfile

import setuptools
from setuptools.command.build_ext import build_ext as _build_ext

__all__ = ["build_ext"]

def has_flag(compiler, flagname):
    """Return a boolean indicating whether a flag name is supported on
    the specified compiler.
    """
    with tempfile.NamedTemporaryFile('w', suffix='.cpp') as f:
        f.write('int main (int argc, char **argv) { return 0; }')
        try:
            compiler.compile([f.name], extra_postargs=[flagname])
        except setuptools.distutils.errors.CompileError:
            return False
    return True

def cpp_flag(compiler):
    """Return the -std=c++[11/14] compiler flag.

    The c++14 is prefered over c++11 (when it is available).
    """
    if has_flag(compiler, '-std=c++14'):
        return '-std=c++14'
    elif has_flag(compiler, '-std=c++11'):
        return '-std=c++11'
    else:
        raise RuntimeError('Unsupported compiler -- at least C++11 support '
                           'is needed!')

class build_ext(_build_ext):
    """
    A custom extension builder that finds the include directories for Eigen
    before compiling.

    """

    c_opts = {
        'msvc': ['/EHsc'],
        'unix': [],
    }

    if sys.platform == 'darwin':
        c_opts['unix'] += ["-mmacosx-version-min=10.7"]

    def build_extensions(self):
        # Add the pybind11 include directory
        import pybind11
        include_dirs = [
            pybind11.get_include(False),
            pybind11.get_include(True),
        ]
        for ext in self.extensions:
            ext.include_dirs += include_dirs

        # Building on RTDs takes a bit of special care
        if os.environ.get("READTHEDOCS", None) == "True":
            for ext in self.extensions:
                ext.extra_compile_args = ["-std=c++14", "-O0"]
            _build_ext.build_extensions(self)
            return

        # Set up pybind11
        ct = self.compiler.compiler_type
        opts = self.c_opts.get(ct, [])
        if ct == 'unix':
            opts.append('-DVERSION_INFO="{0:s}"'
                        .format(self.distribution.get_version()))
            opts.append(cpp_flag(self.compiler))
            for flag in ["-stdlib=libc++", "-fvisibility=hidden",
                         "-Wno-unused-function", "-Wno-uninitialized",
                         "-Wno-unused-local-typedefs"]:
                if has_flag(self.compiler, flag):
                    opts.append(flag)
        elif ct == 'msvc':
            opts.append('/DVERSION_INFO=\\"{0:s}\\"'
                        .format(self.distribution.get_version()))

        for ext in self.extensions:
            ext.extra_compile_args = opts

        # Run the standard build procedure.
        _build_ext.build_extensions(self)
<|endoftext|>


<commit_before>#!/usr/bin/python
# Copyright (c) 2012 The Native Client Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#

"""Download all Native Client toolchains for this platform.

This module downloads toolchain bz2's and expands them. It requires
gsutil to be in the bin PATH and assumes if building on windows that
cygwin is installed to C:\cygwin
"""

import glob
import optparse
import os
import re
import shutil
import stat
import subprocess
import sys
import tarfile
import tempfile
import time
import urllib

SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))
SRC_DIR = os.path.dirname(SCRIPT_DIR)
NACL_BUILD_DIR = os.path.join(SRC_DIR, 'native_client', 'build')

sys.path.append(NACL_BUILD_DIR)

import cygtar

BOT_GSUTIL = '/b/build/scripts/slave/gsutil'
LOCAL_GSUTIL = 'gsutil'
# For local testing on Windows
#LOCAL_GSUTIL = 'python.exe C:\\bin\\gsutil\\gsutil'

GSTORE = 'http://commondatastorage.googleapis.com/'\
         'nativeclient-mirror/nacl/nacl_sdk/'

def DetermineSdkURL(flavor, base_url, version):
  """Download one Native Client toolchain and extract it.

  Arguments:
    flavor: flavor of the sdk to download
    base_url: base url to download toolchain tarballs from
    version: version directory to select tarballs from
  """
  if os.environ.get('BUILDBOT_BUILDERNAME', ''):
    gsutil = BOT_GSUTIL
  else:
    gsutil = LOCAL_GSUTIL

  path = flavor + '.bz2'

  if version == 'latest':
    print 'Looking for latest SDK upload...'
    # Resolve wildcards and pick the highest version
    p = subprocess.Popen(
        ' '.join(gsutil.split() + ['ls',
        base_url + 'trunk.*/' + path]),
        stdout=subprocess.PIPE,
        shell=True)
    (p_stdout, _) = p.communicate()
    assert p.returncode == 0
    versions = p_stdout.splitlines()
    newest = None
    for version in versions:
      m = re.match(base_url.replace(':', '\:').replace('/', '\/') +
                   'trunk\.([01-9]+)/' + path, version)
      if m:
        rev = int(m.group(1))
        if not newest or rev > newest:
          newest = rev
    assert newest
    version = newest

  return GSTORE + 'trunk.' + str(version) + '/' + path

def DownloadAndInstallSDK(url):
  # Pick target directory.
  target_dir = os.path.join(SRC_DIR, 'pepper_XX')

  bz2_dir = SCRIPT_DIR
  bz2_filename = os.path.join(bz2_dir, url.split('/')[-1])

  # Drop old versions.
  old_sdks = glob.glob(os.path.join(bz2_dir, 'pepper_*'))
  if len(old_sdks) > 0:
    print 'Cleaning up old SDKs...'
    if sys.platform in ['win32', 'cygwin']:
      cmd = r'C:\cygwin\bin\rm.exe -rf "%s"' % '" "'.join(old_sdks)
    else:
      cmd = 'rm -rf "%s"' % '" "'.join(old_sdks)
    p = subprocess.Popen(cmd, shell=True)
    p.communicate()
    assert p.returncode == 0

  print 'Downloading "%s" to "%s"...' % (url, bz2_filename)
  sys.stdout.flush()

  # Download it.
  urlret = urllib.urlretrieve(url, bz2_filename)
  assert urlret[0] == bz2_filename

  # Extract toolchain.
  old_cwd = os.getcwd()
  os.chdir(bz2_dir)
  tar_file = None
  try:
    print 'Unpacking tarball...'
    tar_file = cygtar.CygTar(bz2_filename, 'r:bz2')
    names = tar_file.tar.getnames()
    pepper_dir = os.path.commonprefix(names)
    tar_file.Extract()
  except Exception, err:
    print 'Error unpacking %s' % str(err)
    sys.exit(1)
  finally:
    if tar_file:
      tar_file.Close()

  os.chdir(old_cwd)

  actual_dir = os.path.join(bz2_dir, pepper_dir)
  print 'Create toolchain symlink "%s" -> "%s"' % (actual_dir, target_dir)
  if sys.platform in ['win32', 'cygwin']:
    cmd = (r'C:\cygwin\bin\rm.exe -rf ' + target_dir + ' && ' +
           r'C:\cygwin\bin\ln.exe -fsn ' + actual_dir + ' ' +  target_dir)
  else:
    cmd = ('rm -rf ' + target_dir + ' && ' +
           'ln -fsn ' + actual_dir + ' ' + target_dir)

  p = subprocess.Popen(cmd, shell=True)
  p.communicate()
  assert p.returncode == 0

  # Clean up: remove the sdk bz2.
  time.sleep(2)  # Wait for windows.
  os.remove(bz2_filename)

  print 'Install complete.'

PLATFORM_COLLAPSE = {
    'win32': 'win',
    'cygwin': 'win',
    'linux': 'linux',
    'linux2': 'linux',
    'darwin': 'mac',
}

def main(argv):
  parser = optparse.OptionParser()
  parser.add_option(
      '-v', '--version', dest='version',
      default='latest',
      help='which version of the toolchain to download')
  parser.add_option(
      '-f', '--flavor', dest='flavor',
      default='auto',
      help='which flavor of the toolchain to download, e.g. naclsdk_linux or ' +
           'pnaclsdk_linux')
  options, args = parser.parse_args(argv)
  if args:
    parser.print_help()
    print 'ERROR: invalid argument'
    sys.exit(1)

  flavor = options.flavor
  if flavor == 'auto':
    flavor = 'naclsdk_' + PLATFORM_COLLAPSE[sys.platform]

  url = DetermineSdkURL(flavor,
                        base_url='gs://nativeclient-mirror/nacl/nacl_sdk/',
                        version=options.version)
  print 'SDK URL is "%s"' % url

  DownloadAndInstallSDK(url)

if __name__ == '__main__':
  main(sys.argv[1:])
<commit_msg>Update naclports's latest sdk finder to look for .tar.bz2 instead of .bz2, since the newer sdks use extension .tar.bz2.<commit_after>#!/usr/bin/python
# Copyright (c) 2012 The Native Client Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#

"""Download all Native Client toolchains for this platform.

This module downloads toolchain bz2's and expands them. It requires
gsutil to be in the bin PATH and assumes if building on windows that
cygwin is installed to C:\cygwin
"""

import glob
import optparse
import os
import re
import shutil
import stat
import subprocess
import sys
import tarfile
import tempfile
import time
import urllib

SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))
SRC_DIR = os.path.dirname(SCRIPT_DIR)
NACL_BUILD_DIR = os.path.join(SRC_DIR, 'native_client', 'build')

sys.path.append(NACL_BUILD_DIR)

import cygtar

BOT_GSUTIL = '/b/build/scripts/slave/gsutil'
LOCAL_GSUTIL = 'gsutil'
# For local testing on Windows
#LOCAL_GSUTIL = 'python.exe C:\\bin\\gsutil\\gsutil'

GSTORE = 'http://commondatastorage.googleapis.com/'\
         'nativeclient-mirror/nacl/nacl_sdk/'

def DetermineSdkURL(flavor, base_url, version):
  """Download one Native Client toolchain and extract it.

  Arguments:
    flavor: flavor of the sdk to download
    base_url: base url to download toolchain tarballs from
    version: version directory to select tarballs from
  """
  if os.environ.get('BUILDBOT_BUILDERNAME', ''):
    gsutil = BOT_GSUTIL
  else:
    gsutil = LOCAL_GSUTIL

  path = flavor + '.tar.bz2'

  if version == 'latest':
    print 'Looking for latest SDK upload...'
    # Resolve wildcards and pick the highest version
    p = subprocess.Popen(
        ' '.join(gsutil.split() + ['ls',
        base_url + 'trunk.*/' + path]),
        stdout=subprocess.PIPE,
        shell=True)
    (p_stdout, _) = p.communicate()
    assert p.returncode == 0
    versions = p_stdout.splitlines()
    newest = None
    for version in versions:
      m = re.match(base_url.replace(':', '\:').replace('/', '\/') +
                   'trunk\.([01-9]+)/' + path, version)
      if m:
        rev = int(m.group(1))
        if not newest or rev > newest:
          newest = rev
    assert newest
    version = newest

  return GSTORE + 'trunk.' + str(version) + '/' + path

def DownloadAndInstallSDK(url):
  # Pick target directory.
  target_dir = os.path.join(SRC_DIR, 'pepper_XX')

  bz2_dir = SCRIPT_DIR
  bz2_filename = os.path.join(bz2_dir, url.split('/')[-1])

  # Drop old versions.
  old_sdks = glob.glob(os.path.join(bz2_dir, 'pepper_*'))
  if len(old_sdks) > 0:
    print 'Cleaning up old SDKs...'
    if sys.platform in ['win32', 'cygwin']:
      cmd = r'C:\cygwin\bin\rm.exe -rf "%s"' % '" "'.join(old_sdks)
    else:
      cmd = 'rm -rf "%s"' % '" "'.join(old_sdks)
    p = subprocess.Popen(cmd, shell=True)
    p.communicate()
    assert p.returncode == 0

  print 'Downloading "%s" to "%s"...' % (url, bz2_filename)
  sys.stdout.flush()

  # Download it.
  urlret = urllib.urlretrieve(url, bz2_filename)
  assert urlret[0] == bz2_filename

  # Extract toolchain.
  old_cwd = os.getcwd()
  os.chdir(bz2_dir)
  tar_file = None
  try:
    print 'Unpacking tarball...'
    tar_file = cygtar.CygTar(bz2_filename, 'r:bz2')
    names = tar_file.tar.getnames()
    pepper_dir = os.path.commonprefix(names)
    tar_file.Extract()
  except Exception, err:
    print 'Error unpacking %s' % str(err)
    sys.exit(1)
  finally:
    if tar_file:
      tar_file.Close()

  os.chdir(old_cwd)

  actual_dir = os.path.join(bz2_dir, pepper_dir)
  print 'Create toolchain symlink "%s" -> "%s"' % (actual_dir, target_dir)
  if sys.platform in ['win32', 'cygwin']:
    cmd = (r'C:\cygwin\bin\rm.exe -rf ' + target_dir + ' && ' +
           r'C:\cygwin\bin\ln.exe -fsn ' + actual_dir + ' ' +  target_dir)
  else:
    cmd = ('rm -rf ' + target_dir + ' && ' +
           'ln -fsn ' + actual_dir + ' ' + target_dir)

  p = subprocess.Popen(cmd, shell=True)
  p.communicate()
  assert p.returncode == 0

  # Clean up: remove the sdk bz2.
  time.sleep(2)  # Wait for windows.
  os.remove(bz2_filename)

  print 'Install complete.'

PLATFORM_COLLAPSE = {
    'win32': 'win',
    'cygwin': 'win',
    'linux': 'linux',
    'linux2': 'linux',
    'darwin': 'mac',
}

def main(argv):
  parser = optparse.OptionParser()
  parser.add_option(
      '-v', '--version', dest='version',
      default='latest',
      help='which version of the toolchain to download')
  parser.add_option(
      '-f', '--flavor', dest='flavor',
      default='auto',
      help='which flavor of the toolchain to download, e.g. naclsdk_linux or ' +
           'pnaclsdk_linux')
  options, args = parser.parse_args(argv)
  if args:
    parser.print_help()
    print 'ERROR: invalid argument'
    sys.exit(1)

  flavor = options.flavor
  if flavor == 'auto':
    flavor = 'naclsdk_' + PLATFORM_COLLAPSE[sys.platform]

  url = DetermineSdkURL(flavor,
                        base_url='gs://nativeclient-mirror/nacl/nacl_sdk/',
                        version=options.version)
  print 'SDK URL is "%s"' % url

  DownloadAndInstallSDK(url)

if __name__ == '__main__':
  main(sys.argv[1:])
<|endoftext|>


<commit_before>#
# Copyright 2013, Couchbase, Inc.
# 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.
#
from __future__ import absolute_import

import select
import time

import couchbase._libcouchbase as LCB
from couchbase._libcouchbase import (
    Event, TimerEvent, IOEvent,
    LCB_READ_EVENT, LCB_WRITE_EVENT, LCB_RW_EVENT,
    PYCBC_EVSTATE_ACTIVE,
    PYCBC_EVACTION_WATCH,
    PYCBC_EVACTION_UNWATCH
)

class SelectTimer(TimerEvent):
    def __init__(self):
        super(SelectTimer, self).__init__()
        self.pydata = 0

    @property
    def exptime(self):
        return self.pydata

    @exptime.setter
    def exptime(self, val):
        self.pydata = val

    def activate(self, usecs):
        self.exptime = time.time() + usecs / 1000000

    def deactivate(self):
        pass

    @property
    def active(self):
        return self.state == PYCBC_EVSTATE_ACTIVE

    def __cmp__(self, other):
        return cmp(self.exptime, other.exptime)

class SelectIOPS(object):
    def __init__(self):
        self._do_watch = False
        self._ioevents = set()
        self._timers = []

        # Active readers and writers
        self._evwr = set()
        self._evrd = set()

    def _unregister_timer(self, timer):
        timer.deactivate()
        if timer in self._timers:
            self._timers.remove(timer)

    def _unregister_event(self, event):
        try:
            self._evrd.remove(event)
        except KeyError:
            pass
        try:
            self._evwr.remove(event)
        except KeyError:
            pass
        try:
            self._ioevents.remove(event)
        except KeyError:
            pass

    def update_timer(self, timer, action, usecs):
        if action == PYCBC_EVACTION_UNWATCH:
            self._unregister_timer(timer)
            return

        if timer.active:
            self._unregister_timer(timer)
        timer.activate(usecs)
        self._timers.append(timer)

    def update_event(self, event, action, flags, fd=None):
        if action == PYCBC_EVACTION_UNWATCH:
            self._unregister_event(event)
            return

        elif action == PYCBC_EVACTION_WATCH:
            if flags & LCB_READ_EVENT:
                self._evrd.add(event)
            else:
                try:
                    self._evrd.remove(event)
                except KeyError:
                    pass

            if flags & LCB_WRITE_EVENT:
                self._evwr.add(event)
            else:
                try:
                    self._evwr.remove(event)
                except KeyError:
                    pass

    def _poll(self):
        rin = self._evrd
        win = self._evwr
        ein = list(rin) + list(win)

        self._timers.sort()
        mintime = self._timers[0].exptime - time.time()
        if mintime < 0:
            mintime = 0

        if not (rin or win or ein):
            time.sleep(mintime)
            rout = tuple()
            wout = tuple()
            eout = tuple()
        else:
            rout, wout, eout = select.select(rin, win, ein, mintime)

        now = time.time()

        ready_events = {}
        for ev in rout:
            ready_events[ev] = LCB_READ_EVENT

        for ev in wout:
            if ev in ready_events:
                ready_events[ev] |= LCB_WRITE_EVENT
            else:
                ready_events[ev] = LCB_WRITE_EVENT

        for ev in eout:
            ready_events[ev] = LCB_RW_EVENT

        for ev, flags in ready_events.items():
            if ev.state == PYCBC_EVSTATE_ACTIVE:
                ev.ready(flags)

        for timer in self._timers[:]:
            if not timer.active:
                continue

            if timer.exptime > now:
                continue

            timer.ready(0)

    def start_watching(self):
        if self._do_watch:
            return

        self._do_watch = True
        while self._do_watch:
            self._poll()

    def stop_watching(self):
        self._do_watch = False

    def timer_event_factory(self):
        return SelectTimer()
<commit_msg>PYCBC-243: iops/select: Use rich comparison methods rather than __cmp__<commit_after>#
# Copyright 2013, Couchbase, Inc.
# 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.
#
from __future__ import absolute_import

import select
import time

import couchbase._libcouchbase as LCB
from couchbase._libcouchbase import (
    Event, TimerEvent, IOEvent,
    LCB_READ_EVENT, LCB_WRITE_EVENT, LCB_RW_EVENT,
    PYCBC_EVSTATE_ACTIVE,
    PYCBC_EVACTION_WATCH,
    PYCBC_EVACTION_UNWATCH
)

class SelectTimer(TimerEvent):
    def __init__(self):
        super(SelectTimer, self).__init__()
        self.pydata = 0

    @property
    def exptime(self):
        return self.pydata

    @exptime.setter
    def exptime(self, val):
        self.pydata = val

    def activate(self, usecs):
        self.exptime = time.time() + usecs / 1000000

    def deactivate(self):
        pass

    @property
    def active(self):
        return self.state == PYCBC_EVSTATE_ACTIVE

    # Rich comparison operators implemented - __cmp__ not used in Py3
    def __lt__(self, other): return self.exptime < other.exptime
    def __le__(self, other): return self.exptime <= other.exptime
    def __gt__(self, other): return self.exptime > other.exptime
    def __ge__(self, other): return self.exptime >= other.exptime
    def __ne__(self, other): return self.exptime != other.exptime
    def __eq__(self, other): return self.exptime == other.exptime

class SelectIOPS(object):
    def __init__(self):
        self._do_watch = False
        self._ioevents = set()
        self._timers = []

        # Active readers and writers
        self._evwr = set()
        self._evrd = set()

    def _unregister_timer(self, timer):
        timer.deactivate()
        if timer in self._timers:
            self._timers.remove(timer)

    def _unregister_event(self, event):
        try:
            self._evrd.remove(event)
        except KeyError:
            pass
        try:
            self._evwr.remove(event)
        except KeyError:
            pass
        try:
            self._ioevents.remove(event)
        except KeyError:
            pass

    def update_timer(self, timer, action, usecs):
        if action == PYCBC_EVACTION_UNWATCH:
            self._unregister_timer(timer)
            return

        if timer.active:
            self._unregister_timer(timer)
        timer.activate(usecs)
        self._timers.append(timer)

    def update_event(self, event, action, flags, fd=None):
        if action == PYCBC_EVACTION_UNWATCH:
            self._unregister_event(event)
            return

        elif action == PYCBC_EVACTION_WATCH:
            if flags & LCB_READ_EVENT:
                self._evrd.add(event)
            else:
                try:
                    self._evrd.remove(event)
                except KeyError:
                    pass

            if flags & LCB_WRITE_EVENT:
                self._evwr.add(event)
            else:
                try:
                    self._evwr.remove(event)
                except KeyError:
                    pass

    def _poll(self):
        rin = self._evrd
        win = self._evwr
        ein = list(rin) + list(win)

        self._timers.sort()
        mintime = self._timers[0].exptime - time.time()
        if mintime < 0:
            mintime = 0

        if not (rin or win or ein):
            time.sleep(mintime)
            rout = tuple()
            wout = tuple()
            eout = tuple()
        else:
            rout, wout, eout = select.select(rin, win, ein, mintime)

        now = time.time()

        ready_events = {}
        for ev in rout:
            ready_events[ev] = LCB_READ_EVENT

        for ev in wout:
            if ev in ready_events:
                ready_events[ev] |= LCB_WRITE_EVENT
            else:
                ready_events[ev] = LCB_WRITE_EVENT

        for ev in eout:
            ready_events[ev] = LCB_RW_EVENT

        for ev, flags in ready_events.items():
            if ev.state == PYCBC_EVSTATE_ACTIVE:
                ev.ready(flags)

        for timer in self._timers[:]:
            if not timer.active:
                continue

            if timer.exptime > now:
                continue

            timer.ready(0)

    def start_watching(self):
        if self._do_watch:
            return

        self._do_watch = True
        while self._do_watch:
            self._poll()

    def stop_watching(self):
        self._do_watch = False

    def timer_event_factory(self):
        return SelectTimer()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding=utf8

from six.moves import map

class StoreDecorator(object):
    """Base class for store decorators.

    The default implementation will use :func:`getattr` to pass through all
    attribute/method requests to an underlying object stored as
    :attr:`_dstore`. It will also pass through the :attr:`__getattr__` and
    :attr:`__contains__` python special methods.
    """
    def __init__(self, store):
        self._dstore = store

    def __getattr__(self, attr):
        return getattr(self._dstore, attr)

    def __contains__(self, *args, **kwargs):
        return self._dstore.__contains__(*args, **kwargs)

    def __iter__(self, *args, **kwargs):
        return self._dstore.__iter__(*args, **kwargs)

class KeyTransformingDecorator(StoreDecorator):
    # currently undocumented (== not advertised as a feature)
    def _map_key(self, key):
        return key

    def _unmap_key(self, key):
        return key

    def __contains__(self, key):
        return self._map_key(key) in self._dstore

    def __iter__(self):
        return map(self._unmap_key, iter(self._dstore))

    def delete(self, key):
        return self._dstore.delete(self._map_key(key))

    def get(self, key, *args, **kwargs):
        return self._dstore.get(self._map_key(key), *args, **kwargs)

    def get_file(self, key, *args, **kwargs):
        return self._dstore.get_file(self._map_key(key), *args, **kwargs)

    def iter_keys(self):
        return map(self._unmap_key, self._dstore.iter_keys())

    def keys(self):
        return list(map(self._unmap_key, self._dstore.keys()))

    def open(self, key):
        return self._dstore.open(self._map_key(key))

    def put(self, key, *args, **kwargs):
        return self._unmap_key(
            self._dstore.put(self._map_key(key), *args, **kwargs)
        )

    def put_file(self, key, *args, **kwargs):
        return self._unmap_key(
            self._dstore.put_file(self._map_key(key), *args, **kwargs)
        )

    # support for UrlMixin
    def url_for(self, key, *args, **kwargs):
        return self._dstore.url_for(self._map_key(key), *args, **kwargs)

class PrefixDecorator(KeyTransformingDecorator):
    """Prefixes any key with a string before passing it on the decorated
    store. Automatically strips the prefix upon key retrieval.

    :param store: The store to pass keys on to.
    :param prefix: Prefix to add.
    """
    def __init__(self, prefix, store):
        super(PrefixDecorator, self).__init__(store)
        self.prefix = prefix

    def _map_key(self, key):
        self._check_valid_key(key)
        return self.prefix + key

    def _unmap_key(self, key):
        assert key.startswith(self.prefix)

        return key[len(self.prefix):]
<commit_msg>Cherry-pick b2db07afe441a49560ed2b330bf5f10979bf4326 from fmarczin.<commit_after>#!/usr/bin/env python
# coding=utf8

from six.moves import map

class StoreDecorator(object):
    """Base class for store decorators.

    The default implementation will use :func:`getattr` to pass through all
    attribute/method requests to an underlying object stored as
    :attr:`_dstore`. It will also pass through the :attr:`__getattr__` and
    :attr:`__contains__` python special methods.
    """

    def __init__(self, store):
        self._dstore = store

    def __getattr__(self, attr):
        return getattr(self._dstore, attr)

    def __contains__(self, *args, **kwargs):
        return self._dstore.__contains__(*args, **kwargs)

    def __iter__(self, *args, **kwargs):
        return self._dstore.__iter__(*args, **kwargs)

class KeyTransformingDecorator(StoreDecorator):
    # currently undocumented (== not advertised as a feature)
    def _map_key(self, key):
        return key

    def _unmap_key(self, key):
        return key

    def _filter(self, key):
        return True

    def __contains__(self, key):
        return self._map_key(key) in self._dstore

    def __iter__(self):
        return map(self._unmap_key, filter(self._filter, iter(self._dstore)))

    def delete(self, key):
        return self._dstore.delete(self._map_key(key))

    def get(self, key, *args, **kwargs):
        return self._dstore.get(self._map_key(key), *args, **kwargs)

    def get_file(self, key, *args, **kwargs):
        return self._dstore.get_file(self._map_key(key), *args, **kwargs)

    def iter_keys(self):
        return map(self._unmap_key, filter(self._filter,
                                           self._dstore.iter_keys()))

    def keys(self):
        """Return a list of keys currently in store, in any order

        :raises IOError: If there was an error accessing the store.
        """
        return list(self.iter_keys())

    def open(self, key):
        return self._dstore.open(self._map_key(key))

    def put(self, key, *args, **kwargs):
        return self._unmap_key(
            self._dstore.put(self._map_key(key), *args, **kwargs))

    def put_file(self, key, *args, **kwargs):
        return self._unmap_key(
            self._dstore.put_file(self._map_key(key), *args, **kwargs))

    # support for UrlMixin
    def url_for(self, key, *args, **kwargs):
        return self._dstore.url_for(self._map_key(key), *args, **kwargs)

class PrefixDecorator(KeyTransformingDecorator):
    """Prefixes any key with a string before passing it on the decorated
    store. Automatically strips the prefix upon key retrieval.

    :param store: The store to pass keys on to.
    :param prefix: Prefix to add.
    """

    def __init__(self, prefix, store):
        super(PrefixDecorator, self).__init__(store)
        self.prefix = prefix

    def _filter(self, key):
        return key.startswith(self.prefix)

    def _map_key(self, key):
        self._check_valid_key(key)
        return self.prefix + key

    def _unmap_key(self, key):
        assert key.startswith(self.prefix)

        return key[len(self.prefix):]
<|endoftext|>