<commit_before># -*- coding: utf8 -*-

import os
import numpy as np
import tensorflow as tf

from agents.agent import Agent
from memory.memory import Memory
from misc.network_ops import linear

class SAC(Agent):
    def __init__(self, env, monitor_path: str, **usercfg) -> None:
        super(SAC, self).__init__(**usercfg)
        self.env = env
        self.monitor_path: str = monitor_path

        self.config.update(
            n_episodes=100000,
            n_timesteps=env.spec.tags.get("wrapper_config.TimeLimit.max_episode_steps"),
            actor_learning_rate=3e-4,
            softq_learning_rate=3e-4,
            value_learning_rate=3e-4,
            gamma=0.99,
            batch_size=64,
            tau=0.01,
            l2_loss_coef=1e-2,
            n_actor_layers=2,
            logprob_epsilon=1e-6, # For numerical stability when computing tf.log
            n_hidden_units=128,
            n_train_steps=4, # Number of parameter update steps per iteration
            replay_buffer_size=1e6,
            replay_start_size=10000  # Required number of replay buffer entries to start training
        )
        self.config.update(usercfg)

        self.state_shape: list = list(env.observation_space.shape)
        self.n_actions: int = env.action_space.shape[0]
        self.states = tf.placeholder(tf.float32, [None] + self.state_shape, name="states")
        self.actions_taken = tf.placeholder(tf.float32, [None, self.n_actions], name="actions_taken")
        self.softq_target = tf.placeholder(tf.float32, [None, 1], name="softq_target")
        self.value_target = tf.placeholder(tf.float32, [None, 1], name="value_target")
        self.is_training = tf.placeholder(tf.bool, name="is_training")

        # Make networks
        # action_output are the squashed actions and action_original those straight from the normal distribution
        self.action_output, self.action_original, self.action_logprob, self.actor_vars = self.build_actor_network()
        self.softq_output, self.softq_vars = self.build_softq_network()
        self.value_output, self.value_vars = self.build_value_network()
        self.target_value_output, self.value_target_update = self.build_target_value_network(self.value_vars)

        # Make losses
        # softq_target = r(s_t, a_t) - gamma * target_value_output(s_{t+1})
        self.softq_target = tf.Print(self.softq_target, [self.softq_target], message="softq=", first_n=5)
        self.softq_loss = tf.reduce_mean(self.softq_output * (self.softq_output - self.softq_target))
        value_target = self.softq_output + self.action_logprob
        self.value_loss = tf.reduce_mean(self.value_output * (self.value_output - value_target))
        self.actor_loss = tf.reduce_mean(self.action_logprob * \
        (self.action_logprob - self.softq_output + self.value_output))

        # Make train ops
        softq_train_op = tf.train.AdamOptimizer(
            self.config["softq_learning_rate"],
            name="softq_optimizer").minimize(self.softq_loss)
        value_train_op = tf.train.AdamOptimizer(
            self.config["value_learning_rate"],
            name="value_optimizer").minimize(self.value_loss)
        actor_train_op = tf.train.AdamOptimizer(
            self.config["actor_learning_rate"],
            name="actor_optimizer").minimize(self.actor_loss)

        self.train_op = tf.group(softq_train_op, value_train_op, actor_train_op, name="train_op")

        summaries = []
        for v in self.actor_vars + self.softq_vars + self.value_vars:
            summaries.append(tf.summary.histogram(v.name, v))
        self.model_summary_op = tf.summary.merge(summaries)

        self.init_op = tf.global_variables_initializer()

        self.replay_buffer = Memory(int(self.config["replay_buffer_size"]))

        self.n_updates = 0

        self.summary_writer = tf.summary.FileWriter(os.path.join(
            self.monitor_path, "summaries"), tf.get_default_graph())

    def build_actor_network(self):
        w_bound = 3e-3
        x = self.states
        with tf.variable_scope("actor"):
            x = linear(x, self.config["n_hidden_units"], "L1", tf.random_uniform_initializer(-w_bound, w_bound))
            x = tf.nn.relu(x)

            mean = linear(x, self.n_actions, "mean", tf.random_uniform_initializer(-w_bound, w_bound))
            log_std = linear(x, self.n_actions, "log_std", tf.random_uniform_initializer(-w_bound, w_bound))

            normal_dist = tf.distributions.Normal(mean, tf.exp(log_std))
            actions = normal_dist.sample()
            squashed_actions = tf.tanh(actions) # Squash output between [-1, 1]

            logprob = normal_dist.log_prob(actions) - \
            tf.log(1.0 - tf.pow(self.actions_taken, 2) + self.config["logprob_epsilon"])

            actor_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)

        return squashed_actions, actions, logprob, actor_vars

    def build_target_value_network(self, value_vars: list):
        ema = tf.train.ExponentialMovingAverage(decay=1 - self.config["tau"])
        target_update = ema.apply(value_vars)
        target_net = [ema.average(v) for v in value_vars]

        x = self.states
        x = tf.nn.relu(tf.nn.xw_plus_b(x, target_net[0], target_net[1]))
        x = tf.nn.relu(tf.nn.xw_plus_b(x, target_net[2], target_net[3]))
        value = tf.nn.xw_plus_b(x, target_net[4], target_net[5])

        return value, target_update

    def build_softq_network(self):
        x = tf.concat([self.states, self.actions_taken], 1)
        with tf.variable_scope("softq"):
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L1"))
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L2"))
            x = linear(x, 1, "softq", tf.random_uniform_initializer(-3e-3, 3e-3))

            softq_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)

        return x, softq_vars

    def build_value_network(self):
        x = self.states
        with tf.variable_scope("value"):
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L1"))
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L2"))
            x = linear(x, 1, "value", tf.random_uniform_initializer(-3e-3, 3e-3))

            value_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)

        return x, value_vars

    def target_value(self, states: np.ndarray):
        return tf.get_default_session().run(self.target_value_output, feed_dict={
            self.states: states,
            self.is_training: False
        })

    def softq_value(self, states: np.ndarray, actions: np.ndarray):
        return tf.get_default_session().run(self.softq_output, feed_dict={
            self.states: states,
            self.actions_taken: actions,
            self.is_training: False
            })

    def actions(self, states: np.ndarray) -> np.ndarray:
        """Get the actions for a batch of states."""
        return tf.get_default_session().run(self.action_output, feed_dict={
            self.states: states,
            self.is_training: True
        })

    def action(self, state: np.ndarray) -> np.ndarray:
        """Get the action for a single state."""
        return tf.get_default_session().run(self.action_output, feed_dict={
            self.states: [state],
            self.is_training: False
        })[0]

    def train(self):
        sample = self.replay_buffer.get_batch(self.config["batch_size"])

        # for n_actions = 1
        action_batch = np.resize(sample["actions"], [self.config["batch_size"], self.n_actions])

        # Calculate critic targets
        next_value_batch = self.target_value(sample["states1"])
        softq_targets = sample["rewards"] + (1 - sample["terminals1"]) * \
            self.config["gamma"] * next_value_batch.squeeze()
        softq_targets = np.resize(softq_targets, [self.config["batch_size"], 1]).astype(np.float32)
        # Update actor weights
        fetches = [self.softq_output, self.softq_loss, self.train_op]
        predicted_q, softq_loss, _ = tf.get_default_session().run(fetches, feed_dict={
            self.softq_target: softq_targets,
            self.states: sample["states0"],
            self.actions_taken: action_batch,
            self.is_training: True
        })

        summary = tf.Summary()
        summary.value.add(tag="model/softq_loss", simple_value=float(softq_loss))
        summary.value.add(tag="model/predicted_softq_mean", simple_value=np.mean(predicted_q))
        summary.value.add(tag="model/predicted_softq_std", simple_value=np.std(predicted_q))
        self.summary_writer.add_summary(summary, self.n_updates)

        # Update the target networks
        tf.get_default_session().run([self.value_target_update, self.model_summary_op])
        self.n_updates += 1

    def learn(self):
        max_action = self.env.action_space.high
        with tf.Session() as sess, sess.as_default():
            sess.run(self.init_op)
            for episode in range(self.config["n_episodes"]):
                state = self.env.reset()
                episode_reward = 0
                episode_length = 0
                for _ in range(self.config["n_timesteps"]):
                    action = self.action(state)
                    new_state, reward, done, _ = self.env.step(action * max_action)
                    episode_length += 1
                    episode_reward += reward
                    self.replay_buffer.add(state, action, reward, new_state, done)
                    if self.replay_buffer.n_entries > self.config["replay_start_size"]:
                        for _ in range(self.config["n_train_steps"]):
                            self.train()
                    state = new_state
                    if done:
                        summary = tf.Summary()
                        summary.value.add(tag="global/Episode_length",
                                          simple_value=float(episode_length))
                        summary.value.add(tag="global/Reward",
                                          simple_value=float(episode_reward))
                        self.summary_writer.add_summary(summary, episode)
                        self.summary_writer.flush()
                        break
<commit_msg>(Partially) fixed losses<commit_after># -*- coding: utf8 -*-

import os
import numpy as np
import tensorflow as tf

from agents.agent import Agent
from memory.memory import Memory
from misc.network_ops import linear

