<commit_before># Copyright (c) 2021, Google Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of Google Inc. nor the names of its contributors
#    may be used to endorse or promote products derived from this software without
#    specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
r"""Model distillation training binary using a custom training loop.

Distillation attempts to train a smaller student model that mimics the larger
teacher model.

Currently only transformer_learn_values_v2_distill config is
supported for model training.

Example usage:

CONFIG="//learning/genomics/deepconsensus/models/model_configs.py:transformer_learn_values_v2_distill+ccs"
TEACHER_MODEL_DIR=""
OUT_DIR=/tmp

time blaze run -c opt \
//learning/genomics/deepconsensus/models:model_distillation -- \
  --teacher_model_dir ${TEACHER_MODEL_DIR} \
  --params ${CONFIG} \
  --out_dir ${OUT_DIR} \
  --alsologtostderr
"""

import logging
import os
import random
from typing import Optional, Dict

from absl import app
from absl import flags
import ml_collections
from ml_collections.config_flags import config_flags
import tensorflow as tf

from deepconsensus.models import data_providers
from deepconsensus.models import losses_and_metrics
from deepconsensus.models import model_utils

# pylint: disable=unused-import

FLAGS = flags.FLAGS
config_flags.DEFINE_config_file('params', None, 'Training configuration.')
flags.DEFINE_string('teacher_model_dir', None,
                    'Path to the teacher model checkpoint.')
flags.DEFINE_string('out_dir', None,
                    'Output path for logs and model checkpoints.')
flags.DEFINE_string(
    'tpu', None, 'Name of the TPU to use. This gets '
    'populated automatically when using XManager.')
flags.DEFINE_string('tpu_topology', None, 'Tpu topology.')
flags.DEFINE_bool('debug', False,
                  'Enables dumping debug info for TensorBoard Debugger V2.')
flags.DEFINE_bool(
    'write_checkpoint_metrics', False,
    'Whether to write eval metrics for each checkpoint during training.')
flags.DEFINE_bool(
    'eval_and_log_every_step', False, 'Eval and log after every step. '
    'Use this e.g. for testing training and inspecting metrics locally.')

def get_teacher_model(checkpoint_path: str,
                      strategy: tf.distribute.Strategy) -> tf.keras.Model:
  """Get teacher model with an existing checkpoint."""
  params = model_utils.read_params_from_json(checkpoint_path=checkpoint_path)
  with strategy.scope():
    logging.info('Using checkpoint: %s.', checkpoint_path)
    model = model_utils.get_model(params)
    checkpoint = tf.train.Checkpoint(model=model)
    # Note that the `print_model_summary` is necessary because we need to run a
    # forward pass with the model in order for assert_existing_objects_matched
    # to work as expected.
    # If you don't do this, then  assert_existing_objects_matched will not
    # raise an error even if the wrong checkpoint is used.
    # Some context here: b/148023980.
    row_size = data_providers.get_total_rows(params.max_passes)
    input_shape = (1, row_size, params.max_length, params.num_channels)
    model_utils.print_model_summary(model, input_shape)
    checkpoint.restore(
        checkpoint_path).expect_partial().assert_existing_objects_matched()
  return model

def train_model(teacher_model: tf.keras.Model, out_dir: str,
                params: ml_collections.ConfigDict,
                strategy: tf.distribute.Strategy,
                write_checkpoint_metrics: bool) -> None:
  """Trains the model under the given strategy and params."""
  # Freeze config dict here to ensure it is hashable.
  params = ml_collections.FrozenConfigDict(params)
  model_utils.save_params_as_json(out_dir, params)
  train_dataset, eval_dataset = model_utils.get_datasets(params, strategy)
  steps_per_epoch, steps_per_eval = model_utils.get_step_counts(
      params, FLAGS.eval_and_log_every_step)

  with strategy.scope():
    logging.info('Building model.')
    model = model_utils.get_model(params)
    logging.info('Done building model.')
    optimizer = tf.keras.optimizers.Adam(learning_rate=params.learning_rate)
    train_loss = tf.keras.metrics.Mean(name='loss')
    train_metrics = model_utils.get_deepconsensus_metrics(name_prefix='')
    eval_loss = tf.keras.metrics.Mean(name='loss')
    eval_metrics = model_utils.get_deepconsensus_metrics(name_prefix='')
    student_loss_object = model_utils.get_deepconsensus_loss(
        params, reduction=tf.keras.losses.Reduction.NONE)
    distillation_loss_object = losses_and_metrics.DistillationLoss(
        temperature=params.temperature,
        reduction=tf.keras.losses.Reduction.NONE)

    def compute_all_replica_loss(
        per_example_loss: tf.keras.losses.Loss,
        params: ml_collections.ConfigDict) -> tf.Tensor:
      # We divide per-replica losses by global batch size and sum this value
      # across all replicas to compute average loss scaled by global batch size.
      return tf.nn.compute_average_loss(
          per_example_loss, global_batch_size=params.batch_size)

    def compute_loss(labels: tf.Tensor, student_preds: tf.Tensor,
                     student_logits: tf.Tensor,
                     teacher_logits: tf.Tensor) -> Dict[str, tf.Tensor]:
      per_example_student_loss = student_loss_object(labels, student_preds)
      per_example_distill_loss = distillation_loss_object(
          teacher_logits, student_logits)
      per_example_loss = (
          params.student_alpha * per_example_student_loss +
          params.distill_alpha * per_example_distill_loss)

      losses_dict = {}
      losses_dict['total_loss'] = compute_all_replica_loss(
          per_example_loss, params)
      losses_dict['student_loss'] = compute_all_replica_loss(
          per_example_student_loss, params)
      losses_dict['distill_loss'] = compute_all_replica_loss(
          per_example_distill_loss, params)
      return losses_dict

    # model, optimizer, and checkpoint must be created under `strategy.scope`.
    checkpoint, initial_epoch = model_utils.get_checkpoint_and_initial_epoch(
        model, optimizer, out_dir, steps_per_epoch)  # pytype: disable=wrong-arg-types  # typed-keras

  # Create summary writers
  train_writer = tf.summary.create_file_writer(os.path.join(out_dir, 'train'))
  eval_writer = tf.summary.create_file_writer(os.path.join(out_dir, 'eval'))

  def train_step(inputs):
    """Training StepFn."""
    features, labels = inputs
    # Get logits from the teacher model.
    teacher_logits = teacher_model.get_logits(features, training=False)

    with tf.GradientTape() as tape:
      student_logits = model.get_logits(features, training=True)
      student_preds = tf.nn.softmax(student_logits)
      train_losses_dict = compute_loss(labels, student_preds, student_logits,
                                       teacher_logits)
      loss = train_losses_dict['total_loss']
    # Compute gradients.
    grads = tape.gradient(loss, model.trainable_variables)
    # Update weights.
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    train_loss.update_state(loss)
    for metric in train_metrics:
      metric.update_state(labels, student_preds)
    return train_losses_dict

  def eval_step(inputs):
    """Eval StepFn."""
    features, labels = inputs
    # Get logits from the teacher model.
    teacher_logits = teacher_model.get_logits(features, training=False)

    student_logits = model.get_logits(features, training=False)
    student_preds = tf.nn.softmax(student_logits)
    eval_losses_dict = compute_loss(labels, student_preds, student_logits,
                                    teacher_logits)

    eval_loss.update_state(eval_losses_dict['total_loss'])
    for metric in eval_metrics:
      metric.update_state(labels, student_preds)
    return eval_losses_dict

  @tf.function
  def distributed_train_step(iterator):
    per_replica_losses_dict = strategy.run(train_step, args=(next(iterator),))
    reduced_train_losses_dict = {}
    for loss_name in per_replica_losses_dict.keys():
      reduced_train_losses_dict[loss_name] = strategy.reduce(
          tf.distribute.ReduceOp.SUM,
          per_replica_losses_dict[loss_name],
          axis=None)
    return reduced_train_losses_dict

  @tf.function
  def distributed_eval_step(iterator):
    per_replica_losses_dict = strategy.run(eval_step, args=(next(iterator),))
    reduced_eval_losses_dict = {}
    for loss_name in per_replica_losses_dict.keys():
      reduced_eval_losses_dict[loss_name] = strategy.reduce(
          tf.distribute.ReduceOp.SUM,
          per_replica_losses_dict[loss_name],
          axis=None)
    return reduced_eval_losses_dict

  log_steps = 100
  if FLAGS.eval_and_log_every_step:
    log_steps = 1
  train_iterator = iter(train_dataset)
  eval_iterator = iter(eval_dataset)
  min_eval_loss = 1e6
  for epoch in range(initial_epoch, params['num_epochs']):
    logging.info('Starting to run epoch: %s', epoch)
    with train_writer.as_default():
      for step in range(steps_per_epoch):
        reduced_train_losses = distributed_train_step(train_iterator)
        if step % log_steps == 0:
          model_utils.log_and_save_metrics(
              epoch=epoch,
              step=step,
              total_steps=steps_per_epoch,
              optimizer=optimizer,
              losses_dict=reduced_train_losses,
              metrics=train_metrics,
              training=True)
    with eval_writer.as_default():
      for step in range(steps_per_eval):
        reduced_eval_losses = distributed_eval_step(eval_iterator)
      model_utils.log_and_save_metrics(
          epoch=epoch,
          step=step,
          total_steps=steps_per_eval,
          optimizer=optimizer,
          losses_dict=reduced_eval_losses,
          metrics=eval_metrics,
          training=False)
    checkpoint_name = model_utils.save_checkpoint(checkpoint, out_dir,
                                                  train_metrics, eval_metrics,
                                                  write_checkpoint_metrics)
    if min_eval_loss > float(eval_loss.result()):
      min_eval_loss = float(eval_loss.result())
      with tf.io.gfile.GFile(os.path.join(out_dir, 'best_checkpoint.txt'),
                             'w') as f:
        f.write(os.path.basename(checkpoint_name))
    model_utils.reset_all_metrics([train_loss, eval_loss] + train_metrics +
                                  eval_metrics)

