<commit_before>import tensorflow as tf
import math
from pprint import pprint
from settings import FLAGS, model

from models import loss_functions
import data_prep.model_input as input

class Model:
  def __init__(self):

    self.learning_rate = tf.placeholder_with_default(FLAGS.learning_rate, ())
    self.iter_num = tf.placeholder_with_default(FLAGS.num_iterations, ())
    self.summaries = []
    self.noise_std = tf.placeholder_with_default(FLAGS.noise_std, ())
    self.opt = tf.train.AdamOptimizer(self.learning_rate)
    self.model_name = model.__file__

    assert FLAGS.image_range_start + FLAGS.encoder_length + FLAGS.decoder_future_length <= FLAGS.overall_images_count and FLAGS.image_range_start >= 0, \
            "settings for encoder/decoder lengths along with starting range exceed number of available images"
    assert FLAGS.encoder_length >= FLAGS.decoder_reconst_length, "encoder must be at least as long as reconstructer"

  def add_image_summary(self, summary_prefix, frames, encoder_length, decoder_future_length, decoder_reconst_length):
    for i in range(decoder_future_length):
      self.summaries.append(tf.summary.image(summary_prefix + '_future_gen_' + str(i + 1),
                                             self.frames_pred[i], max_outputs=1))
      self.summaries.append(tf.summary.image(summary_prefix + '_future_orig_' + str(i + 1),
                                             frames[:, encoder_length + i, :, :, :], max_outputs=1))
    for i in range(decoder_reconst_length):
      self.summaries.append(tf.summary.image(summary_prefix + '_reconst_gen_' + str(i + 1),
                                             self.frames_reconst[i], max_outputs=1))
      self.summaries.append(tf.summary.image(summary_prefix + '_reconst_orig_' + str(i + 1),
                                             frames[:, i, :, :, :], max_outputs=1))

class TrainModel(Model):
  def __init__(self, summary_prefix, scope_name='train_model'):
    print("Constructing TrainModel")

    with tf.variable_scope(scope_name, reuse=None) as training_scope:
      Model.__init__(self)
      self.scope = training_scope

      tower_grads = []
      tower_losses = []
      for i in range(FLAGS.num_gpus):
        train_batch, _, _ = input.create_batch(FLAGS.tf_records_dir, 'train', FLAGS.batch_size,
                                               int(math.ceil(
                                                 FLAGS.num_iterations / (FLAGS.batch_size * 20))),
                                               False)
        train_batch = tf.cast(train_batch, tf.float32)
        with tf.device('/gpu:%d' % i):
          with tf.name_scope('%s_%d' % ('tower', i)):
            tower_loss, _, _, _ = tower_operations(train_batch[:,FLAGS.image_range_start:,:,:,:], train=True)
            tower_losses.append(tower_loss)

            # Reuse variables for the next tower.
            tf.get_variable_scope().reuse_variables()

            if FLAGS.fine_tuning_weights_list is not None:
              train_vars = []
              for scope_i in FLAGS.fine_tuning_weights_list:
                train_vars += tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope_i)
              pprint('Finetuning. Training only specified weights: %s' % (FLAGS.fine_tuning_weights_list))
              grads = self.opt.compute_gradients(tower_loss, var_list=train_vars)
            else:
              grads = self.opt.compute_gradients(tower_loss)
            tower_grads.append(grads)

      with tf.device('/cpu:0'):
        #copmute average loss
        self.loss = average_losses(tower_losses)

        #compute average over gradients of all towers
        grads = average_gradients(tower_grads)

        # Apply the gradients to adjust the shared variables.
        self.train_op= self.opt.apply_gradients(grads)

      #measure batch time
      self.elapsed_time = tf.placeholder(tf.float32, [])
      self.summaries.append(tf.summary.scalar('batch_duration', self.elapsed_time))

      self.summaries.append(tf.summary.scalar(summary_prefix + '_loss', self.loss))
      self.sum_op = tf.summary.merge(self.summaries)

class ValidationModel(Model):
  def __init__(self, summary_prefix, scope_name='valid_model', reuse_scope=None):
    print("Constructing ValidationModel")

    with tf.variable_scope(scope_name, reuse=None):
      Model.__init__(self)

      assert reuse_scope is not None

      with tf.variable_scope(reuse_scope, reuse=True):
        tower_losses, frames_pred_list, frames_reconst_list, hidden_repr_list, label_batch_list, metadata_batch_list, val_batch_list = [], [], [], [], [], [], []

        for i in range(FLAGS.num_gpus):
          val_batch, label_batch, metadata_batch = input.create_batch(FLAGS.tf_records_dir, 'valid', FLAGS.valid_batch_size, int(
            math.ceil(FLAGS.num_iterations / (FLAGS.batch_size * 20))), False)

          val_batch = tf.cast(val_batch, tf.float32)
          self.val_batch = val_batch

          with tf.device('/gpu:%d' % i):
            with tf.name_scope('%s_%d' % ('tower', i)):
              tower_loss, frames_pred, frames_reconst, hidden_repr = tower_operations(
                val_batch[:, FLAGS.image_range_start:, :, :, :], train=False)
              tower_losses.append(tower_loss)

              frames_pred_list.append(tf.pack(frames_pred))
              frames_reconst_list.append(tf.pack(frames_reconst))
              hidden_repr_list.append(hidden_repr)

              val_batch_list.append(val_batch)
              label_batch_list.append(label_batch)
              metadata_batch_list.append(metadata_batch)
              # Reuse variables for the next tower.
              tf.get_variable_scope().reuse_variables()

        with tf.device('/cpu:0'):
          # compute average loss
          self.loss = average_losses(tower_losses)
          # concatenate outputs of towers to one large tensor each
          self.frames_pred = tf.unstack(tf.concat(1, frames_pred_list))
          self.frames_reconst = tf.unstack(tf.concat(1, frames_reconst_list))
          self.hidden_repr = tf.concat(0, hidden_repr_list)
          self.label = tf.concat(0, label_batch_list)
          self.metadata = tf.concat(0, metadata_batch_list)
          val_set = tf.concat(0, val_batch_list)

      self.add_image_summary(summary_prefix, val_set, FLAGS.encoder_length, FLAGS.decoder_future_length,
                             FLAGS.decoder_reconst_length)

      # evaluate frame predictions for storing on disk
      self.output_frames = self.frames_reconst + self.frames_pred  # join arrays of tensors

      self.summaries.append(tf.summary.scalar(summary_prefix + '_loss', self.loss))
      self.sum_op = tf.summary.merge(self.summaries)

class FeedingValidationModel(Model):
  def __init__(self, scope_name='feeding_model', reuse_scope=None):
    print("Constructing FeedingModel")
    with tf.variable_scope(scope_name, reuse=None):
      Model.__init__(self)

      assert reuse_scope is not None

      with tf.variable_scope(reuse_scope, reuse=True):
        "5D array of batch with videos - shape(batch_size, num_frames, frame_width, frame_higth, num_channels)"
        self.feed_batch = tf.placeholder(tf.float32, shape=(1, FLAGS.encoder_length, 128, 128, 3), name='feed_batch')

        self.frames_pred, self.frames_reconst, self.hidden_repr = \
          tower_operations(self.feed_batch[:, FLAGS.image_range_start:, :, :, :], train=False, compute_loss=False)

def tower_operations(video_batch, train=True, compute_loss=True):
  """
  Build the computation graph from input frame sequences till loss of batch
  :param device number for assining queue runner to CPU
  :param train: boolean that indicates whether train or validation mode
  :param compute_loss: boolean that specifies whether loss should be computed (for feed mode / production compute_loss might be disabled)
  :return batch loss (scalar)
  """
  #only dropout in train mode
  keep_prob_dropout = FLAGS.keep_prob_dopout if train else 1.0

  frames_pred, frames_reconst, hidden_repr = model.composite_model(video_batch, FLAGS.encoder_length,
                                                                   FLAGS.decoder_future_length,
                                                                   FLAGS.decoder_reconst_length,
                                                                   keep_prob_dropout=keep_prob_dropout,
                                                                   noise_std=FLAGS.noise_std,
                                                                   uniform_init=FLAGS.uniform_init,
                                                                   num_channels=FLAGS.num_channels,
                                                                   fc_conv_layer=FLAGS.fc_layer)

  if compute_loss:
    tower_loss = loss_functions.composite_loss(video_batch, frames_pred, frames_reconst, loss_fun=FLAGS.loss_function,
                                  encoder_length=FLAGS.encoder_length,
                                  decoder_future_length=FLAGS.decoder_future_length,
                                  decoder_reconst_length=FLAGS.decoder_reconst_length)
    return tower_loss, frames_pred, frames_reconst, hidden_repr
  else:
    return frames_pred, frames_reconst, hidden_repr

def average_gradients(tower_grads):
  """Calculate the average gradient for each shared variable across all towers.
  Note that this function provides a synchronization point across all towers.
  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over individual gradients. The inner list is over the gradient
      calculation for each tower.
  Returns:
     List of pairs of (gradient, variable) where the gradient has been averaged
     across all towers.
  """

  average_grads = []

  for grad_and_vars in zip(*tower_grads):
    # Note that each grad_and_vars looks like the following:
    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
    grads = []
    for g, _ in grad_and_vars:
      # Add 0 dimension to the gradients to represent the tower.
      expanded_g = tf.expand_dims(g, 0)

      # Append on a 'tower' dimension which we will average over below.
      grads.append(expanded_g)

    # Average over the 'tower' dimension.
    grad = tf.concat(0, grads)
    grad = tf.reduce_mean(grad, 0)

    # Keep in mind that the Variables are redundant because they are shared
    # across towers. So .. we will just return the first tower's pointer to
    # the Variable.
    v = grad_and_vars[0][1]
    grad_and_var = (grad, v)
    average_grads.append(grad_and_var)
  return average_grads

def average_losses(tower_losses):
  """Calculate the average loss among all towers
  Args:
    tower_losses: List of tf.Tensor skalars denoting the loss at each tower.
  Returns:
     loss: tf.Tensor skalar which is the mean over all losses
  """
  losses = []
  for l in tower_losses:
    # Add 0 dimension to the gradients to represent the tower.
    expanded_l = tf.expand_dims(l, 0)

    # Append on a 'tower' dimension which we will average over below.
    losses.append(expanded_l)

  # Average over the 'tower' dimension.
  loss = tf.concat(0, losses)
  loss = tf.reduce_mean(loss, 0)
  return loss

def create_model(mode=None, train_model_scope=None):
  model = None

  if mode is "train":
    model = TrainModel('train', scope_name='train_model')
  elif mode is 'valid':
    assert train_model_scope is not None, "train_model_scope is None, valid mode requires a train scope"
    model = ValidationModel('valid', reuse_scope=train_model_scope)
  elif mode is 'feeding':
    assert train_model_scope is not None, "train_model_scope is None, valid mode requires a train scope"
    model = FeedingValidationModel(reuse_scope=train_model_scope)

  assert model is not None

  return model<commit_msg>adaptive num_channels in feed mode placeholder<commit_after>import tensorflow as tf
import math
from pprint import pprint
from settings import FLAGS, model

from models import loss_functions
import data_prep.model_input as input

class Model:
  def __init__(self):

    self.learning_rate = tf.placeholder_with_default(FLAGS.learning_rate, ())
    self.iter_num = tf.placeholder_with_default(FLAGS.num_iterations, ())
    self.summaries = []
    self.noise_std = tf.placeholder_with_default(FLAGS.noise_std, ())
    self.opt = tf.train.AdamOptimizer(self.learning_rate)
    self.model_name = model.__file__

    assert FLAGS.image_range_start + FLAGS.encoder_length + FLAGS.decoder_future_length <= FLAGS.overall_images_count and FLAGS.image_range_start >= 0, \
            "settings for encoder/decoder lengths along with starting range exceed number of available images"
    assert FLAGS.encoder_length >= FLAGS.decoder_reconst_length, "encoder must be at least as long as reconstructer"

  def add_image_summary(self, summary_prefix, frames, encoder_length, decoder_future_length, decoder_reconst_length):
    for i in range(decoder_future_length):
      self.summaries.append(tf.summary.image(summary_prefix + '_future_gen_' + str(i + 1),
                                             self.frames_pred[i], max_outputs=1))
      self.summaries.append(tf.summary.image(summary_prefix + '_future_orig_' + str(i + 1),
                                             frames[:, encoder_length + i, :, :, :], max_outputs=1))
    for i in range(decoder_reconst_length):
      self.summaries.append(tf.summary.image(summary_prefix + '_reconst_gen_' + str(i + 1),
                                             self.frames_reconst[i], max_outputs=1))
      self.summaries.append(tf.summary.image(summary_prefix + '_reconst_orig_' + str(i + 1),
                                             frames[:, i, :, :, :], max_outputs=1))

