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

# Lint as: python2, python3
r"""Train and Eval SAC.

All hyperparameters come from the SAC paper
https://arxiv.org/pdf/1812.05905.pdf

To run:

```bash
tensorboard --logdir $HOME/tmp/sac/gym/HalfCheetah-v2/ --port 2223 &

python tf_agents/agents/sac/examples/v2/train_eval.py \
  --root_dir=$HOME/tmp/sac/gym/HalfCheetah-v2/ \
  --alsologtostderr
```
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import time

from absl import app
from absl import flags
from absl import logging

import gin
from six.moves import range
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import

from tf_agents.agents.ddpg import critic_network
from tf_agents.agents.sac import sac_agent
from tf_agents.agents.sac import tanh_normal_projection_network
from tf_agents.drivers import dynamic_step_driver
from tf_agents.environments import suite_mujoco
from tf_agents.environments import tf_py_environment
from tf_agents.eval import metric_utils
from tf_agents.metrics import tf_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.policies import greedy_policy
from tf_agents.policies import random_tf_policy
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.utils import common

flags.DEFINE_string('root_dir', os.getenv('TEST_UNDECLARED_OUTPUTS_DIR'),
                    'Root directory for writing logs/summaries/checkpoints.')
flags.DEFINE_multi_string('gin_file', None, 'Path to the trainer config files.')
flags.DEFINE_multi_string('gin_param', None, 'Gin binding to pass through.')

FLAGS = flags.FLAGS

@gin.configurable
def train_eval(
    root_dir,
    env_name='HalfCheetah-v2',
    eval_env_name=None,
    env_load_fn=suite_mujoco.load,
    # The SAC paper reported:
    # Hopper and Cartpole results up to 1000000 iters,
    # Humanoid results up to 10000000 iters,
    # Other mujoco tasks up to 3000000 iters.
    num_iterations=3000000,
    actor_fc_layers=(256, 256),
    critic_obs_fc_layers=None,
    critic_action_fc_layers=None,
    critic_joint_fc_layers=(256, 256),
    # Params for collect
    # Follow https://github.com/haarnoja/sac/blob/master/examples/variants.py
    # HalfCheetah and Ant take 10000 initial collection steps.
    # Other mujoco tasks take 1000.
    # Different choices roughly keep the initial episodes about the same.
    initial_collect_steps=10000,
    collect_steps_per_iteration=1,
    replay_buffer_capacity=1000000,
    # Params for target update
    target_update_tau=0.005,
    target_update_period=1,
    # Params for train
    train_steps_per_iteration=1,
    batch_size=256,
    actor_learning_rate=3e-4,
    critic_learning_rate=3e-4,
    alpha_learning_rate=3e-4,
    td_errors_loss_fn=tf.math.squared_difference,
    gamma=0.99,
    reward_scale_factor=0.1,
    gradient_clipping=None,
    use_tf_functions=True,
    # Params for eval
    num_eval_episodes=30,
    eval_interval=10000,
    # Params for summaries and logging
    train_checkpoint_interval=10000,
    policy_checkpoint_interval=5000,
    rb_checkpoint_interval=50000,
    log_interval=1000,
    summary_interval=1000,
    summaries_flush_secs=10,
    debug_summaries=False,
    summarize_grads_and_vars=False,
    eval_metrics_callback=None):
  """A simple train and eval for SAC."""
  root_dir = os.path.expanduser(root_dir)
  train_dir = os.path.join(root_dir, 'train')
  eval_dir = os.path.join(root_dir, 'eval')

  train_summary_writer = tf.compat.v2.summary.create_file_writer(
      train_dir, flush_millis=summaries_flush_secs * 1000)
  train_summary_writer.set_as_default()

  eval_summary_writer = tf.compat.v2.summary.create_file_writer(
      eval_dir, flush_millis=summaries_flush_secs * 1000)
  eval_metrics = [
      tf_metrics.AverageReturnMetric(buffer_size=num_eval_episodes),
      tf_metrics.AverageEpisodeLengthMetric(buffer_size=num_eval_episodes)
  ]

  global_step = tf.compat.v1.train.get_or_create_global_step()
  with tf.compat.v2.summary.record_if(
      lambda: tf.math.equal(global_step % summary_interval, 0)):
    tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
    eval_env_name = eval_env_name or env_name
    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(eval_env_name))

    time_step_spec = tf_env.time_step_spec()
    observation_spec = time_step_spec.observation
    action_spec = tf_env.action_spec()

    actor_net = actor_distribution_network.ActorDistributionNetwork(
        observation_spec,
        action_spec,
        fc_layer_params=actor_fc_layers,
        continuous_projection_net=tanh_normal_projection_network
        .TanhNormalProjectionNetwork)
    critic_net = critic_network.CriticNetwork(
        (observation_spec, action_spec),
        observation_fc_layer_params=critic_obs_fc_layers,
        action_fc_layer_params=critic_action_fc_layers,
        joint_fc_layer_params=critic_joint_fc_layers,
        kernel_initializer='glorot_uniform',
        last_kernel_initializer='glorot_uniform')

    tf_agent = sac_agent.SacAgent(
        time_step_spec,
        action_spec,
        actor_network=actor_net,
        critic_network=critic_net,
        actor_optimizer=tf.compat.v1.train.AdamOptimizer(
            learning_rate=actor_learning_rate),
        critic_optimizer=tf.compat.v1.train.AdamOptimizer(
            learning_rate=critic_learning_rate),
        alpha_optimizer=tf.compat.v1.train.AdamOptimizer(
            learning_rate=alpha_learning_rate),
        target_update_tau=target_update_tau,
        target_update_period=target_update_period,
        td_errors_loss_fn=td_errors_loss_fn,
        gamma=gamma,
        reward_scale_factor=reward_scale_factor,
        gradient_clipping=gradient_clipping,
        debug_summaries=debug_summaries,
        summarize_grads_and_vars=summarize_grads_and_vars,
        train_step_counter=global_step)
    tf_agent.initialize()

    # Make the replay buffer.
    replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
        data_spec=tf_agent.collect_data_spec,
        batch_size=1,
        max_length=replay_buffer_capacity)
    replay_observer = [replay_buffer.add_batch]

    train_metrics = [
        tf_metrics.NumberOfEpisodes(),
        tf_metrics.EnvironmentSteps(),
        tf_metrics.AverageReturnMetric(
            buffer_size=num_eval_episodes, batch_size=tf_env.batch_size),
        tf_metrics.AverageEpisodeLengthMetric(
            buffer_size=num_eval_episodes, batch_size=tf_env.batch_size),
    ]

    eval_policy = greedy_policy.GreedyPolicy(tf_agent.policy)
    initial_collect_policy = random_tf_policy.RandomTFPolicy(
        tf_env.time_step_spec(), tf_env.action_spec())
    collect_policy = tf_agent.collect_policy

    train_checkpointer = common.Checkpointer(
        ckpt_dir=train_dir,
        agent=tf_agent,
        global_step=global_step,
        metrics=metric_utils.MetricsGroup(train_metrics, 'train_metrics'))
    policy_checkpointer = common.Checkpointer(
        ckpt_dir=os.path.join(train_dir, 'policy'),
        policy=eval_policy,
        global_step=global_step)
    rb_checkpointer = common.Checkpointer(
        ckpt_dir=os.path.join(train_dir, 'replay_buffer'),
        max_to_keep=1,
        replay_buffer=replay_buffer)

    train_checkpointer.initialize_or_restore()
    rb_checkpointer.initialize_or_restore()

    if replay_buffer.num_frames() == 0:
      initial_collect_driver = dynamic_step_driver.DynamicStepDriver(
          tf_env,
          initial_collect_policy,
          observers=replay_observer + train_metrics,
          num_steps=initial_collect_steps)

    collect_driver = dynamic_step_driver.DynamicStepDriver(
        tf_env,
        collect_policy,
        observers=replay_observer + train_metrics,
        num_steps=collect_steps_per_iteration)

    if use_tf_functions:
      initial_collect_driver.run = common.function(initial_collect_driver.run)
      collect_driver.run = common.function(collect_driver.run)
      tf_agent.train = common.function(tf_agent.train)

    # Collect initial replay data.
    logging.info(
        'Initializing replay buffer by collecting experience for %d steps with '
        'a random policy.', initial_collect_steps)
    initial_collect_driver.run()

    results = metric_utils.eager_compute(
        eval_metrics,
        eval_tf_env,
        eval_policy,
        num_episodes=num_eval_episodes,
        train_step=global_step,
        summary_writer=eval_summary_writer,
        summary_prefix='Metrics',
    )
    if eval_metrics_callback is not None:
      eval_metrics_callback(results, global_step.numpy())
    metric_utils.log_metrics(eval_metrics)

    time_step = None
    policy_state = collect_policy.get_initial_state(tf_env.batch_size)

    timed_at_step = global_step.numpy()
    time_acc = 0

    # Prepare replay buffer as dataset with invalid transitions filtered.
    def _filter_invalid_transition(trajectories, unused_arg1):
      return ~trajectories.is_boundary()[0]
    dataset = replay_buffer.as_dataset(
        sample_batch_size=batch_size,
        num_steps=2).unbatch().filter(
            _filter_invalid_transition).batch(batch_size).prefetch(5)
    # Dataset generates trajectories with shape [Bx2x...]
    iterator = iter(dataset)

    def train_step():
      experience, _ = next(iterator)
      return tf_agent.train(experience)

    if use_tf_functions:
      train_step = common.function(train_step)

    for _ in range(num_iterations):
      start_time = time.time()
      time_step, policy_state = collect_driver.run(
          time_step=time_step,
          policy_state=policy_state,
      )
      for _ in range(train_steps_per_iteration):
        train_loss = train_step()
      time_acc += time.time() - start_time

      global_step_val = global_step.numpy()

      if global_step_val % log_interval == 0:
        logging.info('step = %d, loss = %f', global_step_val,
                     train_loss.loss)
        steps_per_sec = (global_step_val - timed_at_step) / time_acc
        logging.info('%.3f steps/sec', steps_per_sec)
        tf.compat.v2.summary.scalar(
            name='global_steps_per_sec', data=steps_per_sec, step=global_step)
        timed_at_step = global_step_val
        time_acc = 0

      for train_metric in train_metrics:
        train_metric.tf_summaries(
            train_step=global_step, step_metrics=train_metrics[:2])

      if global_step_val % eval_interval == 0:
        results = metric_utils.eager_compute(
            eval_metrics,
            eval_tf_env,
            eval_policy,
            num_episodes=num_eval_episodes,
            train_step=global_step,
            summary_writer=eval_summary_writer,
            summary_prefix='Metrics',
        )
        if eval_metrics_callback is not None:
          eval_metrics_callback(results, global_step_val)
        metric_utils.log_metrics(eval_metrics)

      if global_step_val % train_checkpoint_interval == 0:
        train_checkpointer.save(global_step=global_step_val)

      if global_step_val % policy_checkpoint_interval == 0:
        policy_checkpointer.save(global_step=global_step_val)

      if global_step_val % rb_checkpoint_interval == 0:
        rb_checkpointer.save(global_step=global_step_val)
    return train_loss

def main(_):
  tf.compat.v1.enable_v2_behavior()
  logging.set_verbosity(logging.INFO)
  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)
  train_eval(FLAGS.root_dir)

if __name__ == '__main__':
  flags.mark_flag_as_required('root_dir')
  app.run(main)
<commit_msg>Fixed bug where initial_collect_driver is undefined when restoring from checkpoint.<commit_after># coding=utf-8
# Copyright 2018 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python2, python3
r"""Train and Eval SAC.