class SAC(Agent):
    def __init__(self, env, monitor_path: str, **usercfg) -> None:
        super(SAC, self).__init__(**usercfg)
        self.env = env
        self.monitor_path: str = monitor_path

        self.config.update(
            n_episodes=100000,
            n_timesteps=env.spec.tags.get("wrapper_config.TimeLimit.max_episode_steps"),
            actor_learning_rate=3e-4,
            softq_learning_rate=3e-4,
            value_learning_rate=3e-4,
            gamma=0.99,
            batch_size=128,
            tau=0.01,
            l2_loss_coef=1e-2,
            n_actor_layers=2,
            logprob_epsilon=1e-6, # For numerical stability when computing tf.log
            n_hidden_units=128,
            n_train_steps=4, # Number of parameter update steps per iteration
            replay_buffer_size=1e6,
            replay_start_size=128  # Required number of replay buffer entries to start training
        )
        self.config.update(usercfg)

        self.state_shape: list = list(env.observation_space.shape)
        self.n_actions: int = env.action_space.shape[0]
        self.states = tf.placeholder(tf.float32, [None] + self.state_shape, name="states")
        self.actions_taken = tf.placeholder(tf.float32, [None, self.n_actions], name="actions_taken")
        self.softq_target = tf.placeholder(tf.float32, [None, 1], name="softq_target")
        self.value_target = tf.placeholder(tf.float32, [None, 1], name="value_target")
        self.advantage = tf.placeholder(tf.float32, [None, 1], name="advantage")
        self.is_training = tf.placeholder(tf.bool, name="is_training")

        # Make networks
        # action_output are the squashed actions and action_original those straight from the normal distribution
        self.action_output, self.action_original, self.action_logprob, self.actor_vars = self.build_actor_network()
        self.softq_output, self.softq_vars = self.build_softq_network()
        self.value_output, self.value_vars = self.build_value_network()
        self.target_value_output, self.value_target_update = self.build_target_value_network(self.value_vars)

        # Make losses
        # softq_target = r(s_t, a_t) - gamma * target_value_output(s_{t+1})
        self.softq_target = tf.Print(self.softq_target, [self.softq_target], message="softq=", first_n=5)
        self.softq_loss = tf.reduce_mean(0.5 * tf.square(self.softq_output - self.softq_target))
        self.value_loss = tf.reduce_mean(0.5 * tf.square(self.value_output - self.value_target))
        self.actor_loss = tf.reduce_mean(self.action_logprob * self.advantage)

        # Make train ops
        softq_train_op = tf.train.AdamOptimizer(
            self.config["softq_learning_rate"],
            name="softq_optimizer").minimize(self.softq_loss)
        value_train_op = tf.train.AdamOptimizer(
            self.config["value_learning_rate"],
            name="value_optimizer").minimize(self.value_loss)
        actor_train_op = tf.train.AdamOptimizer(
            self.config["actor_learning_rate"],
            name="actor_optimizer").minimize(self.actor_loss)

        self.train_op = tf.group(softq_train_op, value_train_op, actor_train_op, name="train_op")

        summaries = []
        for v in self.actor_vars + self.softq_vars + self.value_vars:
            summaries.append(tf.summary.histogram(v.name, v))
        self.model_summary_op = tf.summary.merge(summaries)

        self.init_op = tf.global_variables_initializer()

        self.replay_buffer = Memory(int(self.config["replay_buffer_size"]))

        self.n_updates = 0

        self.summary_writer = tf.summary.FileWriter(os.path.join(
            self.monitor_path, "summaries"), tf.get_default_graph())

    def build_actor_network(self):
        w_bound = 3e-3
        x = self.states
        with tf.variable_scope("actor"):
            x = linear(x, self.config["n_hidden_units"], "L1", tf.random_uniform_initializer(-w_bound, w_bound))
            x = tf.nn.relu(x)

            mean = linear(x, self.n_actions, "mean", tf.random_uniform_initializer(-w_bound, w_bound))
            log_std = linear(x, self.n_actions, "log_std", tf.random_uniform_initializer(-w_bound, w_bound))
            log_std = tf.clip_by_value(log_std, -20, 2)

            normal_dist = tf.distributions.Normal(mean, tf.exp(log_std))
            actions = normal_dist.sample()
            squashed_actions = tf.tanh(actions) # Squash output between [-1, 1]

            logprob = normal_dist.log_prob(actions) - \
            tf.log(1.0 - tf.pow(self.actions_taken, 2) + self.config["logprob_epsilon"])

            actor_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)

        return squashed_actions, actions, logprob, actor_vars

    def build_target_value_network(self, value_vars: list):
        ema = tf.train.ExponentialMovingAverage(decay=1 - self.config["tau"])
        target_update = ema.apply(value_vars)
        target_net = [ema.average(v) for v in value_vars]

        x = self.states
        x = tf.nn.relu(tf.nn.xw_plus_b(x, target_net[0], target_net[1]))
        x = tf.nn.relu(tf.nn.xw_plus_b(x, target_net[2], target_net[3]))
        value = tf.nn.xw_plus_b(x, target_net[4], target_net[5])

        return value, target_update

    def build_softq_network(self):
        x = tf.concat([self.states, self.actions_taken], 1)
        with tf.variable_scope("softq"):
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L1"))
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L2"))
            x = linear(x, 1, "softq", tf.random_uniform_initializer(-3e-3, 3e-3))

            softq_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)

        return x, softq_vars

    def build_value_network(self):
        x = self.states
        with tf.variable_scope("value"):
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L1"))
            x = tf.nn.relu(linear(x, self.config["n_hidden_units"], "L2"))
            x = linear(x, 1, "value", tf.random_uniform_initializer(-3e-3, 3e-3))

            value_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)

        return x, value_vars

    def value(self, states: np.ndarray):
        return tf.get_default_session().run(self.value_output, feed_dict={
            self.states: states,
            self.is_training: False
        })

    def target_value(self, states: np.ndarray):
        return tf.get_default_session().run(self.target_value_output, feed_dict={
            self.states: states,
            self.is_training: False
        })

    def softq_value(self, states: np.ndarray, actions: np.ndarray):
        return tf.get_default_session().run(self.softq_output, feed_dict={
            self.states: states,
            self.actions_taken: actions,
            self.is_training: False
            })

    def actions(self, states: np.ndarray) -> np.ndarray:
        """Get the actions for a batch of states."""
        return tf.get_default_session().run(self.action_output, feed_dict={
            self.states: states,
            self.is_training: False
        })

    def action(self, state: np.ndarray) -> np.ndarray:
        """Get the action for a single state."""
        return tf.get_default_session().run(self.action_output, feed_dict={
            self.states: [state],
            self.is_training: False
        })[0]

    def train(self):
        sample = self.replay_buffer.get_batch(self.config["batch_size"])

        # for n_actions = 1
        action_batch = np.resize(sample["actions"], [self.config["batch_size"], self.n_actions])

        # Calculate critic targets
        next_value_batch = self.target_value(sample["states1"])
        softq_targets = sample["rewards"] + (1 - sample["terminals1"]) * \
            self.config["gamma"] * next_value_batch.squeeze()
        softq_targets = np.resize(softq_targets, [self.config["batch_size"], 1]).astype(np.float32)

        next_actions = self.actions(sample["states0"])
        next_q_value_batch = self.softq_value(sample["states0"], next_actions)
        value_batch = self.value(sample["states0"])

        log_prob, next_q_value_batch, value_batch = tf.get_default_session().run(
            [self.action_logprob, self.softq_output, self.value_output],
            feed_dict={
                self.states: sample["states0"],
                self.actions_taken: next_actions
            }
        )
        advantage_batch = log_prob - (next_q_value_batch - value_batch)

        value_targets = next_q_value_batch - log_prob
        # Update actor weights
        fetches = [self.softq_output, self.softq_loss, self.actor_loss, self.value_loss, self.train_op]
        predicted_q, softq_loss, actor_loss, value_loss, _ = tf.get_default_session().run(fetches, feed_dict={
            self.softq_target: softq_targets,
            self.states: sample["states0"],
            self.actions_taken: action_batch,
            self.advantage: advantage_batch,
            self.value_target: value_targets,
            self.is_training: True
        })

        summary = tf.Summary()
        summary.value.add(tag="model/softq_loss", simple_value=float(softq_loss))
        summary.value.add(tag="model/value_loss", simple_value=float(value_loss))
        summary.value.add(tag="model/actor_loss", simple_value=float(actor_loss))
        summary.value.add(tag="model/predicted_softq_mean", simple_value=np.mean(predicted_q))
        summary.value.add(tag="model/predicted_softq_std", simple_value=np.std(predicted_q))
        self.summary_writer.add_summary(summary, self.n_updates)

        # Update the target networks
        tf.get_default_session().run([self.value_target_update, self.model_summary_op])
        self.n_updates += 1

    def learn(self):
        max_action = self.env.action_space.high
        timestep_limit = self.env.spec.tags.get('wrapper_config.TimeLimit.max_episode_steps')
        with tf.Session() as sess, sess.as_default():
            sess.run(self.init_op)
            for episode in range(self.config["n_episodes"]):
                state = self.env.reset()
                episode_reward = 0
                episode_length = 0
                for _ in range(self.config["n_timesteps"]):
                    action = self.action(state)
                    new_state, reward, done, _ = self.env.step(action * max_action)
                    episode_length += 1
                    episode_reward += reward
                    self.replay_buffer.add(state, action, reward, new_state, done)
                    if self.replay_buffer.n_entries > self.config["replay_start_size"]:
                        for _ in range(self.config["n_train_steps"]):
                            self.train()
                    state = new_state
                    if done or episode_length >= timestep_limit:
                        summary = tf.Summary()
                        summary.value.add(tag="global/Episode_length",
                                          simple_value=float(episode_length))
                        summary.value.add(tag="global/Reward",
                                          simple_value=float(episode_reward))
                        self.summary_writer.add_summary(summary, episode)
                        self.summary_writer.flush()
                        break