def train(teacher_model_dir: str,
          out_dir: str,
          params: ml_collections.ConfigDict,
          tpu: Optional[str],
          tpu_topology: Optional[str],
          write_checkpoint_metrics: bool,
          debug: Optional[bool] = False):
  """Run the model training and return evaluation output."""
  model_utils.modify_params(params, tpu=tpu, tpu_topology=tpu_topology)
  random.seed(params.seed)
  tf.random.set_seed(params.seed)
  os.environ['TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE'] = 'False'
  while True:
    try:
      if tpu is not None:
        resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu)
        tf.config.experimental_connect_to_cluster(resolver)
        tf.tpu.experimental.initialize_tpu_system(resolver)
        strategy = tf.distribute.TPUStrategy(resolver)
      elif debug:
        strategy = tf.distribute.OneDeviceStrategy(device='/cpu:0')
      else:
        strategy = tf.distribute.MirroredStrategy()
      teacher_model = get_teacher_model(teacher_model_dir, strategy=strategy)
      train_model(teacher_model, out_dir, params, strategy,
                  write_checkpoint_metrics)
      break
    except tf.errors.UnavailableError:
      continue

def main(unused_args=None):
  train(FLAGS.teacher_model_dir, FLAGS.out_dir, FLAGS.params, FLAGS.tpu,
        FLAGS.tpu_topology, FLAGS.write_checkpoint_metrics, FLAGS.debug)

if __name__ == '__main__':
  flags.mark_flags_as_required([
      'teacher_model_dir',
      'params',
      'out_dir',
  ])
  app.run(main)
<commit_msg>Enable initialization of student model weight from teacher model.<commit_after># Copyright (c) 2021, Google Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of Google Inc. nor the names of its contributors
#    may be used to endorse or promote products derived from this software without
#    specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
r"""Model distillation training binary using a custom training loop.

Distillation attempts to train a smaller student model that mimics the larger
teacher model.

Currently only transformer_learn_values_v2_distill config is
supported for model training.

Example usage:

CONFIG="//learning/genomics/deepconsensus/models/model_configs.py:transformer_learn_values_v2_distill+ccs"
TEACHER_MODEL_DIR=""
OUT_DIR=/tmp

time blaze run -c opt \
//learning/genomics/deepconsensus/models:model_distillation -- \
  --teacher_model_dir ${TEACHER_MODEL_DIR} \
  --params ${CONFIG} \
  --out_dir ${OUT_DIR} \
  --alsologtostderr
"""

import logging
import os
import random
from typing import Optional, Dict

from absl import app
from absl import flags
import ml_collections
from ml_collections.config_flags import config_flags
import tensorflow as tf

from deepconsensus.models import data_providers
from deepconsensus.models import losses_and_metrics
from deepconsensus.models import model_utils

# pylint: disable=unused-import

FLAGS = flags.FLAGS
config_flags.DEFINE_config_file('params', None, 'Training configuration.')
flags.DEFINE_string('teacher_model_dir', None,
                    'Path to the teacher model checkpoint.')
flags.DEFINE_string('out_dir', None,
                    'Output path for logs and model checkpoints.')
flags.DEFINE_string(
    'tpu', None, 'Name of the TPU to use. This gets '
    'populated automatically when using XManager.')
flags.DEFINE_string('tpu_topology', None, 'Tpu topology.')
flags.DEFINE_bool('debug', False,
                  'Enables dumping debug info for TensorBoard Debugger V2.')
flags.DEFINE_bool(
    'write_checkpoint_metrics', False,
    'Whether to write eval metrics for each checkpoint during training.')
flags.DEFINE_bool(
    'eval_and_log_every_step', False, 'Eval and log after every step. '
    'Use this e.g. for testing training and inspecting metrics locally.')

def init_student_from_teacher(
    student_model: tf.keras.Model, teacher_model: tf.keras.Model,
    params: ml_collections.ConfigDict) -> tf.keras.Model:
  """Initialize student model using teacher model weights based on params."""
  row_size = data_providers.get_total_rows(params.max_passes)
  input_shape = (1, row_size, params.max_length, params.num_channels)
  model_utils.print_model_summary(teacher_model, input_shape)
  if params.init_encoder_stack:
    teacher2student_encoder_map = dict(
        zip(params.teacher_encoder_layers, params.student_encoder_layers))

    for teacher_layer_id in teacher2student_encoder_map:
      student_layer_id = teacher2student_encoder_map[teacher_layer_id]
      # Copy attention layer.
      teacher_weights = teacher_model.encoder_stack.layers[teacher_layer_id][
          0].layer.get_weights()
      student_model.encoder_stack.layers[student_layer_id][0].layer.set_weights(
          teacher_weights)
      # Copy ffn layer.
      teacher_weights = teacher_model.encoder_stack.layers[teacher_layer_id][
          1].layer.get_weights()
      student_model.encoder_stack.layers[student_layer_id][1].layer.set_weights(
          teacher_weights)

  if params.init_nonencoder_layers:
    # Get layers with weights that are not in the encoder stack.
    layer_ind_for_copy = []
    for ind, layer in enumerate(teacher_model.layers):
      if (layer.trainable_weights) and ('encoder_stack' not in layer.name):
        layer_ind_for_copy.append(ind)

    for layer_ind in layer_ind_for_copy:
      teacher_weights = teacher_model.get_layer(index=layer_ind).get_weights()
      student_model.get_layer(index=layer_ind).set_weights(teacher_weights)
  return student_model

def get_teacher_model(checkpoint_path: str,
                      strategy: tf.distribute.Strategy) -> tf.keras.Model:
  """Get teacher model with an existing checkpoint."""
  params = model_utils.read_params_from_json(checkpoint_path=checkpoint_path)
  with strategy.scope():
    logging.info('Using checkpoint: %s.', checkpoint_path)
    model = model_utils.get_model(params)
    checkpoint = tf.train.Checkpoint(model=model)
    # Note that the `print_model_summary` is necessary because we need to run a
    # forward pass with the model in order for assert_existing_objects_matched
    # to work as expected.
    # If you don't do this, then  assert_existing_objects_matched will not
    # raise an error even if the wrong checkpoint is used.
    # Some context here: b/148023980.
    row_size = data_providers.get_total_rows(params.max_passes)
    input_shape = (1, row_size, params.max_length, params.num_channels)
    model_utils.print_model_summary(model, input_shape)
    checkpoint.restore(
        checkpoint_path).expect_partial().assert_existing_objects_matched()
  return model

def train_model(teacher_model: tf.keras.Model, out_dir: str,
                params: ml_collections.ConfigDict,
                strategy: tf.distribute.Strategy,
                write_checkpoint_metrics: bool) -> None:
  """Trains the model under the given strategy and params."""
  # Freeze config dict here to ensure it is hashable.
  params = ml_collections.FrozenConfigDict(params)
  model_utils.save_params_as_json(out_dir, params)
  train_dataset, eval_dataset = model_utils.get_datasets(params, strategy)
  steps_per_epoch, steps_per_eval = model_utils.get_step_counts(
      params, FLAGS.eval_and_log_every_step)

  with strategy.scope():
    logging.info('Building model.')
    model = model_utils.get_model(params)
    # Note that the `print_model_summary` is necessary because we need to run a
    # forward pass with the model to be able to initialize student from teacher.
    row_size = data_providers.get_total_rows(params.max_passes)
    input_shape = (1, row_size, params.max_length, params.num_channels)
    model_utils.print_model_summary(model, input_shape)
    logging.info('Done building model.')
    # Initialize student model from teacher based on model params.
    model = init_student_from_teacher(model, teacher_model, params)
    optimizer = tf.keras.optimizers.Adam(learning_rate=params.learning_rate)
    train_loss = tf.keras.metrics.Mean(name='loss')
    train_metrics = model_utils.get_deepconsensus_metrics(name_prefix='')
    eval_loss = tf.keras.metrics.Mean(name='loss')
    eval_metrics = model_utils.get_deepconsensus_metrics(name_prefix='')
    student_loss_object = model_utils.get_deepconsensus_loss(
        params, reduction=tf.keras.losses.Reduction.NONE)
    distillation_loss_object = losses_and_metrics.DistillationLoss(
        temperature=params.temperature,
        reduction=tf.keras.losses.Reduction.NONE)

    def compute_all_replica_loss(
        per_example_loss: tf.keras.losses.Loss,
        params: ml_collections.ConfigDict) -> tf.Tensor:
      # We divide per-replica losses by global batch size and sum this value
      # across all replicas to compute average loss scaled by global batch size.
      return tf.nn.compute_average_loss(
          per_example_loss, global_batch_size=params.batch_size)

    def compute_loss(labels: tf.Tensor, student_preds: tf.Tensor,
                     student_logits: tf.Tensor,
                     teacher_logits: tf.Tensor) -> Dict[str, tf.Tensor]:
      per_example_student_loss = student_loss_object(labels, student_preds)
      per_example_distill_loss = distillation_loss_object(
          teacher_logits, student_logits)
      per_example_loss = (
          params.student_alpha * per_example_student_loss +
          params.distill_alpha * per_example_distill_loss)

      losses_dict = {}
      losses_dict['total_loss'] = compute_all_replica_loss(
          per_example_loss, params)
      losses_dict['student_loss'] = compute_all_replica_loss(
          per_example_student_loss, params)
      losses_dict['distill_loss'] = compute_all_replica_loss(
          per_example_distill_loss, params)
      return losses_dict

    # model, optimizer, and checkpoint must be created under `strategy.scope`.
    checkpoint, initial_epoch = model_utils.get_checkpoint_and_initial_epoch(
        model, optimizer, out_dir, steps_per_epoch)  # pytype: disable=wrong-arg-types  # typed-keras

  # Create summary writers
  train_writer = tf.summary.create_file_writer(os.path.join(out_dir, 'train'))
  eval_writer = tf.summary.create_file_writer(os.path.join(out_dir, 'eval'))

  def train_step(inputs):
    """Training StepFn."""
    features, labels = inputs
    # Get logits from the teacher model.
    teacher_logits = teacher_model.get_logits(features, training=False)

    with tf.GradientTape() as tape:
      student_logits = model.get_logits(features, training=True)
      student_preds = tf.nn.softmax(student_logits)
      train_losses_dict = compute_loss(labels, student_preds, student_logits,
                                       teacher_logits)
      loss = train_losses_dict['total_loss']
    # Compute gradients.
    grads = tape.gradient(loss, model.trainable_variables)
    # Update weights.
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    train_loss.update_state(loss)
    for metric in train_metrics:
      metric.update_state(labels, student_preds)
    return train_losses_dict

  def eval_step(inputs):
    """Eval StepFn."""
    features, labels = inputs
    # Get logits from the teacher model.
    teacher_logits = teacher_model.get_logits(features, training=False)

    student_logits = model.get_logits(features, training=False)
    student_preds = tf.nn.softmax(student_logits)
    eval_losses_dict = compute_loss(labels, student_preds, student_logits,
                                    teacher_logits)

    eval_loss.update_state(eval_losses_dict['total_loss'])
    for metric in eval_metrics:
      metric.update_state(labels, student_preds)
    return eval_losses_dict

  @tf.function
  def distributed_train_step(iterator):
    per_replica_losses_dict = strategy.run(train_step, args=(next(iterator),))
    reduced_train_losses_dict = {}
    for loss_name in per_replica_losses_dict.keys():
      reduced_train_losses_dict[loss_name] = strategy.reduce(
          tf.distribute.ReduceOp.SUM,
          per_replica_losses_dict[loss_name],
          axis=None)
    return reduced_train_losses_dict

  @tf.function
  def distributed_eval_step(iterator):
    per_replica_losses_dict = strategy.run(eval_step, args=(next(iterator),))
    reduced_eval_losses_dict = {}
    for loss_name in per_replica_losses_dict.keys():
      reduced_eval_losses_dict[loss_name] = strategy.reduce(
          tf.distribute.ReduceOp.SUM,
          per_replica_losses_dict[loss_name],
          axis=None)
    return reduced_eval_losses_dict

  log_steps = 100
  if FLAGS.eval_and_log_every_step:
    log_steps = 1
  train_iterator = iter(train_dataset)
  eval_iterator = iter(eval_dataset)
  min_eval_loss = 1e6
  for epoch in range(initial_epoch, params['num_epochs']):
    logging.info('Starting to run epoch: %s', epoch)
    with train_writer.as_default():
      for step in range(steps_per_epoch):
        reduced_train_losses = distributed_train_step(train_iterator)
        if step % log_steps == 0:
          model_utils.log_and_save_metrics(
              epoch=epoch,
              step=step,
              total_steps=steps_per_epoch,
              optimizer=optimizer,
              losses_dict=reduced_train_losses,
              metrics=train_metrics,
              training=True)
    with eval_writer.as_default():
      for step in range(steps_per_eval):
        reduced_eval_losses = distributed_eval_step(eval_iterator)
      model_utils.log_and_save_metrics(
          epoch=epoch,
          step=step,
          total_steps=steps_per_eval,
          optimizer=optimizer,
          losses_dict=reduced_eval_losses,
          metrics=eval_metrics,
          training=False)
    checkpoint_name = model_utils.save_checkpoint(checkpoint, out_dir,
                                                  train_metrics, eval_metrics,
                                                  write_checkpoint_metrics)
    if min_eval_loss > float(eval_loss.result()):
      min_eval_loss = float(eval_loss.result())
      with tf.io.gfile.GFile(os.path.join(out_dir, 'best_checkpoint.txt'),
                             'w') as f:
        f.write(os.path.basename(checkpoint_name))
    model_utils.reset_all_metrics([train_loss, eval_loss] + train_metrics +
                                  eval_metrics)