class TrainModel(Model):
  def __init__(self, summary_prefix, scope_name='train_model'):
    print("Constructing TrainModel")

    with tf.variable_scope(scope_name, reuse=None) as training_scope:
      Model.__init__(self)
      self.scope = training_scope

      tower_grads = []
      tower_losses = []
      for i in range(FLAGS.num_gpus):
        train_batch, _, _ = input.create_batch(FLAGS.tf_records_dir, 'train', FLAGS.batch_size,
                                               int(math.ceil(
                                                 FLAGS.num_iterations / (FLAGS.batch_size * 20))),
                                               False)
        train_batch = tf.cast(train_batch, tf.float32)
        with tf.device('/gpu:%d' % i):
          with tf.name_scope('%s_%d' % ('tower', i)):
            tower_loss, _, _, _ = tower_operations(train_batch[:,FLAGS.image_range_start:,:,:,:], train=True)
            tower_losses.append(tower_loss)

            # Reuse variables for the next tower.
            tf.get_variable_scope().reuse_variables()

            if FLAGS.fine_tuning_weights_list is not None:
              train_vars = []
              for scope_i in FLAGS.fine_tuning_weights_list:
                train_vars += tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope_i)
              pprint('Finetuning. Training only specified weights: %s' % (FLAGS.fine_tuning_weights_list))
              grads = self.opt.compute_gradients(tower_loss, var_list=train_vars)
            else:
              grads = self.opt.compute_gradients(tower_loss)
            tower_grads.append(grads)

      with tf.device('/cpu:0'):
        #copmute average loss
        self.loss = average_losses(tower_losses)

        #compute average over gradients of all towers
        grads = average_gradients(tower_grads)

        # Apply the gradients to adjust the shared variables.
        self.train_op= self.opt.apply_gradients(grads)

      #measure batch time
      self.elapsed_time = tf.placeholder(tf.float32, [])
      self.summaries.append(tf.summary.scalar('batch_duration', self.elapsed_time))

      self.summaries.append(tf.summary.scalar(summary_prefix + '_loss', self.loss))
      self.sum_op = tf.summary.merge(self.summaries)

class ValidationModel(Model):
  def __init__(self, summary_prefix, scope_name='valid_model', reuse_scope=None):
    print("Constructing ValidationModel")

    with tf.variable_scope(scope_name, reuse=None):
      Model.__init__(self)

      assert reuse_scope is not None

      with tf.variable_scope(reuse_scope, reuse=True):
        tower_losses, frames_pred_list, frames_reconst_list, hidden_repr_list, label_batch_list, metadata_batch_list, val_batch_list = [], [], [], [], [], [], []

        for i in range(FLAGS.num_gpus):
          val_batch, label_batch, metadata_batch = input.create_batch(FLAGS.tf_records_dir, 'valid', FLAGS.valid_batch_size, int(
            math.ceil(FLAGS.num_iterations / (FLAGS.batch_size * 20))), False)

          val_batch = tf.cast(val_batch, tf.float32)
          self.val_batch = val_batch

          with tf.device('/gpu:%d' % i):
            with tf.name_scope('%s_%d' % ('tower', i)):
              tower_loss, frames_pred, frames_reconst, hidden_repr = tower_operations(
                val_batch[:, FLAGS.image_range_start:, :, :, :], train=False)
              tower_losses.append(tower_loss)

              frames_pred_list.append(tf.pack(frames_pred))
              frames_reconst_list.append(tf.pack(frames_reconst))
              hidden_repr_list.append(hidden_repr)

              val_batch_list.append(val_batch)
              label_batch_list.append(label_batch)
              metadata_batch_list.append(metadata_batch)
              # Reuse variables for the next tower.
              tf.get_variable_scope().reuse_variables()

        with tf.device('/cpu:0'):
          # compute average loss
          self.loss = average_losses(tower_losses)
          # concatenate outputs of towers to one large tensor each
          self.frames_pred = tf.unstack(tf.concat(1, frames_pred_list))
          self.frames_reconst = tf.unstack(tf.concat(1, frames_reconst_list))
          self.hidden_repr = tf.concat(0, hidden_repr_list)
          self.label = tf.concat(0, label_batch_list)
          self.metadata = tf.concat(0, metadata_batch_list)
          val_set = tf.concat(0, val_batch_list)

      self.add_image_summary(summary_prefix, val_set, FLAGS.encoder_length, FLAGS.decoder_future_length,
                             FLAGS.decoder_reconst_length)

      # evaluate frame predictions for storing on disk
      self.output_frames = self.frames_reconst + self.frames_pred  # join arrays of tensors

      self.summaries.append(tf.summary.scalar(summary_prefix + '_loss', self.loss))
      self.sum_op = tf.summary.merge(self.summaries)

class FeedingValidationModel(Model):
  def __init__(self, scope_name='feeding_model', reuse_scope=None):
    print("Constructing FeedingModel")
    with tf.variable_scope(scope_name, reuse=None):
      Model.__init__(self)

      assert reuse_scope is not None

      with tf.variable_scope(reuse_scope, reuse=True):
        "5D array of batch with videos - shape(batch_size, num_frames, frame_width, frame_higth, num_channels)"
        self.feed_batch = tf.placeholder(tf.float32, shape=(1, FLAGS.encoder_length, FLAGS.height, FLAGS.width, FLAGS.num_channels), name='feed_batch')

        self.frames_pred, self.frames_reconst, self.hidden_repr = \
          tower_operations(self.feed_batch[:, FLAGS.image_range_start:, :, :, :], train=False, compute_loss=False)

def tower_operations(video_batch, train=True, compute_loss=True):
  """
  Build the computation graph from input frame sequences till loss of batch
  :param device number for assining queue runner to CPU
  :param train: boolean that indicates whether train or validation mode
  :param compute_loss: boolean that specifies whether loss should be computed (for feed mode / production compute_loss might be disabled)
  :return batch loss (scalar)
  """
  #only dropout in train mode
  keep_prob_dropout = FLAGS.keep_prob_dopout if train else 1.0

  frames_pred, frames_reconst, hidden_repr = model.composite_model(video_batch, FLAGS.encoder_length,
                                                                   FLAGS.decoder_future_length,
                                                                   FLAGS.decoder_reconst_length,
                                                                   keep_prob_dropout=keep_prob_dropout,
                                                                   noise_std=FLAGS.noise_std,
                                                                   uniform_init=FLAGS.uniform_init,
                                                                   num_channels=FLAGS.num_channels,
                                                                   fc_conv_layer=FLAGS.fc_layer)

  if compute_loss:
    tower_loss = loss_functions.composite_loss(video_batch, frames_pred, frames_reconst, loss_fun=FLAGS.loss_function,
                                  encoder_length=FLAGS.encoder_length,
                                  decoder_future_length=FLAGS.decoder_future_length,
                                  decoder_reconst_length=FLAGS.decoder_reconst_length)
    return tower_loss, frames_pred, frames_reconst, hidden_repr
  else:
    return frames_pred, frames_reconst, hidden_repr

def average_gradients(tower_grads):
  """Calculate the average gradient for each shared variable across all towers.
  Note that this function provides a synchronization point across all towers.
  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over individual gradients. The inner list is over the gradient
      calculation for each tower.
  Returns:
     List of pairs of (gradient, variable) where the gradient has been averaged
     across all towers.
  """

  average_grads = []

  for grad_and_vars in zip(*tower_grads):
    # Note that each grad_and_vars looks like the following:
    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
    grads = []
    for g, _ in grad_and_vars:
      # Add 0 dimension to the gradients to represent the tower.
      expanded_g = tf.expand_dims(g, 0)

      # Append on a 'tower' dimension which we will average over below.
      grads.append(expanded_g)

    # Average over the 'tower' dimension.
    grad = tf.concat(0, grads)
    grad = tf.reduce_mean(grad, 0)

    # Keep in mind that the Variables are redundant because they are shared
    # across towers. So .. we will just return the first tower's pointer to
    # the Variable.
    v = grad_and_vars[0][1]
    grad_and_var = (grad, v)
    average_grads.append(grad_and_var)
  return average_grads

def average_losses(tower_losses):
  """Calculate the average loss among all towers
  Args:
    tower_losses: List of tf.Tensor skalars denoting the loss at each tower.
  Returns:
     loss: tf.Tensor skalar which is the mean over all losses
  """
  losses = []
  for l in tower_losses:
    # Add 0 dimension to the gradients to represent the tower.
    expanded_l = tf.expand_dims(l, 0)

    # Append on a 'tower' dimension which we will average over below.
    losses.append(expanded_l)

  # Average over the 'tower' dimension.
  loss = tf.concat(0, losses)
  loss = tf.reduce_mean(loss, 0)
  return loss

def create_model(mode=None, train_model_scope=None):
  model = None

  if mode is "train":
    model = TrainModel('train', scope_name='train_model')
  elif mode is 'valid':
    assert train_model_scope is not None, "train_model_scope is None, valid mode requires a train scope"
    model = ValidationModel('valid', reuse_scope=train_model_scope)
  elif mode is 'feeding':
    assert train_model_scope is not None, "train_model_scope is None, valid mode requires a train scope"
    model = FeedingValidationModel(reuse_scope=train_model_scope)

  assert model is not None

  return model<|endoftext|>


<commit_before># coding: utf-8

import sys

k_bit_rate_num_bits = [ 0, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 32 ]

k_highest_bit_rate = 18
k_lowest_bit_rate = 1
k_num_bit_rates = 19
k_invalid_bit_rate = 255

# This code assumes that rotations, translations, and scales are packed on 3 components (e.g. quat drop w)

if __name__ == "__main__":
	if sys.version_info < (3, 4):
		print('Python 3.4 or higher needed to run this script')
		sys.exit(1)

	permutation_tries = []
	permutation_tries_no_scale = []

	for rotation_bit_rate in range(k_num_bit_rates):
		for translation_bit_rate in range(k_num_bit_rates):
			transform_size = k_bit_rate_num_bits[rotation_bit_rate] * 3 + k_bit_rate_num_bits[translation_bit_rate] * 3
			permutation_tries_no_scale.append((transform_size, rotation_bit_rate, translation_bit_rate))

			for scale_bit_rate in range(k_num_bit_rates):
				transform_size = k_bit_rate_num_bits[rotation_bit_rate] * 3 + k_bit_rate_num_bits[translation_bit_rate] * 3 + k_bit_rate_num_bits[scale_bit_rate] * 3
				permutation_tries.append((transform_size, rotation_bit_rate, translation_bit_rate, scale_bit_rate))

	# Sort by transform size, then by each bit rate
	permutation_tries.sort()
	permutation_tries_no_scale.sort()

	print('constexpr uint8_t k_local_bit_rate_permutations_no_scale[{}][2] ='.format(len(permutation_tries_no_scale)))
	print('{')
	for transform_size, rotation_bit_rate, translation_bit_rate in permutation_tries_no_scale:
		print('\t{{ {}, {} }},\t\t// {} bits per transform'.format(rotation_bit_rate, translation_bit_rate, transform_size))
	print('};')
	print()
	print('constexpr uint8_t k_local_bit_rate_permutations[{}][3] ='.format(len(permutation_tries)))
	print('{')
	for transform_size, rotation_bit_rate, translation_bit_rate, scale_bit_rate in permutation_tries:
		print('\t{{ {}, {}, {} }},\t\t// {} bits per transform'.format(rotation_bit_rate, translation_bit_rate, scale_bit_rate, transform_size))
	print('};')
<commit_msg>fix(tools): use the array length instead of hardcoding<commit_after># coding: utf-8