<|endoftext|>


<commit_before>import logging
from time import time

log = logging.getLogger(__name__)

# This is used to ensure that metrics with a timestamp older than
# RECENT_POINT_THRESHOLD_DEFAULT seconds (or the value passed in to
# the MetricsAggregator constructor) get discarded rather than being
# input into the incorrect bucket. Currently, the MetricsAggregator
# does not support submitting values for the past, and all values get
# submitted for the timestamp passed into the flush() function.
RECENT_POINT_THRESHOLD_DEFAULT = 30

class Infinity(Exception): pass
class UnknownValue(Exception): pass

class Metric(object):
    """
    A base metric class that accepts points, slices them into time intervals
    and performs roll-ups within those intervals.
    """

    def sample(self, value, sample_rate):
        """ Add a point to the given metric. """
        raise NotImplementedError()

    def flush(self, timestamp, interval):
        """ Flush all metrics up to the given timestamp. """
        raise NotImplementedError()

class Gauge(Metric):
    """ A metric that tracks a value at particular points in time. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.value = None
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name
        self.last_sample_time = None

    def sample(self, value, sample_rate):
        self.value = value
        self.last_sample_time = time()

    def flush(self, timestamp, interval):
        if self.value is not None:
            res = [self.formatter(
                metric=self.name,
                timestamp=timestamp,
                value=self.value,
                tags=self.tags,
                hostname=self.hostname,
                device_name=self.device_name
            )]
            self.value = None
            return res

        return []

class Counter(Metric):
    """ A metric that tracks a counter value. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.value = 0
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name

    def sample(self, value, sample_rate):
        self.value += value * int(1 / sample_rate)
        self.last_sample_time = time()

    def flush(self, timestamp, interval):
        try:
            value = self.value / interval
            return [self.formatter(
                metric=self.name,
                value=value,
                timestamp=timestamp,
                tags=self.tags,
                hostname=self.hostname,
                device_name=self.device_name
            )]
        finally:
            self.value = 0

class Histogram(Metric):
    """ A metric to track the distribution of a set of values. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.count = 0
        self.samples = []
        self.percentiles = [0.95]
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name

    def sample(self, value, sample_rate):
        self.count += int(1 / sample_rate)
        self.samples.append(value)
        self.last_sample_time = time()

    def flush(self, ts, interval):
        if not self.count:
            return []

        self.samples.sort()
        length = len(self.samples)

        max_ = self.samples[-1]
        med = self.samples[int(round(length/2 - 1))]
        avg = sum(self.samples) / float(length)

        metric_aggrs = [
            ('max', max_),
            ('median', med),
            ('avg', avg),
            ('count', self.count/interval)
        ]

        metrics = [self.formatter(
                hostname=self.hostname,
                device_name=self.device_name,
                tags=self.tags,
                metric='%s.%s' % (self.name, suffix),
                value=value,
                timestamp=ts
            ) for suffix, value in metric_aggrs
        ]

        for p in self.percentiles:
            val = self.samples[int(round(p * length - 1))]
            name = '%s.%spercentile' % (self.name, int(p * 100))
            metrics.append(self.formatter(
                hostname=self.hostname,
                tags=self.tags,
                metric=name,
                value=val,
                timestamp=ts
            ))

        # Reset our state.
        self.samples = []
        self.count = 0

        return metrics

class Set(Metric):
    """ A metric to track the number of unique elements in a set. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name
        self.values = set()

    def sample(self, value, sample_rate):
        self.values.add(value)
        self.last_sample_time = time()

    def flush(self, timestamp, interval):
        if not self.values:
            return []
        try:
            return [self.formatter(
                hostname=self.hostname,
                device_name=self.device_name,
                tags=self.tags,
                metric=self.name,
                value=len(self.values),
                timestamp=timestamp
            )]
        finally:
            self.values = set()

class Rate(Metric):
    """ Track the rate of metrics over each flush interval """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name
        self.samples = []

    def sample(self, value, sample_rate):
        ts = time()
        self.samples.append((int(ts), value))
        self.last_sample_time = ts

    def _rate(self, sample1, sample2):
        interval = sample2[0] - sample1[0]
        if interval == 0:
            log.warn('Metric %s has an interval of 0. Not flushing.' % self.name)
            raise Infinity()

        delta = sample2[1] - sample1[1]
        if delta < 0:
            log.warn('Metric %s has a rate < 0. Not flushing.' % self.name)
            raise UnknownValue()

        return (delta / float(interval))

    def flush(self, timestamp, interval):
        if len(self.samples) < 2:
            return []
        try:
            try:
                val = self._rate(self.samples[-2], self.samples[-1])
            except:
                return []

            return [self.formatter(
                hostname=self.hostname,
                device_name=self.device_name,
                tags=self.tags,
                metric=self.name,
                value=val,
                timestamp=timestamp
            )]
        finally:
            self.samples = self.samples[-1:]

class MetricsAggregator(object):
    """
    A metric aggregator class.
    """

    # Types of metrics that allow strings
    ALLOW_STRINGS = ['s', ]

    def __init__(self, hostname, interval=1.0, expiry_seconds=300, formatter=None, recent_point_threshold=None):
        self.metrics = {}
        self.events = []
        self.total_count = 0
        self.count = 0
        self.event_count = 0
        self.metric_type_to_class = {
            'g': Gauge,
            'c': Counter,
            'h': Histogram,
            'ms': Histogram,
            's': Set,
            '_dd-r': Rate,
        }
        self.hostname = hostname
        self.expiry_seconds = expiry_seconds
        self.formatter = formatter or api_formatter
        self.interval = float(interval)

        recent_point_threshold = recent_point_threshold or RECENT_POINT_THRESHOLD_DEFAULT
        self.recent_point_threshold = int(recent_point_threshold)
        self.num_discarded_old_points = 0

    def packets_per_second(self, interval):
        return round(float(self.count)/interval, 2)

    def parse_metric_packet(self, packet):
        name_and_metadata = packet.split(':', 1)

        if len(name_and_metadata) != 2:
            raise Exception('Unparseable metric packet: %s' % packet)

        name = name_and_metadata[0]
        metadata = name_and_metadata[1].split('|')

        if len(metadata) < 2:
            raise Exception('Unparseable metric packet: %s' % packet)
        # Try to cast as an int first to avoid precision issues, then as a
        # float.
        try:
            value = int(metadata[0])
        except ValueError:
            try:
                value = float(metadata[0])
            except ValueError:

                # If the data type is Set, we will allow strings
                if metadata[1] in self.ALLOW_STRINGS:
                    value = metadata[0]
                else:
                    # Otherwise, raise an error saying it must be a number
                    raise Exception('Metric value must be a number: %s, %s' % (name, metadata[0]))

        # Parse the optional values - sample rate & tags.
        sample_rate = 1
        tags = None
        for m in metadata[2:]:
            # Parse the sample rate
            if m[0] == '@':
                sample_rate = float(m[1:])
                assert 0 <= sample_rate <= 1
            elif m[0] == '#':
                tags = tuple(sorted(m[1:].split(',')))

        # Submit the metric
        mtype = metadata[1]

        return name, value, mtype, tags, sample_rate

    def _unescape_event_text(self, string):
        return string.replace('\\n', '\n')

    def parse_event_packet(self, packet):
        try:
            name_and_metadata = packet.split(':', 1)
            if len(name_and_metadata) != 2:
                raise Exception(u'Unparseable event packet: %s' % packet)
            # Event syntax:
            # _e{5,4}:title|body|meta
            name = name_and_metadata[0]
            metadata = unicode(name_and_metadata[1])
            title_length, text_length = name.split(',')
            title_length = int(title_length[3:])
            text_length = int(text_length[:-1])

            event = {
                'title': metadata[:title_length],
                'text': self._unescape_event_text(metadata[title_length+1:title_length+text_length+1])
            }
            meta = metadata[title_length+text_length+1:]
            for m in meta.split('|')[1:]:
                if m[0] == u't':
                    event['alert_type'] = m[2:]
                elif m[0] == u'k':
                    event['aggregation_key'] = m[2:]
                elif m[0] == u's':
                    event['source_type_name'] = m[2:]
                elif m[0] == u'd':
                    event['date_happened'] = int(m[2:])
                elif m[0] == u'p':
                    event['priority'] = m[2:]
                elif m[0] == u'h':
                    event['hostname'] = m[2:]
                elif m[0] == u'#':
                    event['tags'] = sorted(m[1:].split(u','))
            return event
        except IndexError, ValueError:
            raise Exception(u'Unparseable event packet: %s' % packet)

    def submit_packets(self, packets):

        for packet in packets.split("\n"):

            if not packet.strip():
                continue

            if packet.startswith('_e'):
                self.event_count += 1
                event = self.parse_event_packet(packet)
                self.event(**event)
            else:
                self.count += 1
                name, value, mtype, tags, sample_rate = self.parse_metric_packet(packet)
                self.submit_metric(name, value, mtype, tags=tags, sample_rate=sample_rate)

    def submit_metric(self, name, value, mtype, tags=None, hostname=None,
                                device_name=None, timestamp=None, sample_rate=1):
        # Avoid calling extra functions to dedupe tags if there are none
        if tags is None:
            context = (name, tuple(), hostname, device_name)
        else:
            context = (name, tuple(sorted(set(tags))), hostname, device_name)
        if context not in self.metrics:
            metric_class = self.metric_type_to_class[mtype]
            self.metrics[context] = metric_class(self.formatter, name, tags,
                hostname or self.hostname, device_name)
        cur_time = time()
        if timestamp is not None and cur_time - int(timestamp) > self.recent_point_threshold:
            self.num_discarded_old_points += 1
        else:
            self.metrics[context].sample(value, sample_rate)

    def gauge(self, name, value, tags=None, hostname=None, device_name=None, timestamp=None):
        self.submit_metric(name, value, 'g', tags, hostname, device_name, timestamp)

    def increment(self, name, value=1, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 'c', tags, hostname, device_name)

    def decrement(self, name, value=-1, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 'c', tags, hostname, device_name)

    def rate(self, name, value, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, '_dd-r', tags, hostname, device_name)

    def histogram(self, name, value, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 'h', tags, hostname, device_name)

    def set(self, name, value, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 's', tags, hostname, device_name)

    def event(self, title, text, date_happened=int(time()), alert_type=None, aggregation_key=None, source_type_name=None, priority=None, tags=None, hostname=None):
        event = {
            'title': title,
            'text': text,
        }
        if date_happened is not None:
            event['date_happened'] = date_happened
        if alert_type is not None:
            event['alert_type'] = alert_type
        if aggregation_key is not None:
            event['aggregation_key'] = aggregation_key
        if source_type_name is not None:
            event['source_type_name'] = source_type_name
        if priority is not None:
            event['priority'] = priority
        if tags is not None:
            event['tags'] = sorted(tags)
        if hostname is not None:
            event['host'] = hostname
        else:
            event['host'] = self.hostname

        self.events.append(event)

    def flush(self):
        timestamp = time()
        expiry_timestamp = timestamp - self.expiry_seconds

        # Flush points and remove expired metrics. We mutate this dictionary
        # while iterating so don't use an iterator.
        metrics = []
        for context, metric in self.metrics.items():
            if metric.last_sample_time < expiry_timestamp:
                log.debug("%s hasn't been submitted in %ss. Expiring." % (context, self.expiry_seconds))
                del self.metrics[context]
            else:
                metrics += metric.flush(timestamp, self.interval)

        # Log a warning regarding metrics with old timestamps being submitted
        if self.num_discarded_old_points > 0:
            log.warn('%s points were discarded as a result of having an old timestamp' % self.num_discarded_old_points)
            self.num_discarded_old_points = 0

        # Save some stats.
        log.debug("received %s payloads since last flush" % self.count)
        self.total_count += self.count
        self.count = 0
        return metrics

    def flush_events(self):
        events = self.events
        self.events = []

        self.total_count += self.event_count
        self.event_count = 0

        log.debug("Received %d events since last flush" % len(events))

        return events

    def send_packet_count(self, metric_name):
        self.submit_metric(metric_name, self.count, 'g')

