<commit_before>import os
import sys
from unittest.mock import patch

from transformers import BertTokenizer, EncoderDecoderModel
from transformers.file_utils import is_datasets_available
from transformers.testing_utils import TestCasePlus, execute_subprocess_async, get_gpu_count, slow
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed

from .finetune_trainer import Seq2SeqTrainingArguments, main
from .seq2seq_trainer import Seq2SeqTrainer
from .test_seq2seq_examples import MBART_TINY

set_seed(42)
MARIAN_MODEL = "sshleifer/student_marian_en_ro_6_1"

class TestFinetuneTrainer(TestCasePlus):
    def test_finetune_trainer(self):
        output_dir = self.run_trainer(1, "12", MBART_TINY, 1)
        logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history
        eval_metrics = [log for log in logs if "eval_loss" in log.keys()]
        first_step_stats = eval_metrics[0]
        assert "eval_bleu" in first_step_stats

    @slow
    def test_finetune_trainer_slow(self):
        # There is a missing call to __init__process_group somewhere
        output_dir = self.run_trainer(eval_steps=2, max_len="128", model_name=MARIAN_MODEL, num_train_epochs=10)

        # Check metrics
        logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history
        eval_metrics = [log for log in logs if "eval_loss" in log.keys()]
        first_step_stats = eval_metrics[0]
        last_step_stats = eval_metrics[-1]

        assert first_step_stats["eval_bleu"] < last_step_stats["eval_bleu"]  # model learned nothing
        assert isinstance(last_step_stats["eval_bleu"], float)

        # test if do_predict saves generations and metrics
        contents = os.listdir(output_dir)
        contents = {os.path.basename(p) for p in contents}
        assert "test_generations.txt" in contents
        assert "test_results.json" in contents

    @slow
    def test_finetune_bert2bert(self):
        if not is_datasets_available():
            return

        import datasets

        bert2bert = EncoderDecoderModel.from_encoder_decoder_pretrained("prajjwal1/bert-tiny", "prajjwal1/bert-tiny")
        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

        bert2bert.config.vocab_size = bert2bert.config.encoder.vocab_size
        bert2bert.config.eos_token_id = tokenizer.sep_token_id
        bert2bert.config.decoder_start_token_id = tokenizer.cls_token_id
        bert2bert.config.max_length = 128

        train_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="train[:1%]")
        val_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="validation[:1%]")

        train_dataset = train_dataset.select(range(32))
        val_dataset = val_dataset.select(range(16))

        rouge = datasets.load_metric("rouge")

        batch_size = 4

        def _map_to_encoder_decoder_inputs(batch):
            # Tokenizer will automatically set [BOS] <text> [EOS]
            inputs = tokenizer(batch["article"], padding="max_length", truncation=True, max_length=512)
            outputs = tokenizer(batch["highlights"], padding="max_length", truncation=True, max_length=128)
            batch["input_ids"] = inputs.input_ids
            batch["attention_mask"] = inputs.attention_mask

            batch["decoder_input_ids"] = outputs.input_ids
            batch["labels"] = outputs.input_ids.copy()
            batch["labels"] = [
                [-100 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch["labels"]
            ]
            batch["decoder_attention_mask"] = outputs.attention_mask

            assert all([len(x) == 512 for x in inputs.input_ids])
            assert all([len(x) == 128 for x in outputs.input_ids])

            return batch

        def _compute_metrics(pred):
            labels_ids = pred.label_ids
            pred_ids = pred.predictions

            # all unnecessary tokens are removed
            pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
            label_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)

            rouge_output = rouge.compute(predictions=pred_str, references=label_str, rouge_types=["rouge2"])[
                "rouge2"
            ].mid

            return {
                "rouge2_precision": round(rouge_output.precision, 4),
                "rouge2_recall": round(rouge_output.recall, 4),
                "rouge2_fmeasure": round(rouge_output.fmeasure, 4),
            }

        # map train dataset
        train_dataset = train_dataset.map(
            _map_to_encoder_decoder_inputs,
            batched=True,
            batch_size=batch_size,
            remove_columns=["article", "highlights"],
        )
        train_dataset.set_format(
            type="torch",
            columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
        )

        # same for validation dataset
        val_dataset = val_dataset.map(
            _map_to_encoder_decoder_inputs,
            batched=True,
            batch_size=batch_size,
            remove_columns=["article", "highlights"],
        )
        val_dataset.set_format(
            type="torch",
            columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
        )

        output_dir = self.get_auto_remove_tmp_dir()

        training_args = Seq2SeqTrainingArguments(
            output_dir=output_dir,
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=batch_size,
            predict_with_generate=True,
            evaluation_strategy="steps",
            do_train=True,
            do_eval=True,
            warmup_steps=0,
            eval_steps=2,
            logging_steps=2,
        )

        # instantiate trainer
        trainer = Seq2SeqTrainer(
            model=bert2bert,
            args=training_args,
            compute_metrics=_compute_metrics,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
        )

        # start training
        trainer.train()

    def run_trainer(self, eval_steps: int, max_len: str, model_name: str, num_train_epochs: int):
        data_dir = self.examples_dir / "seq2seq/test_data/wmt_en_ro"
        output_dir = self.get_auto_remove_tmp_dir()
        args = f"""
            --model_name_or_path {model_name}
            --data_dir {data_dir}
            --output_dir {output_dir}
            --overwrite_output_dir
            --n_train 8
            --n_val 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --val_max_target_length {max_len}
            --do_train
            --do_eval
            --do_predict
            --num_train_epochs {str(num_train_epochs)}
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --learning_rate 3e-3
            --warmup_steps 8
            --evaluation_strategy steps
            --predict_with_generate
            --logging_steps 0
            --save_steps {str(eval_steps)}
            --eval_steps {str(eval_steps)}
            --sortish_sampler
            --label_smoothing 0.1
            --adafactor
            --task translation
            --tgt_lang ro_RO
            --src_lang en_XX
        """.split()
        # --eval_beams  2

        n_gpu = get_gpu_count()
        if n_gpu > 1:
            distributed_args = f"""
                -m torch.distributed.launch
                --nproc_per_node={n_gpu}
                {self.test_file_dir}/finetune_trainer.py
            """.split()
            cmd = [sys.executable] + distributed_args + args
            execute_subprocess_async(cmd, env=self.get_env())
        else:
            # 0 or 1 gpu
            testargs = ["finetune_trainer.py"] + args
            with patch.object(sys, "argv", testargs):
                main()

        return output_dir
<commit_msg>[s2s] broken test (#8613)<commit_after>import os
import sys
from unittest.mock import patch

from transformers import BertTokenizer, EncoderDecoderModel
from transformers.file_utils import is_datasets_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    require_torch_non_multi_gpu_but_fix_me,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed

from .finetune_trainer import Seq2SeqTrainingArguments, main
from .seq2seq_trainer import Seq2SeqTrainer
from .test_seq2seq_examples import MBART_TINY

set_seed(42)
MARIAN_MODEL = "sshleifer/student_marian_en_ro_6_1"

class TestFinetuneTrainer(TestCasePlus):
    def test_finetune_trainer(self):
        output_dir = self.run_trainer(1, "12", MBART_TINY, 1)
        logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history
        eval_metrics = [log for log in logs if "eval_loss" in log.keys()]
        first_step_stats = eval_metrics[0]
        assert "eval_bleu" in first_step_stats

    @slow
    def test_finetune_trainer_slow(self):
        # There is a missing call to __init__process_group somewhere
        output_dir = self.run_trainer(eval_steps=2, max_len="128", model_name=MARIAN_MODEL, num_train_epochs=10)

        # Check metrics
        logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history
        eval_metrics = [log for log in logs if "eval_loss" in log.keys()]
        first_step_stats = eval_metrics[0]
        last_step_stats = eval_metrics[-1]

        assert first_step_stats["eval_bleu"] < last_step_stats["eval_bleu"]  # model learned nothing
        assert isinstance(last_step_stats["eval_bleu"], float)

        # test if do_predict saves generations and metrics
        contents = os.listdir(output_dir)
        contents = {os.path.basename(p) for p in contents}
        assert "test_generations.txt" in contents
        assert "test_results.json" in contents

    @slow
    @require_torch_non_multi_gpu_but_fix_me
    def test_finetune_bert2bert(self):
        if not is_datasets_available():
            return

        import datasets

        bert2bert = EncoderDecoderModel.from_encoder_decoder_pretrained("prajjwal1/bert-tiny", "prajjwal1/bert-tiny")
        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

        bert2bert.config.vocab_size = bert2bert.config.encoder.vocab_size
        bert2bert.config.eos_token_id = tokenizer.sep_token_id
        bert2bert.config.decoder_start_token_id = tokenizer.cls_token_id
        bert2bert.config.max_length = 128

        train_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="train[:1%]")
        val_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="validation[:1%]")

        train_dataset = train_dataset.select(range(32))
        val_dataset = val_dataset.select(range(16))

        rouge = datasets.load_metric("rouge")

        batch_size = 4

        def _map_to_encoder_decoder_inputs(batch):
            # Tokenizer will automatically set [BOS] <text> [EOS]
            inputs = tokenizer(batch["article"], padding="max_length", truncation=True, max_length=512)
            outputs = tokenizer(batch["highlights"], padding="max_length", truncation=True, max_length=128)
            batch["input_ids"] = inputs.input_ids
            batch["attention_mask"] = inputs.attention_mask

            batch["decoder_input_ids"] = outputs.input_ids
            batch["labels"] = outputs.input_ids.copy()
            batch["labels"] = [
                [-100 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch["labels"]
            ]
            batch["decoder_attention_mask"] = outputs.attention_mask

            assert all([len(x) == 512 for x in inputs.input_ids])
            assert all([len(x) == 128 for x in outputs.input_ids])

            return batch

        def _compute_metrics(pred):
            labels_ids = pred.label_ids
            pred_ids = pred.predictions

            # all unnecessary tokens are removed
            pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
            label_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)

            rouge_output = rouge.compute(predictions=pred_str, references=label_str, rouge_types=["rouge2"])[
                "rouge2"
            ].mid

            return {
                "rouge2_precision": round(rouge_output.precision, 4),
                "rouge2_recall": round(rouge_output.recall, 4),
                "rouge2_fmeasure": round(rouge_output.fmeasure, 4),
            }

        # map train dataset
        train_dataset = train_dataset.map(
            _map_to_encoder_decoder_inputs,
            batched=True,
            batch_size=batch_size,
            remove_columns=["article", "highlights"],
        )
        train_dataset.set_format(
            type="torch",
            columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
        )

        # same for validation dataset
        val_dataset = val_dataset.map(
            _map_to_encoder_decoder_inputs,
            batched=True,
            batch_size=batch_size,
            remove_columns=["article", "highlights"],
        )
        val_dataset.set_format(
            type="torch",
            columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
        )

        output_dir = self.get_auto_remove_tmp_dir()

        training_args = Seq2SeqTrainingArguments(
            output_dir=output_dir,
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=batch_size,
            predict_with_generate=True,
            evaluation_strategy="steps",
            do_train=True,
            do_eval=True,
            warmup_steps=0,
            eval_steps=2,
            logging_steps=2,
        )

        # instantiate trainer
        trainer = Seq2SeqTrainer(
            model=bert2bert,
            args=training_args,
            compute_metrics=_compute_metrics,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
        )

        # start training
        trainer.train()

    def run_trainer(self, eval_steps: int, max_len: str, model_name: str, num_train_epochs: int):
        data_dir = self.examples_dir / "seq2seq/test_data/wmt_en_ro"
        output_dir = self.get_auto_remove_tmp_dir()
        args = f"""
            --model_name_or_path {model_name}
            --data_dir {data_dir}
            --output_dir {output_dir}
            --overwrite_output_dir
            --n_train 8
            --n_val 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --val_max_target_length {max_len}
            --do_train
            --do_eval
            --do_predict
            --num_train_epochs {str(num_train_epochs)}
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --learning_rate 3e-3
            --warmup_steps 8
            --evaluation_strategy steps
            --predict_with_generate
            --logging_steps 0
            --save_steps {str(eval_steps)}
            --eval_steps {str(eval_steps)}
            --sortish_sampler
            --label_smoothing 0.1
            --adafactor
            --task translation
            --tgt_lang ro_RO
            --src_lang en_XX
        """.split()
        # --eval_beams  2

        n_gpu = get_gpu_count()
        if n_gpu > 1:
            distributed_args = f"""
                -m torch.distributed.launch
                --nproc_per_node={n_gpu}
                {self.test_file_dir}/finetune_trainer.py
            """.split()
            cmd = [sys.executable] + distributed_args + args
            execute_subprocess_async(cmd, env=self.get_env())
        else:
            # 0 or 1 gpu
            testargs = ["finetune_trainer.py"] + args
            with patch.object(sys, "argv", testargs):
                main()

        return output_dir