import sys

k_bit_rate_num_bits = [ 0, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 32 ]

k_highest_bit_rate = len(k_bit_rate_num_bits) - 1
k_lowest_bit_rate = 1
k_num_bit_rates = len(k_bit_rate_num_bits)
k_invalid_bit_rate = 255

# This code assumes that rotations, translations, and scales are packed on 3 components (e.g. quat drop w)

if __name__ == "__main__":
	if sys.version_info < (3, 4):
		print('Python 3.4 or higher needed to run this script')
		sys.exit(1)

	permutation_tries = []
	permutation_tries_no_scale = []

	for rotation_bit_rate in range(k_num_bit_rates):
		for translation_bit_rate in range(k_num_bit_rates):
			transform_size = k_bit_rate_num_bits[rotation_bit_rate] * 3 + k_bit_rate_num_bits[translation_bit_rate] * 3
			permutation_tries_no_scale.append((transform_size, rotation_bit_rate, translation_bit_rate))

			for scale_bit_rate in range(k_num_bit_rates):
				transform_size = k_bit_rate_num_bits[rotation_bit_rate] * 3 + k_bit_rate_num_bits[translation_bit_rate] * 3 + k_bit_rate_num_bits[scale_bit_rate] * 3
				permutation_tries.append((transform_size, rotation_bit_rate, translation_bit_rate, scale_bit_rate))

	# Sort by transform size, then by each bit rate
	permutation_tries.sort()
	permutation_tries_no_scale.sort()

	print('constexpr uint8_t k_local_bit_rate_permutations_no_scale[{}][2] ='.format(len(permutation_tries_no_scale)))
	print('{')
	for transform_size, rotation_bit_rate, translation_bit_rate in permutation_tries_no_scale:
		print('\t{{ {}, {} }},\t\t// {} bits per transform'.format(rotation_bit_rate, translation_bit_rate, transform_size))
	print('};')
	print()
	print('constexpr uint8_t k_local_bit_rate_permutations[{}][3] ='.format(len(permutation_tries)))
	print('{')
	for transform_size, rotation_bit_rate, translation_bit_rate, scale_bit_rate in permutation_tries:
		print('\t{{ {}, {}, {} }},\t\t// {} bits per transform'.format(rotation_bit_rate, translation_bit_rate, scale_bit_rate, transform_size))
	print('};')
<|endoftext|>


<commit_before># coding=utf-8
from __future__ import absolute_import

__author__ = "Gina Häußge <gina@octoprint.org>"
__license__ = 'GNU Affero General Public License http://www.gnu.org/licenses/agpl.html'
__copyright__ = "Copyright (C) 2015 The OctoPrint Project - Released under terms of the AGPLv3 License"

import os

import time
import traceback
import octoprint.util
import octoprint.plugin

from octoprint.events import Events
from octoprint.server import admin_permission
from flask.ext.login import current_user

import pushbullet
import flask

import sarge

class PushbulletPlugin(octoprint.plugin.EventHandlerPlugin,
		       octoprint.plugin.ProgressPlugin,
                       octoprint.plugin.SettingsPlugin,
                       octoprint.plugin.StartupPlugin,
                       octoprint.plugin.TemplatePlugin,
                       octoprint.plugin.SimpleApiPlugin,
                       octoprint.plugin.AssetPlugin):

	def __init__(self):
		self._bullet = None
		self._channel = None
		self._sender = None
		self._pc_interval = 5
		self._message_count = 0
		self._quiet_time_sec = 1800
		self._last_message = 0
		self._etl_format   = "{days:d}:{hours:02d}h {minutes:02d}m"
		self._eta_strftime = "%H:%M %d-%m"

	def _connect_bullet(self, apikey, channel_name=""):
		self._bullet, self._sender = self._create_sender(apikey, channel_name)

	#~~ progress message helpers
	def _sanitize_current_data(self, currentData):
		if (currentData["progress"]["printTimeLeft"] == None):
			currentData["progress"]["printTimeLeft"] = currentData["job"]["estimatedPrintTime"]
		if (currentData["progress"]["printTimeLeft"] == None):
			self._logger.info("Messasge: Still got no print time {}".format(currentData["progress"]["printTimeLeft"]))
			currentData["progress"]["printTimeLeft"] = 1000
		if (currentData["progress"]["filepos"] == None):
			currentData["progress"]["filepos"] = 0
		if (currentData["progress"]["printTime"] == None):
			currentData["progress"]["printTime"] = currentData["job"]["estimatedPrintTime"]

		currentData["progress"]["printTimeLeftString"] = "No ETL yet"
		currentData["progress"]["printTimeString"] = "Not started yet"
		currentData["progress"]["ETA"] = "No ETA yet"
		#Add additional data
		try:
			currentData["progress"]["printTimeLeftString"] = self._get_time_from_seconds(currentData["progress"]["printTimeLeft"])
			currentData["progress"]["printTimeString"] = self._get_time_from_seconds(currentData["progress"]["printTime"])
			currentData["progress"]["ETA"] = time.strftime(self._eta_strftime, time.localtime(time.time() + currentData["progress"]["printTimeLeft"]))
		except Exception as e:
			self._logger.warning("Caught an exception trying to parse data: {0}\n Error is: {1}\nTraceback:{2}".format(currentData,e,traceback.format_exc()))

		return currentData

	def _get_time_from_seconds(self, seconds):
		days =0
		hours = 0
		minutes = 0
		if seconds >= 86400:
			days = int(seconds / 86400)
			secionds = seconds % 86400
		if seconds >= 3600:
			hours = int(seconds / 3600)
			seconds = seconds % 3600
		if seconds >= 60:
			minutes = int(seconds / 60)
			seconds = seconds % 60
		return self._etl_format.format(**locals())

	#~~ PrintProgressPlugin
	def on_print_progress(self, storage, path, progress):
		if(progress%self._pc_interval==0):
			self._quiet_time_sec = int(self._settings.get(["quiet_minutes"])) * 60
			if self._quiet_time_sec == 0:
				self._quiet_time_sec = 7* 24 * 3600 # use a week if we don't want messages
			self._logger.info("Progress: {} {} {}".format(storage,path,progress))
			try:
				currentData = self._printer.get_current_data()
				currentData = self._sanitize_current_data(currentData)
			except Exception as e:
				self._logger.info("Caught an exception {0}\nTraceback:{1}".format(e,traceback.format_exc()))

			self._logger.info("Messasge: Remaining {} ({})".format(currentData["progress"]["printTimeLeftString"],currentData["progress"]["printTimeLeftOrigin"]))
			self._logger.info("Messasge: Total Print time {}".format(currentData["progress"]["printTimeString"]))
			self._logger.info("Messasge: Estimated Completion  {}".format(currentData["progress"]["ETA"]))
			# first 3 messages, re-calculate interval (not first, analysis estimate can be inacurate)
			self._message_count += 1
			if self._message_count < 4 and self._message_count > 1:
				total_job = currentData["progress"]["printTime"]+currentData["progress"]["printTimeLeft"]
				self._pc_interval = int(100 * self._quiet_time_sec / total_job)
				if self._pc_interval == 0:
					self._pc_interval = 1
				self._logger.info("Messasge: Set interval to {} percent count {} Total estimated {}".format(self._pc_interval, self._message_count,total_job))

			# Suppress message if print is nearly done or interval calculation is off
			no_message = 0
			if(currentData["progress"]["printTime"] < (self._quiet_time_sec * 0.5  + self._last_message)):
				self._logger.info("Messasge: Skip message {} since print is only been running {} since last at {}".format(self._message_count, currentData["progress"]["printTime"], self._last_message))
				no_message = 1

			if(currentData["progress"]["printTimeLeft"] < self._quiet_time_sec):
				self._logger.info("Messasge: Skip trailing message since print is nearly done {} of {}".format(currentData["progress"]["printTimeLeft"], self._quiet_time_sec))
				no_message = 1

			if no_message == 0:
				self._last_message = currentData["progress"]["printTime"]
				file = currentData["job"]["file"]["path"]
				ETA = currentData["progress"]["ETA"]
				print_time = currentData["progress"]["printTimeString"]
				remaining = currentData["progress"]["printTimeLeftString"]

				title = self._settings.get(["printPartial", "title"]).format(**locals())
				body = self._settings.get(["printPartial", "body"]).format(**locals())
				filename = os.path.splitext(file)[0] + ".jpg"

				self._send_message_with_webcam_image(title, body, filename=filename)

	#~~ StartupPlugin

	def on_after_startup(self):
		self._connect_bullet(self._settings.get(["access_token"]),
		                     self._settings.get(["push_channel"]))

	#~~ SettingsPlugin

	def on_settings_load(self):
		data = octoprint.plugin.SettingsPlugin.on_settings_load(self)

		# only return our restricted settings to admin users - this is only needed for OctoPrint <= 1.2.16
		restricted = ("quiet_minutes", "access_token", "push_channel")
		for r in restricted:
			if r in data and (current_user is None or current_user.is_anonymous() or not current_user.is_admin()):
				data[r] = None

		return data

	def on_settings_save(self, data):
		octoprint.plugin.SettingsPlugin.on_settings_save(self, data)

		import threading
		thread = threading.Thread(target=self._connect_bullet, args=(self._settings.get(["quiet_minutes"]),
									     self._settings.get(["access_token"]),
		                                                             self._settings.get(["push_channel"])))
		thread.daemon = True
		thread.start()

	def get_settings_defaults(self):
		return dict(
			quiet_minute=0,
			access_token=None,
			push_channel=None,
			printDone=dict(
				title="Print job finished",
				body="{file} finished printing in {elapsed_time}"
			),
			printPartial=dict(
				title="Print job {progress}% complete",
				body="{progress}% {file}\n{print_time}\n{remaining} left\n{ETA} finish"
			)
		)

	def get_settings_restricted_paths(self):
		# only used in OctoPrint versions > 1.2.16
		return dict(admin=[["quiet_minutes"], ["access_token"], ["push_channel"]])

	#~~ TemplatePlugin API

	def get_template_configs(self):
		return [
			dict(type="settings", name="Pushbullet", custom_bindings=True)
		]

	#~~ AssetPlugin API

	def get_assets(self):
		return dict(js=["js/octobullet.js"])

	#~~ SimpleApiPlugin

	def get_api_commands(self):
		return dict(test=["token"])

	def on_api_command(self, command, data):
		if not admin_permission.can():
			return flask.make_response("Insufficient rights", 403)

		if not command == "test":
			return

		message = data.get("message", "Testing, 1, 2, 3, 4...")
		token = data["token"]
		channel = data.get("channel", None)

		_, sender = self._create_sender(token, channel)

		result = self._send_message_with_webcam_image("Test from the OctoPrint PushBullet Plugin", message, sender=sender)
		return flask.make_response(flask.jsonify(result=result))

	#~~ EventHandlerPlugin

	def on_event(self, event, payload):

		if event == Events.PRINT_DONE:
			file = os.path.basename(payload["file"])
			elapsed_time_in_seconds = payload["time"]

			import datetime
			import octoprint.util
			elapsed_time = octoprint.util.get_formatted_timedelta(datetime.timedelta(seconds=elapsed_time_in_seconds))

			title = self._settings.get(["printDone", "title"]).format(**locals())
			body = self._settings.get(["printDone", "body"]).format(**locals())
			filename = os.path.splitext(file)[0] + ".jpg"

			self._send_message_with_webcam_image(title, body, filename=filename)

		if event == Events.PRINT_STARTED:
			self._pc_interval = 5
			self.message_count = 0
			self._last_message = 0

	##~~ Softwareupdate hook

	def get_update_information(self):
		return dict(
			octobullet=dict(
				displayName="Pushbullet Plugin",
				displayVersion=self._plugin_version,

				# version check: github repository
				type="github_release",
				user="OctoPrint",
				repo="OctoPrint-Pushbullet",
				current=self._plugin_version,

				# update method: pip
				pip="https://github.com/OctoPrint/OctoPrint-Pushbullet/archive/{target_version}.zip"
			)
		)

	##~~ Internal utility methods

	def _send_message_with_webcam_image(self, title, body, filename=None, sender=None):
		if filename is None:
			import random, string
			filename = "{}.jpg".format(random.choice(string.ascii_letters) * 16)

		if sender is None:
			sender = self._sender

		if not sender:
			return False

		snapshot_url = self._settings.global_get(["webcam", "snapshot"])
		if snapshot_url:
			try:
				import urllib
				snapshot_path, headers = urllib.urlretrieve(snapshot_url)
			except Exception as e:
				self._logger.exception(
					"Exception while fetching snapshot from webcam, sending only a note: {message}".format(
						message=str(e)))
			else:
				# ffmpeg can't guess file type it seems
				os.rename(snapshot_path, snapshot_path + ".jpg")
				snapshot_path += ".jpg"

				# flip or rotate as needed
				self._process_snapshot(snapshot_path)

				if self._send_file(sender, snapshot_path, filename, body):
					return True
				self._logger.warn("Could not send a file message with the webcam image, sending only a note")

		return self._send_note(sender, title, body)

	def _send_note(self, sender, title, body):
		try:
			sender.push_note(title, body)
		except:
			self._logger.exception("Error while pushing a note")
			return False
		return True

	def _send_file(self, sender, path, filename, body):
		try:
			with open(path, "rb") as pic:
				try:
					file_data = self._bullet.upload_file(pic, filename)
				except Exception as e:
					self._logger.exception("Error while uploading snapshot, sending only a note: {}".format(str(e)))
					return False

			sender.push_file(file_data["file_name"], file_data["file_url"], file_data["file_type"], body=body)
			return True
		except Exception as e:
			self._logger.exception("Exception while uploading snapshot to Pushbullet, sending only a note: {message}".format(message=str(e)))
			return False
		finally:
			try:
				os.remove(path)
			except:
				self._logger.exception("Could not remove temporary snapshot file {}".format(path))

	def _create_sender(self, token, channel=None):
		try:
			bullet = pushbullet.PushBullet(token)
			sender = bullet

			# Setup channel object if channel setting is present
			if channel:
				for channel_obj in self._bullet.channels:
					if channel_obj.channel_tag == channel:
						sender = channel_obj
						self._logger.info("Connected to PushBullet on channel {}".format(channel))
						break
				else:
					self._logger.warn("Could not find channel {}, please check your configuration!".format(channel))

			self._logger.info("Connected to PushBullet")
			return bullet, sender
		except:
			self._logger.exception("Error while instantiating PushBullet")
			return None, None

	def _process_snapshot(self, snapshot_path, pixfmt="yuv420p"):
		hflip  = self._settings.global_get_boolean(["webcam", "flipH"])
		vflip  = self._settings.global_get_boolean(["webcam", "flipV"])
		rotate = self._settings.global_get_boolean(["webcam", "rotate90"])
		ffmpeg = self._settings.global_get(["webcam", "ffmpeg"])

		if not ffmpeg or not os.access(ffmpeg, os.X_OK) or (not vflip and not hflip and not rotate):
			return

		ffmpeg_command = [ffmpeg, "-y", "-i", snapshot_path]

		rotate_params = ["format={}".format(pixfmt)] # workaround for foosel/OctoPrint#1317
		if rotate:
			rotate_params.append("transpose=2") # 90 degrees counter clockwise
		if hflip:
			rotate_params.append("hflip") 		# horizontal flip
		if vflip:
			rotate_params.append("vflip")		# vertical flip

		ffmpeg_command += ["-vf", sarge.shell_quote(",".join(rotate_params)), snapshot_path]
		self._logger.info("Running: {}".format(" ".join(ffmpeg_command)))

		p = sarge.run(ffmpeg_command, stdout=sarge.Capture(), stderr=sarge.Capture())
		if p.returncode == 0:
			self._logger.info("Rotated/flipped image with ffmpeg")
		else:
			self._logger.warn("Failed to rotate/flip image with ffmpeg, "
			                  "got return code {}: {}, {}".format(p.returncode,
			                                                      p.stdout.text,
			                                                      p.stderr.text))