All hyperparameters come from the SAC paper
https://arxiv.org/pdf/1812.05905.pdf

To run:

```bash
tensorboard --logdir $HOME/tmp/sac/gym/HalfCheetah-v2/ --port 2223 &

python tf_agents/agents/sac/examples/v2/train_eval.py \
  --root_dir=$HOME/tmp/sac/gym/HalfCheetah-v2/ \
  --alsologtostderr
```
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import time

from absl import app
from absl import flags
from absl import logging

import gin
from six.moves import range
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import

from tf_agents.agents.ddpg import critic_network
from tf_agents.agents.sac import sac_agent
from tf_agents.agents.sac import tanh_normal_projection_network
from tf_agents.drivers import dynamic_step_driver
from tf_agents.environments import suite_mujoco
from tf_agents.environments import tf_py_environment
from tf_agents.eval import metric_utils
from tf_agents.metrics import tf_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.policies import greedy_policy
from tf_agents.policies import random_tf_policy
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.utils import common

flags.DEFINE_string('root_dir', os.getenv('TEST_UNDECLARED_OUTPUTS_DIR'),
                    'Root directory for writing logs/summaries/checkpoints.')
flags.DEFINE_multi_string('gin_file', None, 'Path to the trainer config files.')
flags.DEFINE_multi_string('gin_param', None, 'Gin binding to pass through.')

FLAGS = flags.FLAGS

@gin.configurable
def train_eval(
    root_dir,
    env_name='HalfCheetah-v2',
    eval_env_name=None,
    env_load_fn=suite_mujoco.load,
    # The SAC paper reported:
    # Hopper and Cartpole results up to 1000000 iters,
    # Humanoid results up to 10000000 iters,
    # Other mujoco tasks up to 3000000 iters.
    num_iterations=3000000,
    actor_fc_layers=(256, 256),
    critic_obs_fc_layers=None,
    critic_action_fc_layers=None,
    critic_joint_fc_layers=(256, 256),
    # Params for collect
    # Follow https://github.com/haarnoja/sac/blob/master/examples/variants.py
    # HalfCheetah and Ant take 10000 initial collection steps.
    # Other mujoco tasks take 1000.
    # Different choices roughly keep the initial episodes about the same.
    initial_collect_steps=10000,
    collect_steps_per_iteration=1,
    replay_buffer_capacity=1000000,
    # Params for target update
    target_update_tau=0.005,
    target_update_period=1,
    # Params for train
    train_steps_per_iteration=1,
    batch_size=256,
    actor_learning_rate=3e-4,
    critic_learning_rate=3e-4,
    alpha_learning_rate=3e-4,
    td_errors_loss_fn=tf.math.squared_difference,
    gamma=0.99,
    reward_scale_factor=0.1,
    gradient_clipping=None,
    use_tf_functions=True,
    # Params for eval
    num_eval_episodes=30,
    eval_interval=10000,
    # Params for summaries and logging
    train_checkpoint_interval=10000,
    policy_checkpoint_interval=5000,
    rb_checkpoint_interval=50000,
    log_interval=1000,
    summary_interval=1000,
    summaries_flush_secs=10,
    debug_summaries=False,
    summarize_grads_and_vars=False,
    eval_metrics_callback=None):
  """A simple train and eval for SAC."""
  root_dir = os.path.expanduser(root_dir)
  train_dir = os.path.join(root_dir, 'train')
  eval_dir = os.path.join(root_dir, 'eval')

  train_summary_writer = tf.compat.v2.summary.create_file_writer(
      train_dir, flush_millis=summaries_flush_secs * 1000)
  train_summary_writer.set_as_default()

  eval_summary_writer = tf.compat.v2.summary.create_file_writer(
      eval_dir, flush_millis=summaries_flush_secs * 1000)
  eval_metrics = [
      tf_metrics.AverageReturnMetric(buffer_size=num_eval_episodes),
      tf_metrics.AverageEpisodeLengthMetric(buffer_size=num_eval_episodes)
  ]

  global_step = tf.compat.v1.train.get_or_create_global_step()
  with tf.compat.v2.summary.record_if(
      lambda: tf.math.equal(global_step % summary_interval, 0)):
    tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
    eval_env_name = eval_env_name or env_name
    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(eval_env_name))

    time_step_spec = tf_env.time_step_spec()
    observation_spec = time_step_spec.observation
    action_spec = tf_env.action_spec()

    actor_net = actor_distribution_network.ActorDistributionNetwork(
        observation_spec,
        action_spec,
        fc_layer_params=actor_fc_layers,
        continuous_projection_net=tanh_normal_projection_network
        .TanhNormalProjectionNetwork)
    critic_net = critic_network.CriticNetwork(
        (observation_spec, action_spec),
        observation_fc_layer_params=critic_obs_fc_layers,
        action_fc_layer_params=critic_action_fc_layers,
        joint_fc_layer_params=critic_joint_fc_layers,
        kernel_initializer='glorot_uniform',
        last_kernel_initializer='glorot_uniform')

    tf_agent = sac_agent.SacAgent(
        time_step_spec,
        action_spec,
        actor_network=actor_net,
        critic_network=critic_net,
        actor_optimizer=tf.compat.v1.train.AdamOptimizer(
            learning_rate=actor_learning_rate),
        critic_optimizer=tf.compat.v1.train.AdamOptimizer(
            learning_rate=critic_learning_rate),
        alpha_optimizer=tf.compat.v1.train.AdamOptimizer(
            learning_rate=alpha_learning_rate),
        target_update_tau=target_update_tau,
        target_update_period=target_update_period,
        td_errors_loss_fn=td_errors_loss_fn,
        gamma=gamma,
        reward_scale_factor=reward_scale_factor,
        gradient_clipping=gradient_clipping,
        debug_summaries=debug_summaries,
        summarize_grads_and_vars=summarize_grads_and_vars,
        train_step_counter=global_step)
    tf_agent.initialize()

    # Make the replay buffer.
    replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
        data_spec=tf_agent.collect_data_spec,
        batch_size=1,
        max_length=replay_buffer_capacity)
    replay_observer = [replay_buffer.add_batch]

    train_metrics = [
        tf_metrics.NumberOfEpisodes(),
        tf_metrics.EnvironmentSteps(),
        tf_metrics.AverageReturnMetric(
            buffer_size=num_eval_episodes, batch_size=tf_env.batch_size),
        tf_metrics.AverageEpisodeLengthMetric(
            buffer_size=num_eval_episodes, batch_size=tf_env.batch_size),
    ]

    eval_policy = greedy_policy.GreedyPolicy(tf_agent.policy)
    initial_collect_policy = random_tf_policy.RandomTFPolicy(
        tf_env.time_step_spec(), tf_env.action_spec())
    collect_policy = tf_agent.collect_policy

    train_checkpointer = common.Checkpointer(
        ckpt_dir=train_dir,
        agent=tf_agent,
        global_step=global_step,
        metrics=metric_utils.MetricsGroup(train_metrics, 'train_metrics'))
    policy_checkpointer = common.Checkpointer(
        ckpt_dir=os.path.join(train_dir, 'policy'),
        policy=eval_policy,
        global_step=global_step)
    rb_checkpointer = common.Checkpointer(
        ckpt_dir=os.path.join(train_dir, 'replay_buffer'),
        max_to_keep=1,
        replay_buffer=replay_buffer)

    train_checkpointer.initialize_or_restore()
    rb_checkpointer.initialize_or_restore()

    if replay_buffer.num_frames() == 0:
      initial_collect_driver = dynamic_step_driver.DynamicStepDriver(
          tf_env,
          initial_collect_policy,
          observers=replay_observer + train_metrics,
          num_steps=initial_collect_steps)

    collect_driver = dynamic_step_driver.DynamicStepDriver(
        tf_env,
        collect_policy,
        observers=replay_observer + train_metrics,
        num_steps=collect_steps_per_iteration)

    if use_tf_functions:
      initial_collect_driver.run = common.function(initial_collect_driver.run)
      collect_driver.run = common.function(collect_driver.run)
      tf_agent.train = common.function(tf_agent.train)

    if replay_buffer.num_frames() == 0:
      # Collect initial replay data.
      logging.info(
          'Initializing replay buffer by collecting experience for %d steps '
          'with a random policy.', initial_collect_steps)
      initial_collect_driver.run()

    results = metric_utils.eager_compute(
        eval_metrics,
        eval_tf_env,
        eval_policy,
        num_episodes=num_eval_episodes,
        train_step=global_step,
        summary_writer=eval_summary_writer,
        summary_prefix='Metrics',
    )
    if eval_metrics_callback is not None:
      eval_metrics_callback(results, global_step.numpy())
    metric_utils.log_metrics(eval_metrics)

    time_step = None
    policy_state = collect_policy.get_initial_state(tf_env.batch_size)

    timed_at_step = global_step.numpy()
    time_acc = 0

    # Prepare replay buffer as dataset with invalid transitions filtered.
    def _filter_invalid_transition(trajectories, unused_arg1):
      return ~trajectories.is_boundary()[0]
    dataset = replay_buffer.as_dataset(
        sample_batch_size=batch_size,
        num_steps=2).unbatch().filter(
            _filter_invalid_transition).batch(batch_size).prefetch(5)
    # Dataset generates trajectories with shape [Bx2x...]
    iterator = iter(dataset)

    def train_step():
      experience, _ = next(iterator)
      return tf_agent.train(experience)

    if use_tf_functions:
      train_step = common.function(train_step)

    for _ in range(num_iterations):
      start_time = time.time()
      time_step, policy_state = collect_driver.run(
          time_step=time_step,
          policy_state=policy_state,
      )
      for _ in range(train_steps_per_iteration):
        train_loss = train_step()
      time_acc += time.time() - start_time

      global_step_val = global_step.numpy()

      if global_step_val % log_interval == 0:
        logging.info('step = %d, loss = %f', global_step_val,
                     train_loss.loss)
        steps_per_sec = (global_step_val - timed_at_step) / time_acc
        logging.info('%.3f steps/sec', steps_per_sec)
        tf.compat.v2.summary.scalar(
            name='global_steps_per_sec', data=steps_per_sec, step=global_step)
        timed_at_step = global_step_val
        time_acc = 0

      for train_metric in train_metrics:
        train_metric.tf_summaries(
            train_step=global_step, step_metrics=train_metrics[:2])

      if global_step_val % eval_interval == 0:
        results = metric_utils.eager_compute(
            eval_metrics,
            eval_tf_env,
            eval_policy,
            num_episodes=num_eval_episodes,
            train_step=global_step,
            summary_writer=eval_summary_writer,
            summary_prefix='Metrics',
        )
        if eval_metrics_callback is not None:
          eval_metrics_callback(results, global_step_val)
        metric_utils.log_metrics(eval_metrics)

      if global_step_val % train_checkpoint_interval == 0:
        train_checkpointer.save(global_step=global_step_val)

      if global_step_val % policy_checkpoint_interval == 0:
        policy_checkpointer.save(global_step=global_step_val)

      if global_step_val % rb_checkpoint_interval == 0:
        rb_checkpointer.save(global_step=global_step_val)
    return train_loss