<|endoftext|>


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

import abc
from typing import Iterable, List, Optional, Type

from spanner_orm import condition
from spanner_orm import error
from spanner_orm import field
from spanner_orm import foreign_key_relationship
from spanner_orm import model
from spanner_orm.admin import api
from spanner_orm.admin import index_column
from spanner_orm.admin import metadata

class SchemaUpdate(abc.ABC):
  """Base class for specifying schema updates."""

  @abc.abstractmethod
  def ddl(self) -> str:
    raise NotImplementedError

  def execute(self) -> None:
    self.validate()
    api.spanner_admin_api().update_schema(self.ddl())

  @abc.abstractmethod
  def validate(self) -> None:
    raise NotImplementedError

class CreateTable(SchemaUpdate):
  """Update that allows creating a new table."""

  def __init__(self, model_: Type[model.Model]):
    self._model = model_

  def ddl(self) -> str:
    key_fields = [
        '{} {}'.format(name, field.ddl())
        for name, field in self._model.fields.items()
    ]
    key_fields_ddl = ', '.join(key_fields)
    for relation in self._model.foreign_key_relations.values():
      key_fields_ddl += f', {relation.ddl}'
    index_ddl = 'PRIMARY KEY ({})'.format(', '.join(self._model.primary_keys))
    statement = 'CREATE TABLE {} ({}) {}'.format(self._model.table,
                                                 key_fields_ddl, index_ddl)

    if self._model.interleaved:
      statement += ', INTERLEAVE IN PARENT {parent} ON DELETE CASCADE'.format(
          parent=self._model.interleaved.table)
    return statement

  def validate(self) -> None:
    if not self._model.table:
      raise error.SpannerError('New table has no name')

    existing_model = metadata.SpannerMetadata.model(self._model.table)
    if existing_model:
      raise error.SpannerError('Table {} already exists'.format(
          self._model.table))

    if self._model.interleaved:
      self._validate_parent()

    self._validate_primary_keys()

  def _validate_parent(self) -> None:
    """Verifies that the parent table information is valid."""
    parent_primary_keys = self._model.interleaved.primary_keys
    primary_keys = self._model.primary_keys

    message = 'Table {} is not a child of parent table {}'.format(
        self._model.table, self._model.interleaved.table)
    for parent_key, key in zip(parent_primary_keys, primary_keys):
      if parent_key != key:
        raise error.SpannerError(message)
    if len(parent_primary_keys) > len(primary_keys):
      raise error.SpannerError(message)

  def _validate_primary_keys(self) -> None:
    """Verifies that the primary key data is valid."""
    if not self._model.primary_keys:
      raise error.SpannerError('Table {} has no primary key'.format(
          self._model.table))

    for key in self._model.primary_keys:
      if key not in self._model.fields:
        raise error.SpannerError(
            'Table {} column {} in primary key but not in schema'.format(
                self._model.table, key))

class DropTable(SchemaUpdate):
  """Update for dropping an existing table."""

  def __init__(self, table_name: str):
    self._table = table_name

  def ddl(self) -> str:
    return 'DROP TABLE {}'.format(self._table)

  def validate(self) -> None:
    existing_model = metadata.SpannerMetadata.model(self._table)
    if not existing_model:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    # Model indexes include the primary index
    if len(existing_model.indexes) > 1:
      raise error.SpannerError('Table {} has a secondary index'.format(
          self._table))

    self._validate_not_interleaved(existing_model)

  def _validate_not_interleaved(self,
                                existing_model: Type[model.Model]) -> None:
    for model_ in metadata.SpannerMetadata.models().values():
      if model_.interleaved == existing_model:
        raise error.SpannerError('Table {} has interleaved table {}'.format(
            self._table, model_.table))
      for index in model_.indexes.values():
        if index.parent == self._table:
          raise error.SpannerError('Table {} has interleaved index {}'.format(
              self._table, index.name))

class AddColumn(SchemaUpdate):
  """Update for adding a column to an existing table.

  Only supports adding nullable columns
  """

  def __init__(self, table_name: str, column_name: str, field_: field.Field):
    self._table = table_name
    self._column = column_name
    self._field = field_

  def ddl(self) -> str:
    return 'ALTER TABLE {} ADD COLUMN {} {}'.format(self._table, self._column,
                                                    self._field.ddl())

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))
    if not self._field.nullable():
      raise error.SpannerError('Column {} is not nullable'.format(self._column))
    if self._field.primary_key():
      raise error.SpannerError('Column {} is a primary key'.format(
          self._column))

class DropColumn(SchemaUpdate):
  """Update for dropping a column from an existing table."""

  def __init__(self, table_name: str, column_name: str):
    self._table = table_name
    self._column = column_name

  def ddl(self) -> str:
    return 'ALTER TABLE {} DROP COLUMN {}'.format(self._table, self._column)

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    if self._column not in model_.fields:
      raise error.SpannerError('Column {} does not exist on {}'.format(
          self._column, self._table))

    # Verify no indices exist on the column we're trying to drop
    num_indexed_columns = index_column.IndexColumnSchema.count(
        None, condition.equal_to('column_name', self._column),
        condition.equal_to('table_name', self._table))
    if num_indexed_columns > 0:
      raise error.SpannerError('Column {} is indexed'.format(self._column))

class AlterColumn(SchemaUpdate):
  """Update for altering a column an existing table.

  Only supports changing the nullability of a column
  """

  def __init__(self, table_name: str, column_name: str, field_: field.Field):
    self._table = table_name
    self._column = column_name
    self._field = field_

  def ddl(self) -> str:
    return 'ALTER TABLE {} ALTER COLUMN {} {}'.format(self._table, self._column,
                                                      self._field.ddl())

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    if self._column not in model_.fields:
      raise error.SpannerError('Column {} does not exist on {}'.format(
          self._column, self._table))

    if self._column in model_.primary_keys:
      raise error.SpannerError('Column {} is a primary key on {}'.format(
          self._column, self._table))

    old_field = model_.fields[self._column]
    # Validate that the only alteration is to change column nullability
    if self._field.field_type() != old_field.field_type():
      raise error.SpannerError('Column {} is changing type'.format(
          self._column))
    if self._field.nullable() == old_field.nullable():
      raise error.SpannerError('Column {} has no changes'.format(self._column))

class CreateIndex(SchemaUpdate):
  """Update for creating an index on an existing table."""

  def __init__(self,
               table_name: str,
               index_name: str,
               columns: Iterable[str],
               interleaved: Optional[str] = None,
               storing_columns: Optional[Iterable[str]] = None):
    self._table = table_name
    self._index = index_name
    self._columns = columns
    self._parent_table = interleaved
    self._storing_columns = storing_columns or []

  def ddl(self) -> str:
    statement = 'CREATE INDEX {} ON {} ({})'.format(self._index, self._table,
                                                    ', '.join(self._columns))
    if self._storing_columns:
      statement += 'STORING ({})'.format(', '.join(self._storing_columns))
    if self._parent_table:
      statement += ', INTERLEAVE IN {}'.format(self._parent_table)
    return statement

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    if not self._columns:
      raise error.SpannerError('Index {} has no columns'.format(self._index))

    if self._index in model_.indexes:
      raise error.SpannerError('Index {} already exists'.format(self._index))

    self._validate_columns(model_)

    if self._parent_table:
      self._validate_parent(model_)

  def _validate_columns(self, model_: Type[model.Model]) -> None:
    """Verifies all columns exist and are not part of the primary key."""
    for column in self._columns:
      if column not in model_.columns:
        raise error.SpannerError('Table {} has no column {}'.format(
            self._table, column))

    for column in self._storing_columns:
      if column not in model_.columns:
        raise error.SpannerError('Table {} has no column {}'.format(
            self._table, column))
      if column in model_.primary_keys:
        raise error.SpannerError('{} is part of the primary key for {}'.format(
            column, self._table))

  def _validate_parent(self, model_: Type[model.Model]) -> None:
    """Verifies this index can be interleaved in the parent table."""
    parent = model_.interleaved
    while parent:
      if parent == self._parent_table:
        break
      parent = parent.interleaved

    if not parent:
      raise error.SpannerError('{} is not a parent of table {}'.format(
          self._parent_table, self._table))