__plugin_name__ = "Pushbullet"
def __plugin_load__():
	global __plugin_implementation__
	__plugin_implementation__ = PushbulletPlugin()

	global __plugin_hooks__
	__plugin_hooks__ = {
		"octoprint.plugin.softwareupdate.check_config": __plugin_implementation__.get_update_information
	}

<commit_msg>Simple style cleanups, fix typos, moved messages to debug not info<commit_after># coding=utf-8
from __future__ import absolute_import

__author__ = "Gina Häußge <gina@octoprint.org>"
__license__ = 'GNU Affero General Public License http://www.gnu.org/licenses/agpl.html'
__copyright__ = "Copyright (C) 2015 The OctoPrint Project - Released under terms of the AGPLv3 License"

import os

import time
import traceback
import octoprint.util
import octoprint.plugin

from octoprint.events import Events
from octoprint.server import admin_permission
from flask.ext.login import current_user

import pushbullet
import flask

import sarge

class PushbulletPlugin(octoprint.plugin.EventHandlerPlugin,
		       octoprint.plugin.ProgressPlugin,
                       octoprint.plugin.SettingsPlugin,
                       octoprint.plugin.StartupPlugin,
                       octoprint.plugin.TemplatePlugin,
                       octoprint.plugin.SimpleApiPlugin,
                       octoprint.plugin.AssetPlugin):

	def __init__(self):
		self._bullet = None
		self._channel = None
		self._sender = None
		self._interval_percent = 5
		self._message_count = 0
		self._quiet_time_sec = 1800
		self._last_message = 0
		self._time_remaining_format = "{days:d}d {hours:02d}h {minutes:02d}min"
		self._eta_strftime = "%H:%M %d-%m"

	def _connect_bullet(self, apikey, channel_name=""):
		self._bullet, self._sender = self._create_sender(apikey, channel_name)

	#~~ progress message helpers
	def _sanitize_current_data(self, currentData):
		if (currentData["progress"]["printTimeLeft"] == None):
			currentData["progress"]["printTimeLeft"] = currentData["job"]["estimatedPrintTime"]
		if (currentData["progress"]["printTimeLeft"] == None):
			self._logger.debug("Still got no print time {}".format(currentData["progress"]["printTimeLeft"]))
			currentData["progress"]["printTimeLeft"] = 1000
		if (currentData["progress"]["filepos"] == None):
			currentData["progress"]["filepos"] = 0
		if (currentData["progress"]["printTime"] == None):
			currentData["progress"]["printTime"] = currentData["job"]["estimatedPrintTime"]

		currentData["progress"]["printTimeLeftString"] = "No ETL yet"
		currentData["progress"]["printTimeString"] = "Not started yet"
		currentData["progress"]["ETA"] = "No ETA yet"
		#Add additional data
		try:
			currentData["progress"]["printTimeLeftString"] = self._get_time_from_seconds(currentData["progress"]["printTimeLeft"])
			currentData["progress"]["printTimeString"] = self._get_time_from_seconds(currentData["progress"]["printTime"])
			currentData["progress"]["ETA"] = time.strftime(self._eta_strftime, time.localtime(time.time() + currentData["progress"]["printTimeLeft"]))
		except Exception as e:
			self._logger.warning("Caught an exception trying to parse data: {0}\n Error is: {1}\nTraceback:{2}".format(currentData,e,traceback.format_exc()))

		return currentData

	def _get_time_from_seconds(self, seconds):
		days, seconds = divmod(seconds, 86400)
		hours, seconds = divmod(seconds, 3600)
		minutes, seconds = divmod(seconds, 60)
		return self._time_remaining_format.format(**locals())

	#~~ PrintProgressPlugin
	def on_print_progress(self, storage, path, progress):
		if(progress % self._interval_percent == 0):
			self._quiet_time_sec = int(self._settings.get(["quiet_minutes"])) * 60
			if self._quiet_time_sec == 0:
				self._quiet_time_sec = 7* 24 * 3600 # use a week if we don't want messages
			self._logger.debug("Progress: {} {} {}".format(storage,path,progress))
			try:
				currentData = self._printer.get_current_data()
				currentData = self._sanitize_current_data(currentData)
			except Exception as e:
				self._logger.info("Caught an exception {0}\nTraceback:{1}".format(e,traceback.format_exc()))

			self._logger.debug("Remaining {} ({})".format(currentData["progress"]["printTimeLeftString"],currentData["progress"]["printTimeLeftOrigin"]))
			self._logger.debug("Total Print time {}".format(currentData["progress"]["printTimeString"]))
			self._logger.debug("Estimated Completion  {}".format(currentData["progress"]["ETA"]))
			# first 3 messages, re-calculate interval (not first, analysis estimate can be inacurate)
			self._message_count += 1
			if self._message_count < 4 and self._message_count > 1:
				total_job = currentData["progress"]["printTime"] + currentData["progress"]["printTimeLeft"]
				self._interval_percent = int(100 * self._quiet_time_sec / total_job)
				if self._interval_percent == 0:
					self._interval_percent = 1
				self._logger.debug("Set interval to {} percent count {} Total estimated {}".format(self._interval_percent, self._message_count,total_job))

			# Suppress message if print is nearly done or interval calculation is off
			need_message = True
			if(currentData["progress"]["printTime"] < (self._quiet_time_sec * 0.5  + self._last_message)):
				self._logger.debug("Skip message {} since print is only been running {} since last at {}".format(self._message_count, currentData["progress"]["printTime"], self._last_message))
				need_message = False

			if(currentData["progress"]["printTimeLeft"] < self._quiet_time_sec):
				self._logger.debug("Skip trailing message since print is nearly done {} of {}".format(currentData["progress"]["printTimeLeft"], self._quiet_time_sec))
				need_message = False

			if need_message:
				self._last_message = currentData["progress"]["printTime"]
				path = currentData["job"]["file"]["path"]
				eta = currentData["progress"]["ETA"]
				print_time = currentData["progress"]["printTimeString"]
				remaining = currentData["progress"]["printTimeLeftString"]

				title = self._settings.get(["printProgress", "title"]).format(**locals())
				body = self._settings.get(["printProgress", "body"]).format(**locals())
				filename = os.path.splitext(path)[0] + ".jpg"

				self._send_message_with_webcam_image(title, body, filename=filename)

	#~~ StartupPlugin

	def on_after_startup(self):
		self._connect_bullet(self._settings.get(["access_token"]),
		                     self._settings.get(["push_channel"]))

	#~~ SettingsPlugin

	def on_settings_load(self):
		data = octoprint.plugin.SettingsPlugin.on_settings_load(self)

		# only return our restricted settings to admin users - this is only needed for OctoPrint <= 1.2.16
		restricted = ("access_token", "push_channel")
		for r in restricted:
			if r in data and (current_user is None or current_user.is_anonymous() or not current_user.is_admin()):
				data[r] = None

		return data

	def on_settings_save(self, data):
		octoprint.plugin.SettingsPlugin.on_settings_save(self, data)

		import threading
		thread = threading.Thread(target=self._connect_bullet, args=(self._settings.get(["access_token"]),
		                                                             self._settings.get(["push_channel"])))
		thread.daemon = True
		thread.start()

	def get_settings_defaults(self):
		return dict(
			quiet_minutes = 0,
			access_token=None,
			push_channel=None,
			printDone=dict(
				title="Print job finished",
				body="{file} finished printing in {elapsed_time}"
			),
			printProgress=dict(
				title="Print job {progress}% complete",
				body="{progress}% {path}\n{print_time}\n{remaining} left\n{eta} finish"
			)
		)

	def get_settings_restricted_paths(self):
		# only used in OctoPrint versions > 1.2.16
		return dict(admin=[["access_token"], ["push_channel"]])

	#~~ TemplatePlugin API

	def get_template_configs(self):
		return [
			dict(type="settings", name="Pushbullet", custom_bindings=True)
		]

	#~~ AssetPlugin API

	def get_assets(self):
		return dict(js=["js/octobullet.js"])

	#~~ SimpleApiPlugin

	def get_api_commands(self):
		return dict(test=["token"])

	def on_api_command(self, command, data):
		if not admin_permission.can():
			return flask.make_response("Insufficient rights", 403)

		if not command == "test":
			return

		message = data.get("message", "Testing, 1, 2, 3, 4...")
		token = data["token"]
		channel = data.get("channel", None)

		_, sender = self._create_sender(token, channel)

		result = self._send_message_with_webcam_image("Test from the OctoPrint PushBullet Plugin", message, sender=sender)
		return flask.make_response(flask.jsonify(result=result))

	#~~ EventHandlerPlugin

	def on_event(self, event, payload):

		if event == Events.PRINT_DONE:
			file = os.path.basename(payload["file"])
			elapsed_time_in_seconds = payload["time"]

			import datetime
			import octoprint.util
			elapsed_time = octoprint.util.get_formatted_timedelta(datetime.timedelta(seconds=elapsed_time_in_seconds))

			title = self._settings.get(["printDone", "title"]).format(**locals())
			body = self._settings.get(["printDone", "body"]).format(**locals())
			filename = os.path.splitext(file)[0] + ".jpg"

			self._send_message_with_webcam_image(title, body, filename=filename)

		if event == Events.PRINT_STARTED:
			self._interval_percent = 5
			self._message_count = 0
			self._last_message = 0

	##~~ Softwareupdate hook

	def get_update_information(self):
		return dict(
			octobullet=dict(
				displayName="Pushbullet Plugin",
				displayVersion=self._plugin_version,

				# version check: github repository
				type="github_release",
				user="OctoPrint",
				repo="OctoPrint-Pushbullet",
				current=self._plugin_version,

				# update method: pip
				pip="https://github.com/OctoPrint/OctoPrint-Pushbullet/archive/{target_version}.zip"
			)
		)

	##~~ Internal utility methods

	def _send_message_with_webcam_image(self, title, body, filename=None, sender=None):
		if filename is None:
			import random, string
			filename = "{}.jpg".format(random.choice(string.ascii_letters) * 16)

		if sender is None:
			sender = self._sender

		if not sender:
			return False

		snapshot_url = self._settings.global_get(["webcam", "snapshot"])
		if snapshot_url:
			try:
				import urllib
				snapshot_path, headers = urllib.urlretrieve(snapshot_url)
			except Exception as e:
				self._logger.exception(
					"Exception while fetching snapshot from webcam, sending only a note: {message}".format(
						message=str(e)))
			else:
				# ffmpeg can't guess file type it seems
				os.rename(snapshot_path, snapshot_path + ".jpg")
				snapshot_path += ".jpg"

				# flip or rotate as needed
				self._process_snapshot(snapshot_path)

				if self._send_file(sender, snapshot_path, filename, body):
					return True
				self._logger.warn("Could not send a file message with the webcam image, sending only a note")

		return self._send_note(sender, title, body)

	def _send_note(self, sender, title, body):
		try:
			sender.push_note(title, body)
		except:
			self._logger.exception("Error while pushing a note")
			return False
		return True

	def _send_file(self, sender, path, filename, body):
		try:
			with open(path, "rb") as pic:
				try:
					file_data = self._bullet.upload_file(pic, filename)
				except Exception as e:
					self._logger.exception("Error while uploading snapshot, sending only a note: {}".format(str(e)))
					return False

			sender.push_file(file_data["file_name"], file_data["file_url"], file_data["file_type"], body=body)
			return True
		except Exception as e:
			self._logger.exception("Exception while uploading snapshot to Pushbullet, sending only a note: {message}".format(message=str(e)))
			return False
		finally:
			try:
				os.remove(path)
			except:
				self._logger.exception("Could not remove temporary snapshot file {}".format(path))

	def _create_sender(self, token, channel=None):
		try:
			bullet = pushbullet.PushBullet(token)
			sender = bullet

			# Setup channel object if channel setting is present
			if channel:
				for channel_obj in self._bullet.channels:
					if channel_obj.channel_tag == channel:
						sender = channel_obj
						self._logger.info("Connected to PushBullet on channel {}".format(channel))
						break
				else:
					self._logger.warn("Could not find channel {}, please check your configuration!".format(channel))

			self._logger.info("Connected to PushBullet")
			return bullet, sender
		except:
			self._logger.exception("Error while instantiating PushBullet")
			return None, None

	def _process_snapshot(self, snapshot_path, pixfmt="yuv420p"):
		hflip  = self._settings.global_get_boolean(["webcam", "flipH"])
		vflip  = self._settings.global_get_boolean(["webcam", "flipV"])
		rotate = self._settings.global_get_boolean(["webcam", "rotate90"])
		ffmpeg = self._settings.global_get(["webcam", "ffmpeg"])

		if not ffmpeg or not os.access(ffmpeg, os.X_OK) or (not vflip and not hflip and not rotate):
			return

		ffmpeg_command = [ffmpeg, "-y", "-i", snapshot_path]

		rotate_params = ["format={}".format(pixfmt)] # workaround for foosel/OctoPrint#1317
		if rotate:
			rotate_params.append("transpose=2") # 90 degrees counter clockwise
		if hflip:
			rotate_params.append("hflip") 		# horizontal flip
		if vflip:
			rotate_params.append("vflip")		# vertical flip

		ffmpeg_command += ["-vf", sarge.shell_quote(",".join(rotate_params)), snapshot_path]
		self._logger.info("Running: {}".format(" ".join(ffmpeg_command)))

		p = sarge.run(ffmpeg_command, stdout=sarge.Capture(), stderr=sarge.Capture())
		if p.returncode == 0:
			self._logger.info("Rotated/flipped image with ffmpeg")
		else:
			self._logger.warn("Failed to rotate/flip image with ffmpeg, "
			                  "got return code {}: {}, {}".format(p.returncode,
			                                                      p.stdout.text,
			                                                      p.stderr.text))