def train(teacher_model_dir: str,
          out_dir: str,
          params: ml_collections.ConfigDict,
          tpu: Optional[str],
          tpu_topology: Optional[str],
          write_checkpoint_metrics: bool,
          debug: Optional[bool] = False):
  """Run the model training and return evaluation output."""
  model_utils.modify_params(params, tpu=tpu, tpu_topology=tpu_topology)
  random.seed(params.seed)
  tf.random.set_seed(params.seed)
  os.environ['TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE'] = 'False'
  while True:
    try:
      if tpu is not None:
        resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu)
        tf.config.experimental_connect_to_cluster(resolver)
        tf.tpu.experimental.initialize_tpu_system(resolver)
        strategy = tf.distribute.TPUStrategy(resolver)
      elif debug:
        strategy = tf.distribute.OneDeviceStrategy(device='/cpu:0')
      else:
        strategy = tf.distribute.MirroredStrategy()
      teacher_model = get_teacher_model(teacher_model_dir, strategy=strategy)
      train_model(teacher_model, out_dir, params, strategy,
                  write_checkpoint_metrics)
      break
    except tf.errors.UnavailableError:
      continue

def main(unused_args=None):
  train(FLAGS.teacher_model_dir, FLAGS.out_dir, FLAGS.params, FLAGS.tpu,
        FLAGS.tpu_topology, FLAGS.write_checkpoint_metrics, FLAGS.debug)

if __name__ == '__main__':
  flags.mark_flags_as_required([
      'teacher_model_dir',
      'params',
      'out_dir',
  ])
  app.run(main)
<|endoftext|>


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

"""Functionality for interacting with ResultDB's ResultSink.

ResultSink is a process that accepts test results via HTTP requests for
ingesting into ResultDB.

See go/resultdb and go/resultsink for more details.
"""

import base64
import json
import os
import sys

import requests

from typ import json_results
from typ import expectations_parser

# Valid status taken from the "TestStatus" enum in
# https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/proto/v1/test_result.proto
VALID_STATUSES = {
    'PASS',
    'FAIL',
    'CRASH',
    'ABORT',
    'SKIP',
}
STDOUT_KEY = 'typ_stdout'
STDERR_KEY = 'typ_stderr'

class ResultSinkReporter(object):
    def __init__(self, host, disable=False):
        """Class for interacting with ResultDB's ResultSink.

        Args:
            host: A host.Host or host_fake.FakeHost instance.
            disable: Whether to explicitly disable ResultSink integration.
        """
        self._host = host
        self._sink = None
        self._chromium_src_dir = None
        if disable:
            return

        luci_context_file = self._host.getenv('LUCI_CONTEXT')
        if not luci_context_file:
            return
        self._sink = json.loads(
                self._host.read_text_file(luci_context_file)).get('result_sink')
        if not self._sink:
            return

        self._url = ('http://%s/prpc/luci.resultsink.v1.Sink/ReportTestResults'
                     % self._sink['address'])
        self._headers = {
                'Content-Type': 'application/json',
                'Accept': 'application/json',
                'Authorization': 'ResultSink %s' % self._sink['auth_token']
        }
        self._session = requests.Session()

    @property
    def resultdb_supported(self):
        return self._sink is not None

    def report_individual_test_result(
            self, test_name_prefix, result, artifact_output_dir, expectations,
            test_file_location):
        """Reports typ results for a single test to ResultSink.

        Inputs are typically similar to what is passed to
        json_results.make_full_results(), but for a single test/Result instead
        of multiple tests/a ResultSet.

        Args:
            test_name_prefix: A string containing the prefix that will be added
                    to the test name.
            results: A json_results.Results instance containing the results to
                    report.
            artifact_output_dir: The path to the directory where artifacts test
                    artifacts are saved on disk. If a relative path, will be
                    automatically joined with the cwd. Use '.' instead of '' to
                    point to the cwd.
            expectations: An expectations_parser.TestExpectations instance, or
                    None if one is not available.
            test_file_location: A string containing the path to the file
                    containing the test.

        Returns:
            0 if the result was reported successfully or ResultDB is not
            supported, otherwise 1.
        """
        if not self.resultdb_supported:
            return 0

        expectation_tags = expectations.tags if expectations else []

        test_id = test_name_prefix + result.name
        raw_typ_expected_results = (
                expectations.expectations_for(result.name).raw_results
                if expectations
                else [expectations_parser.RESULT_TAGS[t]
                      for t in result.expected])
        result_is_expected = result.actual in result.expected

        tag_list = [
            ('test_name', test_id),
        ]
        for expectation in result.expected:
            tag_list.append(('typ_expectation', expectation))
        for expectation in raw_typ_expected_results:
            tag_list.append(('raw_typ_expectation', expectation))
        if expectation_tags:
            for tag in expectation_tags:
                tag_list.append(('typ_tag', tag))

        artifacts = {}
        original_artifacts = result.artifacts or {}
        assert STDOUT_KEY not in original_artifacts
        assert STDERR_KEY not in original_artifacts
        if original_artifacts:
            assert artifact_output_dir
            if not os.path.isabs(artifact_output_dir):
                artifact_output_dir = self._host.join(
                        self._host.getcwd(), artifact_output_dir)

        for artifact_name, artifact_filepaths in original_artifacts.items():
            # The typ artifact implementation supports multiple artifacts for
            # a single artifact name due to retries, but ResultDB does not.
            if len(artifact_filepaths) > 1:
                for index, filepath in enumerate(artifact_filepaths):
                    artifacts[artifact_name + '-file%d' % index] = {
                        'filePath': self._host.join(
                                artifact_output_dir, filepath),
                    }
            else:
                artifacts[artifact_name] = {
                    'filePath': self._host.join(
                            artifact_output_dir, artifact_filepaths[0]),
                }

        artifacts[STDOUT_KEY] = {
            'contents': base64.b64encode(result.out.encode('utf-8'))
        }
        artifacts[STDERR_KEY] = {
            'contents': base64.b64encode(result.err.encode('utf-8'))
        }
        html_summary = ('
<text-artifact artifact-id="%s"/>
'
                        '
<text-artifact artifact-id="%s"/>
' % (
                            STDOUT_KEY, STDERR_KEY))

        test_location_in_repo = self._convert_path_to_repo_path(
            test_file_location)
        test_metadata = {
            'name': test_id,
            'location': {
                'repo': 'https://chromium.googlesource.com/chromium/src',
                'fileName': test_location_in_repo
            },
        }

        return self._report_result(
                test_id, result.actual, result_is_expected, artifacts,
                tag_list, html_summary, result.took, test_metadata)

    def _report_result(
            self, test_id, status, expected, artifacts, tag_list, html_summary,
            duration, test_metadata):
        """Reports a single test result to ResultSink.

        Args:
            test_id: A string containing the unique identifier of the test.
            status: A string containing the status of the test. Must be in
                    |VALID_STATUSES|.
            expected: A boolean denoting whether |status| is expected or not.
            artifacts: A dict of artifact names (strings) to dicts, specifying
                    either a filepath or base64-encoded artifact content.
            tag_list: A list of tuples of (str, str), each element being a
                    key/value pair to add as tags to the reported result.
            html_summary: A string containing HTML summarizing the test run.
            duration: How long the test took in seconds.
            test_metadata: A dict containing additional test metadata to upload.

        Returns:
            0 if the result was reported successfully or ResultDB is not
            supported, otherwise 1.
        """
        if not self.resultdb_supported:
            return 0

        # TODO(crbug.com/1104252): Handle testLocation key so that ResultDB can
        # look up the correct component for bug filing.
        test_result = _create_json_test_result(
                test_id, status, expected, artifacts, tag_list, html_summary,
                duration, test_metadata)

        return self._post(json.dumps({'testResults': [test_result]}))

    def _post(self, content):
        """POST to ResultSink.

        Args:
            content: A string containing the content to send in the body of the
                    POST request.

        Returns:
            0 if the POST succeeded, otherwise 1.
        """
        res = self._session.post(
            url=self._url,
            headers=self._headers,
            data=content)
        return 0 if res.ok else 1

    def _convert_path_to_repo_path(self, filepath):
        """Converts an absolute file path to a repo relative file path.

        Args:
            filepath: A string containing an absolute path to a file.

        Returns:
            A string containing the path to |filepath| in the Chromium
            src repo, leading with //.
        """
        chromium_src_dir = self._get_chromium_src_dir()
        assert chromium_src_dir in filepath
        repo_location = filepath.replace(chromium_src_dir, '//', 1)
        repo_location = repo_location.replace(self._host.sep, '/')
        return repo_location

    def _get_chromium_src_dir(self):
        if not self._chromium_src_dir:
          src_dir = self._host.abspath(
                  self._host.join(self._host.dirname(__file__),
                                  '..', '..', '..', '..', '..'))
          if not src_dir.endswith(self._host.sep):
              src_dir += self._host.sep
          self._chromium_src_dir = src_dir
        return self._chromium_src_dir