def main(_):
  tf.compat.v1.enable_v2_behavior()
  logging.set_verbosity(logging.INFO)
  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)
  train_eval(FLAGS.root_dir)

if __name__ == '__main__':
  flags.mark_flag_as_required('root_dir')
  app.run(main)
<|endoftext|>


<commit_before># Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core functionality for the classical PSHA hazard calculator.

The difficult part of this calculator is the management of the logic tree
realizations. Let me explain how it works in a real life case.

We want to perform a SHARE calculation with a single source
model with two tectonic region types (Stable Shallow Crust and
Active Shallow Crust) and a complex GMPEs logic tree with 7
branching points with the following structure:

Active_Shallow: 4 GMPEs, weights 0.35, 0.35, 0.20, 0.10
Stable_Shallow: 5 GMPEs, weights 0.2, 0.2, 0.2, 0.2, 0.2
Shield: 2 GMPEs, weights 0.5, 0.5
Subduction_Interface: 4 GMPEs, weights 0.2, 0.2, 0.2, 0.4
Subduction_InSlab: 4 GMPEs, weights 0.2, 0.2, 0.2, 0.4
Volcanic: 1 GMPE, weight 1
Deep: 2 GMPEs, weights 0.6, 0.4

The number of realizations generated by this logic tree is
4 * 5 * 2 * 4 * 4 * 1 * 2 = 1280 and at the end we will
generate 1280 hazard curve containers. However the independent
hazard curve containers are much less, there are actually only 9 of
them, 4 coming from the Active Shallow Crust tectonic region type
and 5 from the Stable Shallow Crust tectonic region type.
The dependent hazard curves (i.e. hazard curves by realization)
can be obtained from the independent hazard curves by a composition
rule implemented in :method:`openquake.engine.models.HazardCurve.build_data`

NB: notice that 1280 / 9 = 142.22, therefore storing all the hazard
curves takes 140+ times more disk space and resources than actually needed.
The plan for the future is store only the independent curves and to
compose them on-the-fly, when they are looked up for a given realization,
since the composition is pretty fast (there are just numpy multiplications),
faster than reading from the database all the redundant data. The
impact on hazard maps has to be determined.
"""
import time
import operator
import itertools

import numpy

from openquake.hazardlib.imt import from_string
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.hazardlib.geo.utils import get_spherical_bounding_box
from openquake.hazardlib.geo.utils import get_longitudinal_extent
from openquake.hazardlib.geo.geodetic import npoints_between

from openquake.calculators.calc import gen_ruptures

from openquake.engine import writer
from openquake.engine.calculators import calculators
from openquake.engine.calculators.hazard import general
from openquake.engine.db import models
from openquake.engine.utils import tasks

inserter = writer.CacheInserter(models.SourceInfo, 100)

class BoundingBox(object):
    """
    A class to store the bounding box in distances, longitudes and magnitudes,
    given a source model and a site. This is used for disaggregation
    calculations. The goal is to determine the minimum and maximum
    distances of the ruptures generated from the model from the site;
    moreover the maximum and minimum longitudes and magnitudes are stored, by
    taking in account the international date line.
    """
    def __init__(self, lt_model_id, site_id):
        self.lt_model_id = lt_model_id
        self.site_id = site_id
        self.min_dist = self.max_dist = None
        self.east = self.west = self.south = self.north = None

    def update(self, dists, lons, lats):
        """
        Compare the current bounding box with the value in the arrays
        dists, lons, lats and enlarge it if needed.

        :param dists:
            a sequence of distances
        :param lons:
            a sequence of longitudes
        :param lats:
            a sequence of latitudes
        """
        if self.min_dist is not None:
            dists = [self.min_dist, self.max_dist] + dists
        if self.west is not None:
            lons = [self.west, self.east] + lons
        if self.south is not None:
            lats = [self.south, self.north] + lats
        self.min_dist, self.max_dist = min(dists), max(dists)
        self.west, self.east, self.north, self.south = \
            get_spherical_bounding_box(lons, lats)

    def update_bb(self, bb):
        """
        Compare the current bounding box with the given bounding box
        and enlarge it if needed.

        :param bb:
            an instance of :class:
            `openquake.engine.calculators.hazard.classical.core.BoundingBox`
        """
        if bb:  # the given bounding box must be non-empty
            self.update([bb.min_dist, bb.max_dist], [bb.west, bb.east],
                        [bb.south, bb.north])

    def bins_edges(self, dist_bin_width, coord_bin_width):
        """
        Define bin edges for disaggregation histograms, from the bin data
        collected from the ruptures.

        :param dists:
            array of distances from the ruptures
        :param lons:
            array of longitudes from the ruptures
        :param lats:
            array of latitudes from the ruptures
        :param dist_bin_width:
            distance_bin_width from job.ini
        :param coord_bin_width:
            coordinate_bin_width from job.ini
        """
        dist_edges = dist_bin_width * numpy.arange(
            int(self.min_dist / dist_bin_width),
            int(numpy.ceil(self.max_dist / dist_bin_width) + 1))

        west = numpy.floor(self.west / coord_bin_width) * coord_bin_width
        east = numpy.ceil(self.east / coord_bin_width) * coord_bin_width
        lon_extent = get_longitudinal_extent(west, east)

        lon_edges, _, _ = npoints_between(
            west, 0, 0, east, 0, 0,
            numpy.round(lon_extent / coord_bin_width) + 1)

        lat_edges = coord_bin_width * numpy.arange(
            int(numpy.floor(self.south / coord_bin_width)),
            int(numpy.ceil(self.north / coord_bin_width) + 1))

        return dist_edges, lon_edges, lat_edges

    def __nonzero__(self):
        """
        True if the bounding box is non empty.
        """
        return (self.min_dist is not None and self.west is not None and
                self.south is not None)

def _calc_pnes(gsim, r_sites, rupture, imts, imls, truncation_level,
               make_ctxt_mon, calc_poes_mon):
    # compute the probabilities of no exceedence for each IMT
    # for the given gsim and rupture; returns a list of pairs
    # [(imt, pnes), ...]
    with make_ctxt_mon:
        sctx, rctx, dctx = ContextMaker([gsim]).make_contexts(r_sites, rupture)
    with calc_poes_mon:
        for imt, levels in itertools.izip(imts, imls):
            poes = gsim.get_poes(
                sctx, rctx, dctx, imt, levels, truncation_level)
            pnes = rupture.get_probability_no_exceedance(poes)
            yield r_sites.expand(pnes, placeholder=1)

@tasks.oqtask
def compute_hazard_curves(sources, sitecol, gsims_by_trt_id, monitor):
    """
    This task computes R2 * I hazard curves (each one is a
    numpy array of S * L floats) from the given source_ruptures
    pairs.

    :param sources:
        a block of source objects
    :param sitecol:
        a :class:`openquake.hazardlib.site.SiteCollection` instance
    :param gsims_by_trt_id:
        a dictionary trt_id -> gsim instances
    :param monitor:
        monitor of the currently running job
    :returns:
        a dictionary trt_model_id -> (curves_by_gsim, bounding_boxes)
    """
    hc = models.oqparam(monitor.job_id)
    trt_model_id = sources[0].trt_model_id
    total_sites = len(sitecol)
    sitemesh = sitecol.mesh
    sorted_imts = sorted(hc.imtls)
    sorted_imls = [hc.imtls[imt]
                   for imt in sorted_imts]
    sorted_imts = map(from_string, sorted_imts)
    trt_model = models.TrtModel.objects.get(pk=trt_model_id)

    gsims = gsims_by_trt_id[trt_model_id]
    curves = [[numpy.ones([total_sites, len(ls)]) for ls in sorted_imls]
              for gsim in gsims]
    if hc.poes_disagg:  # doing disaggregation
        lt_model_id = trt_model.lt_model.id
        bbs = [BoundingBox(lt_model_id, site_id) for site_id in sitecol.sids]
    else:
        bbs = []
    mon = monitor('getting ruptures', measuremem=False)
    make_ctxt_mon = monitor('making contexts', measuremem=False)
    calc_poes_mon = monitor('computing poes', measuremem=False)

    num_sites = 0
    # NB: rows are namedtuples with fields (source, rupture, rupture_sites)
    for source, rows in itertools.groupby(
            gen_ruptures(sources, sitecol, hc.maximum_distance, mon),
            key=operator.attrgetter('source')):
        t0 = time.time()
        num_ruptures = 0
        for _source, rupture, r_sites in rows:
            num_sites = max(num_sites, len(r_sites))
            num_ruptures += 1
            if hc.poes_disagg:  # doing disaggregation
                jb_dists = rupture.surface.get_joyner_boore_distance(sitemesh)
                closest_points = rupture.surface.get_closest_points(sitemesh)
                for bb, dist, point in itertools.izip(
                        bbs, jb_dists, closest_points):
                    if dist < hc.maximum_distance:
                        # ruptures too far away are ignored
                        bb.update([dist], [point.longitude], [point.latitude])

            # compute probabilities for all realizations
            for gsim, curv in itertools.izip(gsims, curves):
                for i, pnes in enumerate(_calc_pnes(
                        gsim, r_sites, rupture, sorted_imts, sorted_imls,
                        hc.truncation_level, make_ctxt_mon, calc_poes_mon)):
                    curv[i] *= pnes

        inserter.add(
            models.SourceInfo(trt_model_id=trt_model_id,
                              source_id=source.source_id,
                              source_class=source.__class__.__name__,
                              num_sites=num_sites,
                              num_ruptures=num_ruptures,
                              calc_time=time.time() - t0))

    make_ctxt_mon.flush()
    calc_poes_mon.flush()
    inserter.flush()

    # the 0 here is a shortcut for filtered sources giving no contribution;
    # this is essential for performance, we want to avoid returning
    # big arrays of zeros (MS)
    curves_by_gsim = [
        (gsim.__class__.__name__,
         [0 if general.all_equal(c, 1) else 1. - c for c in curv])
        for gsim, curv in zip(gsims, curves)]
    return {trt_model_id: (curves_by_gsim, bbs)}

@calculators.add('classical')
class ClassicalHazardCalculator(general.BaseHazardCalculator):
    """
    Classical PSHA hazard calculator. Computes hazard curves for a given set of
    points.

    For each realization of the calculation, we randomly sample source models
    and GMPEs (Ground Motion Prediction Equations) from logic trees.
    """

    core_calc_task = compute_hazard_curves

    def pre_execute(self):
        """
        Do pre-execution work. At the moment, this work entails:
        parsing and initializing sources, parsing and initializing the
        site model (if there is one), parsing vulnerability and
        exposure files and generating logic tree realizations. (The
        latter piece basically defines the work to be done in the
        `execute` phase.).
        """
        weights = super(ClassicalHazardCalculator, self).pre_execute()
        # a dictionary with the bounding boxes for earch source
        # model and each site, defined only for disaggregation
        # calculations:
        if self.oqparam.poes_disagg:
            lt_models = models.LtSourceModel.objects.filter(
                hazard_calculation=self.job)
            self.bb_dict = dict(
                ((lt_model.id, site.id), BoundingBox(lt_model.id, site.id))
                for site in self.site_collection
                for lt_model in lt_models)
        return weights
<commit_msg>Removed the BoundingBox<commit_after># Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core functionality for the classical PSHA hazard calculator.

The difficult part of this calculator is the management of the logic tree
realizations. Let me explain how it works in a real life case.

We want to perform a SHARE calculation with a single source
model with two tectonic region types (Stable Shallow Crust and
Active Shallow Crust) and a complex GMPEs logic tree with 7
branching points with the following structure:

Active_Shallow: 4 GMPEs, weights 0.35, 0.35, 0.20, 0.10
Stable_Shallow: 5 GMPEs, weights 0.2, 0.2, 0.2, 0.2, 0.2
Shield: 2 GMPEs, weights 0.5, 0.5
Subduction_Interface: 4 GMPEs, weights 0.2, 0.2, 0.2, 0.4
Subduction_InSlab: 4 GMPEs, weights 0.2, 0.2, 0.2, 0.4
Volcanic: 1 GMPE, weight 1
Deep: 2 GMPEs, weights 0.6, 0.4

The number of realizations generated by this logic tree is
4 * 5 * 2 * 4 * 4 * 1 * 2 = 1280 and at the end we will
generate 1280 hazard curve containers. However the independent
hazard curve containers are much less, there are actually only 9 of
them, 4 coming from the Active Shallow Crust tectonic region type
and 5 from the Stable Shallow Crust tectonic region type.
The dependent hazard curves (i.e. hazard curves by realization)
can be obtained from the independent hazard curves by a composition
rule implemented in :method:`openquake.engine.models.HazardCurve.build_data`

NB: notice that 1280 / 9 = 142.22, therefore storing all the hazard
curves takes 140+ times more disk space and resources than actually needed.
The plan for the future is store only the independent curves and to
compose them on-the-fly, when they are looked up for a given realization,
since the composition is pretty fast (there are just numpy multiplications),
faster than reading from the database all the redundant data. The
impact on hazard maps has to be determined.
"""
import time
import operator
import itertools