class DropIndex(SchemaUpdate):
  """Update for dropping a secondary index on an existing table."""

  def __init__(self, table_name: str, index_name: str):
    self._table = table_name
    self._index = index_name

  def ddl(self) -> str:
    return 'DROP INDEX {}'.format(self._index)

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    db_index = model_.indexes.get(self._index)
    if not db_index:
      raise error.SpannerError('Index {} does not exist'.format(self._index))
    if db_index.primary:
      raise error.SpannerError('Index {} is the primary index'.format(
          self._index))

class AddForeignKeyRelationship(SchemaUpdate):
  """Update for adding a column to an existing table.

  Only supports adding nullable columns
  """

  def __init__(
      self,
      referencing_table_name: str,
      referenced_table_name: str,
      column_mapping,
  ):
    self._table = table_name
    self._column = column_name
    self._field = field_

  def ddl(self) -> str:
    return 'ALTER TABLE {} ADD'.format(self._table, self._column,
                                                    self._field.ddl())

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

class DropForeignKeyRelationship(SchemaUpdate):
  """Update for dropping a column from an existing table."""

  def __init__(self, table_name: str, column_name: str):
    self._table = table_name
    self._column = column_name

  def ddl(self) -> str:
    return 'ALTER TABLE {} DROP COLUMN {}'.format(self._table, self._column)

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    if self._column not in model_.fields:
      raise error.SpannerError('Column {} does not exist on {}'.format(
          self._column, self._table))

    # Verify no indices exist on the column we're trying to drop
    num_indexed_columns = index_column.IndexColumnSchema.count(
        None, condition.equal_to('column_name', self._column),
        condition.equal_to('table_name', self._table))
    if num_indexed_columns > 0:
      raise error.SpannerError('Column {} is indexed'.format(self._column))

class NoUpdate(SchemaUpdate):
  """Update that does nothing, for migrations that don't update db schemas."""

  def ddl(self) -> str:
    return ''

  def execute(self) -> None:
    pass

  def validate(self) -> None:
    pass

def model_creation_ddl(model_: Type[model.Model]) -> List[str]:
  """Returns the list of ddl statements needed to create the model's table."""
  ddl_list = [CreateTable(model_).ddl()]

  for model_index in model_.indexes.values():
    if model_index.primary:
      continue
    create_index = CreateIndex(
        model_.table,
        model_index.name,
        model_index.columns,
        interleaved=model_index.parent,
        storing_columns=model_index.storing_columns)
    ddl_list.append(create_index.ddl())

  return ddl_list
<commit_msg>Remove changes that will be added in follow-up PR<commit_after># python3
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Used with SpannerAdminApi to manage Spanner schema updates."""

import abc
from typing import Iterable, List, Optional, Type

from spanner_orm import condition
from spanner_orm import error
from spanner_orm import field
from spanner_orm import foreign_key_relationship
from spanner_orm import model
from spanner_orm.admin import api
from spanner_orm.admin import index_column
from spanner_orm.admin import metadata

class SchemaUpdate(abc.ABC):
  """Base class for specifying schema updates."""

  @abc.abstractmethod
  def ddl(self) -> str:
    raise NotImplementedError

  def execute(self) -> None:
    self.validate()
    api.spanner_admin_api().update_schema(self.ddl())

  @abc.abstractmethod
  def validate(self) -> None:
    raise NotImplementedError

class CreateTable(SchemaUpdate):
  """Update that allows creating a new table."""

  def __init__(self, model_: Type[model.Model]):
    self._model = model_

  def ddl(self) -> str:
    key_fields = [
        '{} {}'.format(name, field.ddl())
        for name, field in self._model.fields.items()
    ]
    key_fields_ddl = ', '.join(key_fields)
    for relation in self._model.foreign_key_relations.values():
      key_fields_ddl += f', {relation.ddl}'
    index_ddl = 'PRIMARY KEY ({})'.format(', '.join(self._model.primary_keys))
    statement = 'CREATE TABLE {} ({}) {}'.format(self._model.table,
                                                 key_fields_ddl, index_ddl)

    if self._model.interleaved:
      statement += ', INTERLEAVE IN PARENT {parent} ON DELETE CASCADE'.format(
          parent=self._model.interleaved.table)
    return statement

  def validate(self) -> None:
    if not self._model.table:
      raise error.SpannerError('New table has no name')

    existing_model = metadata.SpannerMetadata.model(self._model.table)
    if existing_model:
      raise error.SpannerError('Table {} already exists'.format(
          self._model.table))

    if self._model.interleaved:
      self._validate_parent()

    self._validate_primary_keys()

  def _validate_parent(self) -> None:
    """Verifies that the parent table information is valid."""
    parent_primary_keys = self._model.interleaved.primary_keys
    primary_keys = self._model.primary_keys

    message = 'Table {} is not a child of parent table {}'.format(
        self._model.table, self._model.interleaved.table)
    for parent_key, key in zip(parent_primary_keys, primary_keys):
      if parent_key != key:
        raise error.SpannerError(message)
    if len(parent_primary_keys) > len(primary_keys):
      raise error.SpannerError(message)

  def _validate_primary_keys(self) -> None:
    """Verifies that the primary key data is valid."""
    if not self._model.primary_keys:
      raise error.SpannerError('Table {} has no primary key'.format(
          self._model.table))

    for key in self._model.primary_keys:
      if key not in self._model.fields:
        raise error.SpannerError(
            'Table {} column {} in primary key but not in schema'.format(
                self._model.table, key))

class DropTable(SchemaUpdate):
  """Update for dropping an existing table."""

  def __init__(self, table_name: str):
    self._table = table_name

  def ddl(self) -> str:
    return 'DROP TABLE {}'.format(self._table)

  def validate(self) -> None:
    existing_model = metadata.SpannerMetadata.model(self._table)
    if not existing_model:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    # Model indexes include the primary index
    if len(existing_model.indexes) > 1:
      raise error.SpannerError('Table {} has a secondary index'.format(
          self._table))

    self._validate_not_interleaved(existing_model)

  def _validate_not_interleaved(self,
                                existing_model: Type[model.Model]) -> None:
    for model_ in metadata.SpannerMetadata.models().values():
      if model_.interleaved == existing_model:
        raise error.SpannerError('Table {} has interleaved table {}'.format(
            self._table, model_.table))
      for index in model_.indexes.values():
        if index.parent == self._table:
          raise error.SpannerError('Table {} has interleaved index {}'.format(
              self._table, index.name))

class AddColumn(SchemaUpdate):
  """Update for adding a column to an existing table.

  Only supports adding nullable columns
  """

  def __init__(self, table_name: str, column_name: str, field_: field.Field):
    self._table = table_name
    self._column = column_name
    self._field = field_

  def ddl(self) -> str:
    return 'ALTER TABLE {} ADD COLUMN {} {}'.format(self._table, self._column,
                                                    self._field.ddl())

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))
    if not self._field.nullable():
      raise error.SpannerError('Column {} is not nullable'.format(self._column))
    if self._field.primary_key():
      raise error.SpannerError('Column {} is a primary key'.format(
          self._column))

class DropColumn(SchemaUpdate):
  """Update for dropping a column from an existing table."""

  def __init__(self, table_name: str, column_name: str):
    self._table = table_name
    self._column = column_name

  def ddl(self) -> str:
    return 'ALTER TABLE {} DROP COLUMN {}'.format(self._table, self._column)

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    if self._column not in model_.fields:
      raise error.SpannerError('Column {} does not exist on {}'.format(
          self._column, self._table))

    # Verify no indices exist on the column we're trying to drop
    num_indexed_columns = index_column.IndexColumnSchema.count(
        None, condition.equal_to('column_name', self._column),
        condition.equal_to('table_name', self._table))
    if num_indexed_columns > 0:
      raise error.SpannerError('Column {} is indexed'.format(self._column))

class AlterColumn(SchemaUpdate):
  """Update for altering a column an existing table.

  Only supports changing the nullability of a column
  """

  def __init__(self, table_name: str, column_name: str, field_: field.Field):
    self._table = table_name
    self._column = column_name
    self._field = field_

  def ddl(self) -> str:
    return 'ALTER TABLE {} ALTER COLUMN {} {}'.format(self._table, self._column,
                                                      self._field.ddl())

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    if self._column not in model_.fields:
      raise error.SpannerError('Column {} does not exist on {}'.format(
          self._column, self._table))

    if self._column in model_.primary_keys:
      raise error.SpannerError('Column {} is a primary key on {}'.format(
          self._column, self._table))

    old_field = model_.fields[self._column]
    # Validate that the only alteration is to change column nullability
    if self._field.field_type() != old_field.field_type():
      raise error.SpannerError('Column {} is changing type'.format(
          self._column))
    if self._field.nullable() == old_field.nullable():
      raise error.SpannerError('Column {} has no changes'.format(self._column))

class CreateIndex(SchemaUpdate):
  """Update for creating an index on an existing table."""

  def __init__(self,
               table_name: str,
               index_name: str,
               columns: Iterable[str],
               interleaved: Optional[str] = None,
               storing_columns: Optional[Iterable[str]] = None):
    self._table = table_name
    self._index = index_name
    self._columns = columns
    self._parent_table = interleaved
    self._storing_columns = storing_columns or []

  def ddl(self) -> str:
    statement = 'CREATE INDEX {} ON {} ({})'.format(self._index, self._table,
                                                    ', '.join(self._columns))
    if self._storing_columns:
      statement += 'STORING ({})'.format(', '.join(self._storing_columns))
    if self._parent_table:
      statement += ', INTERLEAVE IN {}'.format(self._parent_table)
    return statement

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    if not self._columns:
      raise error.SpannerError('Index {} has no columns'.format(self._index))

    if self._index in model_.indexes:
      raise error.SpannerError('Index {} already exists'.format(self._index))

    self._validate_columns(model_)

    if self._parent_table:
      self._validate_parent(model_)

  def _validate_columns(self, model_: Type[model.Model]) -> None:
    """Verifies all columns exist and are not part of the primary key."""
    for column in self._columns:
      if column not in model_.columns:
        raise error.SpannerError('Table {} has no column {}'.format(
            self._table, column))

    for column in self._storing_columns:
      if column not in model_.columns:
        raise error.SpannerError('Table {} has no column {}'.format(
            self._table, column))
      if column in model_.primary_keys:
        raise error.SpannerError('{} is part of the primary key for {}'.format(
            column, self._table))

  def _validate_parent(self, model_: Type[model.Model]) -> None:
    """Verifies this index can be interleaved in the parent table."""
    parent = model_.interleaved
    while parent:
      if parent == self._parent_table:
        break
      parent = parent.interleaved

    if not parent:
      raise error.SpannerError('{} is not a parent of table {}'.format(
          self._parent_table, self._table))

class DropIndex(SchemaUpdate):
  """Update for dropping a secondary index on an existing table."""

  def __init__(self, table_name: str, index_name: str):
    self._table = table_name
    self._index = index_name

  def ddl(self) -> str:
    return 'DROP INDEX {}'.format(self._index)

  def validate(self) -> None:
    model_ = metadata.SpannerMetadata.model(self._table)
    if not model_:
      raise error.SpannerError('Table {} does not exist'.format(self._table))

    db_index = model_.indexes.get(self._index)
    if not db_index:
      raise error.SpannerError('Index {} does not exist'.format(self._index))
    if db_index.primary:
      raise error.SpannerError('Index {} is the primary index'.format(
          self._index))

class NoUpdate(SchemaUpdate):
  """Update that does nothing, for migrations that don't update db schemas."""

  def ddl(self) -> str:
    return ''

  def execute(self) -> None:
    pass

  def validate(self) -> None:
    pass