def _create_json_test_result(
        test_id, status, expected, artifacts, tag_list, html_summary,
        duration, test_metadata):
    """Formats data to be suitable for sending to ResultSink.

    Args:
        test_id: A string containing the unique identifier of the test.
        status: A string containing the status of the test. Must be in
                |VALID_STATUSES|.
        expected: A boolean denoting whether |status| is expected or not.
        artifacts: A dict of artifact names (strings) to dicts, specifying
                either a filepath or base64-encoded artifact content.
        tag_list: A list of tuples of (str, str), each element being a
                    key/value pair to add as tags to the reported result.
        html_summary: A string containing HTML summarizing the test run.
        duration: How long the test took in seconds.
        test_metadata: A dict containing additional test metadata to upload.

    Returns:
        A dict containing the provided data in a format that is ingestable by
        ResultSink.
    """
    assert status in VALID_STATUSES
    # This is based off the protobuf in
    # https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/sink/proto/v1/test_result.proto
    test_result = {
            'testId': test_id,
            'status': status,
            'expected': expected,
            # If the number is too large or small, python formats the number
            # in scientific notation, but google.protobuf.duration doesn't
            # accept an input formatted in scientific notation.
            #
            # .9fs because nanosecond is the smallest precision that
            # google.protobuf.duration supports.
            'duration': '%.9fs' % duration,
            'summaryHtml': html_summary,
            'artifacts': artifacts,
            'tags': [],
            'testMetadata': test_metadata,
    }
    for (k, v) in tag_list:
        test_result['tags'].append({'key': k, 'value': v})

    return test_result

def result_sink_retcode_from_result_set(result_set):
    """Determines whether any interactions with ResultSink failed.

    Args:
        result_set: A json_results.ResultSet instance.

    Returns:
        1 if any Result in |result_set| failed to interact properly with
        ResultSink, otherwise 0.
    """
    return int(any(r.result_sink_retcode for r in result_set.results))
<commit_msg>typ: Set ResultSinkReporter's host to a Typ Host instance by default<commit_after># Copyright 2020 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Functionality for interacting with ResultDB's ResultSink.

ResultSink is a process that accepts test results via HTTP requests for
ingesting into ResultDB.