__plugin_name__ = "Pushbullet"
def __plugin_load__():
	global __plugin_implementation__
	__plugin_implementation__ = PushbulletPlugin()

	global __plugin_hooks__
	__plugin_hooks__ = {
		"octoprint.plugin.softwareupdate.check_config": __plugin_implementation__.get_update_information
	}

<|endoftext|>


<commit_before># Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *
import re

class Ccache(AutotoolsPackage):
    """ccache is a compiler cache. It speeds up recompilation by caching
    previous compilations and detecting when the same compilation is being done
    again."""

    homepage = "https://ccache.samba.org/"
    url      = "https://github.com/ccache/ccache/releases/download/v3.7.9/ccache-3.7.9.tar.gz"

    executables = ['^ccache$']

    version('3.7.11', sha256='34309a59d4b6b6b33756366aa9d3144a4655587be9f914476b4c0e2d36365f01')
    version('3.7.9', sha256='92838e2133c9e704fdab9ee2608dad86c99021278b9ac47d065aa8ff2ea8ce36')
    version('3.7.1', sha256='e562fcdbe766406b6fe4bf97ce5c001d2be8a17465f33bcddefc9499bbb057d8')
    version('3.3.4', sha256='1348b54e7c35dd2f8d17923389e03c546e599cfbde6459d2f31cf6f1521ec538')
    version('3.3.3', sha256='87a399a2267cfac3f36411fbc12ff8959f408cffd050ad15fe423df88e977e8f')
    version('3.3.2', sha256='bf4a150dea611a206a933e122bd545dd6c5111d319505e0e30fef75f88651847')
    version('3.3.1', sha256='4101f9937cd6e8f50d0a5882f7e9a7312ba42c01ff41e4f359c94ae2c9b87879')
    version('3.3',   sha256='b220fce435fe3d86b8b90097e986a17f6c1f971e0841283dd816adb238c5fd6a')
    version('3.2.9', sha256='1e13961b83a3d215c4013469c149414a79312a22d3c7bf9f946abac9ee33e63f')

    depends_on('gperf')
    depends_on('libxslt')
    depends_on('zlib')

    @classmethod
    def determine_version(cls, exe):
        output = Executable(exe)('--version', output=str, error=str)
        match = re.search(r'ccache.*version\s+(\S+)', output)
        return match.group(1) if match else None