import numpy

from openquake.hazardlib.imt import from_string
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.hazardlib.geo.utils import get_spherical_bounding_box
from openquake.hazardlib.geo.utils import get_longitudinal_extent
from openquake.hazardlib.geo.geodetic import npoints_between

from openquake.calculators.calc import gen_ruptures

from openquake.engine import writer
from openquake.engine.calculators import calculators
from openquake.engine.calculators.hazard import general
from openquake.engine.db import models
from openquake.engine.utils import tasks

inserter = writer.CacheInserter(models.SourceInfo, 100)

def _calc_pnes(gsim, r_sites, rupture, imts, imls, truncation_level,
               make_ctxt_mon, calc_poes_mon):
    # compute the probabilities of no exceedence for each IMT
    # for the given gsim and rupture; returns a list of pairs
    # [(imt, pnes), ...]
    with make_ctxt_mon:
        sctx, rctx, dctx = ContextMaker([gsim]).make_contexts(r_sites, rupture)
    with calc_poes_mon:
        for imt, levels in itertools.izip(imts, imls):
            poes = gsim.get_poes(
                sctx, rctx, dctx, imt, levels, truncation_level)
            pnes = rupture.get_probability_no_exceedance(poes)
            yield r_sites.expand(pnes, placeholder=1)

@tasks.oqtask
def compute_hazard_curves(sources, sitecol, gsims_by_trt_id, monitor):
    """
    This task computes R2 * I hazard curves (each one is a
    numpy array of S * L floats) from the given source_ruptures
    pairs.

    :param sources:
        a block of source objects
    :param sitecol:
        a :class:`openquake.hazardlib.site.SiteCollection` instance
    :param gsims_by_trt_id:
        a dictionary trt_id -> gsim instances
    :param monitor:
        monitor of the currently running job
    :returns:
        a dictionary trt_model_id -> (curves_by_gsim, bounding_boxes)
    """
    hc = models.oqparam(monitor.job_id)
    trt_model_id = sources[0].trt_model_id
    total_sites = len(sitecol)
    sitemesh = sitecol.mesh
    sorted_imts = sorted(hc.imtls)
    sorted_imls = [hc.imtls[imt]
                   for imt in sorted_imts]
    sorted_imts = map(from_string, sorted_imts)

    gsims = gsims_by_trt_id[trt_model_id]
    curves = [[numpy.ones([total_sites, len(ls)]) for ls in sorted_imls]
              for gsim in gsims]
    bbs = []
    mon = monitor('getting ruptures', measuremem=False)
    make_ctxt_mon = monitor('making contexts', measuremem=False)
    calc_poes_mon = monitor('computing poes', measuremem=False)

    num_sites = 0
    # NB: rows are namedtuples with fields (source, rupture, rupture_sites)
    for source, rows in itertools.groupby(
            gen_ruptures(sources, sitecol, hc.maximum_distance, mon),
            key=operator.attrgetter('source')):
        t0 = time.time()
        num_ruptures = 0
        for _source, rupture, r_sites in rows:
            num_sites = max(num_sites, len(r_sites))
            num_ruptures += 1
            if hc.poes_disagg:  # doing disaggregation
                jb_dists = rupture.surface.get_joyner_boore_distance(sitemesh)
                closest_points = rupture.surface.get_closest_points(sitemesh)
                for bb, dist, point in itertools.izip(
                        bbs, jb_dists, closest_points):
                    if dist < hc.maximum_distance:
                        # ruptures too far away are ignored
                        bb.update([dist], [point.longitude], [point.latitude])

            # compute probabilities for all realizations
            for gsim, curv in itertools.izip(gsims, curves):
                for i, pnes in enumerate(_calc_pnes(
                        gsim, r_sites, rupture, sorted_imts, sorted_imls,
                        hc.truncation_level, make_ctxt_mon, calc_poes_mon)):
                    curv[i] *= pnes

        inserter.add(
            models.SourceInfo(trt_model_id=trt_model_id,
                              source_id=source.source_id,
                              source_class=source.__class__.__name__,
                              num_sites=num_sites,
                              num_ruptures=num_ruptures,
                              calc_time=time.time() - t0))

    make_ctxt_mon.flush()
    calc_poes_mon.flush()
    inserter.flush()

    # the 0 here is a shortcut for filtered sources giving no contribution;
    # this is essential for performance, we want to avoid returning
    # big arrays of zeros (MS)
    curves_by_gsim = [
        (gsim.__class__.__name__,
         [0 if general.all_equal(c, 1) else 1. - c for c in curv])
        for gsim, curv in zip(gsims, curves)]
    return {trt_model_id: (curves_by_gsim, bbs)}

@calculators.add('classical')
class ClassicalHazardCalculator(general.BaseHazardCalculator):
    """
    Classical PSHA hazard calculator. Computes hazard curves for a given set of
    points.

    For each realization of the calculation, we randomly sample source models
    and GMPEs (Ground Motion Prediction Equations) from logic trees.
    """

    core_calc_task = compute_hazard_curves
<|endoftext|>


<commit_before>from os.path import expanduser
from os.path import basename

import os
import sys
import requests
import ipfshttpclient as ipfsapi

import re
# Datetime conversion to rfc1123
import locale
import datetime
import logging
import platform

from six.moves.urllib.request import urlopen
import json
from .__init__ import __version__ as ipwbVersion

from pkg_resources import parse_version

# from requests.exceptions import ConnectionError
from ipfshttpclient.exceptions import ConnectionError
from ipfshttpclient.exceptions import AddressError
from multiaddr.exceptions import StringParseError

IPFSAPI_MUTLIADDRESS = '/dns/localhost/tcp/5001/http'
# or '/dns/{host}/tcp/{port}/http'
# or '/ip4/{ipaddress}/tcp/{port}/http'
# or '/ip6/{ipaddress}/tcp/{port}/http

IPWBREPLAY_ADDRESS = 'localhost:5000'

(IPWBREPLAY_HOST, IPWBREPLAY_PORT) = IPWBREPLAY_ADDRESS.split(':')
IPWBREPLAY_PORT = int(IPWBREPLAY_PORT)

INDEX_FILE = os.path.sep.join(['samples', 'indexes', 'salam-home.cdxj'])

log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

dtPattern = re.compile(r"^(\d{4})(\d{2})?(\d{2})?(\d{2})?(\d{2})?(\d{2})?$")

def createIPFSClient(daemonMultiaddr=IPFSAPI_MUTLIADDRESS):
    try:
        return ipfsapi.Client(daemonMultiaddr)
    except (StringParseError, AddressError):
        return None  # Malformed multiaddress for the daemon

def isDaemonAlive(daemonMultiaddr=IPFSAPI_MUTLIADDRESS):
    """Ensure that the IPFS daemon is running via HTTP before proceeding"""
    client = createIPFSClient()
    if client is None:
        print("Error initializing IPFS API client")
        return False

    try:
        # ConnectionError/AttributeError if IPFS daemon not running
        client.id()
        return True
    except (ConnectionError):  # exceptions.AttributeError):
        logError("Daemon is not running at " + daemonMultiaddr)
        return False
    except OSError:
        logError("IPFS is likely not installed. "
                 "See https://ipfs.io/docs/install/")
        sys.exit()
    except Exception as e:
        logError('Unknown error in retrieving daemon status')
        logError(sys.exc_info()[0])

def logError(errIn):
    print(errIn, file=sys.stderr)

def isValidCDXJ(stringIn):  # TODO: Check specific strict syntax
    # Also, be sure to mind the meta headers starting with @/#, etc.
    return True

def isValidCDXJLine(cdxjLine):
    try:
        (surtURI, datetime, jsonData) = cdxjLine.split(' ', 2)

        json.loads(jsonData)
        validDatetime = len(datetime) == 14

        validSURT = True  # TODO: check valid SURT URI

        return validSURT and validDatetime
    except ValueError:  # Not valid JSON
        return False
    except NameError:
        return isCDXJMetadataRecord(cdxjLine)
    except Exception as e:
        return False

# Compare versions of software, <0 if a<b, 0 if ==, >1 if b>a
def compareVersions(versionA, versionB):
    return parse_version(versionA) < parse_version(versionB)

def isCDXJMetadataRecord(cdxjLine):
    return cdxjLine.strip()[:1] == '!'

def isLocalHosty(uri):
    # TODO: check for these SW conditions
    # (*, localhost, *); (*, 127/8, *); (*, ::1/128, *)
    localhosts = ['localhost', '127.0.0.1']
    for lh in localhosts:
        if lh in uri:
            return True
    return False

def setLocale():
    currentOS = platform.system()

    if currentOS == 'Darwin':
        newLocale = 'en_US'
    elif currentOS == 'Windows':
        newLocale = 'english'
    else:  # Assume Linux
        newLocale = 'en_US.utf8'

    try:
        locale.setlocale(locale.LC_TIME, newLocale)
    except locale.Error:
        locale.setlocale(locale.LC_TIME, '')

def digits14ToRFC1123(digits14):
    setLocale()
    d = datetime.datetime.strptime(digits14, '%Y%m%d%H%M%S')
    return d.strftime('%a, %d %b %Y %H:%M:%S GMT')

def rfc1123ToDigits14(rfc1123DateString):
    setLocale()
    d = datetime.datetime.strptime(rfc1123DateString,
                                   '%a, %d %b %Y %H:%M:%S %Z')

    # TODO: Account for conversion if TZ other than GMT not specified

    return d.strftime('%Y%m%d%H%M%S')

def iso8601ToDigits14(iso8601DateString):
    setLocale()
    d = datetime.datetime.strptime(iso8601DateString,
                                   "%Y-%m-%dT%H:%M:%SZ")

    # TODO: Account for conversion if TZ other than GMT not specified

    return d.strftime('%Y%m%d%H%M%S')