def model_creation_ddl(model_: Type[model.Model]) -> List[str]:
  """Returns the list of ddl statements needed to create the model's table."""
  ddl_list = [CreateTable(model_).ddl()]

  for model_index in model_.indexes.values():
    if model_index.primary:
      continue
    create_index = CreateIndex(
        model_.table,
        model_index.name,
        model_index.columns,
        interleaved=model_index.parent,
        storing_columns=model_index.storing_columns)
    ddl_list.append(create_index.ddl())

  return ddl_list
<|endoftext|>


<commit_before>import logging
import os
import typing
from typing import Any, Dict, List, Optional, Text

from rasa_nlu.config import RasaNLUModelConfig
from rasa_nlu.extractors import EntityExtractor
from rasa_nlu.model import Metadata
from rasa_nlu.training_data import Message, TrainingData

logger = logging.getLogger(__name__)

if typing.TYPE_CHECKING:
    import mitie

class MitieEntityExtractor(EntityExtractor):

    provides = ["entities"]

    requires = ["tokens", "mitie_feature_extractor", "mitie_file"]

    def __init__(self,
                 component_config: Dict[Text, Any] = None,
                 ner=None
                 ):
        """Construct a new intent classifier using the sklearn framework."""

        super(MitieEntityExtractor, self).__init__(component_config)
        self.ner = ner

    @classmethod
    def required_packages(cls) -> List[Text]:
        return ["mitie"]

    def extract_entities(self, text, tokens, feature_extractor):
        ents = []
        tokens_strs = [token.text for token in tokens]
        if self.ner:
            entities = self.ner.extract_entities(tokens_strs,
                                                 feature_extractor)
            for e in entities:
                if len(e[0]):
                    start = tokens[e[0][0]].offset
                    end = tokens[e[0][-1]].end

                    ents.append({
                        "entity": e[1],
                        "value": text[start:end],
                        "start": start,
                        "end": end,
                        "confidence": None,
                    })

        return ents

    def train(self,
              training_data: TrainingData,
              config: RasaNLUModelConfig,
              **kwargs: Any) -> None:
        import mitie

        model_file = kwargs.get("mitie_file")
        if not model_file:
            raise Exception("Can not run MITIE entity extractor without a "
                            "language model. Make sure this component is "
                            "preceeded by the 'nlp_mitie' component.")

        trainer = mitie.ner_trainer(model_file)
        trainer.num_threads = kwargs.get("num_threads", 1)
        found_one_entity = False

        # filter out pre-trained entity examples
        filtered_entity_examples = self.filter_trainable_entities(
            training_data.training_examples)

        for example in filtered_entity_examples:
            sample = self._prepare_mitie_sample(example)

            found_one_entity = sample.num_entities > 0 or found_one_entity
            trainer.add(sample)

        # Mitie will fail to train if there is not a single entity tagged
        if found_one_entity:
            self.ner = trainer.train()

    def _prepare_mitie_sample(self, training_example):
        import mitie

        text = training_example.text
        tokens = training_example.get("tokens")
        sample = mitie.ner_training_instance([t.text for t in tokens])
        for ent in training_example.get("entities", []):
            try:
                # if the token is not aligned an exception will be raised
                start, end = MitieEntityExtractor.find_entity(
                    ent, text, tokens)
            except ValueError as e:
                logger.warning("Example skipped: {}".format(str(e)))
                continue
            try:
                # mitie will raise an exception on malicious
                # input - e.g. on overlapping entities
                sample.add_entity(list(range(start, end)), ent["entity"])
            except Exception as e:
                logger.warning("Failed to add entity example "
                               "'{}' of sentence '{}'. Reason: "
                               "{}".format(str(e), str(text), e))
                continue
        return sample

    def process(self, message: Message, **kwargs: Any) -> None:

        mitie_feature_extractor = kwargs.get("mitie_feature_extractor")
        if not mitie_feature_extractor:
            raise Exception("Failed to train 'intent_featurizer_mitie'. "
                            "Missing a proper MITIE feature extractor.")

        ents = self.extract_entities(message.text, message.get("tokens"),
                                     mitie_feature_extractor)
        extracted = self.add_extractor_name(ents)
        message.set("entities", message.get("entities", []) + extracted,
                    add_to_output=True)

    @classmethod
    def load(cls,
             meta: Dict,
             model_dir: Text = None,
             model_metadata: Metadata = None,
             cached_component: Optional['MitieEntityExtractor'] = None,
             **kwargs: Any
             ) -> 'MitieEntityExtractor':
        import mitie

        file_name = meta.get("file")

        if not file_name:
            return cls(meta)

        classifier_file = os.path.join(model_dir, file_name)
        if os.path.exists(classifier_file):
            extractor = mitie.named_entity_extractor(classifier_file)
            return cls(meta, extractor)
        else:
            return cls(meta)

    def persist(self,
                file_name: Text,
                model_dir: Text) -> Optional[Dict[Text, Any]]:

        if self.ner:
            file_name = file_name + ".dat"
            entity_extractor_file = os.path.join(model_dir, file_name)
            self.ner.save_to_disk(entity_extractor_file, pure_model=True)
            return {"file": file_name}
        else:
            return {"file": None}
<commit_msg>Update rasa_nlu/extractors/mitie_entity_extractor.py<commit_after>import logging
import os
import typing
from typing import Any, Dict, List, Optional, Text

from rasa_nlu.config import RasaNLUModelConfig
from rasa_nlu.extractors import EntityExtractor
from rasa_nlu.model import Metadata
from rasa_nlu.training_data import Message, TrainingData

logger = logging.getLogger(__name__)

if typing.TYPE_CHECKING:
    import mitie

class MitieEntityExtractor(EntityExtractor):

    provides = ["entities"]

    requires = ["tokens", "mitie_feature_extractor", "mitie_file"]

    def __init__(self,
                 component_config: Dict[Text, Any] = None,
                 ner=None
                 ):
        """Construct a new intent classifier using the sklearn framework."""

        super(MitieEntityExtractor, self).__init__(component_config)
        self.ner = ner

    @classmethod
    def required_packages(cls) -> List[Text]:
        return ["mitie"]

    def extract_entities(self, text, tokens, feature_extractor):
        ents = []
        tokens_strs = [token.text for token in tokens]
        if self.ner:
            entities = self.ner.extract_entities(tokens_strs,
                                                 feature_extractor)
            for e in entities:
                if len(e[0]):
                    start = tokens[e[0][0]].offset
                    end = tokens[e[0][-1]].end

                    ents.append({
                        "entity": e[1],
                        "value": text[start:end],
                        "start": start,
                        "end": end,
                        "confidence": None,
                    })

        return ents

    def train(self,
              training_data: TrainingData,
              config: RasaNLUModelConfig,
              **kwargs: Any) -> None:
        import mitie

        model_file = kwargs.get("mitie_file")
        if not model_file:
            raise Exception("Can not run MITIE entity extractor without a "
                            "language model. Make sure this component is "
                            "preceeded by the 'nlp_mitie' component.")

        trainer = mitie.ner_trainer(model_file)
        trainer.num_threads = kwargs.get("num_threads", 1)
        found_one_entity = False

        # filter out pre-trained entity examples
        filtered_entity_examples = self.filter_trainable_entities(
            training_data.training_examples)

        for example in filtered_entity_examples:
            sample = self._prepare_mitie_sample(example)

            found_one_entity = sample.num_entities > 0 or found_one_entity
            trainer.add(sample)

        # Mitie will fail to train if there is not a single entity tagged
        if found_one_entity:
            self.ner = trainer.train()

    def _prepare_mitie_sample(self, training_example):
        import mitie

        text = training_example.text
        tokens = training_example.get("tokens")
        sample = mitie.ner_training_instance([t.text for t in tokens])
        for ent in training_example.get("entities", []):
            try:
                # if the token is not aligned an exception will be raised
                start, end = MitieEntityExtractor.find_entity(
                    ent, text, tokens)
            except ValueError as e:
                logger.warning("Example skipped: {}".format(str(e)))
                continue
            try:
                # mitie will raise an exception on malicious
                # input - e.g. on overlapping entities
                sample.add_entity(list(range(start, end)), ent["entity"])
            except Exception as e:
                logger.warning("Failed to add entity example "
                               "'{}' of sentence '{}'. Reason: "
                               "{}".format(str(e), str(text), e))
                continue
        return sample

    def process(self, message: Message, **kwargs: Any) -> None:

        mitie_feature_extractor = kwargs.get("mitie_feature_extractor")
        if not mitie_feature_extractor:
            raise Exception("Failed to train 'intent_featurizer_mitie'. "
                            "Missing a proper MITIE feature extractor.")

        ents = self.extract_entities(message.text, message.get("tokens"),
                                     mitie_feature_extractor)
        extracted = self.add_extractor_name(ents)
        message.set("entities", message.get("entities", []) + extracted,
                    add_to_output=True)

    @classmethod
    def load(cls,
             meta: Dict[Text, Any],
             model_dir: Text = None,
             model_metadata: Metadata = None,
             cached_component: Optional['MitieEntityExtractor'] = None,
             **kwargs: Any
             ) -> 'MitieEntityExtractor':
        import mitie

        file_name = meta.get("file")

        if not file_name:
            return cls(meta)

        classifier_file = os.path.join(model_dir, file_name)
        if os.path.exists(classifier_file):
            extractor = mitie.named_entity_extractor(classifier_file)
            return cls(meta, extractor)
        else:
            return cls(meta)

    def persist(self,
                file_name: Text,
                model_dir: Text) -> Optional[Dict[Text, Any]]:

        if self.ner:
            file_name = file_name + ".dat"
            entity_extractor_file = os.path.join(model_dir, file_name)
            self.ner.save_to_disk(entity_extractor_file, pure_model=True)
            return {"file": file_name}
        else:
            return {"file": None}