See go/resultdb and go/resultsink for more details.
"""

import base64
import json
import os
import sys

import requests

from typ import host
from typ import json_results
from typ import expectations_parser

# Valid status taken from the "TestStatus" enum in
# https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/proto/v1/test_result.proto
VALID_STATUSES = {
    'PASS',
    'FAIL',
    'CRASH',
    'ABORT',
    'SKIP',
}
STDOUT_KEY = 'typ_stdout'
STDERR_KEY = 'typ_stderr'

class ResultSinkReporter(object):
    def __init__(self, host=None, disable=False):
        """Class for interacting with ResultDB's ResultSink.

        Args:
            host: A host.Host or host_fake.FakeHost instance.
            disable: Whether to explicitly disable ResultSink integration.
        """
        self.host = host or host.Host()
        self._sink = None
        self._chromium_src_dir = None
        if disable:
            return

        luci_context_file = self.host.getenv('LUCI_CONTEXT')
        if not luci_context_file:
            return
        self._sink = json.loads(
                self.host.read_text_file(luci_context_file)).get('result_sink')
        if not self._sink:
            return

        self._url = ('http://%s/prpc/luci.resultsink.v1.Sink/ReportTestResults'
                     % self._sink['address'])
        self._headers = {
                'Content-Type': 'application/json',
                'Accept': 'application/json',
                'Authorization': 'ResultSink %s' % self._sink['auth_token']
        }
        self._session = requests.Session()

    @property
    def resultdb_supported(self):
        return self._sink is not None

    def report_individual_test_result(
            self, test_name_prefix, result, artifact_output_dir, expectations,
            test_file_location):
        """Reports typ results for a single test to ResultSink.

        Inputs are typically similar to what is passed to
        json_results.make_full_results(), but for a single test/Result instead
        of multiple tests/a ResultSet.

        Args:
            test_name_prefix: A string containing the prefix that will be added
                    to the test name.
            results: A json_results.Results instance containing the results to
                    report.
            artifact_output_dir: The path to the directory where artifacts test
                    artifacts are saved on disk. If a relative path, will be
                    automatically joined with the cwd. Use '.' instead of '' to
                    point to the cwd.
            expectations: An expectations_parser.TestExpectations instance, or
                    None if one is not available.
            test_file_location: A string containing the path to the file
                    containing the test.

        Returns:
            0 if the result was reported successfully or ResultDB is not
            supported, otherwise 1.
        """
        if not self.resultdb_supported:
            return 0

        expectation_tags = expectations.tags if expectations else []

        test_id = test_name_prefix + result.name
        raw_typ_expected_results = (
                expectations.expectations_for(result.name).raw_results
                if expectations
                else [expectations_parser.RESULT_TAGS[t]
                      for t in result.expected])
        result_is_expected = result.actual in result.expected

        tag_list = [
            ('test_name', test_id),
        ]
        for expectation in result.expected:
            tag_list.append(('typ_expectation', expectation))
        for expectation in raw_typ_expected_results:
            tag_list.append(('raw_typ_expectation', expectation))
        if expectation_tags:
            for tag in expectation_tags:
                tag_list.append(('typ_tag', tag))

        artifacts = {}
        original_artifacts = result.artifacts or {}
        assert STDOUT_KEY not in original_artifacts
        assert STDERR_KEY not in original_artifacts
        if original_artifacts:
            assert artifact_output_dir
            if not os.path.isabs(artifact_output_dir):
                artifact_output_dir = self.host.join(
                        self.host.getcwd(), artifact_output_dir)

        for artifact_name, artifact_filepaths in original_artifacts.items():
            # The typ artifact implementation supports multiple artifacts for
            # a single artifact name due to retries, but ResultDB does not.
            if len(artifact_filepaths) > 1:
                for index, filepath in enumerate(artifact_filepaths):
                    artifacts[artifact_name + '-file%d' % index] = {
                        'filePath': self.host.join(
                                artifact_output_dir, filepath),
                    }
            else:
                artifacts[artifact_name] = {
                    'filePath': self.host.join(
                            artifact_output_dir, artifact_filepaths[0]),
                }

        artifacts[STDOUT_KEY] = {
            'contents': base64.b64encode(result.out.encode('utf-8'))
        }
        artifacts[STDERR_KEY] = {
            'contents': base64.b64encode(result.err.encode('utf-8'))
        }
        html_summary = ('
<text-artifact artifact-id="%s"/>
'
                        '
<text-artifact artifact-id="%s"/>
' % (
                            STDOUT_KEY, STDERR_KEY))

        test_location_in_repo = self._convert_path_to_repo_path(
            test_file_location)
        test_metadata = {
            'name': test_id,
            'location': {
                'repo': 'https://chromium.googlesource.com/chromium/src',
                'fileName': test_location_in_repo
            },
        }

        return self._report_result(
                test_id, result.actual, result_is_expected, artifacts,
                tag_list, html_summary, result.took, test_metadata)

    def _report_result(
            self, test_id, status, expected, artifacts, tag_list, html_summary,
            duration, test_metadata):
        """Reports a single test result to ResultSink.

        Args:
            test_id: A string containing the unique identifier of the test.
            status: A string containing the status of the test. Must be in
                    |VALID_STATUSES|.
            expected: A boolean denoting whether |status| is expected or not.
            artifacts: A dict of artifact names (strings) to dicts, specifying
                    either a filepath or base64-encoded artifact content.
            tag_list: A list of tuples of (str, str), each element being a
                    key/value pair to add as tags to the reported result.
            html_summary: A string containing HTML summarizing the test run.
            duration: How long the test took in seconds.
            test_metadata: A dict containing additional test metadata to upload.

        Returns:
            0 if the result was reported successfully or ResultDB is not
            supported, otherwise 1.
        """
        if not self.resultdb_supported:
            return 0

        # TODO(crbug.com/1104252): Handle testLocation key so that ResultDB can
        # look up the correct component for bug filing.
        test_result = _create_json_test_result(
                test_id, status, expected, artifacts, tag_list, html_summary,
                duration, test_metadata)

        return self._post(json.dumps({'testResults': [test_result]}))

    def _post(self, content):
        """POST to ResultSink.

        Args:
            content: A string containing the content to send in the body of the
                    POST request.

        Returns:
            0 if the POST succeeded, otherwise 1.
        """
        res = self._session.post(
            url=self._url,
            headers=self._headers,
            data=content)
        return 0 if res.ok else 1

    def _convert_path_to_repo_path(self, filepath):
        """Converts an absolute file path to a repo relative file path.

        Args:
            filepath: A string containing an absolute path to a file.

        Returns:
            A string containing the path to |filepath| in the Chromium
            src repo, leading with //.
        """
        chromium_src_dir = self._get_chromium_src_dir()
        assert chromium_src_dir in filepath
        repo_location = filepath.replace(chromium_src_dir, '//', 1)
        repo_location = repo_location.replace(self.host.sep, '/')
        return repo_location

    def _get_chromium_src_dir(self):
        if not self._chromium_src_dir:
          src_dir = self.host.abspath(
                  self.host.join(self.host.dirname(__file__),
                                  '..', '..', '..', '..', '..'))
          if not src_dir.endswith(self.host.sep):
              src_dir += self.host.sep
          self._chromium_src_dir = src_dir
        return self._chromium_src_dir

def _create_json_test_result(
        test_id, status, expected, artifacts, tag_list, html_summary,
        duration, test_metadata):
    """Formats data to be suitable for sending to ResultSink.

    Args:
        test_id: A string containing the unique identifier of the test.
        status: A string containing the status of the test. Must be in
                |VALID_STATUSES|.
        expected: A boolean denoting whether |status| is expected or not.
        artifacts: A dict of artifact names (strings) to dicts, specifying
                either a filepath or base64-encoded artifact content.
        tag_list: A list of tuples of (str, str), each element being a
                    key/value pair to add as tags to the reported result.
        html_summary: A string containing HTML summarizing the test run.
        duration: How long the test took in seconds.
        test_metadata: A dict containing additional test metadata to upload.

    Returns:
        A dict containing the provided data in a format that is ingestable by
        ResultSink.
    """
    assert status in VALID_STATUSES
    # This is based off the protobuf in
    # https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/sink/proto/v1/test_result.proto
    test_result = {
            'testId': test_id,
            'status': status,
            'expected': expected,
            # If the number is too large or small, python formats the number
            # in scientific notation, but google.protobuf.duration doesn't
            # accept an input formatted in scientific notation.
            #
            # .9fs because nanosecond is the smallest precision that
            # google.protobuf.duration supports.
            'duration': '%.9fs' % duration,
            'summaryHtml': html_summary,
            'artifacts': artifacts,
            'tags': [],
            'testMetadata': test_metadata,
    }
    for (k, v) in tag_list:
        test_result['tags'].append({'key': k, 'value': v})

    return test_result

def result_sink_retcode_from_result_set(result_set):
    """Determines whether any interactions with ResultSink failed.

    Args:
        result_set: A json_results.ResultSet instance.

    Returns:
        1 if any Result in |result_set| failed to interact properly with
        ResultSink, otherwise 0.
    """
    return int(any(r.result_sink_retcode for r in result_set.results))
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

def term_frequency(terms):
    tf = {}
    for t in terms:
        if not tf.has_key(t):
            tf[t] = 0
        tf[t] += 1
    return tf

def document_frequency(texts):
    df = {}
    for terms in texts:
        for t in list(set(terms)):
            if not df.has_key(t):
                df[t] = 0
            df[t] += 1
    return df

if __name__ == '__main__':
    texts = [u'Emacs（イーマックス）とは高機能でカスタマイズ性の高いテキストエディタである。',
             u'vi（ヴィーアイ）は、Emacsと共にUNIX環境で人気があるテキストエディタ。']

    import tokenizer

    print 'tokenize'
    docs = []
    for t in texts:
        tokens = tokenizer.tokenize(t.encode('utf-8'))
        terms = map(lambda t: t.basic_form, tokens)
        docs.append(terms)
        print '|'.join(terms)

    print '\ntf'
    for d in docs:
        tf = term_frequency(d)
        print '{' + ', '.join(["'%s': %d" % (t, f) for t, f in tf.items()]) + '}'

    print '\ndf'
    df = document_frequency(docs)
    print '{' + ', '.join(["'%s': %d" % (d, f) for d, f in df.items()]) + '}'
<commit_msg>Change df method<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

def term_frequency(terms):
    tf = {}
    for t in terms:
        if not tf.has_key(t):
            tf[t] = 0
        tf[t] += 1
    return tf

def document_frequency(texts):
    from itertools import chain
    all_terms = list(chain.from_iterable(texts)) # flatten
    df = {}
    for t in all_terms:
        df[t] = [t in text for text in texts].count(True)
    return df

if __name__ == '__main__':
    texts = [u'Emacs（イーマックス）とは高機能でカスタマイズ性の高いテキストエディタである。',
             u'vi（ヴィーアイ）は、Emacsと共にUNIX環境で人気があるテキストエディタ。']

    import tokenizer

    print 'tokenize'
    docs = []
    for t in texts:
        tokens = tokenizer.tokenize(t.encode('utf-8'))
        terms = map(lambda t: t.basic_form, tokens)
        docs.append(terms)
        print '|'.join(terms)

    print '\ntf'
    for d in docs:
        tf = term_frequency(d)
        print '{' + ', '.join(["'%s': %d" % (t, f) for t, f in tf.items()]) + '}'

    print '\ndf'
    df = document_frequency(docs)
    print '{' + ', '.join(["'%s': %d" % (d, f) for d, f in df.items()]) + '}'
<|endoftext|>


<commit_before>import re
import socket
import ssl

class MyNntp:
    def __init__(self, server, port, use_ssl):
        """Constructor

        Pass in the server, port, and ssl usage value for connect.
        """

        # just store the values for now
        self.server = server
        self.port = port
        self.ssl = use_ssl

        # define variables we will use throughout our object
        self.s = None
        self.data = unicode()

        self.code = None
        self.text = None

    def parse(self):
        """Parse response to command from server.

        Break the first line from the data buffer into its component parts (if
        it is formatted properly) and retain the three digit server code as well
        as the ASCII string.

        Server code and ASCII string will be set in self.code and self.text if
        it is successfully parsed. The first line will be removed from the data
        buffer as well.

        If the format is not recognized None will be set it both variables. No
        data will be removed from the buffer.
        """

        # make sure we clear any old response stored
        self.code = None
        self.text = None

        # grab the first line from our relieved data (without lin ending)
        index = self.data.find("\r\n")
        line = self.data[:index]

        # break apart the response code and (optionally) the rest of the line
        match = re.match(r"(\d{3})(?: +(.+))?", line)

        # check for match
        if match:
            # store our code and text
            self.code = match.group(1)
            self.text = match.group(2)

            # remove our processed line (including line endings)
            self.data = self.data[index+2:]

        # we are done
        return

    def fetch(self):
        """Get server response.

        Get the response to a command sent to the NNTP server and parse it.
        """

        # receive data from server
        self.data = self.s.recv(1024)

        # parse server response
        self.parse()

        return

    def connect(self):
        """Connect to NNTP server.

        Using the server address, port, and a flag to use ssl, we will connect to
        the server and parse the response from the server using standard sockets.
        """

        # create a socket object and connect
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if self.ssl:
            self.s = ssl.wrap_socket(self.s)
        self.s.connect((self.server, self.port))

        # get data from the server
        self.fetch()

        # check for success
        if self.code != '200':
            return False
        else:
            return True

    def send(self, command):
        """Send a command to the server and get the response.

        Send the command out the connected socket to the server and parse the
        response.
        """

        # send the command to the server
        self.s.sendall(command)

        # get the response from the server
        self.fetch()

        return

    def login(self, username, password):
        """Login to server.

        Login to the server using a username and password. Check for failure to
        login and intelligently handle server responses.
        """

        # send the username to the server
        self.send("AUTHINFO USER " + username + "\r\n")

        # get code 381 if a password is required
        if self.code != '381':
            return False

        # send the password to the server
        self.send("AUTHINFO PASS " + password + "\r\n")

        # get code 281 if successfully logged in
        if self.code != '281':
            return False

        # all went well, return true
        return True

    def listactive(self, processor=None):
        """List Active

        List the active newsgroups available on the server.
        """

        # send the command to the server
        self.send("LIST ACTIVE\r\n")

        # check for 215 for multi line response
        if self.code != '215':
            return False

        # regex pattern to recognize results
        pattern = re.compile(r"(\S+) +(\S+) +(\S+) +(\S+)")

        # this is our end of transmission flag
        eot = False

        # are we processing results?
        if processor is None:
            results = []

        # keep looping until our transmission is finished
        while not eot:
            # process the data in our buffer
            for line in self.data.splitlines(True):
                # check for a full line
                if line.endswith("\r\n"):
                    #  check for end of multi line response
                    if line == ".\r\n":
                        eot = True
                    else:
                        # apply pattern to line
                        match = pattern.match(line)
                        if match:
                            if processor is None:
                                results.append([match.group(1), match.group(2), match.group(3), match.group(4)])
                            else:
                                processor(match.group(1), match.group(2), match.group(3), match.group(4))
                        else:
                            print("unexpected line in results: %s", line)

                    # remove line
                    line, self.data = self.data.split("\r\n", 1)

            # if we have not finished...
            if not eot:
                # receive more data from server
                self.data += self.s.recv(1024)

        # all went well, return true
        if processor is None:
            return results
        else:
            return True

    def post(self, fromheader, subjectheader, newsgroupsheader, article):
        """Post a binary article to a newsgroup.

        """

        # send the post command to the server
        self.send("POST\r\n")

        # get code 340 if we're ok to post
        if self.code != '340':
            return False

        self.s.sendall("From: " + fromheader + "\r\n")
        self.s.sendall("Subject: " + subjectheader + "\r\n")
        self.s.sendall("Newsgroups: " + newsgroupsheader + "\r\n")
        self.s.sendall("\r\n")
        self.s.sendall(article)

        # send our end of transmission character
        self.send(".\r\n")

        # get code 240 if the server accepted our post
        if self.code != '240':
            return False

        return True<commit_msg>send function now adds trailing \r\n to outgoing messages<commit_after>import re
import socket
import ssl

class MyNntp:
    def __init__(self, server, port, use_ssl):
        """Constructor

        Pass in the server, port, and ssl usage value for connect.
        """

        # just store the values for now
        self.server = server
        self.port = port
        self.ssl = use_ssl

        # define variables we will use throughout our object
        self.s = None
        self.data = unicode()

        self.code = None
        self.text = None

    def parse(self):
        """Parse response to command from server.

        Break the first line from the data buffer into its component parts (if
        it is formatted properly) and retain the three digit server code as well
        as the ASCII string.

        Server code and ASCII string will be set in self.code and self.text if
        it is successfully parsed. The first line will be removed from the data
        buffer as well.

        If the format is not recognized None will be set it both variables. No
        data will be removed from the buffer.
        """

        # make sure we clear any old response stored
        self.code = None
        self.text = None

        # grab the first line from our relieved data (without lin ending)
        index = self.data.find("\r\n")
        line = self.data[:index]

        # break apart the response code and (optionally) the rest of the line
        match = re.match(r"(\d{3})(?: +(.+))?", line)

        # check for match
        if match:
            # store our code and text
            self.code = match.group(1)
            self.text = match.group(2)

            # remove our processed line (including line endings)
            self.data = self.data[index+2:]

        # we are done
        return

    def fetch(self):
        """Get server response.

        Get the response to a command sent to the NNTP server and parse it.
        """

        # receive data from server
        self.data = self.s.recv(1024)

        # parse server response
        self.parse()

        return

    def connect(self):
        """Connect to NNTP server.

        Using the server address, port, and a flag to use ssl, we will connect to
        the server and parse the response from the server using standard sockets.
        """

        # create a socket object and connect
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if self.ssl:
            self.s = ssl.wrap_socket(self.s)
        self.s.connect((self.server, self.port))

        # get data from the server
        self.fetch()

        # check for success
        if self.code != '200':
            return False
        else:
            return True

    def send(self, command):
        """Send a command to the server and get the response.

        Send the command out the connected socket to the server and parse the
        response.
        """

        # send the command to the server
        self.s.sendall(command + "\r\n")

        # get the response from the server
        self.fetch()

        return

    def login(self, username, password):
        """Login to server.

        Login to the server using a username and password. Check for failure to
        login and intelligently handle server responses.
        """

        # send the username to the server
        self.send("AUTHINFO USER " + username)

        # get code 381 if a password is required
        if self.code != '381':
            return False

        # send the password to the server
        self.send("AUTHINFO PASS " + password)

        # get code 281 if successfully logged in
        if self.code != '281':
            return False

        # all went well, return true
        return True

    def listactive(self, processor=None):
        """List Active

        List the active newsgroups available on the server.
        """

        # send the command to the server
        self.send("LIST ACTIVE")

        # check for 215 for multi line response
        if self.code != '215':
            return False

        # regex pattern to recognize results
        pattern = re.compile(r"(\S+) +(\S+) +(\S+) +(\S+)")

        # this is our end of transmission flag
        eot = False

        # are we processing results?
        if processor is None:
            results = []

        # keep looping until our transmission is finished
        while not eot:
            # process the data in our buffer
            for line in self.data.splitlines(True):
                # check for a full line
                if line.endswith("\r\n"):
                    #  check for end of multi line response
                    if line == ".\r\n":
                        eot = True
                    else:
                        # apply pattern to line
                        match = pattern.match(line)
                        if match:
                            if processor is None:
                                results.append([match.group(1), match.group(2), match.group(3), match.group(4)])
                            else:
                                processor(match.group(1), match.group(2), match.group(3), match.group(4))
                        else:
                            print("unexpected line in results: %s", line)

                    # remove line
                    line, self.data = self.data.split("\r\n", 1)

            # if we have not finished...
            if not eot:
                # receive more data from server
                self.data += self.s.recv(1024)

        # all went well, return true
        if processor is None:
            return results
        else:
            return True

    def post(self, fromheader, subjectheader, newsgroupsheader, article):
        """Post a binary article to a newsgroup.

        """

        # send the post command to the server
        self.send("POST")

        # get code 340 if we're ok to post
        if self.code != '340':
            return False

        self.s.sendall("From: " + fromheader + "\r\n")
        self.s.sendall("Subject: " + subjectheader + "\r\n")
        self.s.sendall("Newsgroups: " + newsgroupsheader + "\r\n")
        self.s.sendall("\r\n")
        self.s.sendall(article)

        # send our end of transmission character
        self.send(".")

        # get code 240 if the server accepted our post
        if self.code != '240':
            return False

        return True<|endoftext|>


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

"""Base classes for our unit tests.