def api_formatter(metric, value, timestamp, tags, hostname, device_name=None):

    # Workaround for a bug in minjson serialization
    # (https://github.com/DataDog/dd-agent/issues/422)
    if tags is not None and isinstance(tags, tuple) and len(tags) == 1:
        tags = list(tags)
    return {
        'metric': metric,
        'points': [(timestamp, value)],
        'tags': tags,
        'host': hostname,
        'device_name': device_name
    }
<commit_msg>Fix event date_happened default in dogstatsd<commit_after>import logging
from time import time

log = logging.getLogger(__name__)

# This is used to ensure that metrics with a timestamp older than
# RECENT_POINT_THRESHOLD_DEFAULT seconds (or the value passed in to
# the MetricsAggregator constructor) get discarded rather than being
# input into the incorrect bucket. Currently, the MetricsAggregator
# does not support submitting values for the past, and all values get
# submitted for the timestamp passed into the flush() function.
RECENT_POINT_THRESHOLD_DEFAULT = 30

class Infinity(Exception): pass
class UnknownValue(Exception): pass

class Metric(object):
    """
    A base metric class that accepts points, slices them into time intervals
    and performs roll-ups within those intervals.
    """

    def sample(self, value, sample_rate):
        """ Add a point to the given metric. """
        raise NotImplementedError()

    def flush(self, timestamp, interval):
        """ Flush all metrics up to the given timestamp. """
        raise NotImplementedError()

class Gauge(Metric):
    """ A metric that tracks a value at particular points in time. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.value = None
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name
        self.last_sample_time = None

    def sample(self, value, sample_rate):
        self.value = value
        self.last_sample_time = time()

    def flush(self, timestamp, interval):
        if self.value is not None:
            res = [self.formatter(
                metric=self.name,
                timestamp=timestamp,
                value=self.value,
                tags=self.tags,
                hostname=self.hostname,
                device_name=self.device_name
            )]
            self.value = None
            return res

        return []

class Counter(Metric):
    """ A metric that tracks a counter value. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.value = 0
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name

    def sample(self, value, sample_rate):
        self.value += value * int(1 / sample_rate)
        self.last_sample_time = time()

    def flush(self, timestamp, interval):
        try:
            value = self.value / interval
            return [self.formatter(
                metric=self.name,
                value=value,
                timestamp=timestamp,
                tags=self.tags,
                hostname=self.hostname,
                device_name=self.device_name
            )]
        finally:
            self.value = 0