<|endoftext|>


<commit_before>__author__ = 'magus0219'

import logging

import pymongo
import copy

from config import SERVER_CONFIG as _SERVER_CFG

logger_server = logging.getLogger("DeployServer")

try:
    mongodb_client = pymongo.MongoClient(host=_SERVER_CFG["MONGO_HOST"], port=_SERVER_CFG["MONGO_PORT"])
    mongodb_client['deployment'].authenticate(_SERVER_CFG["MONGO_USER"], _SERVER_CFG["MONGO_PWD"])
except Exception as e:
    logger_server.info("MongoDB cannot connect, Please check...")

def serialize_status(status):
    """Rewrite custom object in status dictionary to string for writing to mongo db

    :param status: dictionary of status
    :return:
    """
    copy_status = copy.deepcopy(status)
    copy_status['task_running'] = repr(copy_status['task_running'])
    for i in range(len(copy_status['task_waiting'])):
        copy_status['task_waiting'][i] = repr(copy_status['task_waiting'][i])

    for i in range(len(copy_status['last_tags'])):
        copy_status['last_tags'][i] = repr(copy_status['last_tags'][i])

    copy_status['last_commit_tag'] = repr(copy_status['last_commit_tag'])

    return copy_status<commit_msg>fix config name<commit_after>__author__ = 'magus0219'

import logging

import pymongo
import copy

from config import SERVER_CONFIG as _SERVER_CFG

logger_server = logging.getLogger("DeployServer")

try:
    mongodb_client = pymongo.MongoClient(host=_SERVER_CFG["MONGO_HOST"], port=_SERVER_CFG["MONGO_PORT"])
    mongodb_client['deployment'].authenticate(_SERVER_CFG["MONGO_USERNAME"], _SERVER_CFG["MONGO_PWD"])
except Exception as e:
    logger_server.info("MongoDB cannot connect, Please check...")

def serialize_status(status):
    """Rewrite custom object in status dictionary to string for writing to mongo db

    :param status: dictionary of status
    :return:
    """
    copy_status = copy.deepcopy(status)
    copy_status['task_running'] = repr(copy_status['task_running'])
    for i in range(len(copy_status['task_waiting'])):
        copy_status['task_waiting'][i] = repr(copy_status['task_waiting'][i])

    for i in range(len(copy_status['last_tags'])):
        copy_status['last_tags'][i] = repr(copy_status['last_tags'][i])

    copy_status['last_commit_tag'] = repr(copy_status['last_commit_tag'])

    return copy_status<|endoftext|>


<commit_before>import flask
from flask.ext.script import Manager
from path import path
from mptracker import models

def configure(app):
    project_root = path(__file__).abspath().parent.parent
    data_dir = project_root / '_data'
    data_dir.mkdir_p()
    db_path = data_dir / 'db.sqlite'
    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + db_path

def create_app():
    app = flask.Flask(__name__)
    configure(app)
    models.db.init_app(app)
    return app

manager = Manager(create_app)

@manager.command
def syncdb():
    models.db.create_all()

@manager.command
def import_people():
    from mpscraper.common import install_requests_cache
    from mpscraper.people import PersonScraper
    install_requests_cache()
    ps = PersonScraper()
    session = models.db.session
    for person_info in ps.fetch_people():
        p = models.Person(**person_info)
        session.add(p)
    session.commit()
<commit_msg>don't add duplicates<commit_after>import flask
from flask.ext.script import Manager
from path import path
from mptracker import models

def configure(app):
    project_root = path(__file__).abspath().parent.parent
    data_dir = project_root / '_data'
    data_dir.mkdir_p()
    db_path = data_dir / 'db.sqlite'
    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + db_path

def create_app():
    app = flask.Flask(__name__)
    configure(app)
    models.db.init_app(app)
    return app

manager = Manager(create_app)

@manager.command
def syncdb():
    models.db.create_all()

@manager.command
def import_people():
    from mpscraper.common import install_requests_cache
    from mpscraper.people import PersonScraper
    install_requests_cache()
    ps = PersonScraper()
    existing_cdep_ids = set(p.cdep_id for p in models.Person.query)
    new_people = 0
    session = models.db.session
    for person_info in ps.fetch_people():
        if person_info['cdep_id'] not in existing_cdep_ids:
            print('adding person:', person_info)
            p = models.Person(**person_info)
            session.add(p)
            existing_cdep_ids.add(p.cdep_id)
            new_people += 1
    print('added', new_people, 'people')
    session.commit()
<|endoftext|>


<commit_before>import sys
import functools
from importlib import import_module
from .exception import DoesNotExist, NotConnected
from .util import deprecated
from .match import OUTGOING, INCOMING, EITHER, rel_helper, Traversal, QueryBuilder, NodeSet

# check origin node is saved and not deleted
def check_origin(fn):
    fn_name = fn.func_name if hasattr(fn, 'func_name') else fn.__name__

    @functools.wraps(fn)
    def checker(self, *args, **kwargs):
        self.origin._pre_action_check(self.name + '.' + fn_name)
        return fn(self, *args, **kwargs)
    return checker

class RelationshipManager(object):
    def __init__(self, definition, origin):
        self.direction = definition['direction']
        self.relation_type = definition['relation_type']
        self.label_map = definition['label_map']
        self.definition = definition
        self.origin = origin

    def __str__(self):
        direction = 'either'
        if self.direction == OUTGOING:
            direction = 'a outgoing'
        elif self.direction == INCOMING:
            direction = 'a incoming'

        return "{0} in {1} direction of type {2} on node ({3}) of class '{4}'".format(
            self.description, direction,
            self.relation_type, self.origin._id, self.origin.__class__.__name__)

    @check_origin
    def __bool__(self):
        return len(self) > 0

    @check_origin
    def __nonzero__(self):
        return len(self) > 0

    @check_origin
    def __len__(self):
        return QueryBuilder(self._traversal)._count()

    @check_origin
    def count(self):
        return self.__len__()

    @check_origin
    def all(self):
        return QueryBuilder(self._traversal)._execute()

    def match(self, **kwargs):
        return QueryBuilder(self._traversal.match(**kwargs))._execute()

    @check_origin
    def get(self, **kwargs):
        result = self.search(**kwargs)
        if len(result) == 1:
            return result[0]
        if len(result) > 1:
            raise Exception("Multiple items returned, use search?")
        if not result:
            raise DoesNotExist("No items exist for the specified arguments")

    @check_origin
    @deprecated("search() is now deprecated please use filter() and exclude()")
    def search(self, **kwargs):
        ns = NodeSet(self._traversal)
        for field, value in kwargs.items():
            ns.filter(**{field: value})
        return QueryBuilder(ns)._execute()

    @check_origin
    def is_connected(self, obj):
        self._check_node(obj)

        rel = rel_helper(lhs='a', rhs='b', ident='r', **self.definition)
        q = "START a=node({self}), b=node({them}) MATCH" + rel + "RETURN count(r)"
        return bool(self.origin.cypher(q, {'them': obj._id})[0][0][0])

    def _check_node(self, obj):
        """check for valid node i.e correct class and is saved"""
        for label, cls in self.label_map.items():
            if obj.__class__ is cls:
                if not hasattr(obj, '_id'):
                    raise ValueError("Can't perform operation on unsaved node " + repr(obj))
                return

        allowed_cls = ", ".join([(tcls if isinstance(tcls, str) else tcls.__name__)
                                 for tcls, _ in self.label_map.items()])
        raise ValueError("Expected nodes of class "
                + allowed_cls + " got " + repr(obj)
                + " see relationship definition in " + self.origin.__class__.__name__)

    @check_origin
    def connect(self, obj, properties=None):
        self._check_node(obj)

        if not self.definition['model'] and properties:
            raise NotImplementedError("Relationship properties without " +
                    "using a relationship model is no longer supported")

        new_rel = rel_helper(lhs='us', rhs='them', ident='r', **self.definition)
        q = "START them=node({them}), us=node({self}) CREATE UNIQUE" + new_rel
        params = {'them': obj._id}

        if not properties and not self.definition['model']:
            self.origin.cypher(q, params)
            return True

        rel_model = self.definition['model']
        # need to generate defaults etc to create fake instance
        tmp = rel_model(**properties) if properties else rel_model()

        for p, v in rel_model.deflate(tmp.__properties__).items():
            params['place_holder_' + p] = v
            q += " SET r." + p + " = {place_holder_" + p + "}"

        rel_ = self.origin.cypher(q + " RETURN r", params)[0][0][0]
        rel_instance = rel_model.inflate(rel_)

        if self.definition['direction'] == INCOMING:
            rel_instance._start_node_class = obj.__class__
            rel_instance._end_node_class = self.origin.__class__
        else:
            rel_instance._start_node_class = self.origin.__class__
            rel_instance._end_node_class = obj.__class__

        self.origin.cypher(q, params)
        return rel_instance

    @check_origin
    def relationship(self, obj):
        """relationship: node"""
        self._check_node(obj)
        if not 'model' in self.definition:
            raise NotImplemented("'relationship' method only available on relationships"
                    + " that have a model defined")

        rel_model = self.definition['model']

        my_rel = rel_helper(lhs='us', rhs='them', ident='r', **self.definition)
        q = "START them=node({them}), us=node({self}) MATCH " + my_rel + " RETURN r"
        rel = self.origin.cypher(q, {'them': obj._id})[0][0][0]
        if not rel:
            return
        rel_instance = rel_model.inflate(rel)

        if self.definition['direction'] == INCOMING:
            rel_instance._start_node_class = obj.__class__
            rel_instance._end_node_class = self.origin.__class__
        else:
            rel_instance._start_node_class = self.origin.__class__
            rel_instance._end_node_class = obj.__class__
        return rel_instance

    @check_origin
    def reconnect(self, old_obj, new_obj):
        """reconnect: old_node, new_node"""
        self._check_node(old_obj)
        self._check_node(new_obj)
        if old_obj._id == new_obj._id:
            return
        old_rel = rel_helper(lhs='us', rhs='old', ident='r', **self.definition)

        # get list of properties on the existing rel
        result, meta = self.origin.cypher("START us=node({self}), old=node({old}) MATCH " + old_rel + " RETURN r",
            {'old': old_obj._id})
        if result:
            existing_properties = result[0][0]._properties.keys()
        else:
            raise NotConnected('reconnect', self.origin, old_obj)

        # remove old relationship and create new one
        new_rel = rel_helper(lhs='us', rhs='new', ident='r2', **self.definition)
        q = "START us=node({self}), old=node({old}), new=node({new}) MATCH " + old_rel
        q += " CREATE UNIQUE" + new_rel

        # copy over properties if we have
        for p in existing_properties:
            q += " SET r2.{} = r.{}".format(p, p)
        q += " WITH r DELETE r"

        self.origin.cypher(q, {'old': old_obj._id, 'new': new_obj._id})

    @check_origin
    def disconnect(self, obj):
        rel = rel_helper(lhs='a', rhs='b', ident='r', **self.definition)
        q = "START a=node({self}), b=node({them}) MATCH " + rel + " DELETE r"
        self.origin.cypher(q, {'them': obj._id})

    @check_origin
    def single(self):
        # TODO need to limit to one result
        nodes = QueryBuilder(self._traversal)._execute()
        return nodes[0] if nodes else None