Allows overriding of flags for use of fakes, and some black magic for
inline callbacks.

"""
import contextlib

import datetime
import eventlet
eventlet.monkey_patch(os=False)

import copy
import inspect
import logging as std_logging
import mock
import os

import fixtures
from oslo_concurrency import lockutils
from oslo_config import cfg
from oslo_config import fixture as config_fixture
from oslo_log.fixture import logging_error as log_fixture
from oslo_log import log as logging
from oslo_serialization import jsonutils
from oslo_utils import timeutils
from oslotest import moxstubout
import six
import testtools

from nova import context
from nova import db
from nova.network import manager as network_manager
from nova import objects
from nova.objects import base as objects_base
from nova.tests import fixtures as nova_fixtures
from nova.tests.unit import conf_fixture
from nova.tests.unit import policy_fixture
from nova import utils

CONF = cfg.CONF
CONF.import_opt('enabled', 'nova.api.openstack', group='osapi_v3')

logging.register_options(CONF)
CONF.set_override('use_stderr', False)
logging.setup(CONF, 'nova')

# NOTE(comstud): Make sure we have all of the objects loaded. We do this
# at module import time, because we may be using mock decorators in our
# tests that run at import time.
objects.register_all()

_TRUE_VALUES = ('True', 'true', '1', 'yes')

if six.PY3:
    @contextlib.contextmanager
    def nested(*contexts):
        with contextlib.ExitStack() as stack:
            yield [stack.enter_context(c) for c in contexts]
else:
    nested = contextlib.nested

class SampleNetworks(fixtures.Fixture):

    """Create sample networks in the database."""

    def __init__(self, host=None):
        self.host = host

    def setUp(self):
        super(SampleNetworks, self).setUp()
        ctxt = context.get_admin_context()
        network = network_manager.VlanManager(host=self.host)
        bridge_interface = CONF.flat_interface or CONF.vlan_interface
        network.create_networks(ctxt,
                                label='test',
                                cidr='10.0.0.0/8',
                                multi_host=CONF.multi_host,
                                num_networks=CONF.num_networks,
                                network_size=CONF.network_size,
                                cidr_v6=CONF.fixed_range_v6,
                                gateway=CONF.gateway,
                                gateway_v6=CONF.gateway_v6,
                                bridge=CONF.flat_network_bridge,
                                bridge_interface=bridge_interface,
                                vpn_start=CONF.vpn_start,
                                vlan_start=CONF.vlan_start,
                                dns1=CONF.flat_network_dns)
        for net in db.network_get_all(ctxt):
            network.set_network_host(ctxt, net)

class TestingException(Exception):
    pass

class NullHandler(std_logging.Handler):
    """custom default NullHandler to attempt to format the record.

    Used in conjunction with
    log_fixture.get_logging_handle_error_fixture to detect formatting errors in
    debug level logs without saving the logs.
    """
    def handle(self, record):
        self.format(record)

    def emit(self, record):
        pass

    def createLock(self):
        self.lock = None

class skipIf(object):
    def __init__(self, condition, reason):
        self.condition = condition
        self.reason = reason

    def __call__(self, func_or_cls):
        condition = self.condition
        reason = self.reason
        if inspect.isfunction(func_or_cls):
            @six.wraps(func_or_cls)
            def wrapped(*args, **kwargs):
                if condition:
                    raise testtools.TestCase.skipException(reason)
                return func_or_cls(*args, **kwargs)

            return wrapped
        elif inspect.isclass(func_or_cls):
            orig_func = getattr(func_or_cls, 'setUp')

            @six.wraps(orig_func)
            def new_func(self, *args, **kwargs):
                if condition:
                    raise testtools.TestCase.skipException(reason)
                orig_func(self, *args, **kwargs)

            func_or_cls.setUp = new_func
            return func_or_cls
        else:
            raise TypeError('skipUnless can be used only with functions or '
                            'classes')

def _patch_mock_to_raise_for_invalid_assert_calls():
    def raise_for_invalid_assert_calls(wrapped):
        def wrapper(_self, name):
            valid_asserts = [
                'assert_called_with',
                'assert_called_once_with',
                'assert_has_calls',
                'assert_any_calls']

            if name.startswith('assert') and name not in valid_asserts:
                raise AttributeError('%s is not a valid mock assert method'
                                     % name)

            return wrapped(_self, name)
        return wrapper
    mock.Mock.__getattr__ = raise_for_invalid_assert_calls(
        mock.Mock.__getattr__)

# NOTE(gibi): needs to be called only once at import time
# to patch the mock lib
_patch_mock_to_raise_for_invalid_assert_calls()

class TestCase(testtools.TestCase):
    """Test case base class for all unit tests.

    Due to the slowness of DB access, please consider deriving from
    `NoDBTestCase` first.
    """
    USES_DB = True
    REQUIRES_LOCKING = False

    TIMEOUT_SCALING_FACTOR = 1

    def setUp(self):
        """Run before each test method to initialize test environment."""
        super(TestCase, self).setUp()
        self.useFixture(nova_fixtures.Timeout(
            os.environ.get('OS_TEST_TIMEOUT', 0),
            self.TIMEOUT_SCALING_FACTOR))

        self.useFixture(fixtures.NestedTempfile())
        self.useFixture(fixtures.TempHomeDir())
        self.useFixture(log_fixture.get_logging_handle_error_fixture())

        self.useFixture(nova_fixtures.OutputStreamCapture())

        self.useFixture(nova_fixtures.StandardLogging())

        # NOTE(sdague): because of the way we were using the lock
        # wrapper we eneded up with a lot of tests that started
        # relying on global external locking being set up for them. We
        # consider all of these to be *bugs*. Tests should not require
        # global external locking, or if they do, they should
        # explicitly set it up themselves.
        #
        # The following REQUIRES_LOCKING class parameter is provided
        # as a bridge to get us there. No new tests should be added
        # that require it, and existing classes and tests should be
        # fixed to not need it.
        if self.REQUIRES_LOCKING:
            lock_path = self.useFixture(fixtures.TempDir()).path
            self.fixture = self.useFixture(
                config_fixture.Config(lockutils.CONF))
            self.fixture.config(lock_path=lock_path,
                                group='oslo_concurrency')

        self.useFixture(conf_fixture.ConfFixture(CONF))
        self.useFixture(nova_fixtures.RPCFixture('nova.test'))

        if self.USES_DB:
            self.useFixture(nova_fixtures.Database())

        # NOTE(blk-u): WarningsFixture must be after the Database fixture
        # because sqlalchemy-migrate messes with the warnings filters.
        self.useFixture(nova_fixtures.WarningsFixture())

        # NOTE(danms): Make sure to reset us back to non-remote objects
        # for each test to avoid interactions. Also, backup the object
        # registry.
        objects_base.NovaObject.indirection_api = None
        self._base_test_obj_backup = copy.copy(
            objects_base.NovaObjectRegistry._registry._obj_classes)
        self.addCleanup(self._restore_obj_registry)

        # NOTE(mnaser): All calls to utils.is_neutron() are cached in
        # nova.utils._IS_NEUTRON.  We set it to None to avoid any
        # caching of that value.
        utils._IS_NEUTRON = None

        mox_fixture = self.useFixture(moxstubout.MoxStubout())
        self.mox = mox_fixture.mox
        self.stubs = mox_fixture.stubs
        self.addCleanup(self._clear_attrs)
        self.useFixture(fixtures.EnvironmentVariable('http_proxy'))
        self.policy = self.useFixture(policy_fixture.PolicyFixture())

        self.useFixture(nova_fixtures.PoisonFunctions())

    def _restore_obj_registry(self):
        objects_base.NovaObjectRegistry._registry._obj_classes = \
                self._base_test_obj_backup

    def _clear_attrs(self):
        # Delete attributes that don't start with _ so they don't pin
        # memory around unnecessarily for the duration of the test
        # suite
        for key in [k for k in self.__dict__.keys() if k[0] != '_']:
            # NOTE(gmann): Skip attribute 'id' because if tests are being
            # generated using testscenarios then, 'id' attribute is being
            # added during cloning the tests. And later that 'id' attribute
            # is being used by test suite to generate the results for each
            # newly generated tests by testscenarios.
            if key != 'id':
                del self.__dict__[key]

    def flags(self, **kw):
        """Override flag variables for a test."""
        group = kw.pop('group', None)
        for k, v in six.iteritems(kw):
            CONF.set_override(k, v, group)

    def start_service(self, name, host=None, **kwargs):
        svc = self.useFixture(
            nova_fixtures.ServiceFixture(name, host, **kwargs))
        return svc.service

    def assertJsonEqual(self, expected, observed):
        if isinstance(expected, six.string_types):
            expected = jsonutils.loads(expected)
        if isinstance(observed, six.string_types):
            observed = jsonutils.loads(observed)

        def sort_key(x):
            if isinstance(x, set) or isinstance(x, datetime.datetime):
                return str(x)
            if isinstance(x, dict):
                items = ((sort_key(key), sort_key(value))
                         for key, value in x.items())
                return sorted(items)
            return x

        def inner(expected, observed):
            if isinstance(expected, dict) and isinstance(observed, dict):
                self.assertEqual(len(expected), len(observed))
                expected_keys = sorted(expected)
                observed_keys = sorted(observed)
                self.assertEqual(expected_keys, observed_keys)

                for key in list(six.iterkeys(expected)):
                    inner(expected[key], observed[key])
            elif (isinstance(expected, (list, tuple, set)) and
                      isinstance(observed, (list, tuple, set))):
                self.assertEqual(len(expected), len(observed))

                expected_values_iter = iter(sorted(expected, key=sort_key))
                observed_values_iter = iter(sorted(observed, key=sort_key))

                for i in range(len(expected)):
                    inner(next(expected_values_iter),
                          next(observed_values_iter))
            else:
                self.assertEqual(expected, observed)

        inner(expected, observed)

    def assertPublicAPISignatures(self, baseinst, inst):
        def get_public_apis(inst):
            methods = {}

            def findmethods(object):
                return inspect.ismethod(object) or inspect.isfunction(object)

            for (name, value) in inspect.getmembers(inst, findmethods):
                if name.startswith("_"):
                    continue
                methods[name] = value
            return methods

        baseclass = baseinst.__class__.__name__
        basemethods = get_public_apis(baseinst)
        implmethods = get_public_apis(inst)

        extranames = []
        for name in sorted(implmethods.keys()):
            if name not in basemethods:
                extranames.append(name)

        self.assertEqual([], extranames,
                         "public APIs not listed in base class %s" %
                         baseclass)

        for name in sorted(implmethods.keys()):
            baseargs = inspect.getargspec(basemethods[name])
            implargs = inspect.getargspec(implmethods[name])

            self.assertEqual(baseargs, implargs,
                             "%s args don't match base class %s" %
                             (name, baseclass))

class APICoverage(object):

    cover_api = None

    def test_api_methods(self):
        self.assertTrue(self.cover_api is not None)
        api_methods = [x for x in dir(self.cover_api)
                       if not x.startswith('_')]
        test_methods = [x[5:] for x in dir(self)
                        if x.startswith('test_')]
        self.assertThat(
            test_methods,
            testtools.matchers.ContainsAll(api_methods))

class TimeOverride(fixtures.Fixture):
    """Fixture to start and remove time override."""

    def setUp(self):
        super(TimeOverride, self).setUp()
        timeutils.set_time_override()
        self.addCleanup(timeutils.clear_time_override)

class NoDBTestCase(TestCase):
    """`NoDBTestCase` differs from TestCase in that DB access is not supported.
    This makes tests run significantly faster. If possible, all new tests
    should derive from this class.
    """
    USES_DB = False

class BaseHookTestCase(NoDBTestCase):
    def assert_has_hook(self, expected_name, func):
        self.assertTrue(hasattr(func, '__hook_name__'))
        self.assertEqual(expected_name, func.__hook_name__)

class MatchType(object):
    """Matches any instance of a specified type

    The MatchType class is a helper for use with the
    mock.assert_called_with() method that lets you
    assert that a particular parameter has a specific
    data type. It enables strict check than the built
    in mock.ANY helper, and is the equivalent of the
    mox.IsA() function from the legacy mox library

    Example usage could be:

      mock_some_method.assert_called_once_with(
            "hello",
            MatchType(objects.Instance),
            mock.ANY,
            "world",
            MatchType(objects.KeyPair))
    """
    def __init__(self, wanttype):
        self.wanttype = wanttype

    def __eq__(self, other):
        return type(other) == self.wanttype

    def __ne__(self, other):
        return type(other) != self.wanttype

    def __repr__(self):
        return "<MatchType:" + str(self.wanttype) + ">"
<commit_msg>shut off warnings<commit_after># Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""Base classes for our unit tests.

Allows overriding of flags for use of fakes, and some black magic for
inline callbacks.

"""
import contextlib