class Histogram(Metric):
    """ A metric to track the distribution of a set of values. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.count = 0
        self.samples = []
        self.percentiles = [0.95]
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name

    def sample(self, value, sample_rate):
        self.count += int(1 / sample_rate)
        self.samples.append(value)
        self.last_sample_time = time()

    def flush(self, ts, interval):
        if not self.count:
            return []

        self.samples.sort()
        length = len(self.samples)

        max_ = self.samples[-1]
        med = self.samples[int(round(length/2 - 1))]
        avg = sum(self.samples) / float(length)

        metric_aggrs = [
            ('max', max_),
            ('median', med),
            ('avg', avg),
            ('count', self.count/interval)
        ]

        metrics = [self.formatter(
                hostname=self.hostname,
                device_name=self.device_name,
                tags=self.tags,
                metric='%s.%s' % (self.name, suffix),
                value=value,
                timestamp=ts
            ) for suffix, value in metric_aggrs
        ]

        for p in self.percentiles:
            val = self.samples[int(round(p * length - 1))]
            name = '%s.%spercentile' % (self.name, int(p * 100))
            metrics.append(self.formatter(
                hostname=self.hostname,
                tags=self.tags,
                metric=name,
                value=val,
                timestamp=ts
            ))

        # Reset our state.
        self.samples = []
        self.count = 0

        return metrics

class Set(Metric):
    """ A metric to track the number of unique elements in a set. """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name
        self.values = set()

    def sample(self, value, sample_rate):
        self.values.add(value)
        self.last_sample_time = time()

    def flush(self, timestamp, interval):
        if not self.values:
            return []
        try:
            return [self.formatter(
                hostname=self.hostname,
                device_name=self.device_name,
                tags=self.tags,
                metric=self.name,
                value=len(self.values),
                timestamp=timestamp
            )]
        finally:
            self.values = set()

class Rate(Metric):
    """ Track the rate of metrics over each flush interval """

    def __init__(self, formatter, name, tags, hostname, device_name):
        self.formatter = formatter
        self.name = name
        self.tags = tags
        self.hostname = hostname
        self.device_name = device_name
        self.samples = []

    def sample(self, value, sample_rate):
        ts = time()
        self.samples.append((int(ts), value))
        self.last_sample_time = ts

    def _rate(self, sample1, sample2):
        interval = sample2[0] - sample1[0]
        if interval == 0:
            log.warn('Metric %s has an interval of 0. Not flushing.' % self.name)
            raise Infinity()

        delta = sample2[1] - sample1[1]
        if delta < 0:
            log.warn('Metric %s has a rate < 0. Not flushing.' % self.name)
            raise UnknownValue()

        return (delta / float(interval))

    def flush(self, timestamp, interval):
        if len(self.samples) < 2:
            return []
        try:
            try:
                val = self._rate(self.samples[-2], self.samples[-1])
            except:
                return []

            return [self.formatter(
                hostname=self.hostname,
                device_name=self.device_name,
                tags=self.tags,
                metric=self.name,
                value=val,
                timestamp=timestamp
            )]
        finally:
            self.samples = self.samples[-1:]

class MetricsAggregator(object):
    """
    A metric aggregator class.
    """

    # Types of metrics that allow strings
    ALLOW_STRINGS = ['s', ]

    def __init__(self, hostname, interval=1.0, expiry_seconds=300, formatter=None, recent_point_threshold=None):
        self.metrics = {}
        self.events = []
        self.total_count = 0
        self.count = 0
        self.event_count = 0
        self.metric_type_to_class = {
            'g': Gauge,
            'c': Counter,
            'h': Histogram,
            'ms': Histogram,
            's': Set,
            '_dd-r': Rate,
        }
        self.hostname = hostname
        self.expiry_seconds = expiry_seconds
        self.formatter = formatter or api_formatter
        self.interval = float(interval)

        recent_point_threshold = recent_point_threshold or RECENT_POINT_THRESHOLD_DEFAULT
        self.recent_point_threshold = int(recent_point_threshold)
        self.num_discarded_old_points = 0

    def packets_per_second(self, interval):
        return round(float(self.count)/interval, 2)

    def parse_metric_packet(self, packet):
        name_and_metadata = packet.split(':', 1)

        if len(name_and_metadata) != 2:
            raise Exception('Unparseable metric packet: %s' % packet)

        name = name_and_metadata[0]
        metadata = name_and_metadata[1].split('|')

        if len(metadata) < 2:
            raise Exception('Unparseable metric packet: %s' % packet)
        # Try to cast as an int first to avoid precision issues, then as a
        # float.
        try:
            value = int(metadata[0])
        except ValueError:
            try:
                value = float(metadata[0])
            except ValueError:

                # If the data type is Set, we will allow strings
                if metadata[1] in self.ALLOW_STRINGS:
                    value = metadata[0]
                else:
                    # Otherwise, raise an error saying it must be a number
                    raise Exception('Metric value must be a number: %s, %s' % (name, metadata[0]))

        # Parse the optional values - sample rate & tags.
        sample_rate = 1
        tags = None
        for m in metadata[2:]:
            # Parse the sample rate
            if m[0] == '@':
                sample_rate = float(m[1:])
                assert 0 <= sample_rate <= 1
            elif m[0] == '#':
                tags = tuple(sorted(m[1:].split(',')))

        # Submit the metric
        mtype = metadata[1]

        return name, value, mtype, tags, sample_rate

    def _unescape_event_text(self, string):
        return string.replace('\\n', '\n')

    def parse_event_packet(self, packet):
        try:
            name_and_metadata = packet.split(':', 1)
            if len(name_and_metadata) != 2:
                raise Exception(u'Unparseable event packet: %s' % packet)
            # Event syntax:
            # _e{5,4}:title|body|meta
            name = name_and_metadata[0]
            metadata = unicode(name_and_metadata[1])
            title_length, text_length = name.split(',')
            title_length = int(title_length[3:])
            text_length = int(text_length[:-1])

            event = {
                'title': metadata[:title_length],
                'text': self._unescape_event_text(metadata[title_length+1:title_length+text_length+1])
            }
            meta = metadata[title_length+text_length+1:]
            for m in meta.split('|')[1:]:
                if m[0] == u't':
                    event['alert_type'] = m[2:]
                elif m[0] == u'k':
                    event['aggregation_key'] = m[2:]
                elif m[0] == u's':
                    event['source_type_name'] = m[2:]
                elif m[0] == u'd':
                    event['date_happened'] = int(m[2:])
                elif m[0] == u'p':
                    event['priority'] = m[2:]
                elif m[0] == u'h':
                    event['hostname'] = m[2:]
                elif m[0] == u'#':
                    event['tags'] = sorted(m[1:].split(u','))
            return event
        except IndexError, ValueError:
            raise Exception(u'Unparseable event packet: %s' % packet)

    def submit_packets(self, packets):

        for packet in packets.split("\n"):

            if not packet.strip():
                continue

            if packet.startswith('_e'):
                self.event_count += 1
                event = self.parse_event_packet(packet)
                self.event(**event)
            else:
                self.count += 1
                name, value, mtype, tags, sample_rate = self.parse_metric_packet(packet)
                self.submit_metric(name, value, mtype, tags=tags, sample_rate=sample_rate)

    def submit_metric(self, name, value, mtype, tags=None, hostname=None,
                                device_name=None, timestamp=None, sample_rate=1):
        # Avoid calling extra functions to dedupe tags if there are none
        if tags is None:
            context = (name, tuple(), hostname, device_name)
        else:
            context = (name, tuple(sorted(set(tags))), hostname, device_name)
        if context not in self.metrics:
            metric_class = self.metric_type_to_class[mtype]
            self.metrics[context] = metric_class(self.formatter, name, tags,
                hostname or self.hostname, device_name)
        cur_time = time()
        if timestamp is not None and cur_time - int(timestamp) > self.recent_point_threshold:
            self.num_discarded_old_points += 1
        else:
            self.metrics[context].sample(value, sample_rate)

    def gauge(self, name, value, tags=None, hostname=None, device_name=None, timestamp=None):
        self.submit_metric(name, value, 'g', tags, hostname, device_name, timestamp)

    def increment(self, name, value=1, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 'c', tags, hostname, device_name)

    def decrement(self, name, value=-1, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 'c', tags, hostname, device_name)

    def rate(self, name, value, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, '_dd-r', tags, hostname, device_name)

    def histogram(self, name, value, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 'h', tags, hostname, device_name)

    def set(self, name, value, tags=None, hostname=None, device_name=None):
        self.submit_metric(name, value, 's', tags, hostname, device_name)

    def event(self, title, text, date_happened=None, alert_type=None, aggregation_key=None, source_type_name=None, priority=None, tags=None, hostname=None):
        event = {
            'title': title,
            'text': text,
        }
        if date_happened is not None:
            event['date_happened'] = date_happened
        else:
            event['date_happened'] = int(time())
        if alert_type is not None:
            event['alert_type'] = alert_type
        if aggregation_key is not None:
            event['aggregation_key'] = aggregation_key
        if source_type_name is not None:
            event['source_type_name'] = source_type_name
        if priority is not None:
            event['priority'] = priority
        if tags is not None:
            event['tags'] = sorted(tags)
        if hostname is not None:
            event['host'] = hostname
        else:
            event['host'] = self.hostname

        self.events.append(event)

    def flush(self):
        timestamp = time()
        expiry_timestamp = timestamp - self.expiry_seconds

        # Flush points and remove expired metrics. We mutate this dictionary
        # while iterating so don't use an iterator.
        metrics = []
        for context, metric in self.metrics.items():
            if metric.last_sample_time < expiry_timestamp:
                log.debug("%s hasn't been submitted in %ss. Expiring." % (context, self.expiry_seconds))
                del self.metrics[context]
            else:
                metrics += metric.flush(timestamp, self.interval)

        # Log a warning regarding metrics with old timestamps being submitted
        if self.num_discarded_old_points > 0:
            log.warn('%s points were discarded as a result of having an old timestamp' % self.num_discarded_old_points)
            self.num_discarded_old_points = 0

        # Save some stats.
        log.debug("received %s payloads since last flush" % self.count)
        self.total_count += self.count
        self.count = 0
        return metrics

    def flush_events(self):
        events = self.events
        self.events = []

        self.total_count += self.event_count
        self.event_count = 0

        log.debug("Received %d events since last flush" % len(events))

        return events

    def send_packet_count(self, metric_name):
        self.submit_metric(metric_name, self.count, 'g')

def api_formatter(metric, value, timestamp, tags, hostname, device_name=None):

    # Workaround for a bug in minjson serialization
    # (https://github.com/DataDog/dd-agent/issues/422)
    if tags is not None and isinstance(tags, tuple) and len(tags) == 1:
        tags = list(tags)
    return {
        'metric': metric,
        'points': [(timestamp, value)],
        'tags': tags,
        'host': hostname,
        'device_name': device_name
    }
<|endoftext|>


<commit_before>#!/usr/bin/env python

from __future__ import print_function
from future.standard_library import install_aliases
install_aliases()

from urllib.parse import urlparse, urlencode
from urllib.request import urlopen, Request
from urllib.error import HTTPError

import json
import os

from flask import Flask
from flask import request
from flask import make_response

# Flask app should start in global layout
app = Flask(__name__)

@app.route('/webhook', methods=['POST'])
def webhook():
    req = request.get_json(silent=True, force=True)

    print("Request:")
    print(json.dumps(req, indent=4))

    res = processRequest(req)

    res = json.dumps(res, indent=4)
    print(res)
    r = make_response(res)
    r.headers['Content-Type'] = 'application/json'
    return r

def processRequest(req):
    if req.get("result").get("action") != "alarmCount":
        return {}
    baseurl = "https://dl.dropboxusercontent.com/s/l3kdftm9jlq479w/db.json?dl=0"
    result = urlopen(baseurl).read().decode()
    data = json.loads(result)
    parameters = req.get("result").get("parameters")
    data = filterResult(data,parameters)
    res = makeWebhookResult(data,parameters)
    return res

def filterResult(data,parameters):
    sourceState = parameters.get("source-state")
    ackState = parameters.get("ack-state")
    priority = parameters.get("priority")
    if sourceState == '' and ackState == '' and priority == '':
        return data
    result = []
    if sourceState != '':
        for alarm in data.get('alarms'):
            if alarm.get("Source State").lower() == sourceState.lower():
                result.append(alarm)
        data['alarms'] = result
    result = []
    if ackState != '':
        for alarm in data.get('alarms'):
            if alarm.get("Ack State").lower() == ackState.lower():
                result.append(alarm)
        data['alarms'] = result
    result = []
    if priority != '':
        for alarm in data.get('alarms'):
            if alarm.get("Alarm Class").split('_')[0].lower().startswith(priority.split()[0].lower()):
                result.append(alarm)
        data['alarms'] = result
    return data

def makeWebhookResult(data,parameters):
    alarms = data.get('alarms')
    if alarms is None:
        return {}
    # print(json.dumps(item, indent=4))
    sourceState = parameters.get("source-state")
    ackState = parameters.get("ack-state")
    priority = parameters.get("priority")
    speech = "There"
    count = len(alarms)
    if count == 1:
        speech+=" is 1 alarm"
    else:
        speech+=" are "+str(count)+" alarms"
    if not sourceState and not ackState and not priority:
        speech+=" in total"
    else:
        speech+=" having "
        joint = " "
        count = 0

        if sourceState:
            if count>0:
                joint=" and "
            speech+=joint+"source state as "+sourceState.upper()
            count+=1
        if ackState:
            if count>0:
                joint=" and "
            speech+=joint+"acknowledgement state as "+ackState.upper()
            count+=1
        if priority:
            if count>0:
                joint=" and "
            speech+=joint+"priority as "+priority.split()[0].upper()
            count+=1

    # print("Response:")
    # print(speech)

    return {
        "speech": speech,
        "displayText": speech,
        # "data": data,
        # "contextOut": [],
        "source": "random JSON data"
    }

if __name__ == '__main__':
    port = int(os.getenv('PORT', 5000))

    print("Starting app on port %d" % port)

    app.run(debug=False, port=port, host='0.0.0.0')
<commit_msg>Update alarmCount.py<commit_after>#!/usr/bin/env python

from __future__ import print_function
from future.standard_library import install_aliases
install_aliases()

from urllib.parse import urlparse, urlencode
from urllib.request import urlopen, Request
from urllib.error import HTTPError

import json
import os

from flask import Flask
from flask import request
from flask import make_response

# Flask app should start in global layout
app = Flask(__name__)

@app.route('/webhook', methods=['POST'])
def webhook():
    req = request.get_json(silent=True, force=True)

    print("Request:")
    print(json.dumps(req, indent=4))
    res = processRequest(req)
    res = json.dumps(res, indent=4)
    print(res)
    r = make_response(res)
    r.headers['Content-Type'] = 'application/json'
    return r

def text2int(textnum, numwords={}):
    if not numwords:
        units = [
        "zero", "one", "two", "three", "four", "five", "six", "seven", "eight",
        "nine", "ten", "eleven", "twelve", "thirteen", "fourteen", "fifteen",
        "sixteen", "seventeen", "eighteen", "nineteen",
      ]
        tens = ["", "", "twenty", "thirty", "forty", "fifty", "sixty", "seventy", "eighty", "ninety"]
        scales = ["hundred", "thousand", "million", "billion", "trillion"]
        numwords["and"] = (1, 0)
        for idx, word in enumerate(units):    numwords[word] = (1, idx)
        for idx, word in enumerate(tens):     numwords[word] = (1, idx * 10)
        for idx, word in enumerate(scales):   numwords[word] = (10 ** (idx * 3 or 2), 0)
    current = result = 0
    for word in textnum.split():
        scale, increment = numwords[word]
        current = current * scale + increment
        if scale > 100:
            result += current
            current = 0
    return result + current

def ord2int(textnum):
    ordinal={
        'first':1,
        'second':2,
        'third':3,
        'fourth':4,
        'fifth':5,
        'sixth':6,
        'seventh':7,
        'eighth':8,
        'ninth':9,
        'tenth':10,
        'eleventh':11,
        'twelfth':12,
        'thirteenth':13,
        'fourteenth':14,
        'fifteenth':15,
        'sixteenth':16,
        'seventeenth':17,
        'eighteenth':18,
        'nineteenth':19,
        'twentieth':20,
    }
    return ordinal(textnum)

def processRequest(req):
    baseurl = "https://dl.dropboxusercontent.com/s/l3kdftm9jlq479w/db.json?dl=0"
    result = urlopen(baseurl).read().decode()
    data = json.loads(result)
    parameters = req.get("result").get("parameters")
    data = filterResult(data,parameters)
    res = makeWebhookResult(data,parameters)
    return res

def filterResult(data,parameters):
    sourceState = parameters.get("source-state")
    ackState = parameters.get("ack-state")
    priority = parameters.get("priority")
    numberWord = parameters.get("number")
    # ordinal = parameters.get("ordinal")
    if sourceState == '' and ackState == '' and priority == '':
        return data
    result = []
    if sourceState != '':
        for alarm in data.get('alarms'):
            if alarm.get("Source State").lower() == sourceState.lower():
                result.append(alarm)
        data['alarms'] = result
    result = []
    if ackState != '':
        for alarm in data.get('alarms'):
            if alarm.get("Ack State").lower() == ackState.lower():
                result.append(alarm)
        data['alarms'] = result
    result = []
    if priority != '':
        for alarm in data.get('alarms'):
            if alarm.get("Alarm Class").split('_')[0].lower().startswith(priority.split()[0].lower()):
                result.append(alarm)
        data['alarms'] = result
    result = []
    if numberWord != '':
        if numberWord.lower() == 'one':
            if ordinal == '' or ordinal == 'first':
                result = data.get('alarms')[0]
            elif ordinal.lower() == 'last':
                result = data.get('alarms')[-1]
            else:
                number = ord2int(ordinal)
                data['alarms'] = data.get('alarms')[0:number]
        else:
            number = text2int(numberWord)
            if ordinal == '':
                if number < len(data):
                    data['alarms'] = data.get('alarms')[number-1]
            elif ordinal.lower() == 'first':
                if number < len(data):
                    data['alarms'] = data.get('alarms')[0:number]
            elif ordinal.lower() == 'last':
                if number < len(data):
                    data['alarms'] = data.get('alarms')[number-1:len(data)]
    elif ordinal != '':
        if ordinal == 'last':
            data['alarms'] = data.get('alarms')[-1]
        else:
            number = ord2int(ordinal)
            data['alarms'] = data.get('alarms')[number-1]
    return data

def makeWebhookResult(data,parameters):
    alarms = data.get('alarms')
    if alarms is None:
        return {}
    # print(json.dumps(item, indent=4))
    sourceState = parameters.get("source-state")
    ackState = parameters.get("ack-state")
    priority = parameters.get("priority")
    numberWord = parameters.get("number")
    ordinal = parameters.get("ordinal")
    if req.get("result").get("action") == "alarmCount":
        speech = "There"
        count = len(alarms)
        if count == 1:
            speech+=" is 1 alarm"
        else:
            speech+=" are "+str(count)+" alarms"
        if not sourceState and not ackState and not priority:
            speech+=" in total"
        else:
            speech+=" having "
            joint = " "
            count = 0
            if sourceState:
                if count>0:
                    joint=" and "
                speech+=joint+"source state as "+sourceState.upper()
                count+=1
            if ackState:
                if count>0:
                    joint=" and "
                speech+=joint+"acknowledgement state as "+ackState.upper()
                count+=1
            if priority:
                if count>0:
                    joint=" and "
                speech+=joint+"priority as "+priority.split()[0].upper()
                count+=1

    elif req.get("result").get("action") == "fixAlarm":
        speech = "The following steps need to be followed to fix the "
        if ordinal != '':
            speech+=ordinal
        speech += " alarm "
        if numberWord != '' :
            speech+=numberWord+" "
        speech+='having '
        joint = " "
            count = 0
            if sourceState:
                if count>0:
                    joint=" and "
                speech+=joint+"source state as "+sourceState.upper()
                count+=1
            if ackState:
                if count>0:
                    joint=" and "
                speech+=joint+"acknowledgement state as "+ackState.upper()
                count+=1
            if priority:
                if count>0:
                    joint=" and "
                speech+=joint+"priority as "+priority.split()[0].upper()
                count+=1
        speech+=" "+data.get("alarms")[0].get("fix")

    # elif req.get("result").get("action") == "getAlarm":
    # elif req.get("result").get("action") == "getLatestAlarm":

    # print("Response:")
    # print(speech)

    return {
        "speech": speech,
        "displayText": speech,
        # "data": data,
        # "contextOut": [],
        "source": "random JSON data"
    }

if __name__ == '__main__':
    port = int(os.getenv('PORT', 5000))

    print("Starting app on port %d" % port)

app.run(debug=False, port=port, host='0.0.0.0')
<|endoftext|>


<commit_before>from PyQt4 import QtCore, QtGui

class ActionListWidget(QtGui.QListWidget):
    returnKeyPressed = QtCore.pyqtSignal()

    def __init__(self, *args):
        QtGui.QListWidget.__init__(self, *args)

    def keyPressEvent(self, key_event):
        if ((key_event.key()==QtCore.Qt.Key_Return) or (key_event.key()==QtCore.Qt.Key_Enter)):
            self.returnKeyPressed.emit()
        else:
            return QtGui.QListWidget.keyPressEvent(self, key_event)<commit_msg>Tiny little whitespace tweak.<commit_after>from PyQt4 import QtCore, QtGui

class ActionListWidget(QtGui.QListWidget):
    returnKeyPressed = QtCore.pyqtSignal()

    def __init__(self, *args):
        QtGui.QListWidget.__init__(self, *args)

    def keyPressEvent(self, key_event):
        if ((key_event.key() == QtCore.Qt.Key_Return) or (key_event.key() == QtCore.Qt.Key_Enter)):
            self.returnKeyPressed.emit()
        else:
            return QtGui.QListWidget.keyPressEvent(self, key_event)<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
import os
import sys
from pathlib import Path
from subprocess import check_call, check_output, run

from rich.console import Console

AIRFLOW_SOURCES_DIR = Path(__file__).parents[3].absolute()
BREEZE_IMAGES_DIR = AIRFLOW_SOURCES_DIR / "images" / "breeze"
BREEZE_SOURCES_DIR = AIRFLOW_SOURCES_DIR / "dev" / "breeze" / "src"

SCREENSHOT_WIDTH = "120"

def get_command_list():
    comp_env = os.environ.copy()
    comp_env['COMP_WORDS'] = ""
    comp_env['COMP_CWORD'] = "0"
    comp_env['_BREEZE_COMPLETE'] = 'bash_complete'
    result = check_output('breeze', env=comp_env, text=True)
    return [x.split(",")[1] for x in result.splitlines(keepends=False)]

def print_help_for_all_commands():
    console = Console(width=int(SCREENSHOT_WIDTH), color_system="standard")
    env = os.environ.copy()
    env['AIRFLOW_SOURCES_ROOT'] = str(AIRFLOW_SOURCES_DIR)
    env['RECORD_BREEZE_WIDTH'] = SCREENSHOT_WIDTH
    env['RECORD_BREEZE_TITLE'] = "Breeze commands"
    env['RECORD_BREEZE_OUTPUT_FILE'] = str(BREEZE_IMAGES_DIR / "output-commands.svg")
    env['TERM'] = "xterm-256color"
    env['PYTHONPATH'] = str(BREEZE_SOURCES_DIR)
    new_hash = check_output(
        [
            sys.executable,
            str(BREEZE_SOURCES_DIR / "airflow_breeze" / "breeze.py"),
            "command-hash-export",
            "-",
        ],
        env=env,
        text=True,
    )
    hash_file_path = BREEZE_IMAGES_DIR / "output-commands-hash.txt"
    try:
        old_hash = hash_file_path.read_text()
    except FileNotFoundError:
        old_hash = ""
    if old_hash == new_hash:
        console.print(f"[bright_blue]Skip generation of SVG images as command hash is unchanged {old_hash}")
        return
    if run(["breeze", "--help"], check=False).returncode != 0:
        console.print("[red]ERROR! You need to install breeze with pipx to run this pre-commit[/]")
        console.print("\n[bright_blue]Run this command:[/]\n")
        console.print("        pip install -e ./dev/breeze --force\n")
        sys.exit(1)
    env = os.environ.copy()
    env['AIRFLOW_SOURCES_ROOT'] = str(AIRFLOW_SOURCES_DIR)
    env['RECORD_BREEZE_WIDTH'] = SCREENSHOT_WIDTH
    env['RECORD_BREEZE_TITLE'] = "Breeze commands"
    env['RECORD_BREEZE_OUTPUT_FILE'] = str(BREEZE_IMAGES_DIR / "output-commands.svg")
    env['TERM'] = "xterm-256color"
    check_call(["breeze", "--help"], env=env)
    for command in get_command_list():
        env = os.environ.copy()
        env['AIRFLOW_SOURCES_ROOT'] = str(AIRFLOW_SOURCES_DIR)
        env['RECORD_BREEZE_WIDTH'] = SCREENSHOT_WIDTH
        env['RECORD_BREEZE_TITLE'] = f"Command: {command}"
        env['RECORD_BREEZE_OUTPUT_FILE'] = str(BREEZE_IMAGES_DIR / f"output-{command}.svg")
        env['TERM'] = "xterm-256color"
        check_call(["breeze", command, "--help"], env=env)
    hash_file_path.write_text(new_hash)

def verify_all_commands_described_in_docs():
    console = Console(width=int(SCREENSHOT_WIDTH), color_system="standard")
    errors = []
    doc_content = (AIRFLOW_SOURCES_DIR / "BREEZE.rst").read_text()
    for file in os.listdir(BREEZE_IMAGES_DIR):
        if file.startswith("output-") and file.endswith(".svg"):
            command = file[len("output-") : -len(".svg")]
            if command == "breeze-commands":
                continue
            if file not in doc_content:
                errors.append(command)
            else:
                console.print(f"[green]OK. The {command} screenshot is embedded in BREEZE.rst.")
    if errors:
        console.print("[red]Some of Breeze commands are not described in BREEZE.rst:[/]")
        for command in errors:
            console.print(f"  * [red]{command}[/]")
        console.print()
        console.print(
            "[bright_yellow]Make sure you describe it and embed ./images/breeze/output-<COMMAND>.svg "
            "screenshot as image in the BREEZE.rst file.[/]"
        )
        sys.exit(1)

if __name__ == '__main__':
    print_help_for_all_commands()
    verify_all_commands_described_in_docs()
<commit_msg>Make Breeze help generation indepdent from having breeze installed (#23612)<commit_after>#!/usr/bin/env python
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
import os
import sys
from pathlib import Path
from subprocess import check_call, check_output, run

from rich.console import Console

AIRFLOW_SOURCES_DIR = Path(__file__).parents[3].absolute()
BREEZE_IMAGES_DIR = AIRFLOW_SOURCES_DIR / "images" / "breeze"
BREEZE_INSTALL_DIR = AIRFLOW_SOURCES_DIR / "dev" / "breeze"
BREEZE_SOURCES_DIR = BREEZE_INSTALL_DIR / "src"

SCREENSHOT_WIDTH = "120"

def get_command_list():
    comp_env = os.environ.copy()
    comp_env['COMP_WORDS'] = ""
    comp_env['COMP_CWORD'] = "0"
    comp_env['_BREEZE_COMPLETE'] = 'bash_complete'
    result = check_output('breeze', env=comp_env, text=True)
    return [x.split(",")[1] for x in result.splitlines(keepends=False)]

def print_help_for_all_commands():
    console = Console(width=int(SCREENSHOT_WIDTH), color_system="standard")
    env = os.environ.copy()
    env['AIRFLOW_SOURCES_ROOT'] = str(AIRFLOW_SOURCES_DIR)
    env['RECORD_BREEZE_WIDTH'] = SCREENSHOT_WIDTH
    env['RECORD_BREEZE_TITLE'] = "Breeze commands"
    env['RECORD_BREEZE_OUTPUT_FILE'] = str(BREEZE_IMAGES_DIR / "output-commands.svg")
    env['TERM'] = "xterm-256color"
    env['PYTHONPATH'] = str(BREEZE_SOURCES_DIR)
    new_hash = check_output(
        [
            sys.executable,
            str(BREEZE_SOURCES_DIR / "airflow_breeze" / "breeze.py"),
            "command-hash-export",
            "-",
        ],
        env=env,
        text=True,
    )
    hash_file_path = BREEZE_IMAGES_DIR / "output-commands-hash.txt"
    try:
        old_hash = hash_file_path.read_text()
    except FileNotFoundError:
        old_hash = ""
    if old_hash == new_hash:
        console.print(f"[bright_blue]Skip generation of SVG images as command hash is unchanged {old_hash}")
        return
    run([sys.executable, "-m", "pip", "install", "--upgrade", "-e", BREEZE_INSTALL_DIR])
    env = os.environ.copy()
    env['AIRFLOW_SOURCES_ROOT'] = str(AIRFLOW_SOURCES_DIR)
    env['RECORD_BREEZE_WIDTH'] = SCREENSHOT_WIDTH
    env['RECORD_BREEZE_TITLE'] = "Breeze commands"
    env['RECORD_BREEZE_OUTPUT_FILE'] = str(BREEZE_IMAGES_DIR / "output-commands.svg")
    env['TERM'] = "xterm-256color"
    check_call(["breeze", "--help"], env=env)
    for command in get_command_list():
        env = os.environ.copy()
        env['AIRFLOW_SOURCES_ROOT'] = str(AIRFLOW_SOURCES_DIR)
        env['RECORD_BREEZE_WIDTH'] = SCREENSHOT_WIDTH
        env['RECORD_BREEZE_TITLE'] = f"Command: {command}"
        env['RECORD_BREEZE_OUTPUT_FILE'] = str(BREEZE_IMAGES_DIR / f"output-{command}.svg")
        env['TERM'] = "xterm-256color"
        check_call(["breeze", command, "--help"], env=env)
    hash_file_path.write_text(new_hash)

def verify_all_commands_described_in_docs():
    console = Console(width=int(SCREENSHOT_WIDTH), color_system="standard")
    errors = []
    doc_content = (AIRFLOW_SOURCES_DIR / "BREEZE.rst").read_text()
    for file in os.listdir(BREEZE_IMAGES_DIR):
        if file.startswith("output-") and file.endswith(".svg"):
            command = file[len("output-") : -len(".svg")]
            if command == "breeze-commands":
                continue
            if file not in doc_content:
                errors.append(command)
            else:
                console.print(f"[green]OK. The {command} screenshot is embedded in BREEZE.rst.")
    if errors:
        console.print("[red]Some of Breeze commands are not described in BREEZE.rst:[/]")
        for command in errors:
            console.print(f"  * [red]{command}[/]")
        console.print()
        console.print(
            "[bright_yellow]Make sure you describe it and embed ./images/breeze/output-<COMMAND>.svg "
            "screenshot as image in the BREEZE.rst file.[/]"
        )
        sys.exit(1)

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


<commit_before>"""
Scan Client Tools