class RelationshipDefinition(object):
    def __init__(self, relation_type, cls_name, direction, manager=RelationshipManager, model=None):
        self.module_name = sys._getframe(4).f_globals['__name__']
        self.module_file = sys._getframe(4).f_globals['__file__']
        self.node_class = cls_name
        self.manager = manager
        self.definition = {}
        self.definition['relation_type'] = relation_type
        self.definition['direction'] = direction
        self.definition['model'] = model

    def _lookup(self, name):
        if name.find('.') == -1:
            module = self.module_name
        else:
            module, _, name = name.rpartition('.')

        if not module in sys.modules:
            # yet another hack to get around python semantics
            # __name__ is the namespace of the parent module for __init__.py files,
            # and the namespace of the current module for other .py files,
            # therefore there's a need to define the namespace differently for
            # these two cases in order for . in relative imports to work correctly
            # (i.e. to mean the same thing for both cases).
            # For example in the comments below, namespace == myapp, always
            if '__init__.py' in self.module_file:
                # e.g. myapp/__init__.py -[__name__]-> myapp
                namespace = self.module_name
            else:
                # e.g. myapp/models.py -[__name__]-> myapp.models
                namespace = self.module_name.rpartition('.')[0]

            # load a module from a namespace (e.g. models from myapp)
            if module:
                module = import_module(module, namespace).__name__
            # load the namespace itself (e.g. myapp)
            # (otherwise it would look like import . from myapp)
            else:
                module = import_module(namespace).__name__
        return getattr(sys.modules[module], name)

    def build_label_map(self):
        # get classes for related nodes
        if isinstance(self.node_class, list):
            node_classes = [self._lookup(cls) if isinstance(cls, (str,)) else cls
                        for cls in self.node_class]
        else:
            node_classes = [self._lookup(self.node_class)
                if isinstance(self.node_class, (str,)) else self.node_class]

        # build label map
        self.definition['label_map'] = dict(zip([c.__label__
                for c in node_classes], node_classes))

    def build_manager(self, origin, name):
        self.build_label_map()
        rel = self.manager(self.definition, origin)
        rel.name = name
        rel._traversal = Traversal(origin, name, self.definition)
        return rel

class ZeroOrMore(RelationshipManager):
    description = "zero or more relationships"

def _relate(cls_name, direction, rel_type, cardinality=None, model=None):
    if not isinstance(cls_name, (str, list, object)):
        raise ValueError('Expected class name or list of class names, got ' + repr(cls_name))
    from .relationship import StructuredRel
    if model and not issubclass(model, (StructuredRel,)):
        raise ValueError('model must be a StructuredRel')
    return RelationshipDefinition(rel_type, cls_name, direction, cardinality, model)

def RelationshipTo(cls_name, rel_type, cardinality=ZeroOrMore, model=None):
    return _relate(cls_name, OUTGOING, rel_type, cardinality, model)

def RelationshipFrom(cls_name, rel_type, cardinality=ZeroOrMore, model=None):
    return _relate(cls_name, INCOMING, rel_type, cardinality, model)

def Relationship(cls_name, rel_type, cardinality=ZeroOrMore, model=None):
    return _relate(cls_name, EITHER, rel_type, cardinality, model)
<commit_msg>Fixes #105<commit_after>import sys
import functools
from importlib import import_module
from .exception import DoesNotExist, NotConnected
from .util import deprecated
from .match import OUTGOING, INCOMING, EITHER, rel_helper, Traversal, QueryBuilder, NodeSet

# check origin node is saved and not deleted
def check_origin(fn):
    fn_name = fn.func_name if hasattr(fn, 'func_name') else fn.__name__

    @functools.wraps(fn)
    def checker(self, *args, **kwargs):
        self.origin._pre_action_check(self.name + '.' + fn_name)
        return fn(self, *args, **kwargs)
    return checker

class RelationshipManager(object):
    def __init__(self, definition, origin):
        self.direction = definition['direction']
        self.relation_type = definition['relation_type']
        self.label_map = definition['label_map']
        self.definition = definition
        self.origin = origin

    def __str__(self):
        direction = 'either'
        if self.direction == OUTGOING:
            direction = 'a outgoing'
        elif self.direction == INCOMING:
            direction = 'a incoming'

        return "{0} in {1} direction of type {2} on node ({3}) of class '{4}'".format(
            self.description, direction,
            self.relation_type, self.origin._id, self.origin.__class__.__name__)

    @check_origin
    def __bool__(self):
        return len(self) > 0

    @check_origin
    def __nonzero__(self):
        return len(self) > 0

    @check_origin
    def __len__(self):
        return QueryBuilder(self._traversal)._count()

    @check_origin
    def count(self):
        return self.__len__()

    @check_origin
    def all(self):
        return QueryBuilder(self._traversal)._execute()

    def match(self, **kwargs):
        return QueryBuilder(self._traversal.match(**kwargs))._execute()

    @check_origin
    def get(self, **kwargs):
        result = self.search(**kwargs)
        if len(result) == 1:
            return result[0]
        if len(result) > 1:
            raise Exception("Multiple items returned, use search?")
        if not result:
            raise DoesNotExist("No items exist for the specified arguments")

    @check_origin
    @deprecated("search() is now deprecated please use filter() and exclude()")
    def search(self, **kwargs):
        ns = NodeSet(self._traversal)
        for field, value in kwargs.items():
            ns.filter(**{field: value})
        return QueryBuilder(ns)._execute()

    @check_origin
    def is_connected(self, obj):
        self._check_node(obj)

        rel = rel_helper(lhs='a', rhs='b', ident='r', **self.definition)
        q = "START a=node({self}), b=node({them}) MATCH" + rel + "RETURN count(r)"
        return bool(self.origin.cypher(q, {'them': obj._id})[0][0][0])

    def _check_node(self, obj):
        """check for valid node i.e correct class and is saved"""
        for label, cls in self.label_map.items():
            if obj.__class__ is cls:
                if not hasattr(obj, '_id'):
                    raise ValueError("Can't perform operation on unsaved node " + repr(obj))
                return

        allowed_cls = ", ".join([(tcls if isinstance(tcls, str) else tcls.__name__)
                                 for tcls, _ in self.label_map.items()])
        raise ValueError("Expected nodes of class "
                + allowed_cls + " got " + repr(obj)
                + " see relationship definition in " + self.origin.__class__.__name__)

    @check_origin
    def connect(self, obj, properties=None):
        self._check_node(obj)

        if not self.definition['model'] and properties:
            raise NotImplementedError("Relationship properties without " +
                    "using a relationship model is no longer supported")

        new_rel = rel_helper(lhs='us', rhs='them', ident='r', **self.definition)
        q = "START them=node({them}), us=node({self}) CREATE UNIQUE" + new_rel
        params = {'them': obj._id}

        if not properties and not self.definition['model']:
            self.origin.cypher(q, params)
            return True

        rel_model = self.definition['model']
        # need to generate defaults etc to create fake instance
        tmp = rel_model(**properties) if properties else rel_model()

        for p, v in rel_model.deflate(tmp.__properties__).items():
            params['place_holder_' + p] = v
            q += " SET r." + p + " = {place_holder_" + p + "}"

        rel_ = self.origin.cypher(q + " RETURN r", params)[0][0][0]
        rel_instance = rel_model.inflate(rel_)

        if self.definition['direction'] == INCOMING:
            rel_instance._start_node_class = obj.__class__
            rel_instance._end_node_class = self.origin.__class__
        else:
            rel_instance._start_node_class = self.origin.__class__
            rel_instance._end_node_class = obj.__class__

        self.origin.cypher(q, params)
        return rel_instance

    @check_origin
    def relationship(self, obj):
        """relationship: node"""
        self._check_node(obj)
        if 'model' not in self.definition:
            raise NotImplemented("'relationship' method only available on relationships"
                    + " that have a model defined")

        rel_model = self.definition['model']

        my_rel = rel_helper(lhs='us', rhs='them', ident='r', **self.definition)
        q = "START them=node({them}), us=node({self}) MATCH " + my_rel + " RETURN r"
        rel = self.origin.cypher(q, {'them': obj._id})[0][0][0]
        if not rel:
            return
        rel_instance = rel_model.inflate(rel)

        if self.definition['direction'] == INCOMING:
            rel_instance._start_node_class = obj.__class__
            rel_instance._end_node_class = self.origin.__class__
        else:
            rel_instance._start_node_class = self.origin.__class__
            rel_instance._end_node_class = obj.__class__
        return rel_instance

    @check_origin
    def reconnect(self, old_obj, new_obj):
        """reconnect: old_node, new_node"""
        self._check_node(old_obj)
        self._check_node(new_obj)
        if old_obj._id == new_obj._id:
            return
        old_rel = rel_helper(lhs='us', rhs='old', ident='r', **self.definition)

        # get list of properties on the existing rel
        result, meta = self.origin.cypher("START us=node({self}), old=node({old}) MATCH " + old_rel + " RETURN r",
            {'old': old_obj._id})
        if result:
            existing_properties = result[0][0]._properties.keys()
        else:
            raise NotConnected('reconnect', self.origin, old_obj)

        # remove old relationship and create new one
        new_rel = rel_helper(lhs='us', rhs='new', ident='r2', **self.definition)
        q = "START us=node({self}), old=node({old}), new=node({new}) MATCH " + old_rel
        q += " CREATE UNIQUE" + new_rel

        # copy over properties if we have
        for p in existing_properties:
            q += " SET r2.{} = r.{}".format(p, p)
        q += " WITH r DELETE r"

        self.origin.cypher(q, {'old': old_obj._id, 'new': new_obj._id})

    @check_origin
    def disconnect(self, obj):
        rel = rel_helper(lhs='a', rhs='b', ident='r', **self.definition)
        q = "START a=node({self}), b=node({them}) MATCH " + rel + " DELETE r"
        self.origin.cypher(q, {'them': obj._id})

    @check_origin
    def single(self):
        # TODO need to limit to one result
        nodes = QueryBuilder(self._traversal)._execute()
        return nodes[0] if nodes else None