import datetime
import eventlet
eventlet.monkey_patch(os=False)
import warnings
warnings.filterwarnings("ignore")

import copy
import inspect
import logging as std_logging
import mock
import os

import fixtures
from oslo_concurrency import lockutils
from oslo_config import cfg
from oslo_config import fixture as config_fixture
from oslo_log.fixture import logging_error as log_fixture
from oslo_log import log as logging
from oslo_serialization import jsonutils
from oslo_utils import timeutils
from oslotest import moxstubout
import six
import testtools

from nova import context
from nova import db
from nova.network import manager as network_manager
from nova import objects
from nova.objects import base as objects_base
from nova.tests import fixtures as nova_fixtures
from nova.tests.unit import conf_fixture
from nova.tests.unit import policy_fixture
from nova import utils

CONF = cfg.CONF
CONF.import_opt('enabled', 'nova.api.openstack', group='osapi_v3')

logging.register_options(CONF)
CONF.set_override('use_stderr', False)
logging.setup(CONF, 'nova')

# NOTE(comstud): Make sure we have all of the objects loaded. We do this
# at module import time, because we may be using mock decorators in our
# tests that run at import time.
objects.register_all()

_TRUE_VALUES = ('True', 'true', '1', 'yes')

if six.PY3:
    @contextlib.contextmanager
    def nested(*contexts):
        with contextlib.ExitStack() as stack:
            yield [stack.enter_context(c) for c in contexts]
else:
    nested = contextlib.nested

class SampleNetworks(fixtures.Fixture):

    """Create sample networks in the database."""

    def __init__(self, host=None):
        self.host = host

    def setUp(self):
        super(SampleNetworks, self).setUp()
        ctxt = context.get_admin_context()
        network = network_manager.VlanManager(host=self.host)
        bridge_interface = CONF.flat_interface or CONF.vlan_interface
        network.create_networks(ctxt,
                                label='test',
                                cidr='10.0.0.0/8',
                                multi_host=CONF.multi_host,
                                num_networks=CONF.num_networks,
                                network_size=CONF.network_size,
                                cidr_v6=CONF.fixed_range_v6,
                                gateway=CONF.gateway,
                                gateway_v6=CONF.gateway_v6,
                                bridge=CONF.flat_network_bridge,
                                bridge_interface=bridge_interface,
                                vpn_start=CONF.vpn_start,
                                vlan_start=CONF.vlan_start,
                                dns1=CONF.flat_network_dns)
        for net in db.network_get_all(ctxt):
            network.set_network_host(ctxt, net)