def isRFC1123Compliant(dtstr):
    try:
        datetime.datetime.strptime(dtstr, '%a, %d %b %Y %H:%M:%S GMT')
        return True
    except ValueError as err:
        return False

def getRFC1123OfNow():
    setLocale()
    d = datetime.datetime.now()
    return d.strftime('%a, %d %b %Y %H:%M:%S GMT')

def padDigits14(dtstr, validate=False):
    '''Pad datetime to make a 14-digit string and optionally validate it'''
    match = dtPattern.match(dtstr)
    if match:
        Y = match.group(1)
        m = match.group(2) or '01'
        d = match.group(3) or '01'
        H = match.group(4) or '00'
        M = match.group(5) or '00'
        S = match.group(6) or '00'
        dtstr = f'{Y}{m}{d}{H}{M}{S}'
    if validate:
        datetime.datetime.strptime(dtstr, '%Y%m%d%H%M%S')
    return dtstr

def fetch_remote_file(path):
    try:
        r = requests.get(path)
        return r.text
    except ConnectionError:
        logError(f'File at {path} is unavailable.')
    except Exception as E:
        logError(f'An unknown error occurred while trying to fetch {path}')
        logError(sys.exc_info()[0])
    return None

# IPFS Config manipulation from here on out.
def readIPFSConfig():
    ipfsConfigPath = os.path.sep.join([expanduser("~"), '.ipfs', 'config'])
    if 'IPFS_PATH' in os.environ:
        ipfsConfigPath = os.path.sep.join(
            [os.environ.get('IPFS_PATH'), 'config'])

    try:
        with open(ipfsConfigPath, 'r') as f:
            return json.load(f)
    except IOError:
        logError("IPFS config not found.")
        logError("Have you installed ipfs and run ipfs init?")
        sys.exit()

def writeIPFSConfig(jsonToWrite):
    ipfsConfigPath = os.path.sep.join([expanduser("~"), '.ipfs', 'config'])
    if 'IPFS_PATH' in os.environ:
        ipfsConfigPath = os.path.sep.join(
            [os.environ.get('IPFS_PATH'), 'config'])

    with open(ipfsConfigPath, 'w') as f:
        f.write(json.dumps(jsonToWrite, indent=4, sort_keys=True))

def getIPFSAPIHostAndPort(ipfsJSON=None):
    if not ipfsJSON:
        ipfsJSON = readIPFSConfig()

    (scheme, host, protocol, port) = (
        ipfsJSON['Addresses']['API'][1:].split('/')
    )
    return host + ':' + port

def getIPWBReplayConfig(ipfsJSON=None):
    if not ipfsJSON:
        ipfsJSON = readIPFSConfig()
    port = None
    if ('Ipwb' in ipfsJSON and 'Replay' in ipfsJSON['Ipwb'] and
       'Port' in ipfsJSON['Ipwb']['Replay']):
        host = ipfsJSON['Ipwb']['Replay']['Host']
        port = ipfsJSON['Ipwb']['Replay']['Port']
        return (host, port)
    else:
        return None

def setIPWBReplayConfig(Host, Port, ipfsJSON=None):
    if not ipfsJSON:
        ipfsJSON = readIPFSConfig()
    ipfsJSON['Ipwb'] = {}
    ipfsJSON['Ipwb']['Replay'] = {
      u'Host': Host,
      u'Port': Port
    }
    writeIPFSConfig(ipfsJSON)

def setIPWBReplayIndexPath(cdxj):
    if cdxj is None:
        cdxj = INDEX_FILE
    ipfsJSON = readIPFSConfig()
    ipfsJSON['Ipwb']['Replay']['Index'] = cdxj
    writeIPFSConfig(ipfsJSON)
    return

def getIPWBReplayIndexPath():
    ipfsJSON = readIPFSConfig()
    if 'Ipwb' not in ipfsJSON:
        setIPWBReplayConfig(IPWBREPLAY_HOST, IPWBREPLAY_PORT)
        ipfsJSON = readIPFSConfig()

    if 'Index' in ipfsJSON['Ipwb']['Replay']:
        return ipfsJSON['Ipwb']['Replay']['Index']
    else:
        return ''

# From pywb 2.0.4
def unsurt(surt):
    try:
        index = surt.index(')/')
        parts = surt[0:index].split(',')
        parts.reverse()
        host = '.'.join(parts)
        host += surt[index+1:]
        return host

    except ValueError:
        # May not be a valid surt
        return surt

def compareCurrentAndLatestIPWBVersions():
    try:
        resp = urlopen('https://pypi.python.org/pypi/ipwb/json')
        jResp = json.loads(resp.read())
        latestVersion = jResp['info']['version']
        currentVersion = re.sub(r'\.0+', '.', ipwbVersion)
        return (currentVersion, latestVersion)
    except Exception as e:
        return (None, None)

def checkForUpdate(_):
    (current, latest) = compareCurrentAndLatestIPWBVersions()

    if current != latest and current is not None:
        print('This version of ipwb is outdated.'
              ' Please run pip install --upgrade ipwb.')
    print(f'* Latest version: {latest}')
    print(f'* Installed version: {current}')
<commit_msg>Update ipwb/util.py<commit_after>from os.path import expanduser
from os.path import basename

import os
import sys
import requests
import ipfshttpclient as ipfsapi

import re
# Datetime conversion to rfc1123
import locale
import datetime
import logging
import platform

from six.moves.urllib.request import urlopen
import json
from .__init__ import __version__ as ipwbVersion

from pkg_resources import parse_version

# from requests.exceptions import ConnectionError
from ipfshttpclient.exceptions import ConnectionError
from ipfshttpclient.exceptions import AddressError
from multiaddr.exceptions import StringParseError

IPFSAPI_MUTLIADDRESS = '/dns/localhost/tcp/5001/http'
# or '/dns/{host}/tcp/{port}/http'
# or '/ip4/{ipaddress}/tcp/{port}/http'
# or '/ip6/{ipaddress}/tcp/{port}/http

IPWBREPLAY_ADDRESS = 'localhost:5000'

(IPWBREPLAY_HOST, IPWBREPLAY_PORT) = IPWBREPLAY_ADDRESS.split(':')
IPWBREPLAY_PORT = int(IPWBREPLAY_PORT)

INDEX_FILE = os.path.sep.join(['samples', 'indexes', 'salam-home.cdxj'])

log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

dtPattern = re.compile(r"^(\d{4})(\d{2})?(\d{2})?(\d{2})?(\d{2})?(\d{2})?$")

def createIPFSClient(daemonMultiaddr=IPFSAPI_MUTLIADDRESS):
    try:
        return ipfsapi.Client(daemonMultiaddr)
    except (StringParseError, AddressError):
        return None  # Malformed multiaddress for the daemon

def isDaemonAlive(daemonMultiaddr=IPFSAPI_MUTLIADDRESS):
    """Ensure that the IPFS daemon is running via HTTP before proceeding"""
    client = createIPFSClient()
    if client is None:
        print("Error initializing IPFS API client")
        return False

    try:
        # ConnectionError/AttributeError if IPFS daemon not running
        client.id()
        return True
    except (ConnectionError):  # exceptions.AttributeError):
        logError("Daemon is not running at " + daemonMultiaddr)
        return False
    except OSError:
        logError("IPFS is likely not installed. "
                 "See https://ipfs.io/docs/install/")
        sys.exit()
    except Exception as e:
        logError('Unknown error in retrieving daemon status')
        logError(sys.exc_info()[0])

def logError(errIn):
    print(errIn, file=sys.stderr)

def isValidCDXJ(stringIn):  # TODO: Check specific strict syntax
    # Also, be sure to mind the meta headers starting with @/#, etc.
    return True

def isValidCDXJLine(cdxjLine):
    try:
        (surtURI, datetime, jsonData) = cdxjLine.split(' ', 2)

        json.loads(jsonData)
        validDatetime = len(datetime) == 14

        validSURT = True  # TODO: check valid SURT URI

        return validSURT and validDatetime
    except ValueError:  # Not valid JSON
        return False
    except NameError:
        return isCDXJMetadataRecord(cdxjLine)
    except Exception as e:
        return False

# Compare versions of software, <0 if a<b, 0 if ==, >1 if b>a
def compareVersions(versionA, versionB):
    return parse_version(versionA) < parse_version(versionB)

def isCDXJMetadataRecord(cdxjLine):
    return cdxjLine.strip()[:1] == '!'

def isLocalHosty(uri):
    # TODO: check for these SW conditions
    # (*, localhost, *); (*, 127/8, *); (*, ::1/128, *)
    localhosts = ['localhost', '127.0.0.1']
    for lh in localhosts:
        if lh in uri:
            return True
    return False

def setLocale():
    currentOS = platform.system()

    if currentOS == 'Darwin':
        newLocale = 'en_US'
    elif currentOS == 'Windows':
        newLocale = 'english'
    else:  # Assume Linux
        newLocale = 'en_US.utf8'

    try:
        locale.setlocale(locale.LC_TIME, newLocale)
    except locale.Error:
        locale.setlocale(locale.LC_TIME, '')

def digits14ToRFC1123(digits14):
    setLocale()
    d = datetime.datetime.strptime(digits14, '%Y%m%d%H%M%S')
    return d.strftime('%a, %d %b %Y %H:%M:%S GMT')

def rfc1123ToDigits14(rfc1123DateString):
    setLocale()
    d = datetime.datetime.strptime(rfc1123DateString,
                                   '%a, %d %b %Y %H:%M:%S %Z')

    # TODO: Account for conversion if TZ other than GMT not specified

    return d.strftime('%Y%m%d%H%M%S')

def iso8601ToDigits14(iso8601DateString):
    setLocale()
    d = datetime.datetime.strptime(iso8601DateString,
                                   "%Y-%m-%dT%H:%M:%SZ")

    # TODO: Account for conversion if TZ other than GMT not specified

    return d.strftime('%Y%m%d%H%M%S')

def isRFC1123Compliant(dtstr):
    try:
        datetime.datetime.strptime(dtstr, '%a, %d %b %Y %H:%M:%S GMT')
        return True
    except ValueError as err:
        return False

def getRFC1123OfNow():
    setLocale()
    d = datetime.datetime.now()
    return d.strftime('%a, %d %b %Y %H:%M:%S GMT')

def padDigits14(dtstr, validate=False):
    '''Pad datetime to make a 14-digit string and optionally validate it'''
    match = dtPattern.match(dtstr)
    if match:
        Y = match.group(1)
        m = match.group(2) or '01'
        d = match.group(3) or '01'
        H = match.group(4) or '00'
        M = match.group(5) or '00'
        S = match.group(6) or '00'
        dtstr = f'{Y}{m}{d}{H}{M}{S}'
    if validate:
        datetime.datetime.strptime(dtstr, '%Y%m%d%H%M%S')
    return dtstr

def fetch_remote_file(path):
    try:
        r = requests.get(path)
        return r.text
    except ConnectionError:
        logError(f'File at {path} is unavailable.')
    except Exception as E:
        logError(f'An unknown error occurred while trying to fetch {path}')
        logError(sys.exc_info()[0])
    return None

# IPFS Config manipulation from here on out.
def readIPFSConfig():
    ipfsConfigPath = os.path.sep.join([expanduser("~"), '.ipfs', 'config'])
    if 'IPFS_PATH' in os.environ:
        ipfsConfigPath = os.path.sep.join(
            [os.environ.get('IPFS_PATH'), 'config'])

    try:
        with open(ipfsConfigPath, 'r') as f:
            return json.load(f)
    except IOError:
        logError("IPFS config not found.")
        logError("Have you installed ipfs and run ipfs init?")
        sys.exit()