Helpers for connecting to the Scan Server,
to assemble commands for a scan,
submit them to the server,
monitor the execution.

Shortcuts for 1D, 2D scans.

This code depends on the basic org.csstudio.scan.*
packages and can be invoked from Jython command
lines outside of CSS.

@author: Kay Kasemir
"""

# -------------------------------------------------------
# Path setup

# For now this is hard coded in here to allow use
# from BOY scripts
import sys
# Path to workspace, when running within IDE
workspace="/Kram/MerurialRepos/cs-studio/products/SNS/plugins/"
sys.path.append(workspace + "org.csstudio.scan/bin")
sys.path.append(workspace + "org.csstudio.scan.client/bin")

# Path to binaries, when running within exported product
install="/Users/Fred/Desktop/CSS/plugins/"
sys.path.append(install + "org.csstudio.scan")
sys.path.append(install + "org.csstudio.scan.client")

# -------------------------------------------------------

import org.csstudio.scan.client.ScanServerConnector as ScanServerConnector
from org.csstudio.scan.command import *
import time

class ScanClient(object):
    """
    Base class for a scan client

    Can submit scans to the server and monitor them
    """
    def __init__(self):
        # Connection to the scan server
        self.server = ScanServerConnector.connect()
        # Scan ID
        self.id = -1

    def submit(self, name, command_sequence):
        """
        Submit a CommandSequence to the server for execution

        @param name  Name of the scan
        @param command_sequence  CommandSequence

        @return Scan ID
        """
        self.id = self.server.submitScan(name, command_sequence.getCommands())
        return self.id

    def waitUntilDone(self, id=-1):
        """
        Wait until a submitted scan has finished

        @param id Scan ID, defaulting to the last submitted scan
        """
        if id == -1:
        	id = self.id
        while True:
            info = self.server.getScanInfo(id)
            print info
            if info.isDone():
                break;
            time.sleep(1.0)

    def __str__(self):
        return "Scan client, connected to %s" % self.server.getInfo()

class ScanNd(ScanClient):
    """
    N-dimensional scan that logs arbitrary number of readings
    based on nested loops.

    Arguments:

    * Optional scan name
    * One or more scan specifications: ('device', start, end[, step])
    * Names of device to log in addition to loop'ed devices

    Examples:

    # Scan 'xpos' from 1 to 10, stepping 1, automatically logging 'xpos'
    scan('My first one', ('xpos', 1, 10) )
    scan( ('xpos', 1, 10) )

    # ... also 'readback'
    scan( ('xpos', 1, 10), 'readback')

    # Scan 'xpos' from 1 to 10, stepping 1,
    # inside that looping 'ypos' from 1 to 5 by 0.2,
    # logging 'readback'

    scan('XY Example', ('xpos', 1, 10), ('ypos', 1, 5, 0.2), 'readback')
    """

    def __init__(self):
        ScanClient.__init__(self)

    def _decodeScan(self, parms):
        """ Check for
                ('device', start, end, step)
             or
                ('device', start, end)
             for a default step size of 1
        """
        if (len(parms) == 4):
            return (parms[0], parms[1], parms[2], parms[3])
        elif (len(parms) == 3):
            return (parms[0], parms[1], parms[2], 1)
        else:
            raise Exception('Scan parameters should be (''device'', start, end, step), not %s' % str(parms))

    def __call__(self, *args):
        """ N-dimensional scan command.
            @return ID of scan that was scheduled on the scan server
        """
        args = list(args)
        if len(args) > 0  and  isinstance(args[0], str):
            name = args[0]
            args.pop(0)
        else:
            name = "Scan"

        # Determine the (nested) scans and the devices to log
        # (doesn't really care if the log devices are listed last)
        scans = []
        log = []
        for arg in args:
            if isinstance(arg, tuple):
                scan = self._decodeScan(arg)
                scans.append(scan)
            else:
                log.append(arg)

        # Wrap by scans, going in reverse from inner loop
        scans.reverse()
        cmds = CommandSequence()
        for scan in scans:
            body = CommandSequence()
            # Add loop variable and log
            log.insert(0, scan[0])
            body.log(log)

            # Add previous cmds
            body.add(cmds)

            # Within loop
            cmds = CommandSequence()
            cmds.loop(scan[0], scan[1], scan[2], scan[3], body)

            log = []

        id = self.submit(name, cmds)
        if __name__ == '__main__':
            cmds.dump()
            self.waitUntilDone()
        return id

# Create 'scan' command
scan = ScanNd()

if __name__ == '__main__':
    print 'Welcome to the scan system'
    # print 'Running in %s' % os.getcwd()
    print 'Connected to %s' % scan.server.getInfo()

    # 'Normal' loops
    #scan('Normal 2D', ('xpos', 1, 10), ('ypos', 1, 10, 0.5), 'readback')

    # 'Reversing' inner loop
    #scan('Reversing 2D', ('xpos', 1, 10), ('ypos', 1, 10, -0.5), 'readback')

<commit_msg>SNS scan system: Update script to XML-submission of scans<commit_after>"""
Scan Client Tools