class TestingException(Exception):
    pass

class NullHandler(std_logging.Handler):
    """custom default NullHandler to attempt to format the record.

    Used in conjunction with
    log_fixture.get_logging_handle_error_fixture to detect formatting errors in
    debug level logs without saving the logs.
    """
    def handle(self, record):
        self.format(record)

    def emit(self, record):
        pass

    def createLock(self):
        self.lock = None

class skipIf(object):
    def __init__(self, condition, reason):
        self.condition = condition
        self.reason = reason

    def __call__(self, func_or_cls):
        condition = self.condition
        reason = self.reason
        if inspect.isfunction(func_or_cls):
            @six.wraps(func_or_cls)
            def wrapped(*args, **kwargs):
                if condition:
                    raise testtools.TestCase.skipException(reason)
                return func_or_cls(*args, **kwargs)

            return wrapped
        elif inspect.isclass(func_or_cls):
            orig_func = getattr(func_or_cls, 'setUp')

            @six.wraps(orig_func)
            def new_func(self, *args, **kwargs):
                if condition:
                    raise testtools.TestCase.skipException(reason)
                orig_func(self, *args, **kwargs)

            func_or_cls.setUp = new_func
            return func_or_cls
        else:
            raise TypeError('skipUnless can be used only with functions or '
                            'classes')

def _patch_mock_to_raise_for_invalid_assert_calls():
    def raise_for_invalid_assert_calls(wrapped):
        def wrapper(_self, name):
            valid_asserts = [
                'assert_called_with',
                'assert_called_once_with',
                'assert_has_calls',
                'assert_any_calls']

            if name.startswith('assert') and name not in valid_asserts:
                raise AttributeError('%s is not a valid mock assert method'
                                     % name)

            return wrapped(_self, name)
        return wrapper
    mock.Mock.__getattr__ = raise_for_invalid_assert_calls(
        mock.Mock.__getattr__)

# NOTE(gibi): needs to be called only once at import time
# to patch the mock lib
_patch_mock_to_raise_for_invalid_assert_calls()

class TestCase(testtools.TestCase):
    """Test case base class for all unit tests.

    Due to the slowness of DB access, please consider deriving from
    `NoDBTestCase` first.
    """
    USES_DB = True
    REQUIRES_LOCKING = False

    TIMEOUT_SCALING_FACTOR = 1

    def setUp(self):
        """Run before each test method to initialize test environment."""
        super(TestCase, self).setUp()
        self.useFixture(nova_fixtures.Timeout(
            os.environ.get('OS_TEST_TIMEOUT', 0),
            self.TIMEOUT_SCALING_FACTOR))

        self.useFixture(fixtures.NestedTempfile())
        self.useFixture(fixtures.TempHomeDir())
        self.useFixture(log_fixture.get_logging_handle_error_fixture())

        self.useFixture(nova_fixtures.OutputStreamCapture())

        self.useFixture(nova_fixtures.StandardLogging())

        # NOTE(sdague): because of the way we were using the lock
        # wrapper we eneded up with a lot of tests that started
        # relying on global external locking being set up for them. We
        # consider all of these to be *bugs*. Tests should not require
        # global external locking, or if they do, they should
        # explicitly set it up themselves.
        #
        # The following REQUIRES_LOCKING class parameter is provided
        # as a bridge to get us there. No new tests should be added
        # that require it, and existing classes and tests should be
        # fixed to not need it.
        if self.REQUIRES_LOCKING:
            lock_path = self.useFixture(fixtures.TempDir()).path
            self.fixture = self.useFixture(
                config_fixture.Config(lockutils.CONF))
            self.fixture.config(lock_path=lock_path,
                                group='oslo_concurrency')

        self.useFixture(conf_fixture.ConfFixture(CONF))
        self.useFixture(nova_fixtures.RPCFixture('nova.test'))

        if self.USES_DB:
            self.useFixture(nova_fixtures.Database())

        # NOTE(blk-u): WarningsFixture must be after the Database fixture
        # because sqlalchemy-migrate messes with the warnings filters.
        self.useFixture(nova_fixtures.WarningsFixture())

        # NOTE(danms): Make sure to reset us back to non-remote objects
        # for each test to avoid interactions. Also, backup the object
        # registry.
        objects_base.NovaObject.indirection_api = None
        self._base_test_obj_backup = copy.copy(
            objects_base.NovaObjectRegistry._registry._obj_classes)
        self.addCleanup(self._restore_obj_registry)

        # NOTE(mnaser): All calls to utils.is_neutron() are cached in
        # nova.utils._IS_NEUTRON.  We set it to None to avoid any
        # caching of that value.
        utils._IS_NEUTRON = None

        mox_fixture = self.useFixture(moxstubout.MoxStubout())
        self.mox = mox_fixture.mox
        self.stubs = mox_fixture.stubs
        self.addCleanup(self._clear_attrs)
        self.useFixture(fixtures.EnvironmentVariable('http_proxy'))
        self.policy = self.useFixture(policy_fixture.PolicyFixture())

        self.useFixture(nova_fixtures.PoisonFunctions())

    def _restore_obj_registry(self):
        objects_base.NovaObjectRegistry._registry._obj_classes = \
                self._base_test_obj_backup

    def _clear_attrs(self):
        # Delete attributes that don't start with _ so they don't pin
        # memory around unnecessarily for the duration of the test
        # suite
        for key in [k for k in self.__dict__.keys() if k[0] != '_']:
            # NOTE(gmann): Skip attribute 'id' because if tests are being
            # generated using testscenarios then, 'id' attribute is being
            # added during cloning the tests. And later that 'id' attribute
            # is being used by test suite to generate the results for each
            # newly generated tests by testscenarios.
            if key != 'id':
                del self.__dict__[key]

    def flags(self, **kw):
        """Override flag variables for a test."""
        group = kw.pop('group', None)
        for k, v in six.iteritems(kw):
            CONF.set_override(k, v, group)

    def start_service(self, name, host=None, **kwargs):
        svc = self.useFixture(
            nova_fixtures.ServiceFixture(name, host, **kwargs))
        return svc.service

    def assertJsonEqual(self, expected, observed):
        if isinstance(expected, six.string_types):
            expected = jsonutils.loads(expected)
        if isinstance(observed, six.string_types):
            observed = jsonutils.loads(observed)

        def sort_key(x):
            if isinstance(x, set) or isinstance(x, datetime.datetime):
                return str(x)
            if isinstance(x, dict):
                items = ((sort_key(key), sort_key(value))
                         for key, value in x.items())
                return sorted(items)
            return x

        def inner(expected, observed):
            if isinstance(expected, dict) and isinstance(observed, dict):
                self.assertEqual(len(expected), len(observed))
                expected_keys = sorted(expected)
                observed_keys = sorted(observed)
                self.assertEqual(expected_keys, observed_keys)

                for key in list(six.iterkeys(expected)):
                    inner(expected[key], observed[key])
            elif (isinstance(expected, (list, tuple, set)) and
                      isinstance(observed, (list, tuple, set))):
                self.assertEqual(len(expected), len(observed))

                expected_values_iter = iter(sorted(expected, key=sort_key))
                observed_values_iter = iter(sorted(observed, key=sort_key))

                for i in range(len(expected)):
                    inner(next(expected_values_iter),
                          next(observed_values_iter))
            else:
                self.assertEqual(expected, observed)

        inner(expected, observed)

    def assertPublicAPISignatures(self, baseinst, inst):
        def get_public_apis(inst):
            methods = {}

            def findmethods(object):
                return inspect.ismethod(object) or inspect.isfunction(object)

            for (name, value) in inspect.getmembers(inst, findmethods):
                if name.startswith("_"):
                    continue
                methods[name] = value
            return methods

        baseclass = baseinst.__class__.__name__
        basemethods = get_public_apis(baseinst)
        implmethods = get_public_apis(inst)

        extranames = []
        for name in sorted(implmethods.keys()):
            if name not in basemethods:
                extranames.append(name)

        self.assertEqual([], extranames,
                         "public APIs not listed in base class %s" %
                         baseclass)

        for name in sorted(implmethods.keys()):
            baseargs = inspect.getargspec(basemethods[name])
            implargs = inspect.getargspec(implmethods[name])

            self.assertEqual(baseargs, implargs,
                             "%s args don't match base class %s" %
                             (name, baseclass))

class APICoverage(object):

    cover_api = None

    def test_api_methods(self):
        self.assertTrue(self.cover_api is not None)
        api_methods = [x for x in dir(self.cover_api)
                       if not x.startswith('_')]
        test_methods = [x[5:] for x in dir(self)
                        if x.startswith('test_')]
        self.assertThat(
            test_methods,
            testtools.matchers.ContainsAll(api_methods))

class TimeOverride(fixtures.Fixture):
    """Fixture to start and remove time override."""

    def setUp(self):
        super(TimeOverride, self).setUp()
        timeutils.set_time_override()
        self.addCleanup(timeutils.clear_time_override)

class NoDBTestCase(TestCase):
    """`NoDBTestCase` differs from TestCase in that DB access is not supported.
    This makes tests run significantly faster. If possible, all new tests
    should derive from this class.
    """
    USES_DB = False

class BaseHookTestCase(NoDBTestCase):
    def assert_has_hook(self, expected_name, func):
        self.assertTrue(hasattr(func, '__hook_name__'))
        self.assertEqual(expected_name, func.__hook_name__)

class MatchType(object):
    """Matches any instance of a specified type

    The MatchType class is a helper for use with the
    mock.assert_called_with() method that lets you
    assert that a particular parameter has a specific
    data type. It enables strict check than the built
    in mock.ANY helper, and is the equivalent of the
    mox.IsA() function from the legacy mox library

    Example usage could be:

      mock_some_method.assert_called_once_with(
            "hello",
            MatchType(objects.Instance),
            mock.ANY,
            "world",
            MatchType(objects.KeyPair))
    """
    def __init__(self, wanttype):
        self.wanttype = wanttype

    def __eq__(self, other):
        return type(other) == self.wanttype

    def __ne__(self, other):
        return type(other) != self.wanttype

    def __repr__(self):
        return "<MatchType:" + str(self.wanttype) + ">"
<|endoftext|>