def writeIPFSConfig(jsonToWrite):
    ipfsConfigPath = os.path.sep.join([expanduser("~"), '.ipfs', 'config'])
    if 'IPFS_PATH' in os.environ:
        ipfsConfigPath = os.path.join(
            os.environ.get('IPFS_PATH'), 'config')

    with open(ipfsConfigPath, 'w') as f:
        f.write(json.dumps(jsonToWrite, indent=4, sort_keys=True))

def getIPFSAPIHostAndPort(ipfsJSON=None):
    if not ipfsJSON:
        ipfsJSON = readIPFSConfig()

    (scheme, host, protocol, port) = (
        ipfsJSON['Addresses']['API'][1:].split('/')
    )
    return host + ':' + port

def getIPWBReplayConfig(ipfsJSON=None):
    if not ipfsJSON:
        ipfsJSON = readIPFSConfig()
    port = None
    if ('Ipwb' in ipfsJSON and 'Replay' in ipfsJSON['Ipwb'] and
       'Port' in ipfsJSON['Ipwb']['Replay']):
        host = ipfsJSON['Ipwb']['Replay']['Host']
        port = ipfsJSON['Ipwb']['Replay']['Port']
        return (host, port)
    else:
        return None

def setIPWBReplayConfig(Host, Port, ipfsJSON=None):
    if not ipfsJSON:
        ipfsJSON = readIPFSConfig()
    ipfsJSON['Ipwb'] = {}
    ipfsJSON['Ipwb']['Replay'] = {
      u'Host': Host,
      u'Port': Port
    }
    writeIPFSConfig(ipfsJSON)

def setIPWBReplayIndexPath(cdxj):
    if cdxj is None:
        cdxj = INDEX_FILE
    ipfsJSON = readIPFSConfig()
    ipfsJSON['Ipwb']['Replay']['Index'] = cdxj
    writeIPFSConfig(ipfsJSON)
    return

def getIPWBReplayIndexPath():
    ipfsJSON = readIPFSConfig()
    if 'Ipwb' not in ipfsJSON:
        setIPWBReplayConfig(IPWBREPLAY_HOST, IPWBREPLAY_PORT)
        ipfsJSON = readIPFSConfig()

    if 'Index' in ipfsJSON['Ipwb']['Replay']:
        return ipfsJSON['Ipwb']['Replay']['Index']
    else:
        return ''

# From pywb 2.0.4
def unsurt(surt):
    try:
        index = surt.index(')/')
        parts = surt[0:index].split(',')
        parts.reverse()
        host = '.'.join(parts)
        host += surt[index+1:]
        return host

    except ValueError:
        # May not be a valid surt
        return surt

def compareCurrentAndLatestIPWBVersions():
    try:
        resp = urlopen('https://pypi.python.org/pypi/ipwb/json')
        jResp = json.loads(resp.read())
        latestVersion = jResp['info']['version']
        currentVersion = re.sub(r'\.0+', '.', ipwbVersion)
        return (currentVersion, latestVersion)
    except Exception as e:
        return (None, None)

def checkForUpdate(_):
    (current, latest) = compareCurrentAndLatestIPWBVersions()

    if current != latest and current is not None:
        print('This version of ipwb is outdated.'
              ' Please run pip install --upgrade ipwb.')
    print(f'* Latest version: {latest}')
    print(f'* Installed version: {current}')
<|endoftext|>


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

##############################################################################
#    Copyright (C) 2012 SICLIC http://siclic.fr
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>
#
#############################################################################

from osv import osv, fields
from datetime import datetime
import re
import unicodedata
import netsvc
import decimal_precision as dp
from tools.translate import _

class crossovered_budget(osv.osv):
    _inherit = "crossovered.budget"
    _name = "crossovered.budget"

    def _calc_amounts(self, cr, uid, ids,  field_name, args, context=None):
        res = {}
        for budget in self.browse(cr, uid, ids, context):
            pract = 0.0
            theo = 0.0
            planned = 0.0
            openstc_practical = 0.0
            #Pour chaque budget, on ajoute les montants de toutes leurs lignes budgétaires
            for line in budget.crossovered_budget_line:
                pract += line.practical_amount
                theo += line.theoritical_amount
                planned += line.planned_amount
                openstc_practical += line.openstc_practical_amount
            res.update({budget.id:{'pract_amount':pract, 'theo_amount':theo, 'planned_amount':planned, 'openstc_practical_amount':openstc_practical}})
        return res

    _columns = {
        'planned_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Plannifié", type='float'),
        'pract_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Pratique", type='float'),
        'theo_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Théorique", type='float'),
        'openstc_practical_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Consommé", type='float'),
        'code_budget_ciril':fields.char('CIRIL Budget Code', size=16),
        'service_id':fields.many2one('openstc.service','Service',required=True),
        }

    _defaults = {
        'code_budget_ciril':lambda *a: '0',
        }
crossovered_budget()

class crossovered_budget_lines(osv.osv):
    _rec_name = "name"

    def name_get(self, cr, uid, ids, context=None):
        ret = []
        for budget in self.browse(cr, uid, ids, context=context):
            budget.name
            val = _(u'%s / %s : %s € (%.2f %% consummed)') %(budget.crossovered_budget_id.name, budget.analytic_account_id.name_get()[0][1],budget.planned_amount - budget.openstc_practical_amount, budget.openstc_erosion)
            ret.append((budget.id,val))
        return ret

    def name_search(self, cr, uid, name='', args=[], operator='ilike', context={}, limit=80):
        #ids = self.search(cr, uid, [('analytic_account_id.complete_name',operator,name)] + args, limit=limit, context=context)
        ids = self.search(cr, uid, ['|',('analytic_account_id.name',operator,name),('crossovered_budget_id.name',operator,name)] + args, limit=limit, context=context)
#        if not ids:
#            ids = self.search(cr, uid, [('analytic_account_id.service_id.code',operator,name)] + args, limit=limit, context=context)
        return self.name_get(cr, uid, ids, context=context)

    #custom field for public account : returns amount with taxes included
    def _openstc_pract(self, cr, uid, ids, name, args, context=None):
        #first, we get engage_lines that matches dates and current budget line analytic account
        ret = {}
        for line in self.browse(cr, uid, ids, context):
            ret.setdefault(line.id,{})
            engage_line_ids = self.pool.get("open.engagement.line").search(cr, uid, [('budget_line_id','=',line.id),
                                                                                ('engage_id.date_engage_validated','<=', line.date_to),
                                                                                ('engage_id.date_engage_validated','>=',line.date_from),
                                                                                ('engage_id.state','not in',('draft','to_validate','check_except'))])
            amount = 0.0
            #next, if we found engage_lines, we compute with them,
            #else, we use default openerp practical amount (for example to display sales amount which doesn't work with engages)
            if engage_line_ids:
                for engage_line in self.pool.get("open.engagement.line").browse(cr, uid, engage_line_ids, context):
                    amount += engage_line.amount
            else:
                amount = line.practical_amount
            ret[line.id]['openstc_practical_amount'] = amount
            if line.planned_amount:
                ret[line.id]['openstc_erosion'] = amount * 100.0 / line.planned_amount
            else:
                ret[line.id]['openstc_erosion'] = 0.0

        return ret

    _inherit = "crossovered.budget.lines"
    _columns = {
            'openstc_practical_amount':fields.function(_openstc_pract, multi="openstc_pract_amount", method=True, string="Montant Consommé", type="float", digits_compute=dp.get_precision('Account')),
            'openstc_erosion':fields.function(_openstc_pract, multi="openstc_pract_amount", method=True, string="Taux d'érosion (%)", type="float", digits_compute=dp.get_precision('Account')),
            'openstc_general_account':fields.many2one('account.account', 'M14 account', help="M14 account corresponding to this budget line"),
            'openstc_code_antenne':fields.char('Antenne Code', size=16, help='Antenne code from CIRIL instance'),
            'name':fields.related('analytic_account_id','complete_name',string='Budget name',type='char',store=True),
        }

    def onchange_openstc_general_account(self, cr, uid, ids, openstc_general_account=False):
        if openstc_general_account:
            #we create an account.budget.post to respect base work of budget, even if we don't use it anymore
            account = self.pool.get("account.account").browse(cr, uid, openstc_general_account)
            post = self.pool.get("account.budget.post").search(cr, uid, [('account_ids','=',openstc_general_account)])
            if post:
                post = post[0]
            else:
                post = self.pool.get("account.budget.post").create(cr, uid, {'code':account.code,
                                                                    'name':account.name,
                                                                    'account_ids':[(6,0,[account.id])]})

        return {'value':{'general_budget_id':post}}

crossovered_budget_lines()

    <commit_msg>updated calculation of consummed_amount of budget_lines (faster and avoid conditions on purchases)<commit_after># -*- coding: utf-8 -*-

##############################################################################
#    Copyright (C) 2012 SICLIC http://siclic.fr
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>
#
#############################################################################

from osv import osv, fields
from datetime import datetime
import re
import unicodedata
import netsvc
import decimal_precision as dp
from tools.translate import _

class crossovered_budget(osv.osv):
    _inherit = "crossovered.budget"
    _name = "crossovered.budget"

    def _calc_amounts(self, cr, uid, ids,  field_name, args, context=None):
        res = {}
        for budget in self.browse(cr, uid, ids, context):
            pract = 0.0
            theo = 0.0
            planned = 0.0
            openstc_practical = 0.0
            #Pour chaque budget, on ajoute les montants de toutes leurs lignes budgétaires
            for line in budget.crossovered_budget_line:
                pract += line.practical_amount
                theo += line.theoritical_amount
                planned += line.planned_amount
                openstc_practical += line.openstc_practical_amount
            res.update({budget.id:{'pract_amount':pract, 'theo_amount':theo, 'planned_amount':planned, 'openstc_practical_amount':openstc_practical}})
        return res

    _columns = {
        'planned_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Plannifié", type='float'),
        'pract_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Pratique", type='float'),
        'theo_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Théorique", type='float'),
        'openstc_practical_amount':fields.function(_calc_amounts, method=True, multi = True, string="Montant Consommé", type='float'),
        'code_budget_ciril':fields.char('CIRIL Budget Code', size=16),
        'service_id':fields.many2one('openstc.service','Service',required=True),
        }

    _defaults = {
        'code_budget_ciril':lambda *a: '0',
        }
crossovered_budget()

class crossovered_budget_lines(osv.osv):
    _rec_name = "name"

    def name_get(self, cr, uid, ids, context=None):
        ret = []
        for budget in self.browse(cr, uid, ids, context=context):
            budget.name
            val = _(u'%s / %s : %s € (%.2f %% consummed)') %(budget.crossovered_budget_id.name, budget.analytic_account_id.name_get()[0][1],budget.planned_amount - budget.openstc_practical_amount, budget.openstc_erosion)
            ret.append((budget.id,val))
        return ret

    def name_search(self, cr, uid, name='', args=[], operator='ilike', context={}, limit=80):
        #ids = self.search(cr, uid, [('analytic_account_id.complete_name',operator,name)] + args, limit=limit, context=context)
        ids = self.search(cr, uid, ['|',('analytic_account_id.name',operator,name),('crossovered_budget_id.name',operator,name)] + args, limit=limit, context=context)