class RelationshipDefinition(object):
    def __init__(self, relation_type, cls_name, direction, manager=RelationshipManager, model=None):
        self.module_name = sys._getframe(4).f_globals['__name__']
        if '__file__' in sys._getframe(4).f_globals:
            self.module_file = sys._getframe(4).f_globals['__file__']
        self.node_class = cls_name
        self.manager = manager
        self.definition = {}
        self.definition['relation_type'] = relation_type
        self.definition['direction'] = direction
        self.definition['model'] = model

    def _lookup(self, name):
        if name.find('.') == -1:
            module = self.module_name
        else:
            module, _, name = name.rpartition('.')

        if module not in sys.modules:
            # yet another hack to get around python semantics
            # __name__ is the namespace of the parent module for __init__.py files,
            # and the namespace of the current module for other .py files,
            # therefore there's a need to define the namespace differently for
            # these two cases in order for . in relative imports to work correctly
            # (i.e. to mean the same thing for both cases).
            # For example in the comments below, namespace == myapp, always
            if not hasattr(self, 'module_file'):
                raise ImportError("Couldn't lookup '{}'".format(name))

            if '__init__.py' in self.module_file:
                # e.g. myapp/__init__.py -[__name__]-> myapp
                namespace = self.module_name
            else:
                # e.g. myapp/models.py -[__name__]-> myapp.models
                namespace = self.module_name.rpartition('.')[0]

            # load a module from a namespace (e.g. models from myapp)
            if module:
                module = import_module(module, namespace).__name__
            # load the namespace itself (e.g. myapp)
            # (otherwise it would look like import . from myapp)
            else:
                module = import_module(namespace).__name__
        return getattr(sys.modules[module], name)

    def build_label_map(self):
        # get classes for related nodes
        if isinstance(self.node_class, list):
            node_classes = [self._lookup(cls) if isinstance(cls, (str,)) else cls
                        for cls in self.node_class]
        else:
            node_classes = [self._lookup(self.node_class)
                if isinstance(self.node_class, (str,)) else self.node_class]

        # build label map
        self.definition['label_map'] = dict(zip([c.__label__
                for c in node_classes], node_classes))

    def build_manager(self, origin, name):
        self.build_label_map()
        rel = self.manager(self.definition, origin)
        rel.name = name
        rel._traversal = Traversal(origin, name, self.definition)
        return rel

class ZeroOrMore(RelationshipManager):
    description = "zero or more relationships"

def _relate(cls_name, direction, rel_type, cardinality=None, model=None):
    if not isinstance(cls_name, (str, list, object)):
        raise ValueError('Expected class name or list of class names, got ' + repr(cls_name))
    from .relationship import StructuredRel
    if model and not issubclass(model, (StructuredRel,)):
        raise ValueError('model must be a StructuredRel')
    return RelationshipDefinition(rel_type, cls_name, direction, cardinality, model)

def RelationshipTo(cls_name, rel_type, cardinality=ZeroOrMore, model=None):
    return _relate(cls_name, OUTGOING, rel_type, cardinality, model)

def RelationshipFrom(cls_name, rel_type, cardinality=ZeroOrMore, model=None):
    return _relate(cls_name, INCOMING, rel_type, cardinality, model)

def Relationship(cls_name, rel_type, cardinality=ZeroOrMore, model=None):
    return _relate(cls_name, EITHER, rel_type, cardinality, model)
<|endoftext|>


<commit_before>from boto.exception import BotoServerError

class CfnSphereException(Exception):
    def __init__(self, message="", boto_exception=None):
        self.message = message
        self.str = self.message
        try:
            self.request_id = boto_exception.request_id
            self.str += " (Request ID: {0})".format(self.request_id)
            self.boto_exception = boto_exception
        except:
            pass

    def __str__(self):
        return self.str

class CfnStackActionFailedException(CfnSphereException):
    pass

class TemplateErrorException(CfnSphereException):
    pass

class NoConfigException(CfnSphereException):
    pass

class BadConfigException(CfnSphereException):
    pass

class CyclicDependencyException(CfnSphereException):
    pass

class InvalidDependencyGraphException(CfnSphereException):
    pass

class InvalidEncryptedValueException(CfnSphereException):
    pass

class CfnSphereBotoErrorException(CfnSphereException):
    def __init__(self, e):
        self.boto_exception = e

    def __str__(self):
        return "{0}: {1} (Request ID: {2})".format(self.boto_exception.error_code,
                                                 self.boto_exception.message,
                                                 self.boto_exception.request_id)
<commit_msg>Always initialize self.boto_exception<commit_after>from boto.exception import BotoServerError

class CfnSphereException(Exception):
    def __init__(self, message="", boto_exception=None):
        self.message = message
        self.str = self.message
        self.boto_exception = boto_exception
        try:
            self.request_id = boto_exception.request_id
            self.str += " (Request ID: {0})".format(self.request_id)
        except:
            pass

    def __str__(self):
        return self.str

class CfnStackActionFailedException(CfnSphereException):
    pass

class TemplateErrorException(CfnSphereException):
    pass

class NoConfigException(CfnSphereException):
    pass

class BadConfigException(CfnSphereException):
    pass

class CyclicDependencyException(CfnSphereException):
    pass

class InvalidDependencyGraphException(CfnSphereException):
    pass

class InvalidEncryptedValueException(CfnSphereException):
    pass

class CfnSphereBotoErrorException(CfnSphereException):
    def __init__(self, e):
        self.boto_exception = e

    def __str__(self):
        return "{0}: {1} (Request ID: {2})".format(self.boto_exception.error_code,
                                                 self.boto_exception.message,
                                                 self.boto_exception.request_id)
<|endoftext|>


<commit_before>from utils import *
from DB import ActiveRecord, db
from User import User
from Project import Project
from Sprint import Sprint
from Group import Group
from Goal import Goal
from inspect import getmembers
from datetime import datetime, date

class Status:
	def getIcon(self): return "/static/images/status-%s.png" % self.name.replace(' ', '-')
	icon = property(getIcon)

	def __init__(self, name, text, color, revisionVerb):
		self.name = name
		self.text = text
		self.color = color
		self.revisionVerb = revisionVerb

statuses = [
	Status('blocked', 'Blocked', '#000', 'Blocked'),
	Status('canceled', 'Canceled', '#000', 'Canceled'),
	Status('complete', 'Complete', '#0F0', 'Completed'),
	Status('deferred', 'Deferred', '#000', 'Deferred'),
	Status('in progress', 'In Progress', '#FF0', 'Started'),
	Status('not started', 'Not Started', '#F00', 'Aborted'),
	Status('split', 'Split', '#000', 'Split'),
]
statuses = dict([(s.name, s) for s in statuses])

statusMenu = [('not started', 'in progress', 'complete'), ('blocked',), ('deferred', 'canceled', 'split')]
assert sum(map(len, statusMenu)) == len(statuses)

class Task(ActiveRecord):
	sprint = ActiveRecord.idObjLink(Sprint, 'sprintid')
	creator = ActiveRecord.idObjLink(User, 'creatorid')
	assigned = ActiveRecord.idObjLink(User, 'assignedid')
	group = ActiveRecord.idObjLink(Group, 'groupid')
	goal = ActiveRecord.idObjLink(Goal, 'goalid')

	def getStatus(self): return statuses[self.status]
	def setStatus(self, stat): self.status = stat.name
	stat = property(getStatus, setStatus)

	def __init__(self, groupid, sprintid, creatorid, assignedid, goalid, name, status, hours, seq = None, timestamp = None, deleted = 0, revision = 1, id = None):
		ActiveRecord.__init__(self)
		self.id = id
		self.revision = revision
		self.groupid = groupid
		self.sprintid = sprintid
		self.creatorid = creatorid
		self.assignedid = assignedid
		self.goalid = goalid
		self.name = name
		self.status = status
		self.hours = hours
		self.timestamp = timestamp if timestamp else dateToTs(datetime.now())
		self.seq = seq if seq else maxOr(task.seq for task in self.group.getTasks())+1
		self.deleted = deleted

	def __str__(self):
		return self.safe.name

	def getRevisions(self):
		rows = db().select("SELECT * FROM %s WHERE id = ? ORDER BY revision" % Task.table(), self.id)
		return map(lambda x: Task(**x), rows)

	def getRevisionAt(self, date):
		timestamp = dateToTs(date)
		rows = db().select("SELECT * FROM %s WHERE id = ? AND timestamp <= ? ORDER BY revision DESC LIMIT 1" % Task.table(), self.id, tsEnd(timestamp))
		rows = [x for x in rows]
		return Task(**rows[0]) if len(rows) > 0 else None

	@classmethod
	def load(cls, *id, **attrs):
		if len(id): # Searching by id
			if len(id) != 1:
				raise ArgumentMismatchError
			cur = db().cursor("SELECT * FROM %s WHERE id = ? ORDER BY revision DESC LIMIT 1" % cls.table(), id[0])
		else: # Searching by attributes
			placeholders = ["%s = ?" % k for k in attrs.keys()]
			vals = attrs.values()
			cur = db().cursor("SELECT * FROM %s WHERE %s ORDER BY revision DESC LIMIT 1" % (cls.table(), ' AND '.join(placeholders)), *vals)

		row = cur.fetchone()
		cur.close()
		if row: # Checking in case rowcount == -1 (unsupported)
			return cls(**row)
		else:
			return None

	@classmethod
	def loadAll(cls, orderby = None, **attrs):
		return super(Task, cls).loadAll(groupby = 'id', orderby = orderby, **attrs)

	def save(self):
		if not self.id:
			# Pre-insert since tasks.id isn't autoincrementing
			rows = db().select("SELECT MAX(id) FROM tasks");
			rows = [x for x in rows]
			self.id = (rows[0]['MAX(id)'] or 0) + 1
			db().update("INSERT INTO tasks(id, revision) VALUES(?, ?)", self.id, 1)

			# Shift everything after this sequence
			db().update("UPDATE tasks SET seq = seq + 1 WHERE groupid = ? AND seq >= ?", self.groupid, self.seq)
		return ActiveRecord.save(self, pks = ['id', 'revision'])

	def move(self, newSeq, newGroup):
		# Remove from current group (shift all later tasks up)
		db().update("UPDATE tasks SET seq = seq - 1 WHERE groupid = ? AND seq > ?", self.groupid, self.seq)

		# Switch group (for all revisions)
		if(self.group != newGroup):
			self.group = newGroup
			db().update("UPDATE tasks SET groupid = ? WHERE id = ?", self.groupid, self.id)

		# Add to new group (shift all later tasks down)
		self.seq = newSeq
		db().update("UPDATE tasks SET seq = seq + 1 WHERE groupid = ? AND seq >= ?", self.groupid, self.seq)
		db().update("UPDATE tasks SET seq = ? WHERE id = ?", self.seq, self.id)

	# @classmethod
	# def loadAll(cls, **attrs):
		# if len(attrs):
			# placeholders = ["%s = ?" % k for k in attrs.keys()]
			# vals = attrs.values()
			# rows = db().select("SELECT * FROM %s WHERE %s GROUP BY id" % (cls.table(), ' AND '.join(placeholders)), *vals)
		# else:
			# rows = db().select("SELECT * FROM %s GROUP BY id" % cls.table())
		# return map(lambda x: cls(**x), rows)

	def revise(self):
		cls = self.__class__

		self.revision += 1 # Bump the revision number

		fields = [x for x in set(cls.fields())]
		vals = dict(getmembers(self))
		boundArgs = [vals[k] for k in fields]

		placeholders = ', '.join(map(lambda x: "?", fields))
		db().update("INSERT INTO %s(%s) VALUES(%s)" % (cls.table(), ', '.join(fields), placeholders), *boundArgs)