<commit_msg>ccache: add v4.2.1 and earlier, zstd support (#23039)<commit_after># Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *
import re

class Ccache(CMakePackage):
    """ccache is a compiler cache. It speeds up recompilation by caching
    previous compilations and detecting when the same compilation is being done
    again."""

    homepage = "https://ccache.dev/"
    url      = "https://github.com/ccache/ccache/releases/download/v4.2.1/ccache-4.2.1.tar.gz"

    executables = ['^ccache$']

    version('4.2.1', sha256='320d2b17d2f76393e5d4bb28c8dee5ca783248e9cd23dff0654694d60f8a4b62')
    version('4.2',   sha256='dbf139ff32031b54cb47f2d7983269f328df14b5a427882f89f7721e5c411b7e')
    version('4.1',   sha256='cdeefb827b3eef3b42b5454858123881a4a90abbd46cc72cf8c20b3bd039deb7')
    version('4.0',   sha256='ac97af86679028ebc8555c99318352588ff50f515fc3a7f8ed21a8ad367e3d45')
    version('3.7.11', sha256='34309a59d4b6b6b33756366aa9d3144a4655587be9f914476b4c0e2d36365f01')
    version('3.7.9', sha256='92838e2133c9e704fdab9ee2608dad86c99021278b9ac47d065aa8ff2ea8ce36')
    version('3.7.1', sha256='e562fcdbe766406b6fe4bf97ce5c001d2be8a17465f33bcddefc9499bbb057d8')
    version('3.3.4', sha256='1348b54e7c35dd2f8d17923389e03c546e599cfbde6459d2f31cf6f1521ec538')
    version('3.3.3', sha256='87a399a2267cfac3f36411fbc12ff8959f408cffd050ad15fe423df88e977e8f')
    version('3.3.2', sha256='bf4a150dea611a206a933e122bd545dd6c5111d319505e0e30fef75f88651847')
    version('3.3.1', sha256='4101f9937cd6e8f50d0a5882f7e9a7312ba42c01ff41e4f359c94ae2c9b87879')
    version('3.3',   sha256='b220fce435fe3d86b8b90097e986a17f6c1f971e0841283dd816adb238c5fd6a')
    version('3.2.9', sha256='1e13961b83a3d215c4013469c149414a79312a22d3c7bf9f946abac9ee33e63f')

    depends_on('zstd', when='@4.0:')

    depends_on('gperf', when='@:3.99')
    depends_on('libxslt', when='@:3.99')
    depends_on('zlib', when='@:3.99')

    # Before 4.0 this was an Autotools package
    @when('@:3.99')
    def cmake(self, spec, prefix):
        configure_args = ["--prefix=" + prefix]
        configure(*configure_args)

    @when('@:3.99')
    def build(self, spec, prefix):
        make()

    @when('@:3.99')
    def install(self, spec, prefix):
        make("install")

    @classmethod
    def determine_version(cls, exe):
        output = Executable(exe)('--version', output=str, error=str)
        match = re.search(r'ccache.*version\s+(\S+)', output)
        return match.group(1) if match else None
<|endoftext|>


<commit_before># Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class Dd4hep(CMakePackage):
    """DD4hep is a software framework for providing a complete solution for
       full detector description (geometry, materials, visualization, readout,
       alignment, calibration, etc.) for the full experiment life cycle
       (detector concept development, detector optimization, construction,
       operation). It offers a consistent description through a single source
       of detector information for simulation, reconstruction, analysis, etc.
       It distributed under the LGPLv3 License."""

    homepage = "https://dd4hep.web.cern.ch/dd4hep/"
    url      = "https://github.com/AIDASoft/DD4hep/archive/v01-12-01.tar.gz"
    git      = "https://github.com/AIDASoft/DD4hep.git"

    maintainers = ['vvolkl', 'drbenmorgan']

    tags = ['hep']

    version('master', branch='master')
    version('1.17', sha256='e56071ce5497517fe56af813828b1a5fab90f0b602b86961a277ca22be215232')
    version('1.16.1', sha256='c8b1312aa88283986f89cc008d317b3476027fd146fdb586f9f1fbbb47763f1a')
    version('1.16', sha256='ea9755cd255cf1b058e0e3cd743101ca9ca5ff79f4c60be89f9ba72b1ae5ec69')
    version('1.15', sha256='992a24bd4b3dfaffecec9d1c09e8cde2c7f89d38756879a47b23208242f4e352')
    version('1.14.1', sha256='5b5742f1e23c2b36d3174cca95f810ce909c0eb66f3d6d7acb0ba657819e6717')
    version('1.14', sha256='b603aa3c0db8dda392253aa71fa4a0f0c3c9715d47df0b895d45c1e8849f4895')
    version('1.13.1', sha256='83fa70cd74ce93b2f52f098388dff58d179f05ace5b50aea3f408bb8abf7cb73')
    version('1.13', sha256='0b1f9d902ebe21a9178c1e41204c066b29f68c8836fd1d03a9ce979811ddb295')
    version('1.12.1', sha256='85e8c775ec03c499ce10911e228342e757c81ce9ef2a9195cb253b85175a2e93')
    version('1.12', sha256='133a1fb8ce0466d2482f3ebb03e60b3bebb9b2d3e33d14ba15c8fbb91706b398')
    version('1.11.2', sha256='96a53dd26cb8df11c6dae54669fbc9cc3c90dd47c67e07b24be9a1341c95abc4')
    version('1.11.1', sha256='d7902dd7f6744bbda92f6e303ad5a3410eec4a0d2195cdc86f6c1167e72893f0')
    version('1.11', sha256='25643296f15f9d11ad4ad550b7c3b92e8974fc56f1ee8e4455501010789ae7b6')
    version('1.10', sha256='1d6b5d1c368dc8bcedd9c61b7c7e1a44bad427f8bd34932516aff47c88a31d95')

    generator = 'Ninja'

    # Workarounds for various TBB issues in DD4hep v1.11
    # See https://github.com/AIDASoft/DD4hep/pull/613 .
    patch('tbb-workarounds.patch', when='@1.11')
    patch('tbb2.patch', when='@1.12.1')
    # Workaround for failing build file generation in some cases
    # See https://github.com/spack/spack/issues/24232
    patch('cmake_language.patch', when='@:1.17')

    variant('xercesc', default=False, description="Enable 'Detector Builders' based on XercesC")
    variant('geant4', default=False, description="Enable the simulation part based on Geant4")
    variant('assimp', default=False, description="Enable CAD interface based on Assimp")
    variant('hepmc3', default=False, description="Enable build with hepmc3")
    variant('lcio', default=False, description="Enable build with lcio")
    variant('edm4hep', default=True, description="Enable build with edm4hep")
    variant('geant4units', default=False, description="Use geant4 units throughout")
    variant('debug', default=False,
            description="Enable debug build flag - adds extra info in"
            " some places in addtion to the debug build type")

    depends_on('cmake @3.12:', type='build')
    depends_on('ninja', type='build')
    depends_on('boost @1.49:')
    depends_on('root @6.08: +gdml +math +opengl +python +x')
    extends('python')
    depends_on('xerces-c', when='+xercesc')
    depends_on('geant4@10.2.2:', when='+geant4')
    depends_on('assimp@5.0.2:', when='+assimp')
    depends_on('hepmc3', when="+hepmc3")
    depends_on('lcio', when="+lcio")
    depends_on('edm4hep', when="+edm4hep")

    # See https://github.com/AIDASoft/DD4hep/pull/771
    conflicts('^cmake@3.16:3.17.0', when='@1.15',
              msg='cmake version with buggy FindPython breaks dd4hep cmake config')

    def cmake_args(self):
        spec = self.spec
        cxxstd = spec['root'].variants['cxxstd'].value
        # root can be built with cxxstd=11, but dd4hep requires 14
        if cxxstd == "11":
            cxxstd = "14"
        args = [
            self.define_from_variant('DD4HEP_USE_EDM4HEP', 'edm4hep'),
            self.define_from_variant('DD4HEP_USE_XERCESC', 'xercesc'),
            self.define_from_variant('DD4HEP_USE_GEANT4', 'geant4'),
            self.define_from_variant('DD4HEP_USE_LCIO', 'lcio'),
            self.define_from_variant('DD4HEP_USE_HEPMC3', 'hepmc3'),
            self.define_from_variant('DD4HEP_USE_GEANT4_UNITS', 'geant4units'),
            self.define_from_variant('DD4HEP_BUILD_DEBUG', 'debug'),
            # Downloads assimp from github and builds it on the fly.
            # However, with spack it is preferrable to have a proper external
            # dependency, so we disable it.
            self.define('DD4HEP_LOAD_ASSIMP', False),
            "-DCMAKE_CXX_STANDARD={0}".format(cxxstd),
            "-DBUILD_TESTING={0}".format(self.run_tests),
            "-DBOOST_ROOT={0}".format(spec['boost'].prefix),
            "-DBoost_NO_BOOST_CMAKE=ON",
            "-DPYTHON_EXECUTABLE={0}".format(spec['python'].command.path),
        ]
        return args

    def setup_run_environment(self, env):
        # used p.ex. in ddsim to find DDDetectors dir
        env.set("DD4hepINSTALL", self.prefix)
        env.set("DD4HEP", self.prefix.examples)
        env.set("DD4hep_DIR", self.prefix)
        env.set("DD4hep_ROOT", self.prefix)

    def url_for_version(self, version):
        # dd4hep releases are dashes and padded with a leading zero
        # the patch version is omitted when 0
        # so for example v01-12-01, v01-12 ...
        base_url = self.url.rsplit('/', 1)[0]
        if len(version) == 1:
            major = version[0]
            minor, patch = 0, 0
        elif len(version) == 2:
            major, minor = version
            patch = 0
        else:
            major, minor, patch = version
        # By now the data is normalized enough to handle it easily depending
        # on the value of the patch version
        if patch == 0:
            version_str = 'v%02d-%02d.tar.gz' % (major, minor)
        else:
            version_str = 'v%02d-%02d-%02d.tar.gz' % (major, minor, patch)

        return base_url + '/' + version_str
<commit_msg>dd4hep: Updated version checksum due to updated tag (#24469)<commit_after># Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class Dd4hep(CMakePackage):
    """DD4hep is a software framework for providing a complete solution for
       full detector description (geometry, materials, visualization, readout,
       alignment, calibration, etc.) for the full experiment life cycle
       (detector concept development, detector optimization, construction,
       operation). It offers a consistent description through a single source
       of detector information for simulation, reconstruction, analysis, etc.
       It distributed under the LGPLv3 License."""

    homepage = "https://dd4hep.web.cern.ch/dd4hep/"
    url      = "https://github.com/AIDASoft/DD4hep/archive/v01-12-01.tar.gz"
    git      = "https://github.com/AIDASoft/DD4hep.git"

    maintainers = ['vvolkl', 'drbenmorgan']

    tags = ['hep']

    version('master', branch='master')
    version('1.17', sha256='036a9908aaf1e13eaf5f2f43b6f5f4a8bdda8183ddc5befa77a4448dbb485826')
    version('1.16.1', sha256='c8b1312aa88283986f89cc008d317b3476027fd146fdb586f9f1fbbb47763f1a')
    version('1.16', sha256='ea9755cd255cf1b058e0e3cd743101ca9ca5ff79f4c60be89f9ba72b1ae5ec69')
    version('1.15', sha256='992a24bd4b3dfaffecec9d1c09e8cde2c7f89d38756879a47b23208242f4e352')
    version('1.14.1', sha256='5b5742f1e23c2b36d3174cca95f810ce909c0eb66f3d6d7acb0ba657819e6717')
    version('1.14', sha256='b603aa3c0db8dda392253aa71fa4a0f0c3c9715d47df0b895d45c1e8849f4895')
    version('1.13.1', sha256='83fa70cd74ce93b2f52f098388dff58d179f05ace5b50aea3f408bb8abf7cb73')
    version('1.13', sha256='0b1f9d902ebe21a9178c1e41204c066b29f68c8836fd1d03a9ce979811ddb295')
    version('1.12.1', sha256='85e8c775ec03c499ce10911e228342e757c81ce9ef2a9195cb253b85175a2e93')
    version('1.12', sha256='133a1fb8ce0466d2482f3ebb03e60b3bebb9b2d3e33d14ba15c8fbb91706b398')
    version('1.11.2', sha256='96a53dd26cb8df11c6dae54669fbc9cc3c90dd47c67e07b24be9a1341c95abc4')
    version('1.11.1', sha256='d7902dd7f6744bbda92f6e303ad5a3410eec4a0d2195cdc86f6c1167e72893f0')
    version('1.11', sha256='25643296f15f9d11ad4ad550b7c3b92e8974fc56f1ee8e4455501010789ae7b6')
    version('1.10', sha256='1d6b5d1c368dc8bcedd9c61b7c7e1a44bad427f8bd34932516aff47c88a31d95')

    generator = 'Ninja'

    # Workarounds for various TBB issues in DD4hep v1.11
    # See https://github.com/AIDASoft/DD4hep/pull/613 .
    patch('tbb-workarounds.patch', when='@1.11')
    patch('tbb2.patch', when='@1.12.1')
    # Workaround for failing build file generation in some cases
    # See https://github.com/spack/spack/issues/24232
    patch('cmake_language.patch', when='@:1.17')

    variant('xercesc', default=False, description="Enable 'Detector Builders' based on XercesC")
    variant('geant4', default=False, description="Enable the simulation part based on Geant4")
    variant('assimp', default=False, description="Enable CAD interface based on Assimp")
    variant('hepmc3', default=False, description="Enable build with hepmc3")
    variant('lcio', default=False, description="Enable build with lcio")
    variant('edm4hep', default=True, description="Enable build with edm4hep")
    variant('geant4units', default=False, description="Use geant4 units throughout")
    variant('debug', default=False,
            description="Enable debug build flag - adds extra info in"
            " some places in addtion to the debug build type")

    depends_on('cmake @3.12:', type='build')
    depends_on('ninja', type='build')
    depends_on('boost @1.49:')
    depends_on('root @6.08: +gdml +math +opengl +python +x')
    extends('python')
    depends_on('xerces-c', when='+xercesc')
    depends_on('geant4@10.2.2:', when='+geant4')
    depends_on('assimp@5.0.2:', when='+assimp')
    depends_on('hepmc3', when="+hepmc3")
    depends_on('lcio', when="+lcio")
    depends_on('edm4hep', when="+edm4hep")

    # See https://github.com/AIDASoft/DD4hep/pull/771
    conflicts('^cmake@3.16:3.17.0', when='@1.15',
              msg='cmake version with buggy FindPython breaks dd4hep cmake config')

    def cmake_args(self):
        spec = self.spec
        cxxstd = spec['root'].variants['cxxstd'].value
        # root can be built with cxxstd=11, but dd4hep requires 14
        if cxxstd == "11":
            cxxstd = "14"
        args = [
            self.define_from_variant('DD4HEP_USE_EDM4HEP', 'edm4hep'),
            self.define_from_variant('DD4HEP_USE_XERCESC', 'xercesc'),
            self.define_from_variant('DD4HEP_USE_GEANT4', 'geant4'),
            self.define_from_variant('DD4HEP_USE_LCIO', 'lcio'),
            self.define_from_variant('DD4HEP_USE_HEPMC3', 'hepmc3'),
            self.define_from_variant('DD4HEP_USE_GEANT4_UNITS', 'geant4units'),
            self.define_from_variant('DD4HEP_BUILD_DEBUG', 'debug'),
            # Downloads assimp from github and builds it on the fly.
            # However, with spack it is preferrable to have a proper external
            # dependency, so we disable it.
            self.define('DD4HEP_LOAD_ASSIMP', False),
            "-DCMAKE_CXX_STANDARD={0}".format(cxxstd),
            "-DBUILD_TESTING={0}".format(self.run_tests),
            "-DBOOST_ROOT={0}".format(spec['boost'].prefix),
            "-DBoost_NO_BOOST_CMAKE=ON",
            "-DPYTHON_EXECUTABLE={0}".format(spec['python'].command.path),
        ]
        return args

    def setup_run_environment(self, env):
        # used p.ex. in ddsim to find DDDetectors dir
        env.set("DD4hepINSTALL", self.prefix)
        env.set("DD4HEP", self.prefix.examples)
        env.set("DD4hep_DIR", self.prefix)
        env.set("DD4hep_ROOT", self.prefix)

    def url_for_version(self, version):
        # dd4hep releases are dashes and padded with a leading zero
        # the patch version is omitted when 0
        # so for example v01-12-01, v01-12 ...
        base_url = self.url.rsplit('/', 1)[0]
        if len(version) == 1:
            major = version[0]
            minor, patch = 0, 0
        elif len(version) == 2:
            major, minor = version
            patch = 0
        else:
            major, minor, patch = version
        # By now the data is normalized enough to handle it easily depending
        # on the value of the patch version
        if patch == 0:
            version_str = 'v%02d-%02d.tar.gz' % (major, minor)
        else:
            version_str = 'v%02d-%02d-%02d.tar.gz' % (major, minor, patch)

        return base_url + '/' + version_str
<|endoftext|>


<commit_before># Copyright 2013-2019 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class Eospac(Package):
    """A collection of C routines that can be used to access the Sesame data
       library.
    """

    homepage = "https://laws.lanl.gov/projects/data/eos.html"
    list_url = "https://laws.lanl.gov/projects/data/eos/eospacReleases.php"

    version('6.4.0', 'ecee3abe9c7a3e4b32d68dc3177b088e',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0_612ea8c9b8ffa6d9175d9118955571d9107f1e3c.tgz")
    version('6.4.0beta.4', 'd0875390d08538a3caa8ad1f43c6d789',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.4_aff6429bb6868de31a980278bafa13487c2ce83f.tgz")
    version('6.4.0beta.3', '672c16e522908281168815a74626ec66',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.3_90ff265f62aa1780bfcd0a62dad807b6be6ed461.tgz")
    version('6.4.0beta.2', '9b6e48090647221d5ffe7ec5f9ea4c71',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.2_69196eadbc77506561eef711f19d2f03b4ab0ffa.tgz")
    version('6.4.0beta.1', 'e4e4beabf946f0b8953532832002afc2',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.1_r20171213193219.tgz")
    version('6.3.1', '549fda008c4169a69b02ec2a9de1e434', preferred=True,
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.3.1_r20161202150449.tgz")

    # This patch allows the use of spack's compile wrapper 'flang'
    patch('flang.patch', when='@:6.4.0beta.2%clang')

    def install(self, spec, prefix):
        with working_dir('Source'):
            make('install',
                 'CC={0}'.format(spack_cc),
                 'CXX={0}'.format(spack_cxx),
                 'F77={0}'.format(spack_f77),
                 'F90={0}'.format(spack_fc),
                 'prefix={0}'.format(prefix),
                 'INSTALLED_LIBRARY_DIR={0}'.format(prefix.lib),
                 'INSTALLED_INCLUDE_DIR={0}'.format(prefix.include),
                 'INSTALLED_EXAMPLE_DIR={0}'.format(prefix.example),
                 'INSTALLED_BIN_DIR={0}'.format(prefix.bin))
<commit_msg>eospac package: replace md5 checksums with sha256 (#11757)<commit_after># Copyright 2013-2019 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class Eospac(Package):
    """A collection of C routines that can be used to access the Sesame data
       library.
    """

    homepage = "https://laws.lanl.gov/projects/data/eos.html"
    list_url = "https://laws.lanl.gov/projects/data/eos/eospacReleases.php"

    version('6.4.0',       sha256='15a953beac735c68431afe86ffe33323d540d0fbbbec03ba79438dd29736051d',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0_612ea8c9b8ffa6d9175d9118955571d9107f1e3c.tgz")
    version('6.4.0beta.4', sha256='0ebfd8badff575ea77444aa978629dbdca3135a0b5eb373b8daba058773d4635',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.4_aff6429bb6868de31a980278bafa13487c2ce83f.tgz")
    version('6.4.0beta.3', sha256='9f387ca5356519494c6f3f27adb0c165cf9f9e15e3355a67bf940a4a92eebdab',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.3_90ff265f62aa1780bfcd0a62dad807b6be6ed461.tgz")
    version('6.4.0beta.2', sha256='f9db46cd6c62a6f83960d802350f3e37675921af102969b293c02eb797558a53',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.2_69196eadbc77506561eef711f19d2f03b4ab0ffa.tgz")
    version('6.4.0beta.1', sha256='14c5c804e5f628f41e8ed80bcee5a80adeb6c6f3d130715421ca99a30c7eb7e2',
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.4.0beta.1_r20171213193219.tgz")
    version('6.3.1',       sha256='aa1112c4251c9c3c2883a7ab2c7f2abff2c339f29dbbf8421ef88b0c9df904f8', preferred=True,
            url="https://laws.lanl.gov/projects/data/eos/get_file.php?package=eospac&filename=eospac_v6.3.1_r20161202150449.tgz")

    # This patch allows the use of spack's compile wrapper 'flang'
    patch('flang.patch', when='@:6.4.0beta.2%clang')

    def install(self, spec, prefix):
        with working_dir('Source'):
            make('install',
                 'CC={0}'.format(spack_cc),
                 'CXX={0}'.format(spack_cxx),
                 'F77={0}'.format(spack_f77),
                 'F90={0}'.format(spack_fc),
                 'prefix={0}'.format(prefix),
                 'INSTALLED_LIBRARY_DIR={0}'.format(prefix.lib),
                 'INSTALLED_INCLUDE_DIR={0}'.format(prefix.include),
                 'INSTALLED_EXAMPLE_DIR={0}'.format(prefix.example),
                 'INSTALLED_BIN_DIR={0}'.format(prefix.bin))
<|endoftext|>


<commit_before>##############################################################################
# Copyright (c) 2013-2016, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Created by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://github.com/llnl/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License (as
# published by the Free Software Foundation) version 2.1, February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
from spack import *

class Fenics(Package):
    """FEniCS is organized as a collection of interoperable components
    that together form the FEniCS Project. These components include
    the problem-solving environment DOLFIN, the form compiler FFC, the
    finite element tabulator FIAT, the just-in-time compiler Instant,
    the code generation interface UFC, the form language UFL and a
    range of additional components."""

    homepage = "http://fenicsproject.org/"
    url      = "https://bitbucket.org/fenics-project/dolfin/downloads/dolfin-1.6.0.tar.gz"

    base_url = "https://bitbucket.org/fenics-project/{pkg}/downloads/{pkg}-{version}.tar.gz"  # NOQA: ignore E501

    variant('hdf5',         default=True,  description='Compile with HDF5')
    variant('parmetis',     default=True,  description='Compile with ParMETIS')
    variant('scotch',       default=True,  description='Compile with Scotch')
    variant('petsc',        default=True,  description='Compile with PETSc')
    variant('slepc',        default=True,  description='Compile with SLEPc')
    variant('trilinos',     default=True,  description='Compile with Trilinos')
    variant('suite-sparse', default=True,  description='Compile with SuiteSparse solvers')
    variant('vtk',          default=False, description='Compile with VTK')
    variant('qt',           default=False, description='Compile with QT')
    variant('mpi',          default=True,  description='Enables the distributed memory support')
    variant('openmp',       default=True,  description='Enables the shared memory support')
    variant('shared',       default=True,  description='Enables the build of shared libraries')
    variant('debug',        default=False, description='Builds a debug version of the libraries')

    # not part of spack list for now
    # variant('petsc4py',     default=True,  description='Uses PETSc4py')
    # variant('slepc4py',     default=True,  description='Uses SLEPc4py')
    # variant('pastix',       default=True,  description='Compile with Pastix')

    extends('python')

    depends_on('py-numpy')
    depends_on('py-ply')
    depends_on('py-six')
    depends_on('py-sphinx@1.0.1:', when='+doc')
    depends_on('eigen@3.2.0:')
    depends_on('boost')
    depends_on('mpi', when='+mpi')
    depends_on('hdf5', when='+hdf5')
    depends_on('parmetis@4.0.2:^metis+real64', when='+parmetis')
    depends_on('scotch~metis', when='+scotch~mpi')
    depends_on('scotch+mpi~metis', when='+scotch+mpi')
    depends_on('petsc@3.4:', when='+petsc')
    depends_on('slepc@3.4:', when='+slepc')
    depends_on('trilinos', when='+trilinos')
    depends_on('vtk', when='+vtk')
    depends_on('suite-sparse', when='+suite-sparse')
    depends_on('qt', when='+qt')

    # This are the build dependencies
    depends_on('py-setuptools')
    depends_on('cmake@2.8.12:')
    depends_on('swig@3.0.3:')

    releases = [
        {
            'version': '1.6.0',
            'md5': '35cb4baf7ab4152a40fb7310b34d5800',
            'resources': {
                'ffc': '358faa3e9da62a1b1a717070217b793e',
                'fiat': 'f4509d05c911fd93cea8d288a78a6c6f',
                'instant': '5f2522eb032a5bebbad6597b6fe0732a',
                'ufl': 'c40c5f04eaa847377ab2323122284016',
            }
        },
        {
            'version': '1.5.0',
            'md5': '9b589a3534299a5e6d22c13c5eb30bb8',
            'resources': {
                'ffc': '343f6d30e7e77d329a400fd8e73e0b63',
                'fiat': 'da3fa4dd8177bb251e7f68ec9c7cf6c5',
                'instant': 'b744023ded27ee9df4a8d8c6698c0d58',
                'ufl': '130d7829cf5a4bd5b52bf6d0955116fd',
            }
        },
    ]

    for release in releases:
        version(release['version'], release['md5'], url=base_url.format(pkg='dolfin', version=release['version']))
        for name, md5 in release['resources'].items():
            resource(name=name,
                     url=base_url.format(pkg=name, **release),
                     md5=md5,
                     destination='depends',
                     when='@{version}'.format(**release),
                     placement=name)

    def cmake_is_on(self, option):
        return 'ON' if option in self.spec else 'OFF'

    def install(self, spec, prefix):
        for package in ['ufl', 'ffc', 'fiat', 'instant']:
            with working_dir(join_path('depends', package)):
                python('setup.py', 'install', '--prefix=%s' % prefix)

        cmake_args.extend([
            '-DCMAKE_BUILD_TYPE:STRING={0}'.format(
                'Debug' if '+debug' in spec else 'RelWithDebInfo'),
            '-DBUILD_SHARED_LIBS:BOOL={0}'.format(
                self.cmake_is_on('+shared')),
            '-DDOLFIN_SKIP_BUILD_TESTS:BOOL=ON',
            '-DDOLFIN_ENABLE_OPENMP:BOOL={0}'.format(
                self.cmake_is_on('+openmp')),
            '-DDOLFIN_ENABLE_CHOLMOD:BOOL={0}'.format(
                self.cmake_is_on('suite-sparse')),
            '-DDOLFIN_ENABLE_HDF5:BOOL={0}'.format(
                self.cmake_is_on('hdf5')),
            '-DDOLFIN_ENABLE_MPI:BOOL={0}'.format(
                self.cmake_is_on('mpi')),
            '-DDOLFIN_ENABLE_PARMETIS:BOOL={0}'.format(
                self.cmake_is_on('parmetis')),
            '-DDOLFIN_ENABLE_PASTIX:BOOL={0}'.format(
                self.cmake_is_on('pastix')),
            '-DDOLFIN_ENABLE_PETSC:BOOL={0}'.format(
                self.cmake_is_on('petsc')),
            '-DDOLFIN_ENABLE_PETSC4PY:BOOL={0}'.format(
                self.cmake_is_on('py-petsc4py')),
            '-DDOLFIN_ENABLE_PYTHON:BOOL={0}'.format(
                self.cmake_is_on('python')),
            '-DDOLFIN_ENABLE_QT:BOOL={0}'.format(
                self.cmake_is_on('qt')),
            '-DDOLFIN_ENABLE_SCOTCH:BOOL={0}'.format(
                self.cmake_is_on('scotch')),
            '-DDOLFIN_ENABLE_SLEPC:BOOL={0}'.format(
                self.cmake_is_on('slepc')),
            '-DDOLFIN_ENABLE_SLEPC4PY:BOOL={0}'.format(
                self.cmake_is_on('py-slepc4py')),
            '-DDOLFIN_ENABLE_SPHINX:BOOL={0}'.format(
                self.cmake_is_on('py-sphinx')),
            '-DDOLFIN_ENABLE_TRILINOS:BOOL={0}'.format(
                self.cmake_is_on('trilinos')),
            '-DDOLFIN_ENABLE_UMFPACK:BOOL={0}'.format(
                self.cmake_is_on('suite-sparse')),
            '-DDOLFIN_ENABLE_VTK:BOOL={0}'.format(
                self.cmake_is_on('vtk')),
            '-DDOLFIN_ENABLE_ZLIB:BOOL={0}'.format(
                self.cmake_is_on('zlib')),
        ])
        cmake_args.extend(std_cmake_args)

        with working_dir('build', create=True):
            cmake('..', *cmake_args)

            make()
            make('install')
<commit_msg>Introduced a bug by correcting flake8 warnings<commit_after>##############################################################################
# Copyright (c) 2013-2016, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Created by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://github.com/llnl/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License (as
# published by the Free Software Foundation) version 2.1, February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
from spack import *

class Fenics(Package):
    """FEniCS is organized as a collection of interoperable components
    that together form the FEniCS Project. These components include
    the problem-solving environment DOLFIN, the form compiler FFC, the
    finite element tabulator FIAT, the just-in-time compiler Instant,
    the code generation interface UFC, the form language UFL and a
    range of additional components."""

    homepage = "http://fenicsproject.org/"
    url      = "https://bitbucket.org/fenics-project/dolfin/downloads/dolfin-1.6.0.tar.gz"

    base_url = "https://bitbucket.org/fenics-project/{pkg}/downloads/{pkg}-{version}.tar.gz"  # NOQA: ignore E501

    variant('hdf5',         default=True,  description='Compile with HDF5')
    variant('parmetis',     default=True,  description='Compile with ParMETIS')
    variant('scotch',       default=True,  description='Compile with Scotch')
    variant('petsc',        default=True,  description='Compile with PETSc')
    variant('slepc',        default=True,  description='Compile with SLEPc')
    variant('trilinos',     default=True,  description='Compile with Trilinos')
    variant('suite-sparse', default=True,  description='Compile with SuiteSparse solvers')
    variant('vtk',          default=False, description='Compile with VTK')
    variant('qt',           default=False, description='Compile with QT')
    variant('mpi',          default=True,  description='Enables the distributed memory support')
    variant('openmp',       default=True,  description='Enables the shared memory support')
    variant('shared',       default=True,  description='Enables the build of shared libraries')
    variant('debug',        default=False, description='Builds a debug version of the libraries')

    # not part of spack list for now
    # variant('petsc4py',     default=True,  description='Uses PETSc4py')
    # variant('slepc4py',     default=True,  description='Uses SLEPc4py')
    # variant('pastix',       default=True,  description='Compile with Pastix')

    extends('python')

    depends_on('py-numpy')
    depends_on('py-ply')
    depends_on('py-six')
    depends_on('py-sphinx@1.0.1:', when='+doc')
    depends_on('eigen@3.2.0:')
    depends_on('boost')
    depends_on('mpi', when='+mpi')
    depends_on('hdf5', when='+hdf5')
    depends_on('parmetis@4.0.2:^metis+real64', when='+parmetis')
    depends_on('scotch~metis', when='+scotch~mpi')
    depends_on('scotch+mpi~metis', when='+scotch+mpi')
    depends_on('petsc@3.4:', when='+petsc')
    depends_on('slepc@3.4:', when='+slepc')
    depends_on('trilinos', when='+trilinos')
    depends_on('vtk', when='+vtk')
    depends_on('suite-sparse', when='+suite-sparse')
    depends_on('qt', when='+qt')

    # This are the build dependencies
    depends_on('py-setuptools')
    depends_on('cmake@2.8.12:')
    depends_on('swig@3.0.3:')

    releases = [
        {
            'version': '1.6.0',
            'md5': '35cb4baf7ab4152a40fb7310b34d5800',
            'resources': {
                'ffc': '358faa3e9da62a1b1a717070217b793e',
                'fiat': 'f4509d05c911fd93cea8d288a78a6c6f',
                'instant': '5f2522eb032a5bebbad6597b6fe0732a',
                'ufl': 'c40c5f04eaa847377ab2323122284016',
            }
        },
        {
            'version': '1.5.0',
            'md5': '9b589a3534299a5e6d22c13c5eb30bb8',
            'resources': {
                'ffc': '343f6d30e7e77d329a400fd8e73e0b63',
                'fiat': 'da3fa4dd8177bb251e7f68ec9c7cf6c5',
                'instant': 'b744023ded27ee9df4a8d8c6698c0d58',
                'ufl': '130d7829cf5a4bd5b52bf6d0955116fd',
            }
        },
    ]

    for release in releases:
        version(release['version'], release['md5'], url=base_url.format(pkg='dolfin', version=release['version']))
        for name, md5 in release['resources'].items():
            resource(name=name,
                     url=base_url.format(pkg=name, **release),
                     md5=md5,
                     destination='depends',
                     when='@{version}'.format(**release),
                     placement=name)

    def cmake_is_on(self, option):
        return 'ON' if option in self.spec else 'OFF'

    def install(self, spec, prefix):
        for package in ['ufl', 'ffc', 'fiat', 'instant']:
            with working_dir(join_path('depends', package)):
                python('setup.py', 'install', '--prefix=%s' % prefix)

        cmake_args = [
            '-DCMAKE_BUILD_TYPE:STRING={0}'.format(
                'Debug' if '+debug' in spec else 'RelWithDebInfo'),
            '-DBUILD_SHARED_LIBS:BOOL={0}'.format(
                self.cmake_is_on('+shared')),
            '-DDOLFIN_SKIP_BUILD_TESTS:BOOL=ON',
            '-DDOLFIN_ENABLE_OPENMP:BOOL={0}'.format(
                self.cmake_is_on('+openmp')),
            '-DDOLFIN_ENABLE_CHOLMOD:BOOL={0}'.format(
                self.cmake_is_on('suite-sparse')),
            '-DDOLFIN_ENABLE_HDF5:BOOL={0}'.format(
                self.cmake_is_on('hdf5')),
            '-DDOLFIN_ENABLE_MPI:BOOL={0}'.format(
                self.cmake_is_on('mpi')),
            '-DDOLFIN_ENABLE_PARMETIS:BOOL={0}'.format(
                self.cmake_is_on('parmetis')),
            '-DDOLFIN_ENABLE_PASTIX:BOOL={0}'.format(
                self.cmake_is_on('pastix')),
            '-DDOLFIN_ENABLE_PETSC:BOOL={0}'.format(
                self.cmake_is_on('petsc')),
            '-DDOLFIN_ENABLE_PETSC4PY:BOOL={0}'.format(
                self.cmake_is_on('py-petsc4py')),
            '-DDOLFIN_ENABLE_PYTHON:BOOL={0}'.format(
                self.cmake_is_on('python')),
            '-DDOLFIN_ENABLE_QT:BOOL={0}'.format(
                self.cmake_is_on('qt')),
            '-DDOLFIN_ENABLE_SCOTCH:BOOL={0}'.format(
                self.cmake_is_on('scotch')),
            '-DDOLFIN_ENABLE_SLEPC:BOOL={0}'.format(
                self.cmake_is_on('slepc')),
            '-DDOLFIN_ENABLE_SLEPC4PY:BOOL={0}'.format(
                self.cmake_is_on('py-slepc4py')),
            '-DDOLFIN_ENABLE_SPHINX:BOOL={0}'.format(
                self.cmake_is_on('py-sphinx')),
            '-DDOLFIN_ENABLE_TRILINOS:BOOL={0}'.format(
                self.cmake_is_on('trilinos')),
            '-DDOLFIN_ENABLE_UMFPACK:BOOL={0}'.format(
                self.cmake_is_on('suite-sparse')),
            '-DDOLFIN_ENABLE_VTK:BOOL={0}'.format(
                self.cmake_is_on('vtk')),
            '-DDOLFIN_ENABLE_ZLIB:BOOL={0}'.format(
                self.cmake_is_on('zlib')),
        ]

        cmake_args.extend(std_cmake_args)

        with working_dir('build', create=True):
            cmake('..', *cmake_args)

            make()
            make('install')
<|endoftext|>


<commit_before>"""Utilities for submitting and starting Work Queue workers.

"""
import argparse
import os
import re
import subprocess

from shadho.configuration import ShadhoConfig

def parse_args(args=None):
    p = argparse.ArgumentParser(
        description='Start a Work Queue worker and connect to SHADHO.')

    p.add_argument('-M', '--master', type=str,
        help='name of the Work Queue master to connect to')
    p.add_argument('-u', '--user', type=str, default=os.environ['USER'],
        help='name of the user running the Work Queue master')
    p.add_argument('--timeout', type=int,
        help='amount of time worker idles before exiting')
    p.add_argument('--cores', type=int, default=1,
        help='the number of cores for the worker to use;' +
             ' pass 0 to use all available cores')
    p.add_argument('--feature', type=str, nargs='*', default=[],
        help='user specified feature to advertise, e.g. GPU model name')

    return p.parse_args(args)

def shadho_wq_worker(args=None, config=None):
    """Start a Work Queue worker."""

    if config is None:
        config = ShadhoConfig()

    if args is None:
        cmd_args = ''
    else:
        cmd_args = f'-M {args.master} --cores {args.cores}'
        for feature in args.feature:
            cmd_args += f' --feature {feature}'

    if not re.search(r'(^|[\s])-M([\s]|$)', cmd_args):
        cmd_args = ' '.join([cmd_args, '-M', config.workqueue.name]).strip()

    if not re.search(r'[\s]*-M[\s][\S]*' + args.user + r'.*[\s]*', cmd_args):
        print('Replacing')
        cmd_args = re.sub(r'(^|[\s]*)(.*-M[\s])([\S]+)([\s]*.*$)',
                          r'\1\2\3-' + args.user + r'\4',
                          cmd_args)

    executable = os.path.join(config.shadho_dir, 'bin', 'work_queue_worker')

    print(cmd_args)
    subprocess.run([executable] + cmd_args.split(), stderr=subprocess.STDOUT)

def main():
    args = parse_args()
    shadho_wq_worker(args=args)

if __name__ == '__main__':
    main()
<commit_msg>don't stringify a NoneType master name<commit_after>"""Utilities for submitting and starting Work Queue workers.

"""
import argparse
import os
import re
import subprocess

from shadho.configuration import ShadhoConfig

def parse_args(args=None):
    p = argparse.ArgumentParser(
        description='Start a Work Queue worker and connect to SHADHO.')

    p.add_argument('-M', '--master', type=str, default='',
        help='name of the Work Queue master to connect to')
    p.add_argument('-u', '--user', type=str, default=os.environ['USER'],
        help='name of the user running the Work Queue master')
    p.add_argument('--timeout', type=int,
        help='amount of time worker idles before exiting')
    p.add_argument('--cores', type=int, default=1,
        help='the number of cores for the worker to use;' +
             ' pass 0 to use all available cores')
    p.add_argument('--feature', type=str, nargs='*', default=[],
        help='user specified feature to advertise, e.g. GPU model name')

    return p.parse_args(args)

def shadho_wq_worker(args=None, config=None):
    """Start a Work Queue worker."""

    if config is None:
        config = ShadhoConfig()

    if args is None:
        cmd_args = ''
    else:
        cmd_args = f'{-M if args.master else ""} {args.master} --cores {args.cores}'
        for feature in args.feature:
            cmd_args += f' --feature {feature}'

    if not re.search(r'(^|[\s])-M([\s]|$)', cmd_args):
        cmd_args = ' '.join([cmd_args, '-M', config.workqueue.name]).strip()

    if not re.search(r'[\s]*-M[\s][\S]*' + args.user + r'.*[\s]*', cmd_args):
        print('Replacing')
        cmd_args = re.sub(r'(^|[\s]*)(.*-M[\s])([\S]+)([\s]*.*$)',
                          r'\1\2\3-' + args.user + r'\4',
                          cmd_args)

    executable = os.path.join(config.shadho_dir, 'bin', 'work_queue_worker')

    print(cmd_args)
    subprocess.run([executable] + cmd_args.split(), stderr=subprocess.STDOUT)

def main():
    args = parse_args()
    shadho_wq_worker(args=args)

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