#        if not ids:
#            ids = self.search(cr, uid, [('analytic_account_id.service_id.code',operator,name)] + args, limit=limit, context=context)
        return self.name_get(cr, uid, ids, context=context)

    #custom field for public account : returns amount with taxes included
    def _openstc_pract(self, cr, uid, ids, name, args, context=None):
        #first, we get engage_lines that matches dates and current budget line analytic account
        ret = {}.fromkeys(ids,{'openstc_practical_amount':0.0, 'openstc_erosion':0.0})
        engage_obj = self.pool.get("open.engagement.line")

        cr.execute(''' select budget.id as id, budget.planned_amount, sum(line.amount) as amount
        from open_engagement_line as line, crossovered_budget_lines as budget
        where budget.id = line.budget_line_id
        and line.budget_line_id in {budget_line_ids}
        group by budget.planned_amount, budget.id'''.format(budget_line_ids=tuple(ids)))
        amount = 0.0
        data = cr.fetchall()
        for d in data:
            erosion = 0.0
            if d[1]:
                erosion = d[2] * 100.0 / d[1]
            ret[d[0]] = {'openstc_practical_amount':d[2],
                                 'openstc_erosion':erosion}

        return ret

    _inherit = "crossovered.budget.lines"
    _columns = {
            'openstc_practical_amount':fields.function(_openstc_pract, multi="openstc_pract_amount", method=True, string="Montant Consommé", type="float", digits_compute=dp.get_precision('Account')),
            'openstc_erosion':fields.function(_openstc_pract, multi="openstc_pract_amount", method=True, string="Taux d'érosion (%)", type="float", digits_compute=dp.get_precision('Account')),
            'openstc_general_account':fields.many2one('account.account', 'M14 account', help="M14 account corresponding to this budget line"),
            'openstc_code_antenne':fields.char('Antenne Code', size=16, help='Antenne code from CIRIL instance'),
            'name':fields.related('analytic_account_id','complete_name',string='Budget name',type='char',store=True),
        }

    def onchange_openstc_general_account(self, cr, uid, ids, openstc_general_account=False):
        if openstc_general_account:
            #we create an account.budget.post to respect base work of budget, even if we don't use it anymore
            account = self.pool.get("account.account").browse(cr, uid, openstc_general_account)
            post = self.pool.get("account.budget.post").search(cr, uid, [('account_ids','=',openstc_general_account)])
            if post:
                post = post[0]
            else:
                post = self.pool.get("account.budget.post").create(cr, uid, {'code':account.code,
                                                                    'name':account.name,
                                                                    'account_ids':[(6,0,[account.id])]})

        return {'value':{'general_budget_id':post}}

crossovered_budget_lines()

    <|endoftext|>


<commit_before>import csv
from django.db import connection
from django.conf import settings
from django.utils import dateparse
from django.db import IntegrityError, DataError, ProgrammingError
from django.db.models.loading import get_model
from django.core.management.base import LabelCommand
from calaccess_raw.management.commands import CalAccessCommand

class Command(CalAccessCommand, LabelCommand):
    help = 'Load a cleaned CalAccess file for a model into the database'
    args = '<model name>'
    # Trick for reformating date strings in source data so that they can
    # be gobbled up by MySQL. You'll see how below.
    date_sql = "DATE_FORMAT(str_to_date(@`%s`, '%%c/%%e/%%Y'), '%%Y-%%m-%%d')"

    def handle_label(self, label, **options):
        self.verbosity = options.get("verbosity")
        self.load(label)

    def get_hdrs_and_cnt(self, csv_path):
        """
        Get the headers and the line count
        from a specified csv file
        """
        with open(csv_path) as infile:
            csv_reader = csv.reader(infile)
            headers = csv_reader.next()

        with open(csv_path) as infile:
            csv_count = len(infile.readlines()) - 1

        return headers, csv_count

    def _make_date_case(self, _col):
        """
        This method takes in a column name and generates a
        PostgreSQL "case" for correct insertion into the table.
        """
        return """
        ,CASE
            WHEN "%s" IS NOT NULL AND "%s" != ''
                THEN to_date(substring("%s" from 1 for 10), 'MM/DD/YYYY')
            WHEN "%s" = ''
                THEN to_date('01/01/1900', 'MM/DD/YYYY')
        END AS "%s"\n""" % (_col, _col, _col, _col, _col)

    def _make_int_case(self, _col):
        return """
        ,CASE
            WHEN "%s" = ''
                THEN NULL
            WHEN "%s" = 'Y'
                THEN 1
            WHEN "%s" = 'y'
                THEN 1
            WHEN "%s" = 'X'
                THEN 1
            WHEN "%s" = 'N'
                THEN 0
            WHEN "%s" = 'n'
                THEN 0
            WHEN "%s" IS NOT NULL
                THEN "%s"::int
        END AS "%s"\n""" % (_col, _col, _col, _col, _col, _col, _col, _col, _col)

    def _make_numeric_case(self, _col):
        return """
        ,CASE
            WHEN "%s" = ''
                THEN 0.0
            WHEN "%s" IS NULL
                THEN 0.0
            WHEN "%s" IS NOT NULL
                THEN "%s"::numeric
        END AS "%s"\n""" % (_col, _col, _col, _col, _col)

    def _make_timestamp_case(self, _col):
        # "7/9/2014 12:00:00 AM"
        return """
        ,CASE
            WHEN "%s" IS NOT NULL AND "%s" != ''
                THEN to_timestamp("%s", 'MM/DD/YYYY HH12:MI:SS AM')
            WHEN "%s" = ''
                THEN to_timestamp('01/01/1900 1:00:00 AM', 'MM/DD/YYYY HH12:MI:SS AM')
        END AS "%s"\n""" % (_col, _col, _col, _col, _col)

    def _make_special_not_null_case(self, _col):
        return """
        ,CASE
            WHEN "%s" IS NULL
                THEN ''
        END AS "%s"\n""" % (_col, _col)

    def load_postgresql(self, model, csv_path):
        c = connection.cursor()
        try:
            c.execute('DROP TABLE temporary_table;')
        except ProgrammingError:
            pass
        # clear our table
        c.execute('TRUNCATE TABLE "%s"' % model._meta.db_table)

        headers, csv_count = self.get_hdrs_and_cnt(csv_path)

        #map column names to column types
        name_to_type_map = dict([(col.db_column, col.db_type(connection))
                                for col in model._meta.fields])

        # break out special case column types
        int_columns = []
        numeric_columns = []
        date_columns = []
        time_columns = []
        regular_columns = []

        #fill in those column types
        for col in model._meta.fields:
            if col.db_type(connection).startswith('integer'):
                int_columns.append(col.db_column)
            elif col.db_type(connection).startswith('numeric'):
                numeric_columns.append(col.db_column)
            elif col.db_type(connection).startswith('date'):
                date_columns.append(col.db_column)
            elif col.db_type(connection).startswith('timestamp'):
                time_columns.append(col.db_column)
            else:
                if col.db_column is not None and col.db_column in headers:
                    regular_columns.append(col.db_column)

        file_cols_types = []  # column with its types
        for col in headers:
            if col in regular_columns:
                file_cols_types.append("\"" + col + "\"\t" + name_to_type_map[col])
            else:
                file_cols_types.append("\"" + col + "\"\ttext")

        extra_cols = set([col.db_column for col in model._meta.fields]).difference(set(headers))
        empty_cols = []
        for col in extra_cols:
            if col != None:
                empty_cols.append(col)
                # file_cols_types.append("\"" + col + "\"\t" + name_to_type_map[col])

        # create the temp table w/ columns with types
        c.execute("CREATE TABLE \"temporary_table\" (%s);"
                  % ',\n'.join(file_cols_types))

        temp_insert = """COPY "temporary_table"
            FROM '%s'
            CSV
            HEADER;""" % (csv_path)

        try:
            c.execute(temp_insert)  # insert everything into the temp table
        except DataError as e:
            print "initial insert dataerror error, ", e

        for col in empty_cols:
            c.execute("ALTER TABLE temporary_table ADD COLUMN \"%s\" text" % col)
        # build our insert statement
        insert_statement = "INSERT INTO \"%s\" (\"" % model._meta.db_table
        if not regular_columns:
            c.execute("ALTER TABLE temporary_table ADD COLUMN \"%s\" text" % "DUMMY_COLUMN")
            c.execute("ALTER TABLE \"%s\" ADD COLUMN \"%s\" text" % (model._meta.db_table, "DUMMY_COLUMN"))
            insert_col_list = "\", \"".join(
                ["DUMMY_COLUMN"] + date_columns + int_columns + numeric_columns + time_columns + empty_cols
            )
        else:
            insert_col_list = "\", \"".join(
                regular_columns + date_columns + int_columns + numeric_columns + time_columns + empty_cols
            )

        insert_statement += insert_col_list
        insert_statement += "\")\n"
        # add in the select part for table migration
        select_statement = "SELECT \""
        if not regular_columns:
            select_statement += "\", \"".join(["DUMMY_COLUMN"])
        else:
            select_statement += "\", \"".join(regular_columns)
        select_statement += "\"\n"
        # add in special formatting
        for dcol in date_columns:
            select_statement += self._make_date_case(dcol)
        for icol in int_columns:
            select_statement += self._make_int_case(icol)
        for ncol in numeric_columns:
            select_statement += self._make_numeric_case(ncol)
        for tcol in time_columns:
            select_statement += self._make_timestamp_case(tcol)
        for ecol in empty_cols:
            select_statement += self._make_special_not_null_case(ecol)
        # finalize from statement
        select_statement += "FROM temporary_table;"

        try:
            c.execute(insert_statement + select_statement)
        except DataError as e:
            print insert_statement + select_statement
            # print int_columns
            print "dataerror error, ", e
        except ProgrammingError as e:
            # print insert_statement + select_statement
            # print int_columns
            print "programming error, ", e
        except IntegrityError as e:
            # print insert_statement + select_statement
            # print int_columns
            print "IntegrityError error, ", e

        c.execute('DROP TABLE temporary_table;')
        if not regular_columns:
            c.execute("ALTER TABLE \"%s\" DROP COLUMN \"%s\"" % (model._meta.db_table, "DUMMY_COLUMN"))
        if self.verbosity:
            model_count = model.objects.count()
            if model_count == csv_count:
                self.success("  Table record count matches CSV")
            else:
                msg = '  Table Record count doesn\'t match CSV. \
Table: %s\tCSV: %s'
                self.failure(msg % (model_count, csv_count))

    def load_mysql(self, model, csv_path):
        c = connection.cursor()
        # flush
        # c.execute('TRUNCATE TABLE %s' % model._meta.db_table)

        # Build the MySQL LOAD DATA INFILE command
        bulk_sql_load_part_1 = '''
            LOAD DATA LOCAL INFILE '%s'
            INTO TABLE %s
            FIELDS TERMINATED BY ','
            OPTIONALLY ENCLOSED BY '"'
            LINES TERMINATED BY '\\r\\n'
            IGNORE 1 LINES
            (
        ''' % (csv_path, model._meta.db_table)

        infile = open(csv_path)
        csv_reader = csv.reader(infile)
        headers = csv_reader.next()
        infile.close()

        infile = open(csv_path)
        csv_record_cnt = len(infile.readlines()) - 1
        infile.close()

        header_sql_list = []
        date_set_list = []
        for h in headers:
            # If it is a date field, we need to reformat the data
            # so that MySQL will properly parse it on the way in.
            if h in model.DATE_FIELDS:
                header_sql_list.append('@`%s`' % h)
                date_set_list.append(
                    "`%s` =  %s" % (h, self.date_sql % h)
                )
            else:
                header_sql_list.append('`%s`' % h)

        bulk_sql_load = bulk_sql_load_part_1 + ','.join(header_sql_list) + ')'
        if date_set_list:
            bulk_sql_load += " set %s" % ",".join(date_set_list)

        # Run the query
        # cnt = c.execute(bulk_sql_load)
        # Report back on how we did
        print bulk_sql_load
        if self.verbosity:
            if cnt != csv_record_cnt:
                msg = '  Table Record count doesn\'t match CSV. \
Table: %s\tCSV: %s'
                self.failure(msg % (
                    cnt,
                    csv_record_cnt,
                ))

    def load(self, model_name):
        """
        Loads the source CSV for the provided model.
        """
        if self.verbosity:
            self.log(" Loading %s" % model_name)

        model = get_model("calaccess_raw", model_name)
        csv_path = model.objects.get_csv_path()

        # Flush
        engine = settings.DATABASES['default']['ENGINE']
        if engine == 'django.db.backends.mysql':
            self.load_mysql(model, csv_path)
        elif engine == 'django.db.backends.postgresql_psycopg2':
            self.load_postgresql(model, csv_path)
        else:
            self.failure("Sorry that database is not supported")
