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

"""Optimization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports

import numpy as np

from tensor2tensor.utils import yellowfin

import tensorflow as tf

from tensorflow.python.framework import dtypes

def optimize(loss, learning_rate, hparams, use_tpu=False):
  """Minimize loss."""
  loss = weight_decay_and_noise(loss, hparams, learning_rate)
  loss = tf.identity(loss, name="total_loss")
  log_variable_sizes()
  diet_vars = [
      v for v in tf.global_variables() if v.dtype == dtypes.float16_ref
  ]
  log_variable_sizes(diet_vars, "Diet Variables")
  opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams)
  if use_tpu:
    opt = tf.contrib.tpu.CrossShardOptimizer(opt)

  tf.summary.scalar("learning_rate", learning_rate)
  opt_summaries = ["loss"]
  if hparams.summarize_grads:
    opt_summaries.extend(["gradients", "gradient_norm", "global_gradient_norm"])

  train_op = tf.contrib.layers.optimize_loss(
      name="training",
      loss=loss,
      global_step=tf.train.get_or_create_global_step(),
      learning_rate=learning_rate,
      clip_gradients=hparams.clip_grad_norm or None,
      gradient_noise_scale=hparams.grad_noise_scale or None,
      optimizer=opt,
      summaries=opt_summaries,
      colocate_gradients_with_ops=True)
  return train_op

class ConditionalOptimizer(tf.train.Optimizer):
  """Conditional optimizer."""

  def __init__(self, optimizer_name, lr, hparams):
    if optimizer_name == "Adam":
      # We change the default epsilon for Adam and re-scale lr.
      # Using LazyAdam as it's much faster for large vocabulary embeddings.
      self._opt = tf.contrib.opt.LazyAdamOptimizer(
          lr / 500.0,
          beta1=hparams.optimizer_adam_beta1,
          beta2=hparams.optimizer_adam_beta2,
          epsilon=hparams.optimizer_adam_epsilon)
    elif optimizer_name == "Momentum":
      self._opt = tf.train.MomentumOptimizer(
          lr, momentum=hparams.optimizer_momentum_momentum)
    elif optimizer_name == "YellowFin":
      tf.logging.info("Init YellowFin Optimizer.")
      self._opt = yellowfin.YellowFinOptimizer(
          learning_rate=lr, momentum=hparams.optimizer_momentum_momentum)
    elif optimizer_name == "TrueAdam":
      self._opt = tf.train.AdamOptimizer(
          lr / 500.0,
          beta1=hparams.optimizer_adam_beta1,
          beta2=hparams.optimizer_adam_beta2,
          epsilon=hparams.optimizer_adam_epsilon)
    else:
      self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)

  def compute_gradients(self, loss, var_list=None, **kwargs):
    return self._opt.compute_gradients(loss, var_list, **kwargs)

  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    return self._opt.apply_gradients(
        grads_and_vars, global_step=global_step, name=name)

def _sqrt_decay(step):
  """Decay like 1 / sqrt(step), multiplied by 500 to normalize."""
  return 500.0 / tf.sqrt(tf.maximum(step, 1.0))

def _exp_decay_after(step, rate, from_which_step):
  """Decay exponentially by rate (per step) starting at from_which_step."""
  return tf.cond(
      step < from_which_step,
      lambda: tf.constant(1.0),
      lambda: rate**(step - from_which_step),
      name="exponential_decay_step_cond")

def learning_rate_decay(hparams, num_worker_replicas=1, num_train_steps=1):
  """Inverse-decay learning rate until warmup_steps, then decay."""
  warmup_steps = tf.to_float(
      hparams.learning_rate_warmup_steps * num_worker_replicas)
  step = tf.to_float(tf.train.get_or_create_global_step())
  if hparams.learning_rate_decay_scheme == "noam":
    return 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
        (step + 1) * warmup_steps**-1.5, (step + 1)**-0.5)
  elif hparams.learning_rate_decay_scheme == "exp100k":
    return 0.94**(step // 100000)
  elif hparams.learning_rate_decay_scheme == "cosine":
    cycle_steps = hparams.learning_rate_cosine_cycle_steps
    return 0.5 * (1 + tf.cos(np.pi * (step % cycle_steps) / cycle_steps))
  elif hparams.learning_rate_decay_scheme == "cyclelinear10x":
    # Cycle the rate linearly by 10x every warmup_steps, up and down.
    cycle_steps = hparams.learning_rate_warmup_steps
    cycle_position = step % (2 * cycle_steps)
    cycle_position = tf.to_float(  # Normalize to the interval [-1, 1].
        cycle_position - cycle_steps) / float(cycle_steps)
    cycle_position = 1.0 - tf.abs(cycle_position)  # 0 to 1 and back to 0.
    return (cycle_position + 0.1) * 3.0  # 10x difference each cycle (0.3-3).

  inv_base = tf.exp(tf.log(0.01) / warmup_steps)
  inv_decay = inv_base**(warmup_steps - step)
  if hparams.learning_rate_decay_scheme == "sqrt":
    decay = _sqrt_decay(step - warmup_steps)
  elif hparams.learning_rate_decay_scheme == "exp10k":
    decay = _exp_decay_after(step - warmup_steps, 0.9995,
                             num_train_steps - warmup_steps - 10000)
  elif hparams.learning_rate_decay_scheme == "exp50k":
    decay = _exp_decay_after(step - warmup_steps, 0.99995,
                             num_train_steps - warmup_steps - 50000)
  elif hparams.learning_rate_decay_scheme == "exp500k":
    decay = _exp_decay_after(step - warmup_steps, 0.9999955,
                             num_train_steps - warmup_steps - 500000)
  elif hparams.learning_rate_decay_scheme == "none":
    decay = tf.constant(1.0)
  else:
    raise ValueError("Unrecognized learning rate decay scheme: %s" %
                     hparams.learning_rate_decay_scheme)
  return tf.where(step < warmup_steps, inv_decay, decay)

def weight_decay_and_noise(loss, hparams, learning_rate, var_list=None):
  """Apply weight decay and weight noise."""
  if var_list is None:
    var_list = tf.trainable_variables()

  decay_vars = [v for v in var_list if len(v.shape.as_list()) > 1]
  noise_vars = [v for v in var_list if "/body/" in v.name]

  weight_decay_loss = weight_decay(hparams.weight_decay, decay_vars)
  tf.summary.scalar("losses/weight_decay", weight_decay_loss)
  weight_noise_ops = weight_noise(hparams.weight_noise, learning_rate,
                                  noise_vars)

  with tf.control_dependencies(weight_noise_ops):
    loss = tf.identity(loss)

  loss += weight_decay_loss
  return loss

def weight_noise(noise_rate, learning_rate, var_list):
  """Apply weight noise to vars in var_list."""
  if not noise_rate:
    return [tf.no_op()]

  noise_ops = []

  for v in var_list:
    with tf.device(v._ref().device):  # pylint: disable=protected-access
      scale = noise_rate * learning_rate * 0.001
      tf.summary.scalar("weight_noise_scale", scale)
      noise = tf.truncated_normal(v.shape) * scale
      noise_op = v.assign_add(noise)
      noise_ops.append(noise_op)

  return noise_ops

def weight_decay(decay_rate, var_list):
  """Apply weight decay to vars in var_list."""
  if not decay_rate:
    return 0.

  weight_decays = []
  for v in var_list:
    v_size = int(np.prod(np.array(v.shape.as_list())))

    # Weight decay
    is_bias = len(v.shape.as_list()) <= 1
    if not is_bias:
      with tf.device(v._ref().device):  # pylint: disable=protected-access
        v_loss = tf.nn.l2_loss(v) / v_size
      weight_decays.append(v_loss)

  return tf.add_n(weight_decays) * decay_rate

def log_variable_sizes(var_list=None, tag=None):
  """Log the sizes and shapes of variables, and the total size.

  Args:
    var_list: a list of varaibles; defaults to trainable_variables
    tag: a string; defaults to "Trainable Variables"
  """
  if var_list is None:
    var_list = tf.trainable_variables()
  if tag is None:
    tag = "Trainable Variables"

  if not var_list:
    return

  name_to_var = {v.name: v for v in var_list}
  total_size = 0
  for v_name in sorted(list(name_to_var)):
    v = name_to_var[v_name]
    v_size = int(np.prod(np.array(v.shape.as_list())))
    tf.logging.info("Weight    %s\tshape    %s\tsize    %d",
                    v.name[:-2].ljust(80),
                    str(v.shape).ljust(20), v_size)
    total_size += v_size
  tf.logging.info("%s Total size: %d", tag, total_size)

def get_variable_initializer(hparams):
  """Get variable initializer from hparams."""
  if hparams.initializer == "orthogonal":
    return tf.orthogonal_initializer(gain=hparams.initializer_gain)
  elif hparams.initializer == "uniform":
    max_val = 0.1 * hparams.initializer_gain
    return tf.random_uniform_initializer(-max_val, max_val)
  elif hparams.initializer == "normal_unit_scaling":
    return tf.variance_scaling_initializer(
        hparams.initializer_gain, mode="fan_avg", distribution="normal")
  elif hparams.initializer == "uniform_unit_scaling":
    return tf.variance_scaling_initializer(
        hparams.initializer_gain, mode="fan_avg", distribution="uniform")
  else:
    raise ValueError("Unrecognized initializer: %s" % hparams.initializer)
<commit_msg>New "Adafactor" optimizer - more memory efficient than Adam.<commit_after># coding=utf-8
# Copyright 2017 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Optimization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports

import numpy as np

from tensor2tensor.utils import yellowfin

import tensorflow as tf

from tensorflow.python.framework import dtypes

def optimize(loss, learning_rate, hparams, use_tpu=False):
  """Minimize loss."""
  loss = weight_decay_and_noise(loss, hparams, learning_rate)
  loss = tf.identity(loss, name="total_loss")
  log_variable_sizes()
  diet_vars = [
      v for v in tf.global_variables() if v.dtype == dtypes.float16_ref
  ]
  log_variable_sizes(diet_vars, "Diet Variables")
  opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams)
  if use_tpu:
    opt = tf.contrib.tpu.CrossShardOptimizer(opt)

  tf.summary.scalar("learning_rate", learning_rate)
  opt_summaries = ["loss"]
  if hparams.summarize_grads:
    opt_summaries.extend(["gradients", "gradient_norm", "global_gradient_norm"])

  train_op = tf.contrib.layers.optimize_loss(
      name="training",
      loss=loss,
      global_step=tf.train.get_or_create_global_step(),
      learning_rate=learning_rate,
      clip_gradients=hparams.clip_grad_norm or None,
      gradient_noise_scale=hparams.grad_noise_scale or None,
      optimizer=opt,
      summaries=opt_summaries,
      colocate_gradients_with_ops=True)
  return train_op

class ConditionalOptimizer(tf.train.Optimizer):
  """Conditional optimizer."""

  def __init__(self, optimizer_name, lr, hparams):
    if optimizer_name == "Adam":
      # We change the default epsilon for Adam and re-scale lr.
      # Using LazyAdam as it's much faster for large vocabulary embeddings.
      self._opt = tf.contrib.opt.LazyAdamOptimizer(
          lr / 500.0,
          beta1=hparams.optimizer_adam_beta1,
          beta2=hparams.optimizer_adam_beta2,
          epsilon=hparams.optimizer_adam_epsilon)
    elif optimizer_name == "Momentum":
      self._opt = tf.train.MomentumOptimizer(
          lr, momentum=hparams.optimizer_momentum_momentum)
    elif optimizer_name == "YellowFin":
      tf.logging.info("Init YellowFin Optimizer.")
      self._opt = yellowfin.YellowFinOptimizer(
          learning_rate=lr, momentum=hparams.optimizer_momentum_momentum)
    elif optimizer_name == "TrueAdam":
      self._opt = tf.train.AdamOptimizer(
          lr / 500.0,
          beta1=hparams.optimizer_adam_beta1,
          beta2=hparams.optimizer_adam_beta2,
          epsilon=hparams.optimizer_adam_epsilon)
    elif optimizer_name == "Adafactor":
      self._opt = AdafactorOptimizer(
          lr / 500.0, epsilon=hparams.optimizer_adam_epsilon)
    else:
      self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)

  def compute_gradients(self, loss, var_list=None, **kwargs):
    return self._opt.compute_gradients(loss, var_list, **kwargs)

  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    return self._opt.apply_gradients(
        grads_and_vars, global_step=global_step, name=name)

def _sqrt_decay(step):
  """Decay like 1 / sqrt(step), multiplied by 500 to normalize."""
  return 500.0 / tf.sqrt(tf.maximum(step, 1.0))

def _exp_decay_after(step, rate, from_which_step):
  """Decay exponentially by rate (per step) starting at from_which_step."""
  return tf.cond(
      step < from_which_step,
      lambda: tf.constant(1.0),
      lambda: rate**(step - from_which_step),
      name="exponential_decay_step_cond")

def learning_rate_decay(hparams, num_worker_replicas=1, num_train_steps=1):
  """Inverse-decay learning rate until warmup_steps, then decay."""
  warmup_steps = tf.to_float(
      hparams.learning_rate_warmup_steps * num_worker_replicas)
  step = tf.to_float(tf.train.get_or_create_global_step())
  if hparams.learning_rate_decay_scheme == "noam":
    return 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
        (step + 1) * warmup_steps**-1.5, (step + 1)**-0.5)
  elif hparams.learning_rate_decay_scheme == "exp100k":
    return 0.94**(step // 100000)
  elif hparams.learning_rate_decay_scheme == "cosine":
    cycle_steps = hparams.learning_rate_cosine_cycle_steps
    return 0.5 * (1 + tf.cos(np.pi * (step % cycle_steps) / cycle_steps))
  elif hparams.learning_rate_decay_scheme == "cyclelinear10x":
    # Cycle the rate linearly by 10x every warmup_steps, up and down.
    cycle_steps = hparams.learning_rate_warmup_steps
    cycle_position = step % (2 * cycle_steps)
    cycle_position = tf.to_float(  # Normalize to the interval [-1, 1].
        cycle_position - cycle_steps) / float(cycle_steps)
    cycle_position = 1.0 - tf.abs(cycle_position)  # 0 to 1 and back to 0.
    return (cycle_position + 0.1) * 3.0  # 10x difference each cycle (0.3-3).

  inv_base = tf.exp(tf.log(0.01) / warmup_steps)
  inv_decay = inv_base**(warmup_steps - step)
  if hparams.learning_rate_decay_scheme == "sqrt":
    decay = _sqrt_decay(step - warmup_steps)
  elif hparams.learning_rate_decay_scheme == "exp10k":
    decay = _exp_decay_after(step - warmup_steps, 0.9995,
                             num_train_steps - warmup_steps - 10000)
  elif hparams.learning_rate_decay_scheme == "exp50k":
    decay = _exp_decay_after(step - warmup_steps, 0.99995,
                             num_train_steps - warmup_steps - 50000)
  elif hparams.learning_rate_decay_scheme == "exp500k":
    decay = _exp_decay_after(step - warmup_steps, 0.9999955,
                             num_train_steps - warmup_steps - 500000)
  elif hparams.learning_rate_decay_scheme == "none":
    decay = tf.constant(1.0)
  else:
    raise ValueError("Unrecognized learning rate decay scheme: %s" %
                     hparams.learning_rate_decay_scheme)
  return tf.where(step < warmup_steps, inv_decay, decay)

def weight_decay_and_noise(loss, hparams, learning_rate, var_list=None):
  """Apply weight decay and weight noise."""
  if var_list is None:
    var_list = tf.trainable_variables()

  decay_vars = [v for v in var_list if len(v.shape.as_list()) > 1]
  noise_vars = [v for v in var_list if "/body/" in v.name]

  weight_decay_loss = weight_decay(hparams.weight_decay, decay_vars)
  tf.summary.scalar("losses/weight_decay", weight_decay_loss)
  weight_noise_ops = weight_noise(hparams.weight_noise, learning_rate,
                                  noise_vars)

  with tf.control_dependencies(weight_noise_ops):
    loss = tf.identity(loss)

  loss += weight_decay_loss
  return loss

def weight_noise(noise_rate, learning_rate, var_list):
  """Apply weight noise to vars in var_list."""
  if not noise_rate:
    return [tf.no_op()]

  noise_ops = []

  for v in var_list:
    with tf.device(v._ref().device):  # pylint: disable=protected-access
      scale = noise_rate * learning_rate * 0.001
      tf.summary.scalar("weight_noise_scale", scale)
      noise = tf.truncated_normal(v.shape) * scale
      noise_op = v.assign_add(noise)
      noise_ops.append(noise_op)

  return noise_ops

def weight_decay(decay_rate, var_list):
  """Apply weight decay to vars in var_list."""
  if not decay_rate:
    return 0.

  weight_decays = []
  for v in var_list:
    v_size = int(np.prod(np.array(v.shape.as_list())))

    # Weight decay
    is_bias = len(v.shape.as_list()) <= 1
    if not is_bias:
      with tf.device(v._ref().device):  # pylint: disable=protected-access
        v_loss = tf.nn.l2_loss(v) / v_size
      weight_decays.append(v_loss)

  return tf.add_n(weight_decays) * decay_rate

def log_variable_sizes(var_list=None, tag=None):
  """Log the sizes and shapes of variables, and the total size.

  Args:
    var_list: a list of varaibles; defaults to trainable_variables
    tag: a string; defaults to "Trainable Variables"
  """
  if var_list is None:
    var_list = tf.trainable_variables()
  if tag is None:
    tag = "Trainable Variables"

  if not var_list:
    return

  name_to_var = {v.name: v for v in var_list}
  total_size = 0
  for v_name in sorted(list(name_to_var)):
    v = name_to_var[v_name]
    v_size = int(np.prod(np.array(v.shape.as_list())))
    tf.logging.info("Weight    %s\tshape    %s\tsize    %d",
                    v.name[:-2].ljust(80),
                    str(v.shape).ljust(20), v_size)
    total_size += v_size
  tf.logging.info("%s Total size: %d", tag, total_size)

def get_variable_initializer(hparams):
  """Get variable initializer from hparams."""
  if hparams.initializer == "orthogonal":
    return tf.orthogonal_initializer(gain=hparams.initializer_gain)
  elif hparams.initializer == "uniform":
    max_val = 0.1 * hparams.initializer_gain
    return tf.random_uniform_initializer(-max_val, max_val)
  elif hparams.initializer == "normal_unit_scaling":
    return tf.variance_scaling_initializer(
        hparams.initializer_gain, mode="fan_avg", distribution="normal")
  elif hparams.initializer == "uniform_unit_scaling":
    return tf.variance_scaling_initializer(
        hparams.initializer_gain, mode="fan_avg", distribution="uniform")
  else:
    raise ValueError("Unrecognized initializer: %s" % hparams.initializer)

class AdafactorOptimizer(tf.train.Optimizer):
  """Optimizer that implements the Adafactor algorithm.

  Adafactor is similar to Adam, but seeks to reduce the memory
  requirements due to the moment estimates.  The auxiliary memory
  requirements for an `AxB` weight matrix are `A+B` for Adafactor,
  versus `2AB` for Adam.

  Adam is described in [Kingma et al., 2014](http://arxiv.org/abs/1412.6980)
  ([pdf](http://arxiv.org/pdf/1412.6980.pdf)).

  The differences are as follows:

  1. No momentum - this removes the first-moment estimate.
  2. For an AxB weight matrix, instead of keeping a full AxB second-moment
     estimate matrix, Adafactor keeps only the row and column means of that
     estimate matrix, and estimate the full second-moment estimate matrix
     from on the fly, based on the means.
  3. Adafactor uses a variable decay rate for the second-moment estaimtes -
     faster decay at the start of training and slower decay later. This
     elimnates the awkwardness in Adam related to having biased moment
     estimates at the start of training.

  For non-2d variables:
    We initialize
    ```
    t <- 0
    v <- zeros(shape(var))
    ```

    The update rule is as follows:
    ```
    t <- t + 1
    decay_horizon = min(t, t * relative_decay_horizon + absolute_decay_horizon)
    decay_rate = 1 - 1 / decay_horizon
    v <- decay_rate * v + (1 - decay_rate) * grad^2
    var <- var - lr * grad / (sqrt(v) + epsilon)
    ```

  For 2d variables:
    We initialize
    ```
    t <- 0
    v_r <- zeros([num_rows])
    v_c <- zeros([num_cols])
    ```

    The update rule is as follows:
    ```
    t <- t + 1
    decay_horizon = min(t, t * relative_decay_horizon + absolute_decay_horizon)
    decay_rate = 1 - 1 / decay_horizon
    v_r <- decay_rate * v_r + (1 - decay_rate) * reduce_mean(grad^2, 1)
    v_c <- decay_rate * v_c + (1 - decay_rate) * reduce_mean(grad^2, 0)
    approx_v = expand_dims(v_r, 1) * expand_dims(v_c, 0) / reduce_mean(v_c)
    var <- var - lr * grad / (sqrt(approx_v) + epsilon)
    ```

  TODO(noam): write a paper.
  TODO(noam): we should also apply the 2d logic to the two final dimensions.
    of >2d convolutional kernels.
  """

  def __init__(self,
               learning_rate=0.001,
               epsilon=1e-8,
               relative_decay_horizon=0.2,
               absolute_decay_horizon=100.0,
               use_locking=False,
               name="Adafactor"):
    """Construct a new Adafactor optimizer.

    See class comment.

    Args:
      learning_rate: A Tensor or a floating point value.  The learning rate.
      epsilon: A small constant for numerical stability.
      relative_decay_horizon: a floating point value <= 1
      absolute_decay_horizon: a floating point value (representing a step count)
      use_locking: If True use locks for update operations.
      name: Optional name for the operations created when applying gradients.
        Defaults to "AdafactorOptimizer".
    """
    super(AdafactorOptimizer, self).__init__(use_locking, name)
    self._lr = learning_rate
    self._relative_decay_horizon = relative_decay_horizon
    self._absolute_decay_horizon = absolute_decay_horizon
    self._epsilon = epsilon

  def _prepare(self):
    global_step = tf.to_float(tf.train.get_or_create_global_step()) + 1.0
    decay_horizon = tf.minimum(global_step,
                               global_step * self._relative_decay_horizon +
                               self._absolute_decay_horizon)
    self._mixing_rate = 1.0 / decay_horizon
    self._decay_rate = 1.0 - self._mixing_rate
    self._epsilon = tf.to_float(self._epsilon)
    self._lr = tf.to_float(self._lr)

  def _should_use_factored_second_moment_estimate(self, shape):
    """Should we use a factored second moment estimator.

    Based on the shape of the variable.

    Args:
      shape: a list of integers
    Returns:
      a boolean
    """
    return len(shape) == 2

  def _create_slots(self, var_list):
    for v in var_list:
      shape = v.get_shape().as_list()
      if self._should_use_factored_second_moment_estimate(shape):
        r_val = tf.zeros([shape[0]], dtype=tf.float32)
        c_val = tf.zeros([shape[1]], dtype=tf.float32)
        self._get_or_make_slot(v, r_val, "vr", self._name)
        self._get_or_make_slot(v, c_val, "vc", self._name)
      else:
        self._zeros_slot(v, "v", self._name)

  def _apply_dense(self, grad, var):
    return self._resource_apply_dense(grad, var)

  def _resource_apply_dense(self, grad, var):
    shape = var.get_shape().as_list()
    grad_squared = tf.square(grad)
    updates = []
    if self._should_use_factored_second_moment_estimate(shape):
      vr = self.get_slot(var, "vr")
      new_vr = (self._decay_rate * vr +
                self._mixing_rate * tf.reduce_mean(grad_squared, 1))
      vc = self.get_slot(var, "vc")
      new_vc = (self._decay_rate * vc +
                self._mixing_rate * tf.reduce_mean(grad_squared, 0))
      vr_update = tf.assign(vr, new_vr, use_locking=self._use_locking)
      vc_update = tf.assign(vc, new_vc, use_locking=self._use_locking)
      updates = [vr_update, vc_update]
      vr = tf.sqrt(new_vr) + self._epsilon
      vc = tf.sqrt(new_vc) + self._epsilon
      vc /= tf.reduce_mean(vc)
      denom = tf.expand_dims(vr, 1) * tf.expand_dims(vc, 0)
    else:
      v = self.get_slot(var, "v")
      new_v = (self._decay_rate * v + self._mixing_rate * grad_squared)
      v_update = tf.assign(v, new_v, use_locking=self._use_locking)
      updates = [v_update]
      denom = tf.sqrt(new_v) + self._epsilon
    subtrahend = self._lr * grad / denom
    var_update = tf.assign_sub(var, subtrahend, use_locking=self._use_locking)
    updates = [var_update] + updates
    return tf.group(*updates)
<|endoftext|>


<commit_before>from collections import deque, defaultdict

class Node:
    def __init__(self, name, level, parent):
        self.name = name
        self.level = level
        self.ancestor = [parent] if parent is not None else []
        self.fill_ancestor()

    def fill_ancestor(self):
        number_of_link = self.level.bit_length()
        for i in range(1, number_of_link):
            self.ancestor.append(self.ancestor[i-1].kth_ancestor(2**(i-1)))

    def kth_ancestor(self, k):
        if k == 0:
            return self

        next_link = k.bit_length() - 1
        if next_link < len(self.ancestor):
            return self.ancestor[next_link].kth_ancestor(k - 2**next_link)
        else:
            return None

    def __repr__(self):
        return "{{name: {}, level: {}, ancestor: {}}}".format(self.name, self.level, [a.name for a in self.ancestor if a is not None])

def build_tree(edges):
    root = Node(edges.pop(0).pop(), 0, None)
    tree = {root.name: root}
    next_queue = deque([root])
    while next_queue:
        current = next_queue.popleft()
        for child in edges[current.name]:
            node = Node(child, current.level+1, current)
            tree[child] = node
            next_queue.append(node)
    return tree

def interpret(tree, *args):
    if args[0] == 0:
        tree[args[2]] = Node(args[2], tree[args[1]].level+1, tree[args[1]])
    elif args[0] == 1:
        tree.pop(args[1])
    elif args[0] == 2:
        ancestor = tree[args[1]].kth_ancestor(args[2]) if args[1] in tree.keys() else None
        print(0 if ancestor is None else ancestor.name)

def main():
    T = int(input())
    for _ in range(T):
        P = int(input())

        edges = defaultdict(list)
        for _ in range(P):
            X, Y = tuple(int(v) for v in input().split())
            edges[Y].append(X)

        tree = build_tree(edges)

        Q = int(input())
        for _ in range(Q):
            interpret(tree, *tuple(int(v) for v in input().split()))

if __name__ == "__main__":
    main()
<commit_msg>HACKER: Kth Ancestor (partial)<commit_after>class Node:
    def __init__(self, name, level, parent):
        self.name = name
        self.level = level
        self.ancestor = [] if parent is None else [parent]
        self.fill_ancestor()

    def fill_ancestor(self):
        number_of_link = self.level.bit_length()
        for i in range(1, number_of_link):
            self.ancestor.append(self.ancestor[i-1].kth_ancestor(2**(i-1)))

    def kth_ancestor(self, k):
        if k == 0:
            return self

        next_link = k.bit_length() - 1
        if next_link < len(self.ancestor):
            return self.ancestor[next_link].kth_ancestor(k - 2**next_link)
        else:
            return None

#    def __repr__(self):
#        return "{{name:{},level:{},ready:{},ancestor:{}}}".format(self.name,
#                                                                  self.level,
#                                                                  self.ready,
#                                                                  [a.name for a in self.ancestor if a is not None])

def main():
    T = int(input())
    for _ in range(T):
        P = int(input())

        X, Y = tuple(int(v) for v in input().split())
        root = Node(X, 0, None)
        tree = {root.name: root}
        for _ in range(P-1):
            X, Y = tuple(int(v) for v in input().split())
            node = Node(X, tree[Y].level+1, tree[Y])
            tree[X] = node
        Q = int(input())
        for _ in range(Q):
            cmd = tuple(int(v) for v in input().split())
            if cmd[0] == 0:
                tree[cmd[2]] = Node(cmd[2], tree[cmd[1]].level+1, tree[cmd[1]])
            elif cmd[0] == 1:
                tree.pop(cmd[1])
            elif cmd[0] == 2:
                ancestor = tree[cmd[1]].kth_ancestor(cmd[2]) if cmd[1] in tree.keys() else None
                print(0 if ancestor is None else ancestor.name)

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


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

from berg import git_util, logger
from berg.configuration import config, BERG_REPO_ROOT
from berg.util import check_call

def cp(src_path, dest_path):
  if logger.level < logging.DEBUG:
    check_call('gsutil -q cp %s %s' % (src_path, dest_path))
  else:
    check_call('gsutil cp %s %s' % (src_path, dest_path))

def rsync(src_dir, dest_dir):
  dest_dir = 'gs://' + dest_dir
  if logger.level < logging.DEBUG:
    check_call('gsutil -q -m rsync -r %s %s' % (src_dir, dest_dir))
  else:
    check_call('gsutil -m rsync -r %s %s' % (src_dir, dest_dir))

def upload_repo(repo_name):
  rsync(git_util.current_repo_path(),
        os.path.join(config.gcs_repo_root, repo_name))

def upload_berg_repo_for_self_update():
  rsync(BERG_REPO_ROOT, os.path.join(config.gcs_repo_root, 'berg'))

def download_repo(repo_name):
  local_repo = os.path.join('/root/code', repo_name)
  os.makedirs(local_repo, exist_ok=True)
  rsync(os.path.join(config.gcs_repo_root, repo_name), local_repo)
<commit_msg>Update gsutil.py<commit_after># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import os

from berg import git_util, logger
from berg.configuration import config, BERG_REPO_ROOT
from berg.util import check_call

def cp(src_path, dest_path):
  if logger.level < logging.DEBUG:
    check_call('gsutil -q cp %s %s' % (src_path, dest_path))
  else:
    check_call('gsutil cp %s %s' % (src_path, dest_path))

def rsync(src_dir, dest_dir):
  if logger.level < logging.DEBUG:
    check_call('gsutil -q -m rsync -r %s %s' % (src_dir, dest_dir))
  else:
    check_call('gsutil -m rsync -r %s %s' % (src_dir, dest_dir))

def upload_repo(repo_name):
  rsync(git_util.current_repo_path(),
        os.path.join(config.gcs_repo_root, repo_name))

def upload_berg_repo_for_self_update():
  rsync(BERG_REPO_ROOT, os.path.join(config.gcs_repo_root, 'berg'))

def download_repo(repo_name):
  local_repo = os.path.join('/root/code', repo_name)
  os.makedirs(local_repo, exist_ok=True)
  rsync(os.path.join(config.gcs_repo_root, repo_name), local_repo)
<|endoftext|>


<commit_before># This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# 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
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Container for all the import related functions."""

try:
    import simplejson as json
except ImportError:
    import json

import bson
import datetime
import glob
import os
import types

import models
import models.defconfig as mdefconfig
import models.job as mjob
import utils
import utils.db

def import_and_save_job(json_obj, db_options, base_path=utils.BASE_PATH):
    """Wrapper function to be used as an external task.

    This function should only be called by Celery or other task managers.

    :param json_obj: The JSON object with the values to be parsed.
    :type json_obj: dict
    :param db_options: The mongodb database connection parameters.
    :type db_options: dict
    :return The ID of the job document created.
    """
    database = utils.db.get_db_connection(db_options)
    docs, job_id = import_job_from_json(json_obj, database, base_path)

    if docs:
        utils.LOG.info(
            "Importing %d documents with job ID: %s", len(docs), job_id)
        utils.db.save_all(database, docs, manipulate=True)
    else:
        utils.LOG.info("No jobs to save")

    return job_id

def import_job_from_json(json_obj, database, base_path=utils.BASE_PATH):
    """Import a job based on the provided JSON object.

    The provided JSON object, a dict-like object, should contain at least the
    `job` and `kernel` keys. These keys will be used to traverse the directory
    structure found at the `/job/kernel` path (starting from the default
    location).

    :param json_obj: A dict-like object, that should contain the keys `job` and
        `kernel`.
    :param base_path: The base path where to start constructing the traverse
        directory. It defaults to: /var/www/images/kernel-ci.
    :return The documents to be saved, and the job document ID.
    """
    job_dir = json_obj[models.JOB_KEY]
    kernel_dir = json_obj[models.KERNEL_KEY]

    return _import_job(job_dir, kernel_dir, database, base_path)

def _import_job(job, kernel, database, base_path=utils.BASE_PATH):
    """Traverse the job dir and create the documenst to save.

    :param job: The name of the job.
    :param kernel: The name of the kernel.
    :param base_path: The base path where to strat the traversing.
    :return The documents to be saved, and the job document ID.
    """
    docs = []
    ret_val = 201
    job_id = None

    job_dir = os.path.join(base_path, job)
    kernel_dir = os.path.join(job_dir, kernel)

    if utils.is_hidden(job) or utils.is_hidden(kernel):
        return docs

    job_name = (
        models.JOB_DOCUMENT_NAME %
        {models.JOB_KEY: job, models.KERNEL_KEY: kernel})

    saved_doc = utils.db.find_one(
        database[models.JOB_COLLECTION], [job_name], field=models.NAME_KEY)
    if saved_doc:
        job_doc = mjob.JobDocument.from_json(saved_doc)
        job_id = job_doc.id
    else:
        job_doc = mjob.JobDocument(job, kernel)
        ret_val, job_id = utils.db.save(database, job_doc, manipulate=True)

    if all([ret_val == 201, job_id is not None]):
        job_doc.id = job_id
        docs = _traverse_kernel_dir(job_doc, kernel_dir, database)
    else:
        utils.LOG.error("Unable to save job document %s", job_name)
        job_id = None

    return docs, job_id

def _traverse_kernel_dir(job_doc, kernel_dir, database):
    """Traverse the kernel directory looking for defconfig dirs.

    :param job_doc: The created `JobDocument`.
    :param kernel_dir: The kernel directory to traverse.
    :param database: The database connection.
    """
    docs = []

    if os.path.isdir(kernel_dir):
        if (os.path.exists(os.path.join(kernel_dir, models.DONE_FILE)) or
                glob.glob(os.path.join(kernel_dir, models.DONE_FILE_PATTERN))):
            job_doc.status = models.PASS_STATUS
        else:
            job_doc.status = models.UNKNOWN_STATUS

        if not job_doc.created_on:
            job_doc.created_on = datetime.datetime.now(tz=bson.tz_util.utc)

        docs.extend(
            [
                _traverse_defconf_dir(
                    job_doc, kernel_dir, defconf_dir, database
                ) for defconf_dir in os.listdir(kernel_dir)
                if os.path.isdir(os.path.join(kernel_dir, defconf_dir))
                if not utils.is_hidden(defconf_dir)
                if not utils.is_lab_dir(defconf_dir)
                if not None
            ]
        )
    else:
        job_doc.status = models.BUILD_STATUS
        job_doc.created_on = datetime.datetime.now(tz=bson.tz_util.utc)

    # Kind of a hack:
    # We want to store some metadata at the job document level as well, like
    # git tree, git commit...
    # Since, at the moment, we do not have the metadata file at the job level
    # we need to pick one from the build documents, and extract some values.
    docs_len = len(docs)
    if docs_len > 1:
        idx = 0
        while idx < docs_len:
            defconf_doc = docs[idx]
            if isinstance(defconf_doc, mjob.JobDocument):
                idx += 1
            elif isinstance(defconf_doc, mdefconfig.DefconfigDocument):
                if (defconf_doc.job == job_doc.job and
                        defconf_doc.kernel == job_doc.kernel):
                    job_doc.git_commit = defconf_doc.git_commit
                    job_doc.git_describe = defconf_doc.git_describe
                    job_doc.git_url = defconf_doc.git_url
                    job_doc.git_branch = defconf_doc.git_branch
                    break
            else:
                idx += 1
        # Save, again, the job-doc or we end up without the git data.
        utils.db.save(database, job_doc)

    return docs

def _traverse_defconf_dir(
        job_doc, kernel_dir, defconfig_dir, database):
    """Traverse the defconfig directory looking for files.

    :param job_doc: The created `JobDocument`.
    :param kernel_dir: The parent directory of this defconfig.
    :param defconfig_dir: The actual defconfig directory to parse.
    :param database: The database connection.
    :return A `DefconfigDocument` instance.
    """
    job = job_doc.job
    kernel = job_doc.kernel
    real_dir = os.path.join(kernel_dir, defconfig_dir)

    utils.LOG.info("Traversing directory '%s'", real_dir)

    defconfig_doc = None
    for dirname, subdirs, files in os.walk(real_dir):
        # Consider only the actual directory and its files.
        subdirs[:] = []

        data_file = os.path.join(dirname, models.BUILD_META_JSON_FILE)

        if os.path.isfile(data_file):
            defconfig_doc = _parse_build_data(
                data_file, job, kernel, dirname)
            defconfig_doc.job_id = job_doc.id
            # Search for previous defconfig doc. This is only useful when
            # re-importing data and we want to have the same ID as before.
            doc_id, c_date = _search_prev_defconfig_doc(
                defconfig_doc, database)
            defconfig_doc.id = doc_id
            if c_date:
                defconfig_doc.created_on = c_date
            else:
                # Give the defconfig doc the same date as the job one.
                # In this way all defconfigs will have the same date regardless
                # of when they were saved on the file system.
                defconfig_doc.created_on = job_doc.created_on
        else:
            utils.LOG.warn("No build data file found in '%s'", real_dir)

    return defconfig_doc

def _search_prev_defconfig_doc(defconfig_doc, database):
    """Search for a similar defconfig document in the database.

    Search for an already imported defconfig/build document in the database
    and return its object ID and creation date. This is done to make sure
    we do not create double documents when re-importing the same data or
    updating it.

    :param defconfig_doc: The new defconfig document.
    :param database: The db connection.
    :return The previous doc ID and its creation date, or None.
    """
    doc_id = None
    c_date = None

    if all([defconfig_doc, database]):
        spec = {
            models.JOB_KEY: defconfig_doc.job,
            models.KERNEL_KEY: defconfig_doc.kernel,
            models.DEFCONFIG_KEY: defconfig_doc.defconfig,
            models.DEFCONFIG_FULL_KEY: defconfig_doc.defconfig_full,
            models.ARCHITECTURE_KEY: defconfig_doc.arch
        }
        prev_doc = utils.db.find(
            database[models.DEFCONFIG_COLLECTION],
            10,
            0,
            spec=spec,
            fields=[models.ID_KEY, models.CREATED_KEY]
        )

        prev_doc_count = prev_doc.count()
        if prev_doc_count > 0:
            if prev_doc_count == 1:
                doc_id = prev_doc[0].get(models.ID_KEY, None)
                c_date = prev_doc[0].get(models.CREATED_KEY, None)
            else:
                utils.LOG.warn(
                    "Found multiple defconfig docs matching: %s",
                    spec)
                utils.LOG.error(
                    "Cannot keep old document ID, don't know which one to "
                    "use!")

    return doc_id, c_date

def _parse_build_data(data_file, job, kernel, defconfig_dir):
    """Parse the metadata file contained in thie build directory.

    :param data_file: The path to the metadata file.
    :param defconf_doc: The `DefconfigDocument` whose metadata will be updated.
    """
    build_data = None
    with open(data_file) as data:
        build_data = json.load(data)

    defconfig_doc = None

    if all([build_data, isinstance(build_data, types.DictionaryType)]):
        data_pop = build_data.pop

        try:
            defconfig = data_pop(models.DEFCONFIG_KEY)
            defconfig_full = data_pop(models.DEFCONFIG_FULL_KEY, None)
            kconfig_fragments = data_pop(models.KCONFIG_FRAGMENTS_KEY, None)

            defconfig_full = utils.get_defconfig_full(
                defconfig_dir, defconfig, defconfig_full, kconfig_fragments)

            # Err on the safe side.
            job = data_pop(models.JOB_KEY, None) or job
            kernel = data_pop(models.KERNEL_KEY, None) or kernel

            defconfig_doc = mdefconfig.DefconfigDocument(
                job, kernel, defconfig, defconfig_full)

            defconfig_doc.dirname = defconfig_dir

            defconfig_doc.arch = data_pop(models.ARCHITECTURE_KEY, None)
            defconfig_doc.build_log = data_pop(models.BUILD_LOG_KEY, None)
            defconfig_doc.build_platform = data_pop(
                models.BUILD_PLATFORM_KEY, [])
            defconfig_doc.build_time = data_pop(models.BUILD_TIME_KEY, 0)
            defconfig_doc.dtb_dir = data_pop(models.DTB_DIR_KEY, None)
            defconfig_doc.errors = data_pop(models.BUILD_ERRORS_KEY, 0)
            defconfig_doc.file_server_resource = data_pop(
                models.FILE_SERVER_RESOURCE_KEY, None)
            defconfig_doc.file_server_url = data_pop(
                models.FILE_SERVER_URL_KEY, None)
            defconfig_doc.git_branch = data_pop(models.GIT_BRANCH_KEY, None)
            defconfig_doc.git_commit = data_pop(models.GIT_COMMIT_KEY, None)
            defconfig_doc.git_describe = data_pop(
                models.GIT_DESCRIBE_KEY, None)
            defconfig_doc.git_url = data_pop(models.GIT_URL_KEY, None)
            defconfig_doc.kconfig_fragments = kconfig_fragments
            defconfig_doc.kernel_config = data_pop(
                models.KERNEL_CONFIG_KEY, None)
            defconfig_doc.kernel_image = data_pop(
                models.KERNEL_IMAGE_KEY, None)
            defconfig_doc.modules = data_pop(models.MODULES_KEY, None)
            defconfig_doc.modules_dir = data_pop(models.MODULES_DIR_KEY, None)
            defconfig_doc.status = data_pop(
                models.BUILD_RESULT_KEY, models.UNKNOWN_STATUS)
            defconfig_doc.system_map = data_pop(models.SYSTEM_MAP_KEY, None)
            defconfig_doc.text_offset = data_pop(models.TEXT_OFFSET_KEY, None)
            defconfig_doc.version = data_pop(models.VERSION_KEY, "1.0")
            defconfig_doc.warnings = data_pop(models.BUILD_WARNINGS_KEY, 0)

            defconfig_doc.metadata = build_data
        except KeyError, ex:
            utils.LOG.exception(ex)
            utils.LOG.error(
                "Missing mandatory key in build data file '%s'",
                data_file)

    return defconfig_doc
<commit_msg>Catch JSON error if the file is empty.<commit_after># This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# 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
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Container for all the import related functions."""

try:
    import simplejson as json
except ImportError:
    import json

import bson
import datetime
import glob
import os
import types

import models
import models.defconfig as mdefconfig
import models.job as mjob
import utils
import utils.db

def import_and_save_job(json_obj, db_options, base_path=utils.BASE_PATH):
    """Wrapper function to be used as an external task.

    This function should only be called by Celery or other task managers.

    :param json_obj: The JSON object with the values to be parsed.
    :type json_obj: dict
    :param db_options: The mongodb database connection parameters.
    :type db_options: dict
    :return The ID of the job document created.
    """
    database = utils.db.get_db_connection(db_options)
    docs, job_id = import_job_from_json(json_obj, database, base_path)

    if docs:
        utils.LOG.info(
            "Importing %d documents with job ID: %s", len(docs), job_id)
        utils.db.save_all(database, docs, manipulate=True)
    else:
        utils.LOG.info("No jobs to save")

    return job_id

def import_job_from_json(json_obj, database, base_path=utils.BASE_PATH):
    """Import a job based on the provided JSON object.

    The provided JSON object, a dict-like object, should contain at least the
    `job` and `kernel` keys. These keys will be used to traverse the directory
    structure found at the `/job/kernel` path (starting from the default
    location).

    :param json_obj: A dict-like object, that should contain the keys `job` and
        `kernel`.
    :param base_path: The base path where to start constructing the traverse
        directory. It defaults to: /var/www/images/kernel-ci.
    :return The documents to be saved, and the job document ID.
    """
    job_dir = json_obj[models.JOB_KEY]
    kernel_dir = json_obj[models.KERNEL_KEY]

    return _import_job(job_dir, kernel_dir, database, base_path)

def _import_job(job, kernel, database, base_path=utils.BASE_PATH):
    """Traverse the job dir and create the documenst to save.

    :param job: The name of the job.
    :param kernel: The name of the kernel.
    :param base_path: The base path where to strat the traversing.
    :return The documents to be saved, and the job document ID.
    """
    docs = []
    ret_val = 201
    job_id = None

    job_dir = os.path.join(base_path, job)
    kernel_dir = os.path.join(job_dir, kernel)

    if utils.is_hidden(job) or utils.is_hidden(kernel):
        return docs

    job_name = (
        models.JOB_DOCUMENT_NAME %
        {models.JOB_KEY: job, models.KERNEL_KEY: kernel})

    saved_doc = utils.db.find_one(
        database[models.JOB_COLLECTION], [job_name], field=models.NAME_KEY)
    if saved_doc:
        job_doc = mjob.JobDocument.from_json(saved_doc)
        job_id = job_doc.id
    else:
        job_doc = mjob.JobDocument(job, kernel)
        ret_val, job_id = utils.db.save(database, job_doc, manipulate=True)

    if all([ret_val == 201, job_id is not None]):
        job_doc.id = job_id
        docs = _traverse_kernel_dir(job_doc, kernel_dir, database)
    else:
        utils.LOG.error("Unable to save job document %s", job_name)
        job_id = None

    return docs, job_id

def _traverse_kernel_dir(job_doc, kernel_dir, database):
    """Traverse the kernel directory looking for defconfig dirs.

    :param job_doc: The created `JobDocument`.
    :param kernel_dir: The kernel directory to traverse.
    :param database: The database connection.
    """
    docs = []

    if os.path.isdir(kernel_dir):
        if (os.path.exists(os.path.join(kernel_dir, models.DONE_FILE)) or
                glob.glob(os.path.join(kernel_dir, models.DONE_FILE_PATTERN))):
            job_doc.status = models.PASS_STATUS
        else:
            job_doc.status = models.UNKNOWN_STATUS

        if not job_doc.created_on:
            job_doc.created_on = datetime.datetime.now(tz=bson.tz_util.utc)

        docs.extend(
            [
                _traverse_defconf_dir(
                    job_doc, kernel_dir, defconf_dir, database
                ) for defconf_dir in os.listdir(kernel_dir)
                if os.path.isdir(os.path.join(kernel_dir, defconf_dir))
                if not utils.is_hidden(defconf_dir)
                if not utils.is_lab_dir(defconf_dir)
                if not None
            ]
        )
    else:
        job_doc.status = models.BUILD_STATUS
        job_doc.created_on = datetime.datetime.now(tz=bson.tz_util.utc)

    # Kind of a hack:
    # We want to store some metadata at the job document level as well, like
    # git tree, git commit...
    # Since, at the moment, we do not have the metadata file at the job level
    # we need to pick one from the build documents, and extract some values.
    docs_len = len(docs)
    if docs_len > 1:
        idx = 0
        while idx < docs_len:
            defconf_doc = docs[idx]
            if isinstance(defconf_doc, mjob.JobDocument):
                idx += 1
            elif isinstance(defconf_doc, mdefconfig.DefconfigDocument):
                if (defconf_doc.job == job_doc.job and
                        defconf_doc.kernel == job_doc.kernel):
                    job_doc.git_commit = defconf_doc.git_commit
                    job_doc.git_describe = defconf_doc.git_describe
                    job_doc.git_url = defconf_doc.git_url
                    job_doc.git_branch = defconf_doc.git_branch
                    break
            else:
                idx += 1
        # Save, again, the job-doc or we end up without the git data.
        utils.db.save(database, job_doc)

    return docs

def _traverse_defconf_dir(
        job_doc, kernel_dir, defconfig_dir, database):
    """Traverse the defconfig directory looking for files.

    :param job_doc: The created `JobDocument`.
    :param kernel_dir: The parent directory of this defconfig.
    :param defconfig_dir: The actual defconfig directory to parse.
    :param database: The database connection.
    :return A `DefconfigDocument` instance.
    """
    job = job_doc.job
    kernel = job_doc.kernel
    real_dir = os.path.join(kernel_dir, defconfig_dir)

    utils.LOG.info("Traversing directory '%s'", real_dir)

    defconfig_doc = None
    for dirname, subdirs, files in os.walk(real_dir):
        # Consider only the actual directory and its files.
        subdirs[:] = []

        data_file = os.path.join(dirname, models.BUILD_META_JSON_FILE)

        if os.path.isfile(data_file):
            defconfig_doc = _parse_build_data(
                data_file, job, kernel, dirname)
            defconfig_doc.job_id = job_doc.id
            # Search for previous defconfig doc. This is only useful when
            # re-importing data and we want to have the same ID as before.
            doc_id, c_date = _search_prev_defconfig_doc(
                defconfig_doc, database)
            defconfig_doc.id = doc_id
            if c_date:
                defconfig_doc.created_on = c_date
            else:
                # Give the defconfig doc the same date as the job one.
                # In this way all defconfigs will have the same date regardless
                # of when they were saved on the file system.
                defconfig_doc.created_on = job_doc.created_on
        else:
            utils.LOG.warn("No build data file found in '%s'", real_dir)

    return defconfig_doc

def _search_prev_defconfig_doc(defconfig_doc, database):
    """Search for a similar defconfig document in the database.

    Search for an already imported defconfig/build document in the database
    and return its object ID and creation date. This is done to make sure
    we do not create double documents when re-importing the same data or
    updating it.

    :param defconfig_doc: The new defconfig document.
    :param database: The db connection.
    :return The previous doc ID and its creation date, or None.
    """
    doc_id = None
    c_date = None

    if all([defconfig_doc, database]):
        spec = {
            models.JOB_KEY: defconfig_doc.job,
            models.KERNEL_KEY: defconfig_doc.kernel,
            models.DEFCONFIG_KEY: defconfig_doc.defconfig,
            models.DEFCONFIG_FULL_KEY: defconfig_doc.defconfig_full,
            models.ARCHITECTURE_KEY: defconfig_doc.arch
        }
        prev_doc = utils.db.find(
            database[models.DEFCONFIG_COLLECTION],
            10,
            0,
            spec=spec,
            fields=[models.ID_KEY, models.CREATED_KEY]
        )

        prev_doc_count = prev_doc.count()
        if prev_doc_count > 0:
            if prev_doc_count == 1:
                doc_id = prev_doc[0].get(models.ID_KEY, None)
                c_date = prev_doc[0].get(models.CREATED_KEY, None)
            else:
                utils.LOG.warn(
                    "Found multiple defconfig docs matching: %s",
                    spec)
                utils.LOG.error(
                    "Cannot keep old document ID, don't know which one to "
                    "use!")

    return doc_id, c_date

def _parse_build_data(data_file, job, kernel, defconfig_dir):
    """Parse the metadata file contained in thie build directory.

    :param data_file: The path to the metadata file.
    :param defconf_doc: The `DefconfigDocument` whose metadata will be updated.
    """
    build_data = None
    try:
        with open(data_file) as data:
            build_data = json.load(data)
    except json.JSONDecodeError, ex:
        utils.LOG.exception(ex)
        utils.LOG.error("Error loading JSON data from '%s'", defconfig_dir)

    defconfig_doc = None

    if all([build_data, isinstance(build_data, types.DictionaryType)]):
        data_pop = build_data.pop

        try:
            defconfig = data_pop(models.DEFCONFIG_KEY)
            defconfig_full = data_pop(models.DEFCONFIG_FULL_KEY, None)
            kconfig_fragments = data_pop(models.KCONFIG_FRAGMENTS_KEY, None)

            defconfig_full = utils.get_defconfig_full(
                defconfig_dir, defconfig, defconfig_full, kconfig_fragments)

            # Err on the safe side.
            job = data_pop(models.JOB_KEY, None) or job
            kernel = data_pop(models.KERNEL_KEY, None) or kernel

            defconfig_doc = mdefconfig.DefconfigDocument(
                job, kernel, defconfig, defconfig_full)

            defconfig_doc.dirname = defconfig_dir

            defconfig_doc.arch = data_pop(models.ARCHITECTURE_KEY, None)
            defconfig_doc.build_log = data_pop(models.BUILD_LOG_KEY, None)
            defconfig_doc.build_platform = data_pop(
                models.BUILD_PLATFORM_KEY, [])
            defconfig_doc.build_time = data_pop(models.BUILD_TIME_KEY, 0)
            defconfig_doc.dtb_dir = data_pop(models.DTB_DIR_KEY, None)
            defconfig_doc.errors = data_pop(models.BUILD_ERRORS_KEY, 0)
            defconfig_doc.file_server_resource = data_pop(
                models.FILE_SERVER_RESOURCE_KEY, None)
            defconfig_doc.file_server_url = data_pop(
                models.FILE_SERVER_URL_KEY, None)
            defconfig_doc.git_branch = data_pop(models.GIT_BRANCH_KEY, None)
            defconfig_doc.git_commit = data_pop(models.GIT_COMMIT_KEY, None)
            defconfig_doc.git_describe = data_pop(
                models.GIT_DESCRIBE_KEY, None)
            defconfig_doc.git_url = data_pop(models.GIT_URL_KEY, None)
            defconfig_doc.kconfig_fragments = kconfig_fragments
            defconfig_doc.kernel_config = data_pop(
                models.KERNEL_CONFIG_KEY, None)
            defconfig_doc.kernel_image = data_pop(
                models.KERNEL_IMAGE_KEY, None)
            defconfig_doc.modules = data_pop(models.MODULES_KEY, None)
            defconfig_doc.modules_dir = data_pop(models.MODULES_DIR_KEY, None)
            defconfig_doc.status = data_pop(
                models.BUILD_RESULT_KEY, models.UNKNOWN_STATUS)
            defconfig_doc.system_map = data_pop(models.SYSTEM_MAP_KEY, None)
            defconfig_doc.text_offset = data_pop(models.TEXT_OFFSET_KEY, None)
            defconfig_doc.version = data_pop(models.VERSION_KEY, "1.0")
            defconfig_doc.warnings = data_pop(models.BUILD_WARNINGS_KEY, 0)

            defconfig_doc.metadata = build_data
        except KeyError, ex:
            utils.LOG.exception(ex)
            utils.LOG.error(
                "Missing mandatory key in build data file '%s'",
                data_file)

    return defconfig_doc
<|endoftext|>


<commit_before># encoding: utf-8
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Contact: Kyle Lahnakoski (kyle@lahnakoski.com)
#

from __future__ import absolute_import, division, unicode_literals

import json
import sys

PY3 = sys.version_info[0] == 3
PY2 = sys.version_info[0] == 2

PYPY = False
try:
    import __pypy__ as _
    PYPY=True
except Exception:
    PYPY=False

none_type = type(None)
boolean_type = type(True)

if PY3:
    import itertools
    import collections
    from collections import Callable
    from functools import cmp_to_key, reduce, update_wrapper
    from configparser import ConfigParser
    from itertools import zip_longest
    import builtins as __builtin__
    from builtins import input

    izip = zip
    zip_longest = itertools.zip_longest

    text = str
    text_type = str
    string_types = str
    binary_type = bytes
    integer_types = (int, )
    number_types = (int, float)
    long = int
    unichr = chr

    xrange = range
    def _gen():
        yield

    generator_types = (
        type(_gen()),
        type(filter(lambda x: True, [])),
        type({}.items()),
        type({}.values()),
        type(map(lambda: 0, []))
    )
    unichr = chr

    round = round
    from html.parser import HTMLParser
    from urllib.parse import urlparse
    from io import StringIO
    from io import BytesIO
    from _thread import allocate_lock, get_ident, start_new_thread, interrupt_main

    def items(d):
        return list(d.items())

    def iteritems(d):
        return d.items()

    def transpose(*args):
        return list(zip(*args))

    def get_function_name(func):
        return func.__name__

    def get_function_arguments(func):
        return func.__code__.co_varnames[:func.__code__.co_argcount]

    def get_function_code(func):
        return func.__code__

    def get_function_defaults(func):
        return func.__defaults__

    def sort_using_cmp(data, cmp):
        return sorted(data, key=cmp_to_key(cmp))

    def sort_using_key(data, key):
        return sorted(data, key=key)

    def first(values):
        try:
            return iter(values).__next__()
        except StopIteration:
            return None

    def is_text(t):
        return t.__class__ is str

    def is_binary(b):
        return b.__class__ is bytes

    utf8_json_encoder = json.JSONEncoder(
        skipkeys=False,
        ensure_ascii=False,  # DIFF FROM DEFAULTS
        check_circular=True,
        allow_nan=True,
        indent=None,
        separators=(',', ':'),
        default=None,
        sort_keys=True   # <-- IMPORTANT!  sort_keys==True
    ).encode

    UserDict = collections.UserDict

else:
    import collections
    from collections import Callable
    from functools import cmp_to_key, reduce, update_wrapper

    import __builtin__
    from types import GeneratorType
    from ConfigParser import ConfigParser
    from itertools import izip_longest as zip_longest
    from __builtin__ import zip as transpose
    from itertools import izip
    from __builtin__ import raw_input as input

    reduce = __builtin__.reduce
    text = __builtin__.unicode
    text_type = __builtin__.unicode
    string_types = (str, unicode)
    binary_type = str
    integer_types = (int, long)
    number_types = (int, long, float)
    long = __builtin__.long
    unichr = __builtin__.unichr

    xrange = __builtin__.xrange
    generator_types = (GeneratorType, type(reversed([])))
    unichr = __builtin__.unichr

    round = __builtin__.round
    import HTMLParser
    from urlparse import urlparse
    from StringIO import StringIO
    from io import BytesIO
    from thread import allocate_lock, get_ident, start_new_thread, interrupt_main

    def items(d):
        return d.items()

    def iteritems(d):
        return d.iteritems()

    def get_function_name(func):
        return func.func_name

    def get_function_arguments(func):
        return func.func_code.co_varnames[:func.func_code.co_argcount]

    def get_function_code(func):
        return func.func_code

    def get_function_defaults(func):
        return func.func_defaults

    def sort_using_cmp(data, cmp):
        return sorted(data, cmp=cmp)

    def sort_using_key(data, key):
        return sorted(data, key=key)
        # return sorted(
        #     ((key(d), d) for d in data),
        #     lambda a, b: (1 if (a[0]>b[0]) else (-1 if (a[0]<b[0]) else 0))
        # )

    def first(values):
        try:
            return iter(values).next()
        except StopIteration:
            return None

    def is_text(t):
        return t.__class__ is unicode

    def is_binary(b):
        return b.__class__ is str

    utf8_json_encoder = json.JSONEncoder(
        skipkeys=False,
        ensure_ascii=False,  # DIFF FROM DEFAULTS
        check_circular=True,
        allow_nan=True,
        indent=None,
        separators=(',', ':'),
        encoding='utf-8',  # DIFF FROM DEFAULTS
        default=None,
        sort_keys=True   # <-- IMPORTANT!  sort_keys==True
    ).encode

    # COPIED FROM Python's collections.UserDict (copied July 2018)
    class UserDict(collections.MutableMapping):

        # Start by filling-out the abstract methods
        def __init__(*args, **kwargs):
            if not args:
                raise TypeError("descriptor '__init__' of 'UserDict' object "
                                "needs an argument")
            self, args = args[0], args[1:]
            if len(args) > 1:
                raise TypeError('expected at most 1 arguments, got %d' % len(args))
            if args:
                dict = args[0]
            elif 'dict' in kwargs:
                dict = kwargs.pop('dict')
                import warnings
                warnings.warn("Passing 'dict' as keyword argument is deprecated",
                              DeprecationWarning, stacklevel=2)
            else:
                dict = None
            self.data = {}
            if dict is not None:
                self.update(dict)
            if len(kwargs):
                self.update(kwargs)
        def __len__(self): return len(self.data)
        def __getitem__(self, key):
            if key in self.data:
                return self.data[key]
            if hasattr(self.__class__, "__missing__"):
                return self.__class__.__missing__(self, key)
            raise KeyError(key)
        def __setitem__(self, key, item): self.data[key] = item
        def __delitem__(self, key): del self.data[key]
        def __iter__(self):
            return iter(self.data)

        # Modify __contains__ to work correctly when __missing__ is present
        def __contains__(self, key):
            return key in self.data

        # Now, add the methods in dicts but not in MutableMapping
        def __repr__(self): return repr(self.data)
        def copy(self):
            if self.__class__ is UserDict:
                return UserDict(self.data.copy())
            import copy
            data = self.data
            try:
                self.data = {}
                c = copy.copy(self)
            finally:
                self.data = data
            c.update(self)
            return c
        @classmethod
        def fromkeys(cls, iterable, value=None):
            d = cls()
            for key in iterable:
                d[key] = value
            return d

class decorate(object):
    def __init__(self, func):
        self.func = func

    def __call__(self, caller):
        """
        :param caller: A METHOD THAT IS EXPECTED TO CALL func
        :return: caller, BUT WITH SIGNATURE OF  self.func
        """
        return update_wrapper(caller, self.func)

_keep_imports = (ConfigParser, zip_longest, reduce, transpose, izip, HTMLParser, urlparse, StringIO, BytesIO, allocate_lock, get_ident, start_new_thread, interrupt_main)
<commit_msg>more generators<commit_after># encoding: utf-8
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Contact: Kyle Lahnakoski (kyle@lahnakoski.com)
#

from __future__ import absolute_import, division, unicode_literals

import json
import sys

PY3 = sys.version_info[0] == 3
PY2 = sys.version_info[0] == 2

PYPY = False
try:
    import __pypy__ as _
    PYPY=True
except Exception:
    PYPY=False

none_type = type(None)
boolean_type = type(True)

if PY3:
    import itertools
    import collections
    from collections import Callable
    from functools import cmp_to_key, reduce, update_wrapper
    from configparser import ConfigParser
    from itertools import zip_longest
    import builtins as __builtin__
    from builtins import input

    izip = zip
    zip_longest = itertools.zip_longest

    text = str
    text_type = str
    string_types = str
    binary_type = bytes
    integer_types = (int, )
    number_types = (int, float)
    long = int
    unichr = chr

    xrange = range
    def _gen():
        yield

    generator_types = (
        type(_gen()),
        type(filter(lambda x: True, [])),
        type({}.items()),
        type({}.values()),
        type(map(lambda: 0, [])),
        type(reversed([])),
        type(sorted([]))
    )
    unichr = chr

    round = round
    from html.parser import HTMLParser
    from urllib.parse import urlparse
    from io import StringIO
    from io import BytesIO
    from _thread import allocate_lock, get_ident, start_new_thread, interrupt_main

    def items(d):
        return list(d.items())

    def iteritems(d):
        return d.items()

    def transpose(*args):
        return list(zip(*args))

    def get_function_name(func):
        return func.__name__

    def get_function_arguments(func):
        return func.__code__.co_varnames[:func.__code__.co_argcount]

    def get_function_code(func):
        return func.__code__

    def get_function_defaults(func):
        return func.__defaults__

    def sort_using_cmp(data, cmp):
        return sorted(data, key=cmp_to_key(cmp))

    def sort_using_key(data, key):
        return sorted(data, key=key)

    def first(values):
        try:
            return iter(values).__next__()
        except StopIteration:
            return None

    def is_text(t):
        return t.__class__ is str

    def is_binary(b):
        return b.__class__ is bytes

    utf8_json_encoder = json.JSONEncoder(
        skipkeys=False,
        ensure_ascii=False,  # DIFF FROM DEFAULTS
        check_circular=True,
        allow_nan=True,
        indent=None,
        separators=(',', ':'),
        default=None,
        sort_keys=True   # <-- IMPORTANT!  sort_keys==True
    ).encode

    UserDict = collections.UserDict

else:
    import collections
    from collections import Callable
    from functools import cmp_to_key, reduce, update_wrapper

    import __builtin__
    from types import GeneratorType
    from ConfigParser import ConfigParser
    from itertools import izip_longest as zip_longest
    from __builtin__ import zip as transpose
    from itertools import izip
    from __builtin__ import raw_input as input

    reduce = __builtin__.reduce
    text = __builtin__.unicode
    text_type = __builtin__.unicode
    string_types = (str, unicode)
    binary_type = str
    integer_types = (int, long)
    number_types = (int, long, float)
    long = __builtin__.long
    unichr = __builtin__.unichr

    xrange = __builtin__.xrange
    generator_types = (GeneratorType, type(reversed([])))
    unichr = __builtin__.unichr

    round = __builtin__.round
    import HTMLParser
    from urlparse import urlparse
    from StringIO import StringIO
    from io import BytesIO
    from thread import allocate_lock, get_ident, start_new_thread, interrupt_main

    def items(d):
        return d.items()

    def iteritems(d):
        return d.iteritems()

    def get_function_name(func):
        return func.func_name

    def get_function_arguments(func):
        return func.func_code.co_varnames[:func.func_code.co_argcount]

    def get_function_code(func):
        return func.func_code

    def get_function_defaults(func):
        return func.func_defaults

    def sort_using_cmp(data, cmp):
        return sorted(data, cmp=cmp)

    def sort_using_key(data, key):
        return sorted(data, key=key)
        # return sorted(
        #     ((key(d), d) for d in data),
        #     lambda a, b: (1 if (a[0]>b[0]) else (-1 if (a[0]<b[0]) else 0))
        # )

    def first(values):
        try:
            return iter(values).next()
        except StopIteration:
            return None

    def is_text(t):
        return t.__class__ is unicode

    def is_binary(b):
        return b.__class__ is str

    utf8_json_encoder = json.JSONEncoder(
        skipkeys=False,
        ensure_ascii=False,  # DIFF FROM DEFAULTS
        check_circular=True,
        allow_nan=True,
        indent=None,
        separators=(',', ':'),
        encoding='utf-8',  # DIFF FROM DEFAULTS
        default=None,
        sort_keys=True   # <-- IMPORTANT!  sort_keys==True
    ).encode

    # COPIED FROM Python's collections.UserDict (copied July 2018)
    class UserDict(collections.MutableMapping):

        # Start by filling-out the abstract methods
        def __init__(*args, **kwargs):
            if not args:
                raise TypeError("descriptor '__init__' of 'UserDict' object "
                                "needs an argument")
            self, args = args[0], args[1:]
            if len(args) > 1:
                raise TypeError('expected at most 1 arguments, got %d' % len(args))
            if args:
                dict = args[0]
            elif 'dict' in kwargs:
                dict = kwargs.pop('dict')
                import warnings
                warnings.warn("Passing 'dict' as keyword argument is deprecated",
                              DeprecationWarning, stacklevel=2)
            else:
                dict = None
            self.data = {}
            if dict is not None:
                self.update(dict)
            if len(kwargs):
                self.update(kwargs)
        def __len__(self): return len(self.data)
        def __getitem__(self, key):
            if key in self.data:
                return self.data[key]
            if hasattr(self.__class__, "__missing__"):
                return self.__class__.__missing__(self, key)
            raise KeyError(key)
        def __setitem__(self, key, item): self.data[key] = item
        def __delitem__(self, key): del self.data[key]
        def __iter__(self):
            return iter(self.data)

        # Modify __contains__ to work correctly when __missing__ is present
        def __contains__(self, key):
            return key in self.data

        # Now, add the methods in dicts but not in MutableMapping
        def __repr__(self): return repr(self.data)
        def copy(self):
            if self.__class__ is UserDict:
                return UserDict(self.data.copy())
            import copy
            data = self.data
            try:
                self.data = {}
                c = copy.copy(self)
            finally:
                self.data = data
            c.update(self)
            return c
        @classmethod
        def fromkeys(cls, iterable, value=None):
            d = cls()
            for key in iterable:
                d[key] = value
            return d

class decorate(object):
    def __init__(self, func):
        self.func = func

    def __call__(self, caller):
        """
        :param caller: A METHOD THAT IS EXPECTED TO CALL func
        :return: caller, BUT WITH SIGNATURE OF  self.func
        """
        return update_wrapper(caller, self.func)

_keep_imports = (ConfigParser, zip_longest, reduce, transpose, izip, HTMLParser, urlparse, StringIO, BytesIO, allocate_lock, get_ident, start_new_thread, interrupt_main)
<|endoftext|>


<commit_before>#!/usr/bin/python

# This script does package renamings on source code (and .xml files and build files, etc), as part of migrating code to Jetpack.
# For example, this script may replace the text
#   import android.support.annotation.RequiresApi;
# with
#   import androidx.annotation.RequiresApi;

# See also b/74074903
import json
import os.path
import subprocess

HARDCODED_RULES_REVERSE = [
  "s|androidx.core.media.MediaBrowserCompat|android.support.v4.media.MediaBrowserCompat|g\n",
  "s|androidx.core.media.MediaDescriptionCompat|android.support.v4.media.MediaDescriptionCompat|g\n",
  "s|androidx.core.media.MediaMetadataCompat|android.support.v4.media.MediaMetadataCompat|g\n",
  "s|androidx.core.media.RatingCompat|android.support.v4.media.RatingCompat|g\n",
  "s|androidx.core.media.session.MediaControllerCompat|android.support.v4.media.session.MediaControllerCompat|g\n",
  "s|androidx.core.media.session.MediaSessionCompat|android.support.v4.media.session.MediaSessionCompat|g\n",
  "s|androidx.core.media.session.ParcelableVolumeInfo|android.support.v4.media.session.ParcelableVolumeInfo|g\n",
  "s|androidx.core.media.session.PlaybackStateCompat|android.support.v4.media.session.PlaybackStateCompat|g\n",
]

class StringBuilder(object):
  def __init__(self, item=None):
    self.items = []
    if item is not None:
      self.add(item)

  def add(self, item):
    self.items.append(str(item))
    return self

  def __str__(self):
    return "".join(self.items)

class ExecutionConfig(object):
  """Stores configuration about this execution of the package renaming.
     For example, file paths of the source code.
     This config could potentially be affected by command-line arguments.
  """
  def __init__(self, jetifierConfig, sourceRoots, excludeDirs):
    self.jetifierConfig = jetifierConfig
    self.sourceRoots = sourceRoots
    self.excludeDirs = excludeDirs

class SourceRewriteRule(object):
  def __init__(self, fromName, toName):
    self.fromName = fromName
    self.toName = toName

  def serialize(self):
    return self.fromName + ":" + self.toName

class JetifierConfig(object):
  """Stores configuration about the renaming itself, such as package rename rules.
     This config isn't supposed to be affected by command-line arguments.
  """
  @staticmethod
  def parse(filePath):
    with open(filePath) as f:
      lines = f.readlines()
      nonCommentLines = [line for line in lines if not line.strip().startswith("#")]
      parsed = json.loads("".join(nonCommentLines))
      return JetifierConfig(parsed)

  def __init__(self, parsedJson):
    self.json = parsedJson

  def getTypesMap(self):
    rules = []
    for rule in self.json["rules"]:
      fromName = rule["from"].replace("/", ".").replace("(.*)", "")
      toName = rule["to"].replace("/", ".").replace("{0}", "")
      if not toName.startswith("ignore"):
        rules.append(SourceRewriteRule(fromName, toName))

    return rules

def createRewriteCommand(executionConfig):
  # create command to find source files
  finderTextBuilder = StringBuilder("find")
  for sourceRoot in executionConfig.sourceRoots:
    finderTextBuilder.add(" ").add(sourceRoot)
  for exclusion in executionConfig.excludeDirs:
    finderTextBuilder.add(" -name ").add(exclusion).add(" -prune -o")
  finderTextBuilder.add(" -iregex '.*\.java\|.*\.xml\|.*\.cfg\|.*\.flags' -print")

  # create command to rewrite one source
  print("Building sed instructions")
  rewriterTextBuilder = StringBuilder()
  rewriteRules = executionConfig.jetifierConfig.getTypesMap()
  for rule in rewriteRules:
    rewriterTextBuilder.add("s|").add(rule.fromName).add("|").add(rule.toName).add("|g\n")
  for rule in HARDCODED_RULES_REVERSE:
    rewriterTextBuilder.add(rule)
  scriptPath = "/tmp/jetifier-sed-script.txt"
  print("Writing " + scriptPath)
  with open(scriptPath, 'w') as scriptFile:
    scriptFile.write(str(rewriterTextBuilder))

  # create the command to do the rewrites
  fullCommand = "time " + str(finderTextBuilder) + " | xargs -n 1 --no-run-if-empty -P 64 sed -i -f " + scriptPath

  return fullCommand

def processConfig(executionConfig):
  print("Building rewrite command")
  rewriteCommand = createRewriteCommand(executionConfig)
  commandLength = len(rewriteCommand)
  print("""
Will run command:

""" + rewriteCommand + """

""")
  response = raw_input("Ok? [y/n]")
  if response == "y":
    subprocess.check_output(rewriteCommand, shell=True)

def main():
  pathOfThisFile = os.path.realpath(__file__)
  jetifierPath = os.path.abspath(os.path.join(pathOfThisFile, "..", ".."))

  jetifierConfigPath = os.path.join(jetifierPath, "core", "src", "main", "resources", "default.config")
  print("Parsing " + jetifierConfigPath)
  jetifierConfig = JetifierConfig.parse(jetifierConfigPath)

  sourceRoot = os.getcwd()
  excludeDirs = ["out", ".git", ".repo"]

  executionConfig = ExecutionConfig(jetifierConfig, [sourceRoot], excludeDirs)

  processConfig(executionConfig)

main()

<commit_msg>Update package migration script for ResultsReceiver.<commit_after>#!/usr/bin/python

# This script does package renamings on source code (and .xml files and build files, etc), as part of migrating code to Jetpack.
# For example, this script may replace the text
#   import android.support.annotation.RequiresApi;
# with
#   import androidx.annotation.RequiresApi;

# See also b/74074903
import json
import os.path
import subprocess

HARDCODED_RULES_REVERSE = [
  "s|androidx.core.os.ResultReceiver|android.support.v4.os.ResultsReceiver|g\n",
  "s|androidx.core.media.MediaBrowserCompat|android.support.v4.media.MediaBrowserCompat|g\n",
  "s|androidx.core.media.MediaDescriptionCompat|android.support.v4.media.MediaDescriptionCompat|g\n",
  "s|androidx.core.media.MediaMetadataCompat|android.support.v4.media.MediaMetadataCompat|g\n",
  "s|androidx.core.media.RatingCompat|android.support.v4.media.RatingCompat|g\n",
  "s|androidx.core.media.session.MediaControllerCompat|android.support.v4.media.session.MediaControllerCompat|g\n",
  "s|androidx.core.media.session.MediaSessionCompat|android.support.v4.media.session.MediaSessionCompat|g\n",
  "s|androidx.core.media.session.ParcelableVolumeInfo|android.support.v4.media.session.ParcelableVolumeInfo|g\n",
  "s|androidx.core.media.session.PlaybackStateCompat|android.support.v4.media.session.PlaybackStateCompat|g\n",
]

class StringBuilder(object):
  def __init__(self, item=None):
    self.items = []
    if item is not None:
      self.add(item)

  def add(self, item):
    self.items.append(str(item))
    return self

  def __str__(self):
    return "".join(self.items)

class ExecutionConfig(object):
  """Stores configuration about this execution of the package renaming.
     For example, file paths of the source code.
     This config could potentially be affected by command-line arguments.
  """
  def __init__(self, jetifierConfig, sourceRoots, excludeDirs):
    self.jetifierConfig = jetifierConfig
    self.sourceRoots = sourceRoots
    self.excludeDirs = excludeDirs

class SourceRewriteRule(object):
  def __init__(self, fromName, toName):
    self.fromName = fromName
    self.toName = toName

  def serialize(self):
    return self.fromName + ":" + self.toName

class JetifierConfig(object):
  """Stores configuration about the renaming itself, such as package rename rules.
     This config isn't supposed to be affected by command-line arguments.
  """
  @staticmethod
  def parse(filePath):
    with open(filePath) as f:
      lines = f.readlines()
      nonCommentLines = [line for line in lines if not line.strip().startswith("#")]
      parsed = json.loads("".join(nonCommentLines))
      return JetifierConfig(parsed)

  def __init__(self, parsedJson):
    self.json = parsedJson

  def getTypesMap(self):
    rules = []
    for rule in self.json["rules"]:
      fromName = rule["from"].replace("/", ".").replace("(.*)", "")
      toName = rule["to"].replace("/", ".").replace("{0}", "")
      if not toName.startswith("ignore"):
        rules.append(SourceRewriteRule(fromName, toName))

    return rules

def createRewriteCommand(executionConfig):
  # create command to find source files
  finderTextBuilder = StringBuilder("find")
  for sourceRoot in executionConfig.sourceRoots:
    finderTextBuilder.add(" ").add(sourceRoot)
  for exclusion in executionConfig.excludeDirs:
    finderTextBuilder.add(" -name ").add(exclusion).add(" -prune -o")
  finderTextBuilder.add(" -iregex '.*\.java\|.*\.xml\|.*\.cfg\|.*\.flags' -print")

  # create command to rewrite one source
  print("Building sed instructions")
  rewriterTextBuilder = StringBuilder()
  rewriteRules = executionConfig.jetifierConfig.getTypesMap()
  for rule in rewriteRules:
    rewriterTextBuilder.add("s|").add(rule.fromName).add("|").add(rule.toName).add("|g\n")
  for rule in HARDCODED_RULES_REVERSE:
    rewriterTextBuilder.add(rule)
  scriptPath = "/tmp/jetifier-sed-script.txt"
  print("Writing " + scriptPath)
  with open(scriptPath, 'w') as scriptFile:
    scriptFile.write(str(rewriterTextBuilder))

  # create the command to do the rewrites
  fullCommand = "time " + str(finderTextBuilder) + " | xargs -n 1 --no-run-if-empty -P 64 sed -i -f " + scriptPath

  return fullCommand

def processConfig(executionConfig):
  print("Building rewrite command")
  rewriteCommand = createRewriteCommand(executionConfig)
  commandLength = len(rewriteCommand)
  print("""
Will run command:

""" + rewriteCommand + """

""")
  response = raw_input("Ok? [y/n]")
  if response == "y":
    subprocess.check_output(rewriteCommand, shell=True)

def main():
  pathOfThisFile = os.path.realpath(__file__)
  jetifierPath = os.path.abspath(os.path.join(pathOfThisFile, "..", ".."))

  jetifierConfigPath = os.path.join(jetifierPath, "core", "src", "main", "resources", "default.config")
  print("Parsing " + jetifierConfigPath)
  jetifierConfig = JetifierConfig.parse(jetifierConfigPath)

  sourceRoot = os.getcwd()
  excludeDirs = ["out", ".git", ".repo"]

  executionConfig = ExecutionConfig(jetifierConfig, [sourceRoot], excludeDirs)

  processConfig(executionConfig)

main()

<|endoftext|>


<commit_before># coding=utf-8
from __future__ import print_function, unicode_literals

import glob
import imp
import os
import time
import traceback

from bgmi.config import SCRIPT_PATH, MAX_PAGE
from bgmi.lib.download import download_prepare
from bgmi.lib.models import STATUS_UPDATED, STATUS_FOLLOWED
from bgmi.lib.models import Scripts
from bgmi.lib.fetch import DATA_SOURCE_MAP
from bgmi.utils import print_success, print_warning, print_info

class ScriptRunner(object):
    _defined = None
    scripts = []
    download_queue = []

    def __new__(cls, *args, **kwargs):
        if cls._defined is None:

            script_files = glob.glob('{}{}*.py'.format(SCRIPT_PATH, os.path.sep))
            for i in script_files:
                try:
                    s = imp.load_source('script', os.path.join(SCRIPT_PATH, i))
                    script_class = getattr(s, 'Script')()

                    if cls.check(script_class):
                        cls.scripts.append(script_class)
                        print_info('Load script {} successfully.'.format(i))

                except:
                    print_warning('Load script {} failed, ignored'.format(i))
                    if os.getenv('DEBUG_SCRIPT'):  # pragma: no cover
                        traceback.print_exc()
                        # self.scripts = filter(self._check_followed, self.scripts)
                        # self.scripts = filter(self._check_bangumi, self.scripts)

            cls._defined = super(ScriptRunner, cls).__new__(cls, *args, **kwargs)

        return cls._defined

    @classmethod
    def check(cls, script):
        condition = [
            'script.Model().due_date > datetime.datetime.now()',
        ]

        for i in condition:
            try:
                if not eval(i):
                    return False
            except:
                # ignore if error
                if os.getenv('DEBUG_SCRIPT'):  # pragma: no cover
                    traceback.print_exc()

        return True

    def get_model(self, name):
        for script in self.scripts:
            if script.Model.bangumi_name == name:
                return script.Model().obj

    def get_models_dict(self):
        return [dict(script.Model()) for script in self.scripts if script.bangumi_name is not None]

    @staticmethod
    def make_dict(script):
        return [{
            'name': script.bangumi_name,
            'title': '[{}][{}]'.format(script.bangumi_name, k),
            'episode': k,
            'download': v
        } for k, v in script.get_download_url().items()]

    def run(self, return_=True, download=False):
        for script in self.scripts:
            print_info('fetching {} ...'.format(script.bangumi_name))
            download_item = self.make_dict(script)

            script_obj = script.Model().obj

            if not download_item:
                print_info('Got nothing, quit script {}.'.format(script))
                continue

            max_episode = max(download_item, key=lambda d: d['episode'])
            episode = max_episode['episode']
            episode_range = range(script_obj.episode + 1, episode + 1)

            if episode <= script_obj.episode:
                continue

            print_success('{} updated, episode: {}'.format(script.bangumi_name, episode))
            script_obj.episode = episode
            script_obj.status = STATUS_UPDATED
            script_obj.updated_time = int(time.time())
            script_obj.save()

            download_queue = []
            for i in episode_range:
                for e in download_item:
                    if i == e['episode']:
                        download_queue.append(e)

            if return_:
                self.download_queue.extend(download_queue)
                continue

            if download:
                print_success('Start downloading of {}'.format(script))
                download_prepare(download_queue)

        return self.download_queue

    def get_download_cover(self):
        return [script['cover'] for script in self.get_models_dict()]

class ScriptBase(object):
    class Model(object):
        obj = None

        # data
        bangumi_name = None
        cover = None
        update_time = None
        due_date = None

        # source
        source = None

        # fetch_episode_of_bangumi(self, bangumi_id, subtitle_list=None, max_page=MAX_PAGE):
        _bangumi_id = None
        _subtitle_list = []
        _max_page = MAX_PAGE

        def __init__(self):
            if self.bangumi_name is not None:
                s, _ = Scripts.get_or_create(bangumi_name=self.bangumi_name,
                                             defaults={'episode': 0, 'status': STATUS_FOLLOWED})
                self.obj = s

        def __iter__(self):
            for i in ('bangumi_name', 'cover', 'update_time'):
                yield (i, getattr(self, i))

            # patch for cal
            yield ('name', self.bangumi_name)
            yield ('status', self.obj.status)
            yield ('updated_time', self.obj.updated_time)
            yield ('subtitle_group', '')
            yield ('episode', self.obj.episode)

    @property
    def _data(self):
        return {
            'bangumi_id': self.Model._bangumi_id,
            'subtitle_list': self.Model._subtitle_list,
            'max_page': int(self.Model._max_page),
        }

    @property
    def source(self):
        return self.Model.source

    @property
    def name(self):
        return self.Model.bangumi_name

    @property
    def bangumi_name(self):
        return self.Model.bangumi_name

    @property
    def cover(self):
        return self.Model.cover

    @property
    def updated_time(self):
        return self.Model.update_time

    def get_download_url(self):
        """Get the download url, and return a dict of episode and the url.
        Download url also can be magnet link.
        For example:
        ```
            {
                1: 'http://example.com/Bangumi/1/1.mp4'
                2: 'http://example.com/Bangumi/1/2.mp4'
                3: 'http://example.com/Bangumi/1/3.mp4'
            }
        ```
        The keys `1`, `2`, `3` is the episode, the value is the url of bangumi.
        :return: dict
        """
        if self.source is not None:
            source = DATA_SOURCE_MAP.get(self.source, None)()
            if source is None:
                raise Exception('Script data source is invalid, usable sources: {}'
                                .format(', '.join(DATA_SOURCE_MAP.keys())))
            ret = {}
            data = source.fetch_episode_of_bangumi(**self._data)
            for i in data:
                if int(i['episode']) not in data:
                    ret[int(i['episode'])] = i['download']
            return ret
        else:
            return {}

if __name__ == '__main__':
    runner = ScriptRunner()
    runner.run()
<commit_msg>fix bug of bgmi script<commit_after># coding=utf-8
from __future__ import print_function, unicode_literals

import glob
import imp
import os
import time
import datetime
import traceback

from bgmi.config import SCRIPT_PATH, MAX_PAGE
from bgmi.lib.download import download_prepare
from bgmi.lib.models import STATUS_UPDATED, STATUS_FOLLOWED
from bgmi.lib.models import Scripts
from bgmi.lib.fetch import DATA_SOURCE_MAP
from bgmi.utils import print_success, print_warning, print_info

class ScriptRunner(object):
    _defined = None
    scripts = []
    download_queue = []

    def __new__(cls, *args, **kwargs):
        if cls._defined is None:

            script_files = glob.glob('{}{}*.py'.format(SCRIPT_PATH, os.path.sep))
            for i in script_files:
                try:
                    s = imp.load_source('script', os.path.join(SCRIPT_PATH, i))
                    script_class = getattr(s, 'Script')()

                    if cls.check(script_class):
                        cls.scripts.append(script_class)
                        print_info('Load script {} successfully.'.format(i))

                except:
                    print_warning('Load script {} failed, ignored'.format(i))
                    if os.getenv('DEBUG_SCRIPT'):  # pragma: no cover
                        traceback.print_exc()
                        # self.scripts = filter(self._check_followed, self.scripts)
                        # self.scripts = filter(self._check_bangumi, self.scripts)

            cls._defined = super(ScriptRunner, cls).__new__(cls, *args, **kwargs)

        return cls._defined

    @classmethod
    def check(cls, script):
        condition = [
            'script.Model().due_date > datetime.datetime.now()',
        ]

        for i in condition:
            try:
                if not eval(i):
                    return False
            except:
                # ignore if error
                if os.getenv('DEBUG_SCRIPT'):  # pragma: no cover
                    traceback.print_exc()

        return True

    def get_model(self, name):
        for script in self.scripts:
            if script.Model.bangumi_name == name:
                return script.Model().obj

    def get_models_dict(self):
        return [dict(script.Model()) for script in self.scripts if script.bangumi_name is not None]

    @staticmethod
    def make_dict(script):
        return [{
            'name': script.bangumi_name,
            'title': '[{}][{}]'.format(script.bangumi_name, k),
            'episode': k,
            'download': v
        } for k, v in script.get_download_url().items()]

    def run(self, return_=True, download=False):
        for script in self.scripts:
            print_info('fetching {} ...'.format(script.bangumi_name))
            download_item = self.make_dict(script)

            script_obj = script.Model().obj

            if not download_item:
                print_info('Got nothing, quit script {}.'.format(script))
                continue

            max_episode = max(download_item, key=lambda d: d['episode'])
            episode = max_episode['episode']
            episode_range = range(script_obj.episode + 1, episode + 1)

            if episode <= script_obj.episode:
                continue

            print_success('{} updated, episode: {}'.format(script.bangumi_name, episode))
            script_obj.episode = episode
            script_obj.status = STATUS_UPDATED
            script_obj.updated_time = int(time.time())
            script_obj.save()

            download_queue = []
            for i in episode_range:
                for e in download_item:
                    if i == e['episode']:
                        download_queue.append(e)

            if return_:
                self.download_queue.extend(download_queue)
                continue

            if download:
                print_success('Start downloading of {}'.format(script))
                download_prepare(download_queue)

        return self.download_queue

    def get_download_cover(self):
        return [script['cover'] for script in self.get_models_dict()]

class ScriptBase(object):
    class Model(object):
        obj = None

        # data
        bangumi_name = None
        cover = None
        update_time = None
        due_date = None

        # source
        source = None

        # fetch_episode_of_bangumi(self, bangumi_id, subtitle_list=None, max_page=MAX_PAGE):
        _bangumi_id = None
        _subtitle_list = []
        _max_page = MAX_PAGE

        def __init__(self):
            if self.bangumi_name is not None:
                s, _ = Scripts.get_or_create(bangumi_name=self.bangumi_name,
                                             defaults={'episode': 0, 'status': STATUS_FOLLOWED})
                self.obj = s

        def __iter__(self):
            for i in ('bangumi_name', 'cover', 'update_time'):
                yield (i, getattr(self, i))

            # patch for cal
            yield ('name', self.bangumi_name)
            yield ('status', self.obj.status)
            yield ('updated_time', self.obj.updated_time)
            yield ('subtitle_group', '')
            yield ('episode', self.obj.episode)

    @property
    def _data(self):
        return {
            'bangumi_id': self.Model._bangumi_id,
            'subtitle_list': self.Model._subtitle_list,
            'max_page': int(self.Model._max_page),
        }

    @property
    def source(self):
        return self.Model.source

    @property
    def name(self):
        return self.Model.bangumi_name

    @property
    def bangumi_name(self):
        return self.Model.bangumi_name

    @property
    def cover(self):
        return self.Model.cover

    @property
    def updated_time(self):
        return self.Model.update_time

    def get_download_url(self):
        """Get the download url, and return a dict of episode and the url.
        Download url also can be magnet link.
        For example:
        ```
            {
                1: 'http://example.com/Bangumi/1/1.mp4'
                2: 'http://example.com/Bangumi/1/2.mp4'
                3: 'http://example.com/Bangumi/1/3.mp4'
            }
        ```
        The keys `1`, `2`, `3` is the episode, the value is the url of bangumi.
        :return: dict
        """
        if self.source is not None:
            source = DATA_SOURCE_MAP.get(self.source, None)()
            if source is None:
                raise Exception('Script data source is invalid, usable sources: {}'
                                .format(', '.join(DATA_SOURCE_MAP.keys())))
            ret = {}
            data = source.fetch_episode_of_bangumi(**self._data)
            for i in data:
                if int(i['episode']) not in data:
                    ret[int(i['episode'])] = i['download']
            return ret
        else:
            return {}

if __name__ == '__main__':
    runner = ScriptRunner()
    runner.run()
<|endoftext|>