Helpers for connecting to the Scan Server,
to assemble commands for a scan,
submit them to the server,
monitor the execution.

Shortcuts for 1D, 2D scans.

This code depends on the basic org.csstudio.scan.*
packages and can be invoked from Jython command
lines outside of CSS.

@author: Kay Kasemir
"""

# -------------------------------------------------------
# Path setup

# For now this is hard coded in here to allow use
# from BOY scripts
import sys
# Path to workspace, when running within IDE
workspace="/Kram/MerurialRepos/cs-studio/products/SNS/plugins/"
sys.path.append(workspace + "org.csstudio.scan/bin")
sys.path.append(workspace + "org.csstudio.scan.client/bin")

# Path to binaries, when running within exported product
install="/Users/Fred/Desktop/CSS/plugins/"
sys.path.append(install + "org.csstudio.scan")
sys.path.append(install + "org.csstudio.scan.client")

# -------------------------------------------------------

import org.csstudio.scan.client.ScanServerConnector as ScanServerConnector
from org.csstudio.scan.command import *
import time

class ScanClient(object):
    """
    Base class for a scan client

    Can submit scans to the server and monitor them
    """
    def __init__(self):
        # Connection to the scan server
        self.server = ScanServerConnector.connect()
        # Scan ID
        self.id = -1

    def submit(self, name, command_sequence):
        """
        Submit a CommandSequence to the server for execution

        @param name  Name of the scan
        @param command_sequence  CommandSequence

        @return Scan ID
        """
        self.id = self.server.submitScan(name, command_sequence.getXML())
        return self.id

    def waitUntilDone(self, id=-1):
        """
        Wait until a submitted scan has finished

        @param id Scan ID, defaulting to the last submitted scan
        """
        if id == -1:
        	id = self.id
        while True:
            info = self.server.getScanInfo(id)
            print info
            if info.isDone():
                break;
            time.sleep(1.0)

    def __str__(self):
        return "Scan client, connected to %s" % self.server.getInfo()

class ScanNd(ScanClient):
    """
    N-dimensional scan that logs arbitrary number of readings
    based on nested loops.

    Arguments:

    * Optional scan name
    * One or more scan specifications: ('device', start, end[, step])
    * Names of device to log in addition to loop'ed devices

    Examples:

    # Scan 'xpos' from 1 to 10, stepping 1, automatically logging 'xpos'
    scan('My first one', ('xpos', 1, 10) )
    scan( ('xpos', 1, 10) )

    # ... also 'readback'
    scan( ('xpos', 1, 10), 'readback')

    # Scan 'xpos' from 1 to 10, stepping 1,
    # inside that looping 'ypos' from 1 to 5 by 0.2,
    # logging 'readback'

    scan('XY Example', ('xpos', 1, 10), ('ypos', 1, 5, 0.2), 'readback')
    """

    def __init__(self):
        ScanClient.__init__(self)

    def _decodeScan(self, parms):
        """ Check for
                ('device', start, end, step)
             or
                ('device', start, end)
             for a default step size of 1
        """
        if (len(parms) == 4):
            return (parms[0], parms[1], parms[2], parms[3])
        elif (len(parms) == 3):
            return (parms[0], parms[1], parms[2], 1)
        else:
            raise Exception('Scan parameters should be (''device'', start, end, step), not %s' % str(parms))

    def __call__(self, *args):
        """ N-dimensional scan command.
            @return ID of scan that was scheduled on the scan server
        """
        args = list(args)
        if len(args) > 0  and  isinstance(args[0], str):
            name = args[0]
            args.pop(0)
        else:
            name = "Scan"

        # Determine the (nested) scans and the devices to log
        # (doesn't really care if the log devices are listed last)
        scans = []
        log = []
        for arg in args:
            if isinstance(arg, tuple):
                scan = self._decodeScan(arg)
                scans.append(scan)
            else:
                log.append(arg)

        # Wrap by scans, going in reverse from inner loop
        scans.reverse()
        cmds = CommandSequence()
        for scan in scans:
            body = CommandSequence()
            # Add loop variable and log
            log.insert(0, scan[0])
            body.log(log)

            # Add previous cmds
            body.add(cmds)

            # Within loop
            cmds = CommandSequence()
            cmds.loop(scan[0], scan[1], scan[2], scan[3], body)

            log = []

        id = self.submit(name, cmds)
        if __name__ == '__main__':
            cmds.dump()
            self.waitUntilDone()
        return id

# Create 'scan' command
scan = ScanNd()

if __name__ == '__main__':
    print 'Welcome to the scan system'
    # print 'Running in %s' % os.getcwd()
    print 'Connected to %s' % scan.server.getInfo()

    # 'Normal' loops
    #scan('Normal 2D', ('xpos', 1, 10), ('ypos', 1, 10, 0.5), 'readback')

    # 'Reversing' inner loop
    #scan('Reversing 2D', ('xpos', 1, 10), ('ypos', 1, 10, -0.5), 'readback')

<|endoftext|>