<commit_msg>Switch to using the revision icons for now (too hard to differentiate the sprint icons on the projector)<commit_after>from utils import *
from DB import ActiveRecord, db
from User import User
from Project import Project
from Sprint import Sprint
from Group import Group
from Goal import Goal
from inspect import getmembers
from datetime import datetime, date

class Status:
	def getIcon(self): return "/static/images/revision-%s.png" % self.name.replace(' ', '-')
	icon = property(getIcon)

	def __init__(self, name, text, color, revisionVerb):
		self.name = name
		self.text = text
		self.color = color
		self.revisionVerb = revisionVerb

statuses = [
	Status('blocked', 'Blocked', '#000', 'Blocked'),
	Status('canceled', 'Canceled', '#000', 'Canceled'),
	Status('complete', 'Complete', '#0F0', 'Completed'),
	Status('deferred', 'Deferred', '#000', 'Deferred'),
	Status('in progress', 'In Progress', '#FF0', 'Started'),
	Status('not started', 'Not Started', '#F00', 'Aborted'),
	Status('split', 'Split', '#000', 'Split'),
]
statuses = dict([(s.name, s) for s in statuses])

statusMenu = [('not started', 'in progress', 'complete'), ('blocked',), ('deferred', 'canceled', 'split')]
assert sum(map(len, statusMenu)) == len(statuses)

class Task(ActiveRecord):
	sprint = ActiveRecord.idObjLink(Sprint, 'sprintid')
	creator = ActiveRecord.idObjLink(User, 'creatorid')
	assigned = ActiveRecord.idObjLink(User, 'assignedid')
	group = ActiveRecord.idObjLink(Group, 'groupid')
	goal = ActiveRecord.idObjLink(Goal, 'goalid')

	def getStatus(self): return statuses[self.status]
	def setStatus(self, stat): self.status = stat.name
	stat = property(getStatus, setStatus)

	def __init__(self, groupid, sprintid, creatorid, assignedid, goalid, name, status, hours, seq = None, timestamp = None, deleted = 0, revision = 1, id = None):
		ActiveRecord.__init__(self)
		self.id = id
		self.revision = revision
		self.groupid = groupid
		self.sprintid = sprintid
		self.creatorid = creatorid
		self.assignedid = assignedid
		self.goalid = goalid
		self.name = name
		self.status = status
		self.hours = hours
		self.timestamp = timestamp if timestamp else dateToTs(datetime.now())
		self.seq = seq if seq else maxOr(task.seq for task in self.group.getTasks())+1
		self.deleted = deleted

	def __str__(self):
		return self.safe.name

	def getRevisions(self):
		rows = db().select("SELECT * FROM %s WHERE id = ? ORDER BY revision" % Task.table(), self.id)
		return map(lambda x: Task(**x), rows)

	def getRevisionAt(self, date):
		timestamp = dateToTs(date)
		rows = db().select("SELECT * FROM %s WHERE id = ? AND timestamp <= ? ORDER BY revision DESC LIMIT 1" % Task.table(), self.id, tsEnd(timestamp))
		rows = [x for x in rows]
		return Task(**rows[0]) if len(rows) > 0 else None

	@classmethod
	def load(cls, *id, **attrs):
		if len(id): # Searching by id
			if len(id) != 1:
				raise ArgumentMismatchError
			cur = db().cursor("SELECT * FROM %s WHERE id = ? ORDER BY revision DESC LIMIT 1" % cls.table(), id[0])
		else: # Searching by attributes
			placeholders = ["%s = ?" % k for k in attrs.keys()]
			vals = attrs.values()
			cur = db().cursor("SELECT * FROM %s WHERE %s ORDER BY revision DESC LIMIT 1" % (cls.table(), ' AND '.join(placeholders)), *vals)

		row = cur.fetchone()
		cur.close()
		if row: # Checking in case rowcount == -1 (unsupported)
			return cls(**row)
		else:
			return None

	@classmethod
	def loadAll(cls, orderby = None, **attrs):
		return super(Task, cls).loadAll(groupby = 'id', orderby = orderby, **attrs)

	def save(self):
		if not self.id:
			# Pre-insert since tasks.id isn't autoincrementing
			rows = db().select("SELECT MAX(id) FROM tasks");
			rows = [x for x in rows]
			self.id = (rows[0]['MAX(id)'] or 0) + 1
			db().update("INSERT INTO tasks(id, revision) VALUES(?, ?)", self.id, 1)

			# Shift everything after this sequence
			db().update("UPDATE tasks SET seq = seq + 1 WHERE groupid = ? AND seq >= ?", self.groupid, self.seq)
		return ActiveRecord.save(self, pks = ['id', 'revision'])

	def move(self, newSeq, newGroup):
		# Remove from current group (shift all later tasks up)
		db().update("UPDATE tasks SET seq = seq - 1 WHERE groupid = ? AND seq > ?", self.groupid, self.seq)

		# Switch group (for all revisions)
		if(self.group != newGroup):
			self.group = newGroup
			db().update("UPDATE tasks SET groupid = ? WHERE id = ?", self.groupid, self.id)

		# Add to new group (shift all later tasks down)
		self.seq = newSeq
		db().update("UPDATE tasks SET seq = seq + 1 WHERE groupid = ? AND seq >= ?", self.groupid, self.seq)
		db().update("UPDATE tasks SET seq = ? WHERE id = ?", self.seq, self.id)

	# @classmethod
	# def loadAll(cls, **attrs):
		# if len(attrs):
			# placeholders = ["%s = ?" % k for k in attrs.keys()]
			# vals = attrs.values()
			# rows = db().select("SELECT * FROM %s WHERE %s GROUP BY id" % (cls.table(), ' AND '.join(placeholders)), *vals)
		# else:
			# rows = db().select("SELECT * FROM %s GROUP BY id" % cls.table())
		# return map(lambda x: cls(**x), rows)

	def revise(self):
		cls = self.__class__

		self.revision += 1 # Bump the revision number

		fields = [x for x in set(cls.fields())]
		vals = dict(getmembers(self))
		boundArgs = [vals[k] for k in fields]

		placeholders = ', '.join(map(lambda x: "?", fields))
		db().update("INSERT INTO %s(%s) VALUES(%s)" % (cls.table(), ', '.join(fields), placeholders), *boundArgs)
<|endoftext|>


<commit_before>""" Finds and returns list of continuum pixels, as a mask. """

def find_contpix(lambdas, fluxes, ivars):
    """ Find and return continuum pixels

    Parameters
    ----------
    lambdas: numpy ndarray of length npixels
        rest-frame wavelength vector
    fluxes: numpy ndarray of shape (nstars, pixels)
        pixel intensities
    ivars: numpy ndarray of shape (nstars, npixels)
        inverse variances, parallel to fluxes

    Returns
    -------
    contmask: boolean mask of length npixels
        True indicates that the pixel is continuum
    """

<commit_msg>set f_cut and sig_cut by experimentation<commit_after>""" Finds and returns list of continuum pixels, as a mask. """

def get_num_contpix(f_cut, sig_cut, lambdas, fluxes, ivars):
    """ Find and return continuum pixels given the flux and sigma cut

    Parameters
    ----------
    f_cut: float
        the upper limit imposed on the quantity (fbar-1)
    sig_cut: float
        the upper limit imposed on the quantity (f_sig)
    lambdas: numpy ndarray of length npixels
        rest-frame wavelength vector
    fluxes: numpy ndarray of shape (nstars, npixels)
        pixel intensities
    ivars: numpy ndarray of shape nstars, npixels
        inverse variances, parallel to fluxes

    Returns
    -------
    contmask: boolean mask of length npixels
        True indicates that the pixel is continuum
    """
    f_bar = np.median(fluxes, axis=0)
    sigma_f = np.var(fluxes, axis=0)
    cont1 = np.abs(f_bar-1) < f_cut
    cont2 = sigma_f < sigma_cut
    contmask = np.logical_and(cont1, cont2)
    return contmask

def find_contpix(lambdas, fluxes, ivars):
    """ Find and return continuum pixels

    Searches through some flux and sigma cuts such that between 5 and 10%
    of pixels are identified as continuum.

    Parameters
    ----------
    lambdas: numpy ndarray of length npixels
        rest-frame wavelength vector
    fluxes: numpy ndarray of shape (nstars, pixels)
        pixel intensities
    ivars: numpy ndarray of shape (nstars, npixels)
        inverse variances, parallel to fluxes

    Returns
    -------
    contmask: boolean mask of length npixels
        True indicates that the pixel is continuum
    """
    # Apply a cut based on the median and variance vector
    # Adjust the cut levels so that 5-7% of pixels are identified as continuum
    npixels = len(lambdas)
    f_cut = 0.003
    stepsize = 0.0001
    sig_cut = 0.003
    contmask = get_num_contpix(f_cut, sig_cut, lambdas, fluxes, ivars)
    frac = sum(contmask)/npixels
    while frac < 0.05*npixels:
        f_cut += sig_step
        contmask = get_num_contpix(f_cut, sig_cut, lambdas, fluxes, ivars)
        frac = sum(contmask)/npixels
    if frac > 0.10*npixels:
        print("Warning: Over 10% of pixels identified as continuum.")
    print("%s out of %s pixels identified as continuum" %(sum(contmask),
                                                          npixels))
    return contmask
<|endoftext|>