<commit_msg>removed comments, added back mysql code<commit_after>import csv
from django.db import connection
from django.conf import settings
from django.utils import dateparse
from django.db import IntegrityError, DataError, ProgrammingError
from django.db.models.loading import get_model
from django.core.management.base import LabelCommand
from calaccess_raw.management.commands import CalAccessCommand

class Command(CalAccessCommand, LabelCommand):
    help = 'Load a cleaned CalAccess file for a model into the database'
    args = '<model name>'
    # Trick for reformating date strings in source data so that they can
    # be gobbled up by MySQL. You'll see how below.
    date_sql = "DATE_FORMAT(str_to_date(@`%s`, '%%c/%%e/%%Y'), '%%Y-%%m-%%d')"

    def handle_label(self, label, **options):
        self.verbosity = options.get("verbosity")
        self.load(label)

    def get_hdrs_and_cnt(self, csv_path):
        """
        Get the headers and the line count
        from a specified csv file
        """
        with open(csv_path) as infile:
            csv_reader = csv.reader(infile)
            headers = csv_reader.next()

        with open(csv_path) as infile:
            csv_count = len(infile.readlines()) - 1

        return headers, csv_count

    def _make_date_case(self, _col):
        """
        This method takes in a column name and generates a
        PostgreSQL "case" for correct insertion into the table.
        """
        return """
        ,CASE
            WHEN "%s" IS NOT NULL AND "%s" != ''
                THEN to_date(substring("%s" from 1 for 10), 'MM/DD/YYYY')
            WHEN "%s" = ''
                THEN to_date('01/01/1900', 'MM/DD/YYYY')
        END AS "%s"\n""" % (_col, _col, _col, _col, _col)

    def _make_int_case(self, _col):
        return """
        ,CASE
            WHEN "%s" = ''
                THEN NULL
            WHEN "%s" = 'Y'
                THEN 1
            WHEN "%s" = 'y'
                THEN 1
            WHEN "%s" = 'X'
                THEN 1
            WHEN "%s" = 'N'
                THEN 0
            WHEN "%s" = 'n'
                THEN 0
            WHEN "%s" IS NOT NULL
                THEN "%s"::int
        END AS "%s"\n""" % (_col, _col, _col, _col, _col, _col, _col, _col, _col)

    def _make_numeric_case(self, _col):
        return """
        ,CASE
            WHEN "%s" = ''
                THEN 0.0
            WHEN "%s" IS NULL
                THEN 0.0
            WHEN "%s" IS NOT NULL
                THEN "%s"::numeric
        END AS "%s"\n""" % (_col, _col, _col, _col, _col)

    def _make_timestamp_case(self, _col):
        # "7/9/2014 12:00:00 AM"
        return """
        ,CASE
            WHEN "%s" IS NOT NULL AND "%s" != ''
                THEN to_timestamp("%s", 'MM/DD/YYYY HH12:MI:SS AM')
            WHEN "%s" = ''
                THEN to_timestamp('01/01/1900 1:00:00 AM', 'MM/DD/YYYY HH12:MI:SS AM')
        END AS "%s"\n""" % (_col, _col, _col, _col, _col)

    def _make_special_not_null_case(self, _col):
        return """
        ,CASE
            WHEN "%s" IS NULL
                THEN ''
        END AS "%s"\n""" % (_col, _col)

    def load_postgresql(self, model, csv_path):
        c = connection.cursor()
        try:
            c.execute('DROP TABLE temporary_table;')
        except ProgrammingError:
            pass
        # clear our table
        c.execute('TRUNCATE TABLE "%s"' % model._meta.db_table)

        headers, csv_count = self.get_hdrs_and_cnt(csv_path)

        #map column names to column types
        name_to_type_map = dict([(col.db_column, col.db_type(connection))
                                for col in model._meta.fields])

        # break out special case column types
        int_columns = []
        numeric_columns = []
        date_columns = []
        time_columns = []
        regular_columns = []

        #fill in those column types
        for col in model._meta.fields:
            if col.db_type(connection).startswith('integer'):
                int_columns.append(col.db_column)
            elif col.db_type(connection).startswith('numeric'):
                numeric_columns.append(col.db_column)
            elif col.db_type(connection).startswith('date'):
                date_columns.append(col.db_column)
            elif col.db_type(connection).startswith('timestamp'):
                time_columns.append(col.db_column)
            else:
                if col.db_column is not None and col.db_column in headers:
                    regular_columns.append(col.db_column)

        file_cols_types = []  # column with its types
        for col in headers:
            if col in regular_columns:
                file_cols_types.append("\"" + col + "\"\t" + name_to_type_map[col])
            else:
                file_cols_types.append("\"" + col + "\"\ttext")

        extra_cols = set([col.db_column for col in model._meta.fields]).difference(set(headers))
        empty_cols = []
        for col in extra_cols:
            if col != None:
                empty_cols.append(col)

        # create the temp table w/ columns with types
        c.execute("CREATE TABLE \"temporary_table\" (%s);"
                  % ',\n'.join(file_cols_types))

        temp_insert = """COPY "temporary_table"
            FROM '%s'
            CSV
            HEADER;""" % (csv_path)

        try:
            c.execute(temp_insert)  # insert everything into the temp table
        except DataError as e:
            print "initial insert dataerror error, ", e

        for col in empty_cols:
            c.execute("ALTER TABLE temporary_table ADD COLUMN \"%s\" text" % col)
        # build our insert statement
        insert_statement = "INSERT INTO \"%s\" (\"" % model._meta.db_table
        if not regular_columns:
            c.execute("ALTER TABLE temporary_table ADD COLUMN \"%s\" text" % "DUMMY_COLUMN")
            c.execute("ALTER TABLE \"%s\" ADD COLUMN \"%s\" text" % (model._meta.db_table, "DUMMY_COLUMN"))
            insert_col_list = "\", \"".join(
                ["DUMMY_COLUMN"] + date_columns + int_columns + numeric_columns + time_columns + empty_cols
            )
        else:
            insert_col_list = "\", \"".join(
                regular_columns + date_columns + int_columns + numeric_columns + time_columns + empty_cols
            )

        insert_statement += insert_col_list
        insert_statement += "\")\n"
        # add in the select part for table migration
        select_statement = "SELECT \""
        if not regular_columns:
            select_statement += "\", \"".join(["DUMMY_COLUMN"])
        else:
            select_statement += "\", \"".join(regular_columns)
        select_statement += "\"\n"
        # add in special formatting
        for dcol in date_columns:
            select_statement += self._make_date_case(dcol)
        for icol in int_columns:
            select_statement += self._make_int_case(icol)
        for ncol in numeric_columns:
            select_statement += self._make_numeric_case(ncol)
        for tcol in time_columns:
            select_statement += self._make_timestamp_case(tcol)
        for ecol in empty_cols:
            select_statement += self._make_special_not_null_case(ecol)
        # finalize from statement
        select_statement += "FROM temporary_table;"

        try:
            c.execute(insert_statement + select_statement)
        except DataError as e:
            print "dataerror error, ", e
        except ProgrammingError as e:
            print "programming error, ", e
        except IntegrityError as e:
            print "IntegrityError error, ", e

        c.execute('DROP TABLE temporary_table;')
        if not regular_columns:
            c.execute("ALTER TABLE \"%s\" DROP COLUMN \"%s\"" % (model._meta.db_table, "DUMMY_COLUMN"))
        if self.verbosity:
            model_count = model.objects.count()
            if model_count == csv_count:
                self.success("  Table record count matches CSV")
            else:
                msg = '  Table Record count doesn\'t match CSV. \
Table: %s\tCSV: %s'
                self.failure(msg % (model_count, csv_count))

    def load_mysql(self, model, csv_path):
        c = connection.cursor()
        # flush
        c.execute('TRUNCATE TABLE %s' % model._meta.db_table)

        # Build the MySQL LOAD DATA INFILE command
        bulk_sql_load_part_1 = '''
            LOAD DATA LOCAL INFILE '%s'
            INTO TABLE %s
            FIELDS TERMINATED BY ','
            OPTIONALLY ENCLOSED BY '"'
            LINES TERMINATED BY '\\r\\n'
            IGNORE 1 LINES
            (
        ''' % (csv_path, model._meta.db_table)

        infile = open(csv_path)
        csv_reader = csv.reader(infile)
        headers = csv_reader.next()
        infile.close()

        infile = open(csv_path)
        csv_record_cnt = len(infile.readlines()) - 1
        infile.close()

        header_sql_list = []
        date_set_list = []
        for h in headers:
            # If it is a date field, we need to reformat the data
            # so that MySQL will properly parse it on the way in.
            if h in model.DATE_FIELDS:
                header_sql_list.append('@`%s`' % h)
                date_set_list.append(
                    "`%s` =  %s" % (h, self.date_sql % h)
                )
            else:
                header_sql_list.append('`%s`' % h)

        bulk_sql_load = bulk_sql_load_part_1 + ','.join(header_sql_list) + ')'
        if date_set_list:
            bulk_sql_load += " set %s" % ",".join(date_set_list)

        # Run the query
        cnt = c.execute(bulk_sql_load)
        # Report back on how we did
        if self.verbosity:
            if cnt != csv_record_cnt:
                msg = '  Table Record count doesn\'t match CSV. \
Table: %s\tCSV: %s'
                self.failure(msg % (
                    cnt,
                    csv_record_cnt,
                ))

    def load(self, model_name):
        """
        Loads the source CSV for the provided model.
        """
        if self.verbosity:
            self.log(" Loading %s" % model_name)

        model = get_model("calaccess_raw", model_name)
        csv_path = model.objects.get_csv_path()

        # Flush
        engine = settings.DATABASES['default']['ENGINE']
        if engine == 'django.db.backends.mysql':
            self.load_mysql(model, csv_path)
        elif engine == 'django.db.backends.postgresql_psycopg2':
            self.load_postgresql(model, csv_path)
        else:
            self.failure("Sorry that database is not supported")
<|endoftext|>