<commit_before>
import tensorflow as tf
from ale_python_interface import ALEInterface
import threading
from skimage.transform import resize
import numpy as np
from collections import deque
from model import build_network, loss
from environment import Environment

ACTIONS = 4
NUM_CONCURRENT = 1
GAMMA = 0.99
NETWORK_UPDATE_FREQUENCY = 5
TARGET_NETWORK_UPDATE_FREQUENCY = 10000

LEARNING_RATE = 0.00025
GRADIENT_MOMENTUM = 0.95
SQUARED_GRADIENT_MOMENTUM = 0.95
MIN_SQUARED_GRADIENT = 0.01

NOOP_MAX = 30
RESIZED_HEIGHT = 84
RESIZED_WIDTH = 84
AGENT_HISTORY_LENGTH = 4

INITIAL_EXPLORATION = 1.0
FINAL_EXPLORATION = 0.1
FINAL_EXPLORATION_FRAME = 1000000

class GlobalThread(object):
  """
  Main global thread responsible for:
    - Initializing shared global network and target network parameters
    - Kicking off actor-learner threads
  """
  def __init__(self, session, graph, ale_io_lock, lock):
    self._session = session
    self._graph = graph
    self._ale_io_lock = ale_io_lock
    self.build_graph()
    self.rng = np.random.RandomState(123456)
    self._session.run(self._reset_target_network_params)
    self.lock = lock

  def build_graph(self):
    """ Placeholder for the function that builds up the model
    """
    # Input Placeholders
    state = tf.placeholder("float", [None, 84, 84, 4]) # Previous agent_history_length frames
    a = tf.placeholder("float", [None, ACTIONS]) # One hot vector representing chosen action at time t
    y = tf.placeholder("float", [None]) # Float holding y_t, the target value for chosen action at time t
    self._state=state
    self._a = a
    self._y = y

    # Set up online network and target network.
    with tf.variable_scope('network_params') as scope:
      network = build_network(state)
    with tf.variable_scope('target_network_params') as scope:
      target_network = build_network(state)
    self._network = network
    self._target_network = target_network

    # Set up loss
    cost = loss(network, a, y)
    self._cost = cost

    # Op for periodically updating target network with online network weights
    network_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                     "network_params")
    target_network_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                        "target_network_params")
    self._reset_target_network_params = [target_network_params[i].assign(network_params[i]) for i in range(len(target_network_params))]

    # One gradients list per thread
    for i in range(NUM_CONCURRENT):
      with tf.variable_scope('grads_'+str(i)) as scope:
        grad_vars = [tf.Variable(np.zeros(var.get_shape().as_list(), dtype=np.float32), trainable=False) for var in network_params]

    # One optimizer per thread
    optimizers = []
    for i in range(NUM_CONCURRENT):
        optimizers.append(tf.train.AdamOptimizer(LEARNING_RATE))
    self._optimizers = optimizers

    # One compute_gradients op per thread
    compute_gradients_ops = []
    for i in range(NUM_CONCURRENT):
        compute_gradients_ops.append(self._optimizers[i].compute_gradients(cost, var_list=network_params))
    self._compute_gradients_ops = compute_gradients_ops
    self._grads_and_vars = compute_gradients_ops

    # One gradient assign add (+=) op per thread
    assign_add_gradients_ops = []
    for i in range(NUM_CONCURRENT):
        namespace = 'grads_'+str(i)
        grad_vars = tf.get_collection(tf.GraphKeys.VARIABLES, namespace)
        assign_add_gradients_ops.append([tf.assign_add(v,g) for v,g in zip(grad_vars, [grad[0] for grad in self._compute_gradients_ops[i]])])
    self._assign_add_gradients_ops = assign_add_gradients_ops

    # "Async apply gradients" consists of two operations:
    # 1) Copy thread-local gradients to global thread's gradient placeholders
    # 2) Have global thread call apply_gradients on it's network_params using the
    #    gradients currently in the placeholders

    # Set up placeholder gradients (threads will copy their local gradients to these placeholders, then call apply)
    placeholder_gradients = []
    for var in network_params:
      variable = tf.Variable(np.zeros(var.get_shape().as_list(), dtype=np.float32), trainable=False)
      tensor = variable.value()
      placeholder_gradients.append((tensor, variable))
    self._placeholder_gradients = placeholder_gradients

    # One copy gradients op per thread
    copy_gradients_ops = []
    for i in range(NUM_CONCURRENT):
        namespace = 'grads_'+str(i)
        grad_vars = tf.get_collection(tf.GraphKeys.VARIABLES, namespace) # list of just variables
        copy_gradients_i = [placeholder_gradients[i][1].assign(grad_vars[i]) for i in range(len(placeholder_gradients))]
        copy_gradients_ops.append(copy_gradients_i)
    self._copy_gradients_ops = copy_gradients_ops

    # Set up global optimizer (just to apply gradients from threads to shared network_params)
    apply_grads_optimizer = tf.train.AdamOptimizer(LEARNING_RATE)
    # Op for applying the gradients
    apply_gradients = apply_grads_optimizer.apply_gradients(placeholder_gradients)
    self._apply_gradients = apply_gradients

    # Zero out gradients
    clear_gradients_ops = []
    for i in range(NUM_CONCURRENT):
        namespace = 'grads_'+str(i)
        grad_vars = tf.get_collection(tf.GraphKeys.VARIABLES, namespace)
        clear_gradients_ops.append([tf.assign(grad,np.zeros(grad.get_shape().as_list(), dtype=np.float32)) for grad in grad_vars])
    self._clear_gradients_ops = clear_gradients_ops

    # Global counter variable
    T = tf.Variable(0)
    self._increment_T = T.assign_add(1)
    self.T = T

    tf.initialize_all_variables().run()
    print("VARIABLES INITIALIZED")

  def async_update_gradients(self, thread_id):
    """
    Called by an actor-learner thread.
    Copies the actor-learner's accumulated gradients to
    the global thread's placeholder gradients, runs the
    apply_gradients op.
    """
    self._session.run(self._copy_gradients_ops[thread_id])
    self._session.run(self._apply_gradients)

  def _actor_learner_thread(self, i):
    """
    Main actor learner thread:
    - Initialize thread-specific environment
    - Run an infinite training loop, periodically
      sending asyncronous gradient updates to main model
    """
    # Initialize this thread's agent's environment
    environment = Environment(self._ale_io_lock)

    # Get initial game state
    state = environment.get_initial_state()

    epsilon = INITIAL_EXPLORATION

    # Main learning loop
    T_max = 1000000
    t = 0
    episode_reward = 0
    episode_max_q_value = 0
    episode_steps = 0
    while self._session.run(self.T) < T_max:
      # Get Q(s,a) values
      Q_s_a = self._session.run(self._network, feed_dict={self._state : state })

      # Take action a according to the e-greedy policy
      if epsilon > FINAL_EXPLORATION:
          epsilon -= (INITIAL_EXPLORATION - FINAL_EXPLORATION) / FINAL_EXPLORATION_FRAME
      if self.rng.rand() < epsilon:
        action = self.rng.randint(0, environment.num_actions())
      else:
        action =  np.argmax(Q_s_a)

      # Execute the chosen action in the environment, observe new state and reward.
      (reward, new_state, terminal) = environment.execute(action)
      reward = np.clip(reward, -1, 1)

      # Compute y using target network
      target_Q_s_a = self._session.run(self._target_network, feed_dict={self._state : state })
      if terminal:
        y = reward
      else:
        y = reward + GAMMA * np.max(target_Q_s_a)

      # Build a one hot action vector
      actions_one_hot = np.zeros(environment.num_actions())
      actions_one_hot[action]=1.0

      # Calculate loss
      loss = self._session.run(self._cost, feed_dict={self._state : state, self._a: [actions_one_hot], self._y: [y]})

      # Accumulate gradients
      self._session.run(self._assign_add_gradients_ops[i], feed_dict={self._state : state, self._a: [actions_one_hot], self._y: [y]})

      # Collect some stats (TODO: put in tensorflow summaries)
      episode_reward += reward
      episode_max_q_value += np.max(Q_s_a)
      episode_steps += 1

      # s = s'
      if not terminal:
        state = new_state
      else:
        episode_average_reward = episode_reward / float(episode_steps)
        episode_average_max_q_value = episode_max_q_value / float(episode_steps)
        print "Reward:", episode_average_reward, "Max Q: ", episode_average_max_q_value
        episode_reward = 0
        episode_max_q_value = 0
        episode_steps = 0
        # Episode ended, so next state is the next episode's initial state
        state = environment.get_initial_state()

      # T += 1
      self._session.run(self._increment_T)
      T = self._session.run(self.T)
      t += 1

      # Update the target network
      if T % TARGET_NETWORK_UPDATE_FREQUENCY == 0:
        self._session.run(self._reset_target_network_params)

      # Perform asynchronous update
      if (t % NETWORK_UPDATE_FREQUENCY == 0) or terminal:
        self.async_update_gradients(i)
        self._session.run(self._clear_gradients_ops[i])

    print("=======================================done==============================")

  def train(self):
    """
    Creates and kicks off num_concurrent actor-learner threads
    """
    workers = []
    for thread_id in xrange(NUM_CONCURRENT):
      t = threading.Thread(target=self._actor_learner_thread, args=(thread_id,))
      t.start()
      workers.append(t)
    for t in workers:
      t.join()

def main(_):
  g = tf.Graph()
  with g.as_default(), tf.Session() as session:
    with tf.device("/cpu:0"):
      ale_io_lock = threading.Lock() # using a lock to avoid race condition on ALE init
      lock = threading.Lock()
      global_thread = GlobalThread(session, g, ale_io_lock, lock)
      global_thread.train()

if __name__ == "__main__":
  tf.app.run()

<commit_msg>actually apply gradients<commit_after>
import tensorflow as tf
from ale_python_interface import ALEInterface
import threading
from skimage.transform import resize
import numpy as np
from collections import deque
from model import build_network, loss
from environment import Environment

ACTIONS = 4
NUM_CONCURRENT = 4 # 16
GAMMA = 0.99
NETWORK_UPDATE_FREQUENCY = 5
TARGET_NETWORK_UPDATE_FREQUENCY = 10000

LEARNING_RATE = 0.00025
GRADIENT_MOMENTUM = 0.95
SQUARED_GRADIENT_MOMENTUM = 0.95
MIN_SQUARED_GRADIENT = 0.01

NOOP_MAX = 30
RESIZED_HEIGHT = 84
RESIZED_WIDTH = 84
AGENT_HISTORY_LENGTH = 4

INITIAL_EXPLORATION = 1.0
FINAL_EXPLORATION = 0.1
FINAL_EXPLORATION_FRAME = 1000000

class GlobalThread(object):
  """
  Main global thread responsible for:
    - Initializing shared global network and target network parameters
    - Kicking off actor-learner threads
  """
  def __init__(self, session, graph, ale_io_lock, lock):
    self._session = session
    self._graph = graph
    self._ale_io_lock = ale_io_lock
    self.build_graph()
    self.rng = np.random.RandomState(123456)
    self._session.run(self._reset_target_network_params)
    self.lock = lock

  def build_graph(self):
    """ Placeholder for the function that builds up the model
    """
    # Input Placeholders
    state = tf.placeholder("float", [None, 84, 84, 4]) # Previous agent_history_length frames
    a = tf.placeholder("float", [None, ACTIONS]) # One hot vector representing chosen action at time t
    y = tf.placeholder("float", [None]) # Float holding y_t, the target value for chosen action at time t
    self._state=state
    self._a = a
    self._y = y

    # Set up online network and target network.
    with tf.variable_scope('network_params') as scope:
      network = build_network(state)
    with tf.variable_scope('target_network_params') as scope:
      target_network = build_network(state)
    self._network = network
    self._target_network = target_network

    # Set up loss
    cost = loss(network, a, y)
    self._cost = cost

    # Op for periodically updating target network with online network weights
    network_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                     "network_params")
    target_network_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                        "target_network_params")
    self._reset_target_network_params = [target_network_params[i].assign(network_params[i]) for i in range(len(target_network_params))]

    # One gradients list per thread
    for i in range(NUM_CONCURRENT):
      with tf.variable_scope('grads_'+str(i)) as scope:
        grad_vars = [tf.Variable(np.zeros(var.get_shape().as_list(), dtype=np.float32), trainable=False) for var in network_params]

    # One optimizer per thread
    optimizers = []
    for i in range(NUM_CONCURRENT):
        optimizers.append(tf.train.AdamOptimizer(LEARNING_RATE))
    self._optimizers = optimizers

    # One gradient assign add (+=) op per thread
    # Assign_add computes gradients using the thread's optimizer,
    # then assigns the computed gradients to the thread's grad_vars
    assign_add_gradients_ops = []
    for i in range(NUM_CONCURRENT):
        namespace = 'grads_'+str(i)
        compute_gradients = self._optimizers[i].compute_gradients(cost, var_list=network_params)
        grad_vars = tf.get_collection(tf.GraphKeys.VARIABLES, namespace)
        assign_add_gradients_ops.append([tf.assign_add(v,g) for v,g in zip(grad_vars, [grad[0] for grad in compute_gradients])])
    self._assign_add_gradients_ops = assign_add_gradients_ops

    # Global optimizer (just to apply threads' gradients)
    apply_grads_optimizer = tf.train.AdamOptimizer(LEARNING_RATE)
    global_grads_and_vars = apply_grads_optimizer.compute_gradients(cost, var_list=network_params)
    self._global_grads_and_vars = global_grads_and_vars

    # Set up placeholder gradients
    placeholder_gradients = []
    for grad_var in global_grads_and_vars:
       placeholder_gradients.append((tf.placeholder('float', shape=grad_var[1].get_shape()) ,grad_var[1]))
    self._placeholder_gradients = placeholder_gradients

    # Op for applying the gradients
    apply_gradients = apply_grads_optimizer.apply_gradients(placeholder_gradients)
    self._async_apply_gradients = apply_gradients

    # Zero out gradients
    clear_gradients_ops = []
    for i in range(NUM_CONCURRENT):
        namespace = 'grads_'+str(i)
        grad_vars = tf.get_collection(tf.GraphKeys.VARIABLES, namespace)
        clear_gradients_ops.append([tf.assign(grad,np.zeros(grad.get_shape().as_list(), dtype=np.float32)) for grad in grad_vars])
    self._clear_gradients_ops = clear_gradients_ops

    # Global counter variable
    T = tf.Variable(0)
    self._increment_T = T.assign_add(1)
    self.T = T

    tf.initialize_all_variables().run()
    print("VARIABLES INITIALIZED")

  def _actor_learner_thread(self, i):
    """
    Main actor learner thread:
    - Initialize thread-specific environment
    - Run an infinite training loop, periodically
      sending asyncronous gradient updates to main model
    """
    # Initialize this thread's agent's environment
    environment = Environment(self._ale_io_lock)

    # Get initial game state
    state = environment.get_initial_state()

    epsilon = INITIAL_EXPLORATION

    # Main learning loop
    T_max = 1000000
    t = 0
    episode_reward = 0
    episode_max_q_value = 0
    episode_steps = 0
    while self._session.run(self.T) < T_max:
      # Get Q(s,a) values
      Q_s_a = self._session.run(self._network, feed_dict={self._state : state })

      # Take action a according to the e-greedy policy
      if epsilon > FINAL_EXPLORATION:
          epsilon -= (INITIAL_EXPLORATION - FINAL_EXPLORATION) / FINAL_EXPLORATION_FRAME
      if self.rng.rand() < epsilon:
        action = self.rng.randint(0, environment.num_actions())
      else:
        action =  np.argmax(Q_s_a)

      # Execute the chosen action in the environment, observe new state and reward.
      (reward, new_state, terminal) = environment.execute(action)
      episode_reward += reward
      reward = np.clip(reward, -1, 1)

      # Compute y using target network
      target_Q_s_a = self._session.run(self._target_network, feed_dict={self._state : state })
      if terminal:
        y = reward
      else:
        y = reward + GAMMA * np.max(target_Q_s_a)

      # Build a one hot action vector
      actions_one_hot = np.zeros(environment.num_actions())
      actions_one_hot[action]=1.0

      # Calculate loss
      # loss = self._session.run(self._cost, feed_dict={self._state : state, self._a: [actions_one_hot], self._y: [y]})

      # Accumulate gradients
      self._session.run(self._assign_add_gradients_ops[i], feed_dict={self._state : state, self._a: [actions_one_hot], self._y: [y]})

      # Collect some stats (TODO: put in tensorflow summaries)
      episode_max_q_value += np.max(Q_s_a)
      episode_steps += 1

      # s = s'
      if not terminal:
        state = new_state
      else:
        episode_average_reward = episode_reward #/ float(episode_steps)
        episode_average_max_q_value = episode_max_q_value / float(episode_steps)
        print "Reward:", round(episode_average_reward, 3), "Max Q: ", round(episode_average_max_q_value,3)
        episode_reward = 0
        episode_max_q_value = 0
        episode_steps = 0
        # Episode ended, so next state is the next episode's initial state
        state = environment.get_initial_state()

      # T += 1
      self._session.run(self._increment_T)
      T = self._session.run(self.T)
      t += 1

      # Update the target network
      if T % TARGET_NETWORK_UPDATE_FREQUENCY == 0:
        self._session.run(self._reset_target_network_params)

      # Perform asynchronous update
      if (t % NETWORK_UPDATE_FREQUENCY == 0) or terminal:
        grad_vals = self._session.run([grad for grad, _ in self._global_grads_and_vars], feed_dict={self._state : state, self._a: [actions_one_hot], self._y: [y]})

        # Build a variable name-to-gradient dictionary.
        var_to_grad = {}
        for j in range(len(grad_vals)):
            var = self._global_grads_and_vars[j][1]
            var_to_grad[var.name] = grad_vals[j]

        feed_dict = {}
        for j, (grad, var) in enumerate(self._global_grads_and_vars):
          feed_dict[self._placeholder_gradients[j][0]] = grad_vals[j]
        self._session.run(self._async_apply_gradients, feed_dict=feed_dict)

        self._session.run(self._clear_gradients_ops[i])

    print("=======================================done==============================")

  def train(self):
    """
    Creates and kicks off num_concurrent actor-learner threads
    """
    workers = []
    for thread_id in xrange(NUM_CONCURRENT):
      t = threading.Thread(target=self._actor_learner_thread, args=(thread_id,))
      t.start()
      workers.append(t)
    for t in workers:
      t.join()

def main(_):
  g = tf.Graph()
  with g.as_default(), tf.Session() as session:
    with tf.device("/cpu:0"):
      ale_io_lock = threading.Lock() # using a lock to avoid race condition on ALE init
      lock = threading.Lock()
      global_thread = GlobalThread(session, g, ale_io_lock, lock)
      global_thread.train()

if __name__ == "__main__":
  tf.app.run()

<|endoftext|>


<commit_before>import os
import contextlib
import nuke
import re
import logging
from collections import OrderedDict

from ..vendor import six

log = logging.getLogger(__name__)

@contextlib.contextmanager
def maintained_selection():
    previous_selection = nuke.selectedNodes()
    try:
        yield
    finally:
        # unselect all selection in case there is some
        current_seletion = nuke.selectedNodes()
        [n['selected'].setValue(False) for n in current_seletion]
        # and select all previously selected nodes
        if previous_selection:
            [n['selected'].setValue(True) for n in previous_selection]

def imprint(node, data, tab=None):
    """Store attributes with value on node

    Parse user data into Node knobs.
    Use `collections.OrderedDict` to ensure knob order.

    Args:
        node(nuke.Node): node object from Nuke
        data(dict): collection of attributes and their value

    Returns:
        None

    Examples:
        ```
        import nuke
        from avalon.nuke import lib

        node = nuke.createNode("NoOp")
        data = {
            # Regular type of attributes
            "myList": ["x", "y", "z"],
            "myBool": True,
            "myFloat": 0.1,
            "myInt": 5,

            # Creating non-default imprint type of knob
            "MyFilePath": lib.Knobby("File_Knob", "/file/path"),
            "divider": lib.Knobby("Text_Knob", ""),

            # Manual nice knob naming
            ("my_knob", "Nice Knob Name"): "some text",

            # dict type will be created as knob group
            "KnobGroup": {
                "knob1": 5,
                "knob2": "hello",
                "knob3": ["a", "b"],
            },

            # Nested dict will be created as tab group
            "TabGroup": {
                "tab1": {"count": 5},
                "tab2": {"isGood": True},
                "tab3": {"direction": ["Left", "Right"]},
            },
        }
        lib.imprint(node, data, tab="Demo")

        ```

    """
    for knob in create_knobs(data, tab):
        node.addKnob(knob)

class Knobby(object):
    """For creating knob which it's type isn't mapped in `create_knobs`

    Args:
        type (string): Nuke knob type name
        value: Value to be set with `Knob.setValue`, put `None` if not required
        *args: Args other than knob name for initializing knob class

    """

    def __init__(self, type, value, *args):
        self.type = type
        self.value = value
        self.args = args

    def create(self, name, nice=None):
        knob_cls = getattr(nuke, self.type)
        knob = knob_cls(name, nice, *self.args)
        if self.value is not None:
            knob.setValue(self.value)
        return knob

def create_knobs(data, tab=None):
    """Create knobs by data

    Depending on the type of each dict value and creates the correct Knob.

    Mapped types:
        bool: nuke.Boolean_Knob
        int: nuke.Int_Knob
        float: nuke.Double_Knob
        list: nuke.Enumeration_Knob
        six.string_types: nuke.String_Knob

        dict: If it's a nested dict (all values are dict), will turn into
            A tabs group. Or just a knobs group.

    Args:
        data (dict): collection of attributes and their value
        tab (string, optional): Knobs' tab name

    Returns:
        list: A list of `nuke.Knob` objects

    """
    def nice_naming(key):
        """Convert camelCase name into UI Display Name"""
        words = re.findall('[A-Z][^A-Z]*', key[0].upper() + key[1:])
        return " ".join(words)

    # Turn key-value pairs into knobs
    knobs = list()

    if tab:
        knobs.append(nuke.Tab_Knob(tab))

    for key, value in data.items():
        # Knob name
        if isinstance(key, tuple):
            name, nice = key
        else:
            name, nice = key, nice_naming(key)

        # Create knob by value type
        if isinstance(value, Knobby):
            knobby = value
            knob = knobby.create(name, nice)

        elif isinstance(value, float):
            knob = nuke.Double_Knob(name, nice)
            knob.setValue(value)

        elif isinstance(value, bool):
            knob = nuke.Boolean_Knob(name, nice)
            knob.setValue(value)
            knob.setFlag(nuke.STARTLINE)

        elif isinstance(value, int):
            knob = nuke.Int_Knob(name, nice)
            knob.setValue(value)

        elif isinstance(value, six.string_types):
            knob = nuke.String_Knob(name, nice)
            knob.setValue(value)

        elif isinstance(value, list):
            knob = nuke.Enumeration_Knob(name, nice, value)

        elif isinstance(value, dict):
            if all(isinstance(v, dict) for v in value.values()):
                # Create a group of tabs
                begain = nuke.BeginTabGroup_Knob()
                end = nuke.EndTabGroup_Knob()
                begain.setName(name)
                end.setName(name + "_End")
                knobs.append(begain)
                for k, v in value.items():
                    knobs += create_knobs(v, tab=k)
                knobs.append(end)
            else:
                # Create a group of knobs
                knobs.append(nuke.Tab_Knob(name, nice, nuke.TABBEGINGROUP))
                knobs += create_knobs(value)
                knobs.append(nuke.Tab_Knob(name, nice, nuke.TABENDGROUP))
            continue

        else:
            raise TypeError("Unsupported type: %r" % type(value))

        knobs.append(knob)

    return knobs

def add_publish_knob(node):
    """Add Publish knob to node

    Arguments:
        node (obj): nuke node to be processed

    Returns:
        node (obj): processed nuke node

    """
    if "publish" not in node.knobs():
        body = OrderedDict()
        body[("divd", "")] = Knobby("Text_Knob", "")
        body["publish"] = True
        imprint(node, body)
    return node

def set_avalon_knob_data(node, data=None, prefix="avalon:"):
    """ Sets data into nodes's avalon knob

    Arguments:
        node (obj): Nuke node to imprint with data,
        data ( dict): Data to be imprinted into AvalonTab
        prefix (str, optional): filtering prefix

    Returns:
        node (obj)

    Examples:
        data = {
            'asset': 'sq020sh0280',
            'family': 'render',
            'subset': 'subsetMain'
        }
    """
    editable = ["asset", "subset", "name", "namespace"]

    data = data or dict()

    group = OrderedDict()
    body = OrderedDict()

    group["avalonDataGroup"] = body
    body[("warn", "")] = Knobby("Text_Knob",
                                "Warning! Do not change following data!")
    body[("divd", "")] = Knobby("Text_Knob", "")

    for key, value in data.items():
        name = prefix + key
        if key in editable:
            body[name] = value
        else:
            body[name] = Knobby("Text_Knob", str(value))

    imprint(node, group, tab="AvalonTab")

    return node

def get_avalon_knob_data(node, prefix="avalon:"):
    """ Get data from nodes's avalon knob

    Arguments:
        node (obj): Nuke node to search for data,
        prefix (str, optional): filtering prefix

    Returns:
        data (dict)
    """
    data = {
        knob[len(prefix):]: node[knob].value()
        for knob in node.knobs().keys()
        if knob.startsWith(prefix)
    }
    return data

def fix_data_for_node_create(data):
    for k, v in data.items():

        data[k] = str(v)

        if "True" in v:
            data[k] = True
        if "False" in v:
            data[k] = False
        if "0x" in v:
            data[k] = int(v, 16)
    return data

def add_write_node(name, **kwarg):
    frame_range = kwarg.get("frame_range", None)

    w = nuke.createNode(
        "Write",
        "name {}".format(name))

    w["file"].setValue(kwarg["file"])

    for k, v in kwarg.items():
        if "frame_range" in k:
            continue
        log.info([k, v])
        try:
            w[k].setValue(v)
        except KeyError as e:
            log.debug(e)
            continue

    if frame_range:
        w["use_limit"].setValue(True)
        w["first"].setValue(frame_range[0])
        w["last"].setValue(frame_range[1])

    log.info(w)
    return w

def get_node_path(path, padding=4):
    """Get filename for the Nuke write with padded number as '#'

    >>> get_frame_path("test.exr")
    ('test', 4, '.exr')

    >>> get_frame_path("filename.#####.tif")
    ('filename.', 5, '.tif')

    >>> get_frame_path("foobar##.tif")
    ('foobar', 2, '.tif')

    >>> get_frame_path("foobar_%08d.tif")
    ('foobar_', 8, '.tif')

    Args:
        path (str): The path to render to.

    Returns:
        tuple: head, padding, tail (extension)

    """
    filename, ext = os.path.splitext(path)

    # Find a final number group
    if '%' in filename:
        match = re.match('.*?(%[0-9]+d)$', filename)
        if match:
            padding = int(match.group(1).replace('%', '').replace('d', ''))
            # remove number from end since fusion
            # will swap it with the frame number
            filename = filename.replace(match.group(1), '')
    elif '#' in filename:
        match = re.match('.*?(#+)$', filename)

        if match:
            padding = len(match.group(1))
            # remove number from end since fusion
            # will swap it with the frame number
            filename = filename.replace(match.group(1), '')

    return filename, padding, ext
<commit_msg>Hide knob name prefix on GUI<commit_after>import os
import contextlib
import nuke
import re
import logging
from collections import OrderedDict

from ..vendor import six

log = logging.getLogger(__name__)

@contextlib.contextmanager
def maintained_selection():
    previous_selection = nuke.selectedNodes()
    try:
        yield
    finally:
        # unselect all selection in case there is some
        current_seletion = nuke.selectedNodes()
        [n['selected'].setValue(False) for n in current_seletion]
        # and select all previously selected nodes
        if previous_selection:
            [n['selected'].setValue(True) for n in previous_selection]

def imprint(node, data, tab=None):
    """Store attributes with value on node

    Parse user data into Node knobs.
    Use `collections.OrderedDict` to ensure knob order.

    Args:
        node(nuke.Node): node object from Nuke
        data(dict): collection of attributes and their value

    Returns:
        None

    Examples:
        ```
        import nuke
        from avalon.nuke import lib

        node = nuke.createNode("NoOp")
        data = {
            # Regular type of attributes
            "myList": ["x", "y", "z"],
            "myBool": True,
            "myFloat": 0.1,
            "myInt": 5,

            # Creating non-default imprint type of knob
            "MyFilePath": lib.Knobby("File_Knob", "/file/path"),
            "divider": lib.Knobby("Text_Knob", ""),

            # Manual nice knob naming
            ("my_knob", "Nice Knob Name"): "some text",

            # dict type will be created as knob group
            "KnobGroup": {
                "knob1": 5,
                "knob2": "hello",
                "knob3": ["a", "b"],
            },

            # Nested dict will be created as tab group
            "TabGroup": {
                "tab1": {"count": 5},
                "tab2": {"isGood": True},
                "tab3": {"direction": ["Left", "Right"]},
            },
        }
        lib.imprint(node, data, tab="Demo")

        ```

    """
    for knob in create_knobs(data, tab):
        node.addKnob(knob)

class Knobby(object):
    """For creating knob which it's type isn't mapped in `create_knobs`

    Args:
        type (string): Nuke knob type name
        value: Value to be set with `Knob.setValue`, put `None` if not required
        *args: Args other than knob name for initializing knob class

    """

    def __init__(self, type, value, *args):
        self.type = type
        self.value = value
        self.args = args

    def create(self, name, nice=None):
        knob_cls = getattr(nuke, self.type)
        knob = knob_cls(name, nice, *self.args)
        if self.value is not None:
            knob.setValue(self.value)
        return knob

def create_knobs(data, tab=None):
    """Create knobs by data

    Depending on the type of each dict value and creates the correct Knob.

    Mapped types:
        bool: nuke.Boolean_Knob
        int: nuke.Int_Knob
        float: nuke.Double_Knob
        list: nuke.Enumeration_Knob
        six.string_types: nuke.String_Knob

        dict: If it's a nested dict (all values are dict), will turn into
            A tabs group. Or just a knobs group.

    Args:
        data (dict): collection of attributes and their value
        tab (string, optional): Knobs' tab name

    Returns:
        list: A list of `nuke.Knob` objects

    """
    def nice_naming(key):
        """Convert camelCase name into UI Display Name"""
        words = re.findall('[A-Z][^A-Z]*', key[0].upper() + key[1:])
        return " ".join(words)

    # Turn key-value pairs into knobs
    knobs = list()

    if tab:
        knobs.append(nuke.Tab_Knob(tab))

    for key, value in data.items():
        # Knob name
        if isinstance(key, tuple):
            name, nice = key
        else:
            name, nice = key, nice_naming(key)

        # Create knob by value type
        if isinstance(value, Knobby):
            knobby = value
            knob = knobby.create(name, nice)

        elif isinstance(value, float):
            knob = nuke.Double_Knob(name, nice)
            knob.setValue(value)

        elif isinstance(value, bool):
            knob = nuke.Boolean_Knob(name, nice)
            knob.setValue(value)
            knob.setFlag(nuke.STARTLINE)

        elif isinstance(value, int):
            knob = nuke.Int_Knob(name, nice)
            knob.setValue(value)

        elif isinstance(value, six.string_types):
            knob = nuke.String_Knob(name, nice)
            knob.setValue(value)

        elif isinstance(value, list):
            knob = nuke.Enumeration_Knob(name, nice, value)

        elif isinstance(value, dict):
            if all(isinstance(v, dict) for v in value.values()):
                # Create a group of tabs
                begain = nuke.BeginTabGroup_Knob()
                end = nuke.EndTabGroup_Knob()
                begain.setName(name)
                end.setName(name + "_End")
                knobs.append(begain)
                for k, v in value.items():
                    knobs += create_knobs(v, tab=k)
                knobs.append(end)
            else:
                # Create a group of knobs
                knobs.append(nuke.Tab_Knob(name, nice, nuke.TABBEGINGROUP))
                knobs += create_knobs(value)
                knobs.append(nuke.Tab_Knob(name, nice, nuke.TABENDGROUP))
            continue

        else:
            raise TypeError("Unsupported type: %r" % type(value))

        knobs.append(knob)

    return knobs

def add_publish_knob(node):
    """Add Publish knob to node

    Arguments:
        node (obj): nuke node to be processed

    Returns:
        node (obj): processed nuke node

    """
    if "publish" not in node.knobs():
        body = OrderedDict()
        body[("divd", "")] = Knobby("Text_Knob", "")
        body["publish"] = True
        imprint(node, body)
    return node

def set_avalon_knob_data(node, data=None, prefix="avalon:"):
    """ Sets data into nodes's avalon knob

    Arguments:
        node (obj): Nuke node to imprint with data,
        data ( dict): Data to be imprinted into AvalonTab
        prefix (str, optional): filtering prefix

    Returns:
        node (obj)

    Examples:
        data = {
            'asset': 'sq020sh0280',
            'family': 'render',
            'subset': 'subsetMain'
        }
    """
    editable = ["asset", "subset", "name", "namespace"]

    data = data or dict()

    group = OrderedDict()
    body = OrderedDict()

    group["avalonDataGroup"] = body
    body[("warn", "")] = Knobby("Text_Knob",
                                "Warning! Do not change following data!")
    body[("divd", "")] = Knobby("Text_Knob", "")

    for key, value in data.items():
        name = (prefix + key, key)  # Hide prefix on GUI
        if key in editable:
            body[name] = value
        else:
            body[name] = Knobby("Text_Knob", str(value))

    imprint(node, group, tab="AvalonTab")

    return node

def get_avalon_knob_data(node, prefix="avalon:"):
    """ Get data from nodes's avalon knob

    Arguments:
        node (obj): Nuke node to search for data,
        prefix (str, optional): filtering prefix

    Returns:
        data (dict)
    """
    data = {
        knob[len(prefix):]: node[knob].value()
        for knob in node.knobs().keys()
        if knob.startsWith(prefix)
    }
    return data

def fix_data_for_node_create(data):
    for k, v in data.items():

        data[k] = str(v)

        if "True" in v:
            data[k] = True
        if "False" in v:
            data[k] = False
        if "0x" in v:
            data[k] = int(v, 16)
    return data

def add_write_node(name, **kwarg):
    frame_range = kwarg.get("frame_range", None)

    w = nuke.createNode(
        "Write",
        "name {}".format(name))

    w["file"].setValue(kwarg["file"])

    for k, v in kwarg.items():
        if "frame_range" in k:
            continue
        log.info([k, v])
        try:
            w[k].setValue(v)
        except KeyError as e:
            log.debug(e)
            continue

    if frame_range:
        w["use_limit"].setValue(True)
        w["first"].setValue(frame_range[0])
        w["last"].setValue(frame_range[1])

    log.info(w)
    return w

def get_node_path(path, padding=4):
    """Get filename for the Nuke write with padded number as '#'

    >>> get_frame_path("test.exr")
    ('test', 4, '.exr')

    >>> get_frame_path("filename.#####.tif")
    ('filename.', 5, '.tif')

    >>> get_frame_path("foobar##.tif")
    ('foobar', 2, '.tif')

    >>> get_frame_path("foobar_%08d.tif")
    ('foobar_', 8, '.tif')

    Args:
        path (str): The path to render to.

    Returns:
        tuple: head, padding, tail (extension)

    """
    filename, ext = os.path.splitext(path)

    # Find a final number group
    if '%' in filename:
        match = re.match('.*?(%[0-9]+d)$', filename)
        if match:
            padding = int(match.group(1).replace('%', '').replace('d', ''))
            # remove number from end since fusion
            # will swap it with the frame number
            filename = filename.replace(match.group(1), '')
    elif '#' in filename:
        match = re.match('.*?(#+)$', filename)

        if match:
            padding = len(match.group(1))
            # remove number from end since fusion
            # will swap it with the frame number
            filename = filename.replace(match.group(1), '')

    return filename, padding, ext
<|endoftext|>


<commit_before>#!/usr/bin/env  python
import subprocess
import io
import cv2
import numpy as np
from signal import signal, SIGPIPE, SIG_DFL
from CVhandles import get_image, pushframe, pushdata
from multiprocessing import Process, Pool
import time
signal(SIGPIPE, SIG_DFL)

if __name__ == "__main__":
    curimage = cv2.imread("/home/zhukov/testimage.jpg")
    while True:
        #curimage = get_image(0)
        data = {"Foo":"https://i.pinimg.com/originals/00/21/52/002152ecac89b72d602059193ebdc161.jpg"}
        time.sleep(0.1)
        # CV stuff goes here
        pushframe(curimage,1)
        pushdata(data,1)
<commit_msg>commenting cvsimulate<commit_after>#!/usr/bin/env  python
import subprocess
import io
import cv2
import numpy as np
from signal import signal, SIGPIPE, SIG_DFL
from CVhandles import get_image, pushframe, pushdata
from multiprocessing import Process, Pool
import time
signal(SIGPIPE, SIG_DFL)

#to use cv simulate, place a test image in your home directory named "testimage.jpg" an example working test image was placed in the pakfront foler
if __name__ == "__main__":
    curimage = cv2.imread("$HOME/testimage.jpg")
    while True:
        #curimage = get_image(0)
        data = {"Foo":"https://i.pinimg.com/originals/00/21/52/002152ecac89b72d602059193ebdc161.jpg"}
        time.sleep(0.1)
        # CV stuff goes here
        pushframe(curimage,1)
        pushdata(data,1)
<|endoftext|>


<commit_before>import click
from pprint import pprint
import re
import requests
import sys
import time

try:
    from api_token import API_TOKEN
except ImportError:
    print('Please store API_TOKEN in api_token.py')
    sys.exit(1)

MESSAGE_ENDPOINT = 'https://slack.com/api/search.messages'
USER_ENDPOINT = 'https://slack.com/api/users.info'
CHANNEL_ENDPOINT = 'https://slack.com/api/channels.info'

def lookup_username(userid):
    '''Looks up username from Slack API based on userid

    Expects:
        userid: a valid userid of form <@XXX>

    Returns:
        string of username in form @YYYYY
    '''
    userid = re.sub('[<@|>]', '', userid)
    payload = {'token': API_TOKEN, 'user': userid}
    results = requests.get(USER_ENDPOINT, params=payload).json()
    if results['ok']:
        return '@' + results['user']['name']
    return 'Unknown User'

def lookup_channel(channelid):
    '''Looks up channel name from Slack API based on channelid

    Expects:
        channelid: a valid userid of form <#XXX>

    Returns:
        string of username in form #YYYYY
    '''
    channelid = re.sub('[<#|>]', '', channelid)
    payload = {'token': API_TOKEN, 'channel': channelid}
    results = requests.get(CHANNEL_ENDPOINT, params=payload).json()
    if results['ok']:
        return '#' + results['channel']['name']
    return 'Unknown Channel'

def filter_by_date(messages, start, end):
    try:
        start = time.mktime(time.strptime(start, '%m-%d-%Y'))
        end = time.mktime(time.strptime(end, '%m-%d-%Y'))
    except:
        raise ValueError('Invalid date')
    return [m for m in messages if end > float(m['ts']) > start]

def format_text(text):
    '''Replaces all userids and channelids with plaintext names

    Expects:
        A string with 0 or more userids and channelids

    Returns:
        The same text with those ids replaced
    '''
    text = re.sub('<@\w*>', lambda x: lookup_username(x.group()), text)
    text = re.sub('<#\w*>', lambda x: lookup_channel(x.group()), text)
    return text

@click.command()
@click.option('--emoji', help=('Text name of the emoji to search, without '
              'surrounding :s'))
@click.option('--outfile', help=('Outputfile to write results to.'))
@click.option('--startdate', help=('Date to begin search. Must be in '
              'MM-DD-YYYY format.'), default='01-01-1999')
@click.option('--enddate', help=('Date to end search. Must be in '
              'MM-DD-YYYY format.'), default='12-12-2222')
def get_messages(emoji, outfile, startdate, enddate):
    if not emoji:
        raise ValueError('Please supply an emoji name to search by')
    print('Querying Slack api')
    matches = []
    paging = {'page': 0, 'pages': 99999}
    while paging['page'] < paging['pages']:
        payload = {'token': API_TOKEN,
                   'query': 'has::' + emoji + ':',
                   'page': paging['page'] + 1}
        response = requests.get('https://slack.com/api/search.messages',
                                params=payload).json()
        if not response['ok']:
            raise Exception('Query Failed!')
        matches += response['messages']['matches']
        paging = response['messages']['paging']
        print('Fetched results page {} of {}'.format(paging['page'],
                                                     paging['pages']))
    print('Filtering by date')
    matches = filter_by_date(matches, startdate, enddate)
    print('Formatting results')
    results = [{'ts': time.ctime(float(m['ts'])),
                'username': m['username'],
                'permalink': m['permalink'],
                'text': format_text(m['text'])}
               for m in matches]
    if outfile:
        print('Writing results to file')
        with open(outfile, 'w') as f:
            pprint(results, stream=f)
    else:
        pprint(results)

if __name__ == '__main__':
    get_messages()
<commit_msg>Added shebang line to script to allow executing with users default python<commit_after>#! /usr/bin/env python

import click
from pprint import pprint
import re
import requests
import sys
import time

try:
    from api_token import API_TOKEN
except ImportError:
    print('Please store API_TOKEN in api_token.py')
    sys.exit(1)

MESSAGE_ENDPOINT = 'https://slack.com/api/search.messages'
USER_ENDPOINT = 'https://slack.com/api/users.info'
CHANNEL_ENDPOINT = 'https://slack.com/api/channels.info'

def lookup_username(userid):
    '''Looks up username from Slack API based on userid

    Expects:
        userid: a valid userid of form <@XXX>

    Returns:
        string of username in form @YYYYY
    '''
    userid = re.sub('[<@|>]', '', userid)
    payload = {'token': API_TOKEN, 'user': userid}
    results = requests.get(USER_ENDPOINT, params=payload).json()
    if results['ok']:
        return '@' + results['user']['name']
    return 'Unknown User'

def lookup_channel(channelid):
    '''Looks up channel name from Slack API based on channelid

    Expects:
        channelid: a valid userid of form <#XXX>

    Returns:
        string of username in form #YYYYY
    '''
    channelid = re.sub('[<#|>]', '', channelid)
    payload = {'token': API_TOKEN, 'channel': channelid}
    results = requests.get(CHANNEL_ENDPOINT, params=payload).json()
    if results['ok']:
        return '#' + results['channel']['name']
    return 'Unknown Channel'

def filter_by_date(messages, start, end):
    try:
        start = time.mktime(time.strptime(start, '%m-%d-%Y'))
        end = time.mktime(time.strptime(end, '%m-%d-%Y'))
    except:
        raise ValueError('Invalid date')
    return [m for m in messages if end > float(m['ts']) > start]

def format_text(text):
    '''Replaces all userids and channelids with plaintext names

    Expects:
        A string with 0 or more userids and channelids

    Returns:
        The same text with those ids replaced
    '''
    text = re.sub('<@\w*>', lambda x: lookup_username(x.group()), text)
    text = re.sub('<#\w*>', lambda x: lookup_channel(x.group()), text)
    return text

@click.command()
@click.option('--emoji', help=('Text name of the emoji to search, without '
              'surrounding :s'))
@click.option('--outfile', help=('Outputfile to write results to.'))
@click.option('--startdate', help=('Date to begin search. Must be in '
              'MM-DD-YYYY format.'), default='01-01-1999')
@click.option('--enddate', help=('Date to end search. Must be in '
              'MM-DD-YYYY format.'), default='12-12-2222')
def get_messages(emoji, outfile, startdate, enddate):
    if not emoji:
        raise ValueError('Please supply an emoji name to search by')
    print('Querying Slack api')
    matches = []
    paging = {'page': 0, 'pages': 99999}
    while paging['page'] < paging['pages']:
        payload = {'token': API_TOKEN,
                   'query': 'has::' + emoji + ':',
                   'page': paging['page'] + 1}
        response = requests.get('https://slack.com/api/search.messages',
                                params=payload).json()
        if not response['ok']:
            raise Exception('Query Failed!')
        matches += response['messages']['matches']
        paging = response['messages']['paging']
        print('Fetched results page {} of {}'.format(paging['page'],
                                                     paging['pages']))
    print('Filtering by date')
    matches = filter_by_date(matches, startdate, enddate)
    print('Formatting results')
    results = [{'ts': time.ctime(float(m['ts'])),
                'username': m['username'],
                'permalink': m['permalink'],
                'text': format_text(m['text'])}
               for m in matches]
    if outfile:
        print('Writing results to file')
        with open(outfile, 'w') as f:
            pprint(results, stream=f)
    else:
        pprint(results)

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


<commit_before>from kivy.logger import Logger
from kivy.uix.scatter import Scatter
from kivy.uix.widget import Widget
from kivy.properties import StringProperty, ListProperty, NumericProperty
from kivy.uix.relativelayout import RelativeLayout
from kivy.clock import Clock
from kivy.graphics import *

from defines import *
from gui import *

import Queue

black_filename = "./images/black_transparent.png"
white_filename = "./images/white_transparent.png"
x_filename = "./images/X_transparent.png"

import pdb

class BoardWidget(RelativeLayout):
    source = StringProperty(None)
    # TODO: Get the names and times hooked up
    black_name = StringProperty("Freddo")
    white_name = StringProperty("Deep Thunk")
    black_time = StringProperty("0:00")
    white_time = StringProperty("0:00")
    black_to_move_marker = StringProperty("*")
    white_to_move_marker = StringProperty("")
    black_captures = StringProperty("0")
    white_captures = StringProperty("0")
    gridlines = ListProperty([])
    # TODO: Only the vertical offset is used so far.
    board_offset = ListProperty([0,80.0])

    def __init__(self, *args, **kwargs):
        self.marker = None
        self.stones_by_board_pos = {}
        self.action_queue = Queue.Queue()

        super(BoardWidget, self).__init__(*args, **kwargs)

    def set_game(self, game):
        self.game = game

        # We must watch what happens to the logical board, and update accordingly
        board = game.get_board()
        board.add_observer(self)

        self.trig = Clock.create_trigger(self.perform)
        self.set_up_grid()

        # TODO: convert to use convention of BLACK = 1
        self.black_name = game.get_player_name(0)
        self.white_name = game.get_player_name(1)

        # start the game
        prompt = game.prompt_for_action(self)
        # self.display_feedback_string(prompt)

    def display_feedback_string(self, message):
        # TODO: Update screen
        print message

    def request_move(self, name):
        # TODO: Update side panel instead, use player object param instead of name
        your_move = "Your move, %s" % name
        return your_move

    def enqueue_action(self, action):
        self.action_queue.put(action)
        self.trig()

    def perform(self, dt):
        if self.action_queue.empty():
            return
        action = self.action_queue.get()
        action.perform(self.game)
        self.update_captures_and_winner()
        prompt = self.game.prompt_for_action(self)
        #self.display_feedback_string(prompt)

    def board_size(self):
        return self.game.size()

    def grid_size(self):
        """ The Grid on the screen allows extra space at the edges """
        return self.game.size() + 1

    def before_set_occ(self, pos, colour):
        pass

    def after_set_occ(self, pos, colour):
        self.make_move_on_the_gui_board(pos, colour)
        self.update_captures_and_winner()

    def update_captures_and_winner(self):
        #pdb.set_trace()
        self.black_captures = str(self.game.get_captured(BLACK))
        self.white_captures = str(self.game.get_captured(WHITE))

        if self.game.finished():
            winner = self.game.winner_name()
            if winner == BLACK:
                self.black_to_move_marker = "won by"
                self.white_to_move_marker = ""
            elif winner == WHITE:
                self.black_to_move_marker = ""
                self.white_to_move_marker = "won by"
            # draws are exceedingly unlikely...
        else:
            # Mark who is to move. TODO: Underline?
            to_move = self.game.to_move_colour()
            if to_move == BLACK:
                self.black_to_move_marker = "*"
                self.white_to_move_marker = ""
            elif to_move == WHITE:
                self.black_to_move_marker = ""
                self.white_to_move_marker = "*"

    def setup_grid_lines(self):
        size_x, size_y = self.size
        # Adjust for the position of the board being shifted so we can have a panel of
        # extra information.
        size_y -= self.board_offset[1]

        lines = []

        # This part is using the relative layout to get the lines in the right place
        # RL does not scale though.
        Color(0, 0, 0) # Black lines
        grid_size_x = float(size_x) / (self.grid_size())
        grid_size_y = float(size_y) / (self.grid_size())
        GS = self.grid_size()
        # horizontal lines
        for y in range(1,GS):
            # Each line is drawn twice, in opposite directions so we don't
            # get diagonal lines drawn as well
            this_end = (grid_size_x, grid_size_y*y)
            lines.extend(this_end)
            lines.extend((size_x-grid_size_x, grid_size_y*y))
            lines.extend(this_end)
        # vertical lines
        for x in range(1,GS):
            this_end = (grid_size_x*x, grid_size_y)
            lines.extend(this_end)
            lines.extend((grid_size_x*x, size_y-grid_size_y))
            lines.extend(this_end)
        return lines

    def set_up_grid(self, _dt=None):
        if hasattr(self, "game"):
            self.gridlines = self.setup_grid_lines()

    def snap_to_grid(self, screen_pos):
        return self.board_to_screen(self.screen_to_board(screen_pos))

    # screen_to_board and board_to_screen do their own compensating for
    # the canvas coordinate system covering the whole screen. It would be nice
    # if the RelativeLayout handled this. So far only the vertical axis has been
    # affected.
    def screen_to_board(self, screen_pos):
        """ Convert a screen position (in pixels) to a board coordinate pair,
            dependant on the size of the board """
        size_x, size_y = self.size
        bsp = screen_pos[0], screen_pos[1]-self.board_offset[1]
        size_y -= self.board_offset[1]
        GS = self.grid_size()
        board_x = int(round(GS * bsp[0] / size_x) - 1)
        board_y = int(round(GS * bsp[1] / size_y) - 1)
        return board_x, board_y

    def board_to_screen(self, board_pos):
        """ Convert a board coordinate pair to a screen position (in pixels),
            dependant on the size of the board """
        size_x, size_y = self.size
        size_y -= self.board_offset[1]

        # Use float() to avoid python int / int = int prob
        GS = float(self.grid_size())
        screen_x = ((board_pos[0] + 1) / GS) * size_x
        screen_y = ((board_pos[1] + 1) / GS) * size_y
        screen_y += self.board_offset[1]
        return screen_x, screen_y

    def on_touch_down(self, touch):
        # Check that it is a human's turn.
        current_player = self.game.get_current_player()
        if current_player.get_type() == "human":
            # Place a marker at the (snapped) cursor position.
            if self.marker == None:
                try:
                    # load the image
                    # TODO: Separate class?
                    self.marker = Piece(self.game, source=x_filename)
                except Exception, e:
                    Logger.exception('Board: Unable to load <%s>' % x_filename)
            self.marker.pos = self.snap_to_grid(touch.pos)
            self.add_widget(self.marker)
        else:
            self.display_feedback_string("It is not your turn!")

    def on_touch_up(self, touch):
        # If there is an active marker,
        # replace the marker with a piece of the appropriate colour
        if self.marker != None:
            self.remove_widget(self.marker)
            self.marker = None
            # Quick hack to get both coloured stones on the board
            board_pos = self.screen_to_board(touch.pos)

            # Make and Queue MoveAction
            ma = MoveAction.create_from_move(board_pos)
            self.enqueue_action(ma)

    def make_move_on_the_gui_board(self, board_pos, colour):
        if self.stones_by_board_pos.has_key(board_pos) or colour == EMPTY:
            # There is a piece there already, remove it.
            assert colour == EMPTY
            current_piece = self.stones_by_board_pos.pop(board_pos)
            self.remove_widget(current_piece)
        else:
            # Nothing there yet, place a stone
            filename = ["", black_filename, white_filename][colour]

            try:
                # load the image
                new_piece = Piece(self.game, source=filename)
                self.stones_by_board_pos[board_pos] = new_piece
                new_piece.pos = self.board_to_screen(board_pos)
                self.add_widget(new_piece)
            except Exception, e:
                Logger.exception('Board: Unable to load <%s>' % filename)

    def on_touch_move(self, touch):
        # TODO: Check for off board? remove marker?
        if self.marker != None:
            # Move the marker position
            self.marker.pos = self.snap_to_grid(touch.pos)

    def on_size(self,*args,**kwargs):
        self.set_up_grid()

    def get_gridlines(self):
        return []

class Piece(Scatter):
    source = StringProperty(None)

    def __init__(self, game, *args, **kwargs):
        game_size = game.size()
        self.scale = 9. / game_size
        super(Piece, self).__init__(*args, **kwargs)

<commit_msg>Fixed winner display<commit_after>from kivy.logger import Logger
from kivy.uix.scatter import Scatter
from kivy.uix.widget import Widget
from kivy.properties import StringProperty, ListProperty, NumericProperty
from kivy.uix.relativelayout import RelativeLayout
from kivy.clock import Clock
from kivy.graphics import *

from defines import *
from gui import *

import Queue

black_filename = "./images/black_transparent.png"
white_filename = "./images/white_transparent.png"
x_filename = "./images/X_transparent.png"

import pdb

class BoardWidget(RelativeLayout):
    source = StringProperty(None)
    # TODO: Get the names and times hooked up
    black_name = StringProperty("Freddo")
    white_name = StringProperty("Deep Thunk")
    black_time = StringProperty("0:00")
    white_time = StringProperty("0:00")
    black_to_move_marker = StringProperty("*")
    white_to_move_marker = StringProperty("")
    black_captures = StringProperty("0")
    white_captures = StringProperty("0")
    gridlines = ListProperty([])
    # TODO: Only the vertical offset is used so far.
    board_offset = ListProperty([0,80.0])

    def __init__(self, *args, **kwargs):
        self.marker = None
        self.stones_by_board_pos = {}
        self.action_queue = Queue.Queue()

        super(BoardWidget, self).__init__(*args, **kwargs)

    def set_game(self, game):
        self.game = game

        # We must watch what happens to the logical board, and update accordingly
        board = game.get_board()
        board.add_observer(self)

        self.trig = Clock.create_trigger(self.perform)
        self.set_up_grid()

        # TODO: convert to use convention of BLACK = 1
        self.black_name = game.get_player_name(0)
        self.white_name = game.get_player_name(1)

        # start the game
        prompt = game.prompt_for_action(self)
        # self.display_feedback_string(prompt)

    def display_feedback_string(self, message):
        # TODO: Update screen
        print message

    def request_move(self, name):
        # TODO: Update side panel instead, use player object param instead of name
        your_move = "Your move, %s" % name
        return your_move

    def enqueue_action(self, action):
        self.action_queue.put(action)
        self.trig()

    def perform(self, dt):
        if self.action_queue.empty():
            return
        action = self.action_queue.get()
        action.perform(self.game)
        self.update_captures_and_winner()
        prompt = self.game.prompt_for_action(self)
        #self.display_feedback_string(prompt)

    def board_size(self):
        return self.game.size()

    def grid_size(self):
        """ The Grid on the screen allows extra space at the edges """
        return self.game.size() + 1

    def before_set_occ(self, pos, colour):
        pass

    def after_set_occ(self, pos, colour):
        self.make_move_on_the_gui_board(pos, colour)
        self.update_captures_and_winner()

    def update_captures_and_winner(self):
        """ Update fields in the panel from changes to the game state """
        #pdb.set_trace()
        self.black_captures = str(self.game.get_captured(BLACK))
        self.white_captures = str(self.game.get_captured(WHITE))

        if self.game.finished():
            winner = self.game.winner()
            if winner == BLACK:
                self.black_to_move_marker = "won by"
                self.white_to_move_marker = ""
            elif winner == WHITE:
                self.black_to_move_marker = ""
                self.white_to_move_marker = "won by"
            # draws are exceedingly unlikely...
        else:
            # Mark who is to move. TODO: Underline?
            to_move = self.game.to_move_colour()
            if to_move == BLACK:
                self.black_to_move_marker = "*"
                self.white_to_move_marker = ""
            elif to_move == WHITE:
                self.black_to_move_marker = ""
                self.white_to_move_marker = "*"

    def setup_grid_lines(self):
        size_x, size_y = self.size
        # Adjust for the position of the board being shifted so we can have a panel of
        # extra information.
        size_y -= self.board_offset[1]

        lines = []

        # This part is using the relative layout to get the lines in the right place
        # RL does not scale though.
        Color(0, 0, 0) # Black lines
        grid_size_x = float(size_x) / (self.grid_size())
        grid_size_y = float(size_y) / (self.grid_size())
        GS = self.grid_size()
        # horizontal lines
        for y in range(1,GS):
            # Each line is drawn twice, in opposite directions so we don't
            # get diagonal lines drawn as well
            this_end = (grid_size_x, grid_size_y*y)
            lines.extend(this_end)
            lines.extend((size_x-grid_size_x, grid_size_y*y))
            lines.extend(this_end)
        # vertical lines
        for x in range(1,GS):
            this_end = (grid_size_x*x, grid_size_y)
            lines.extend(this_end)
            lines.extend((grid_size_x*x, size_y-grid_size_y))
            lines.extend(this_end)
        return lines

    def set_up_grid(self, _dt=None):
        if hasattr(self, "game"):
            self.gridlines = self.setup_grid_lines()

    def snap_to_grid(self, screen_pos):
        return self.board_to_screen(self.screen_to_board(screen_pos))

    # screen_to_board and board_to_screen do their own compensating for
    # the canvas coordinate system covering the whole screen. It would be nice
    # if the RelativeLayout handled this. So far only the vertical axis has been
    # affected.
    def screen_to_board(self, screen_pos):
        """ Convert a screen position (in pixels) to a board coordinate pair,
            dependant on the size of the board """
        size_x, size_y = self.size
        bsp = screen_pos[0], screen_pos[1]-self.board_offset[1]
        size_y -= self.board_offset[1]
        GS = self.grid_size()
        board_x = int(round(GS * bsp[0] / size_x) - 1)
        board_y = int(round(GS * bsp[1] / size_y) - 1)
        return board_x, board_y

    def board_to_screen(self, board_pos):
        """ Convert a board coordinate pair to a screen position (in pixels),
            dependant on the size of the board """
        size_x, size_y = self.size
        size_y -= self.board_offset[1]

        # Use float() to avoid python int / int = int prob
        GS = float(self.grid_size())
        screen_x = ((board_pos[0] + 1) / GS) * size_x
        screen_y = ((board_pos[1] + 1) / GS) * size_y
        screen_y += self.board_offset[1]
        return screen_x, screen_y

    def on_touch_down(self, touch):
        # Check that it is a human's turn.
        current_player = self.game.get_current_player()
        if current_player.get_type() == "human":
            # Place a marker at the (snapped) cursor position.
            if self.marker == None:
                try:
                    # load the image
                    # TODO: Separate class?
                    self.marker = Piece(self.game, source=x_filename)
                except Exception, e:
                    Logger.exception('Board: Unable to load <%s>' % x_filename)
            self.marker.pos = self.snap_to_grid(touch.pos)
            self.add_widget(self.marker)
        else:
            self.display_feedback_string("It is not your turn!")

    def on_touch_up(self, touch):
        # If there is an active marker,
        # replace the marker with a piece of the appropriate colour
        if self.marker != None:
            self.remove_widget(self.marker)
            self.marker = None
            # Quick hack to get both coloured stones on the board
            board_pos = self.screen_to_board(touch.pos)

            # Make and Queue MoveAction
            ma = MoveAction.create_from_move(board_pos)
            self.enqueue_action(ma)

    def make_move_on_the_gui_board(self, board_pos, colour):
        if self.stones_by_board_pos.has_key(board_pos) or colour == EMPTY:
            # There is a piece there already, remove it.
            assert colour == EMPTY
            current_piece = self.stones_by_board_pos.pop(board_pos)
            self.remove_widget(current_piece)
        else:
            # Nothing there yet, place a stone
            filename = ["", black_filename, white_filename][colour]

            try:
                # load and place the appropriate stone image
                new_piece = Piece(self.game, source=filename)
                self.stones_by_board_pos[board_pos] = new_piece
                new_piece.pos = self.board_to_screen(board_pos)
                self.add_widget(new_piece)
            except Exception, e:
                Logger.exception('Board: Unable to load <%s>' % filename)

    def on_touch_move(self, touch):
        # TODO: Check for off board? remove marker?
        if self.marker != None:
            # Move the marker position
            self.marker.pos = self.snap_to_grid(touch.pos)

    def on_size(self,*args,**kwargs):
        self.set_up_grid()

    def get_gridlines(self):
        return []

class Piece(Scatter):
    source = StringProperty(None)

    def __init__(self, game, *args, **kwargs):
        game_size = game.size()
        self.scale = 9. / game_size
        super(Piece, self).__init__(*args, **kwargs)

<|endoftext|>


<commit_before>version = '4.4.0a4'

name = 'awsume'
author = 'Trek10, Inc'
author_email = 'package-management@trek10.com'
description = 'Awsume - A cli that makes using AWS IAM credentials easy'
license = 'MIT'
homepage = 'https://github.com/trek10inc/awsume'

message = 'Thank you for using AWSume! Check us out at https://trek10.com'
<commit_msg>version 4.4.0a5<commit_after>version = '4.4.0a5'

name = 'awsume'
author = 'Trek10, Inc'
author_email = 'package-management@trek10.com'
description = 'Awsume - A cli that makes using AWS IAM credentials easy'
license = 'MIT'
homepage = 'https://github.com/trek10inc/awsume'

message = 'Thank you for using AWSume! Check us out at https://trek10.com'
<|endoftext|>


<commit_before># Copyright 2014, Oliver Nagy <olitheolix@gmail.com>
#
# This file is part of Azrael (https://github.com/olitheolix/azrael)
#
# Azrael 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.
#
# Azrael 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 Azrael. If not, see <http://www.gnu.org/licenses/>.

"""
Database abstractions.
"""
import copy
import logging

import azrael.config as config
from IPython import embed as ipshell
from azrael.aztypes import typecheck, RetVal

# Global database handles.
logit = logging.getLogger('azrael.' + __name__)

# Connect to MongoDB and store the relevant collection handles in the
# 'dbHandles' variables to avoid hard coded collection names in Azrael.
client = config.getMongoClient()
dbName = 'azrael'
dbHandles = {
    'Commands': client[dbName]['Cmd'],
    'Templates': client[dbName]['template'],
    'ObjInstances': client[dbName]['objinstances'],
    'Counters': client[dbName]['Counters'],
    'Constraints': client[dbName]['Constraints'],
}

@typecheck
def init():
    """
    Flush the database.
    """
    # Delete the database.
    client.drop_database(dbName)
    dbHandles['ObjInstances'].ensure_index([('objID', 1)])

@typecheck
def getUniqueObjectIDs(numIDs: int):
    """
    Return ``numIDs`` unique object IDs as a tuple.

    If ``numIDs`` is Zero then return a scalar with the current value.

    The ``numIDs`` is positive then return a tuple with ``numIDs`` entries,
    each of which constitutes a unique ID.

    This function returns an error unless ``numIDs`` is non-negative.

    :param int numIDs: non-negative integer.
    :return tuple or scalar: object IDs (numIDs > 0) or last issued object ID
                            (numIDs = 0)
    """
    # Sanity check.
    if numIDs < 0:
        return RetVal(False, 'numIDs must be non-negative', None)

    # Increment the counter by ``numIDs``.
    fam = dbHandles['Counters'].find_and_modify
    doc = fam({'name': 'objcnt'},
              {'$inc': {'cnt': numIDs}},
              new=True, upsert=True)

    # Error check.
    if doc is None:
        return RetVal(False, 'Cannot determine new object IDs', None)

    # Extract the new counter value (after it was incremented).
    cnt = doc['cnt']

    # Return either the current value or the range of new IDs.
    if numIDs == 0:
        return RetVal(True, None, cnt)
    else:
        newIDs = tuple(range(cnt - numIDs + 1, cnt + 1))
        return RetVal(True, None, newIDs)

class DatabaseInMemory:
    def __init__(self, name: tuple):
        self.dbname = name
        self.reset()

    def reset(self):
        self.content = {}
        return RetVal(True, None, None)

    def count(self):
        return RetVal(True, None, len(self.content))

    def put(self, ops: dict):
        ret = {}
        for aid, op in ops.items():
            exists, data = op['exists'], op['data']
            if exists and aid in self.content:
                self.content[aid] = data
                ret[aid] = True
            elif not exists and aid not in self.content:
                self.content[aid] = data
                ret[aid] = True
            else:
                ret[aid] = False
        return RetVal(True, None, ret)

    def hasKey(self, d, key_hierarchy):
        try:
            tmp = d
            for key in key_hierarchy:
                tmp = tmp[key]

        except (KeyError, TypeError):
            return False
        return True

    def setKey(self, d, key_hierarchy, value):
        tmp = d
        for key in key_hierarchy[:-1]:
            if key not in tmp:
                tmp[key] = {}
            tmp = tmp[key]
        tmp[key_hierarchy[-1]] = value
        return True

    def incKey(self, d, key_hierarchy, value):
        tmp = d
        for key in key_hierarchy[:-1]:
            if key not in tmp:
                tmp[key] = {}
            tmp = tmp[key]

        try:
            tmp[key_hierarchy[-1]] += value
        except TypeError:
            return False
        return True

    def delKey(self, d, key_hierarchy):
        try:
            tmp = d
            for key in key_hierarchy[:-1]:
                tmp = tmp[key]
            del tmp[key_hierarchy[-1]]
        except (KeyError, TypeError):
            return False
        return True

    def getKey(self, d, key_hierarchy):
        tmp = d
        for key in key_hierarchy:
            tmp = tmp[key]
        return tmp

    def mod(self, ops):
        ret = {}
        for aid, op in ops.items():
            # Verify the specified items exist.
            try:
                c = self.content[aid]
                for key, yes in op['exists'].items():
                    assert self.hasKey(c, key) is yes

                for key_hierarchy in op['inc']:
                    tmp = c
                    for key in key_hierarchy:
                        tmp = tmp[key]
                    assert isinstance(tmp, (float, int))
                ret[aid] = True
            except (AssertionError, KeyError):
                ret[aid] = False
                continue

            for key, val in op['inc'].items():
                self.incKey(self.content[aid], key, val)

            for key in op['unset']:
                self.delKey(self.content[aid], key)

            for key, val in op['set'].items():
                self.setKey(self.content[aid], key, val)

        return RetVal(True, None, ret)

    def remove(self, aids: (tuple, list)):
        num_deleted = 0
        for aid in aids:
            try:
                del self.content[aid]
                num_deleted += 1
            except KeyError:
                pass
        return RetVal(True, None, num_deleted)

    def allKeys(self):
        keys = list(self.content.keys())
        return RetVal(True, None, keys)

    def project(self, doc, prj):
        doc = copy.deepcopy(doc)
        out = {}
        for p in prj:
            try:
                self.setKey(out, p, self.getKey(doc, p))
            except KeyError:
                continue
        return out

    def getOne(self, aid, prj=[]):
        doc = self.content.get(aid, None)
        if doc is None:
            return RetVal(False, None, None)
        else:
            if len(prj) > 0:
                doc = self.project(doc, prj)
            return RetVal(True, None, doc)

    def getMulti(self, aids, prj=[]):
        docs = {aid: self.content[aid] for aid in aids if aid in self.content}
        if len(prj) > 0:
            for doc in docs:
                docs[doc] = self.project(docs[doc], prj)

        return RetVal(True, None, docs)

    def getAll(self, prj=[]):
        docs = copy.deepcopy(self.content)
        if len(prj) > 0:
            for doc in docs:
                docs[doc] = self.project(docs[doc], prj)

        return RetVal(True, None, docs)

class DatabaseMongo:
    def __init__(self, name: tuple):
        import pymongo
        self.name_db, self.name_col = name
        client = pymongo.MongoClient()
        self.db = client[self.name_db][self.name_col]
        self.db.ensure_index([('objID', 1)])

    def reset(self):
        self.db.drop()
        return RetVal(True, None, None)

    def count(self):
        return RetVal(True, None, self.db.count())

    def put(self, ops: dict):
        ret = {}
        for aid, op in ops.items():
            exists, data = op['exists'], op['data']
            data = copy.deepcopy(data)
            data['objID'] = aid
            if exists:
                r = self.db.update({'objID': aid}, data, upsert=False)
                if r['updatedExisting']:
                    ret[aid] = True
                else:
                    ret[aid] = False
            else:
                r = self.db.update_one(
                    {'objID': aid},
                    {'$setOnInsert': data},
                    upsert=True
                )
                if r.upserted_id is None:
                    ret[aid] = False
                else:
                    ret[aid] = True
        return RetVal(True, None, ret)

    def mod(self, ops):
        ret = {}
        for aid, op_tmp in ops.items():
            query = {'.'.join(key): {'$exists': yes} for key, yes in op_tmp['exists'].items()}
            query['objID'] = aid

            # Update operations.
            op = {
                '$inc': {'.'.join(key): val for key, val in op_tmp['inc'].items()},
                '$set': {'.'.join(key): val for key, val in op_tmp['set'].items()},
                '$unset': {'.'.join(key): True for key in op_tmp['unset']},
            }
            # Prune update operations.
            op = {k: v for k, v in op.items() if len(v) > 0}

            # If no updates are necessary then skip this object.
            if len(op) == 0:
                continue

            # Issue the database query.
            r = self.db.update_one(query, op, upsert=False)

            # fixme: this used to be acknowledged. That was wrong yet it passed
            # all unit tests. Devise a test to identify this problem.
#            ret[aid] = r.acknowledged
            ret[aid] = (r.matched_count == 1)
        return RetVal(True, None, ret)

    def _removeAID(self, docs):
        docs = {doc['objID']: doc for doc in docs}
        for aid, doc in docs.items():
            del doc['objID']
        return docs

    def remove(self, aids: (tuple, list)):
        ret = self.db.delete_many({'objID': {'$in': aids}})
        return RetVal(True, None, ret.deleted_count)

    def allKeys(self):
        keys = self.db.distinct('objID')
        return RetVal(True, None, keys)

    def getOne(self, aid, prj=[]):
        prj = {'.'.join(_): True for _ in prj}
        prj['_id'] = False

        doc = self.db.find_one({'objID': aid}, prj)
        try:
            del doc['objID']
        except (KeyError, TypeError):
            pass

        if doc is None:
            return RetVal(False, None, None)
        else:
            return RetVal(True, None, doc)

    def getMulti(self, aids, prj=[]):
        prj = {'.'.join(_): True for _ in prj}
        if len(prj) > 0:
            prj['objID'] = True
        prj['_id'] = False
        cursor = self.db.find({'objID': {'$in': aids}}, prj)
        docs = self._removeAID(cursor)
        return RetVal(True, None, docs)

    # fixme: don't use a list default argument
    def getAll(self, prj=[]):
        prj = {'.'.join(_): True for _ in prj}
        if len(prj) > 0:
            prj['objID'] = True
        prj['_id'] = False
        cursor = self.db.find({}, prj)
        docs = self._removeAID(cursor)
        return RetVal(True, None, docs)
<commit_msg>Updated 'database.dbHandles' to reference the new data stores.<commit_after># Copyright 2014, Oliver Nagy <olitheolix@gmail.com>
#
# This file is part of Azrael (https://github.com/olitheolix/azrael)
#
# Azrael 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.
#
# Azrael 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 Azrael. If not, see <http://www.gnu.org/licenses/>.

"""
Database abstractions.
"""
import copy
import logging

import azrael.config as config
from IPython import embed as ipshell
from azrael.aztypes import typecheck, RetVal

# Global database handles.
logit = logging.getLogger('azrael.' + __name__)

@typecheck
def init():
    """
    Flush the database.
    """
    # Delete the database.
    client.drop_database(dbName)

@typecheck
def getUniqueObjectIDs(numIDs: int):
    """
    Return ``numIDs`` unique object IDs as a tuple.

    If ``numIDs`` is Zero then return a scalar with the current value.

    The ``numIDs`` is positive then return a tuple with ``numIDs`` entries,
    each of which constitutes a unique ID.

    This function returns an error unless ``numIDs`` is non-negative.

    :param int numIDs: non-negative integer.
    :return tuple or scalar: object IDs (numIDs > 0) or last issued object ID
                            (numIDs = 0)
    """
    # Sanity check.
    if numIDs < 0:
        return RetVal(False, 'numIDs must be non-negative', None)

    # Increment the counter by ``numIDs``.
    fam = dbHandles['Counters'].find_and_modify
    doc = fam({'name': 'objcnt'},
              {'$inc': {'cnt': numIDs}},
              new=True, upsert=True)

    # Error check.
    if doc is None:
        return RetVal(False, 'Cannot determine new object IDs', None)

    # Extract the new counter value (after it was incremented).
    cnt = doc['cnt']

    # Return either the current value or the range of new IDs.
    if numIDs == 0:
        return RetVal(True, None, cnt)
    else:
        newIDs = tuple(range(cnt - numIDs + 1, cnt + 1))
        return RetVal(True, None, newIDs)

class DatabaseInMemory:
    def __init__(self, name: tuple):
        self.dbname = name
        self.reset()

    def reset(self):
        self.content = {}
        return RetVal(True, None, None)

    def count(self):
        return RetVal(True, None, len(self.content))

    def put(self, ops: dict):
        ret = {}
        for aid, op in ops.items():
            exists, data = op['exists'], op['data']
            if exists and aid in self.content:
                self.content[aid] = data
                ret[aid] = True
            elif not exists and aid not in self.content:
                self.content[aid] = data
                ret[aid] = True
            else:
                ret[aid] = False
        return RetVal(True, None, ret)

    def hasKey(self, d, key_hierarchy):
        try:
            tmp = d
            for key in key_hierarchy:
                tmp = tmp[key]

        except (KeyError, TypeError):
            return False
        return True

    def setKey(self, d, key_hierarchy, value):
        tmp = d
        for key in key_hierarchy[:-1]:
            if key not in tmp:
                tmp[key] = {}
            tmp = tmp[key]
        tmp[key_hierarchy[-1]] = value
        return True

    def incKey(self, d, key_hierarchy, value):
        tmp = d
        for key in key_hierarchy[:-1]:
            if key not in tmp:
                tmp[key] = {}
            tmp = tmp[key]

        try:
            tmp[key_hierarchy[-1]] += value
        except TypeError:
            return False
        return True

    def delKey(self, d, key_hierarchy):
        try:
            tmp = d
            for key in key_hierarchy[:-1]:
                tmp = tmp[key]
            del tmp[key_hierarchy[-1]]
        except (KeyError, TypeError):
            return False
        return True

    def getKey(self, d, key_hierarchy):
        tmp = d
        for key in key_hierarchy:
            tmp = tmp[key]
        return tmp

    def mod(self, ops):
        ret = {}
        for aid, op in ops.items():
            # Verify the specified items exist.
            try:
                c = self.content[aid]
                for key, yes in op['exists'].items():
                    assert self.hasKey(c, key) is yes

                for key_hierarchy in op['inc']:
                    tmp = c
                    for key in key_hierarchy:
                        tmp = tmp[key]
                    assert isinstance(tmp, (float, int))
                ret[aid] = True
            except (AssertionError, KeyError):
                ret[aid] = False
                continue

            for key, val in op['inc'].items():
                self.incKey(self.content[aid], key, val)

            for key in op['unset']:
                self.delKey(self.content[aid], key)

            for key, val in op['set'].items():
                self.setKey(self.content[aid], key, val)

        return RetVal(True, None, ret)

    def remove(self, aids: (tuple, list)):
        num_deleted = 0
        for aid in aids:
            try:
                del self.content[aid]
                num_deleted += 1
            except KeyError:
                pass
        return RetVal(True, None, num_deleted)

    def allKeys(self):
        keys = list(self.content.keys())
        return RetVal(True, None, keys)

    def project(self, doc, prj):
        doc = copy.deepcopy(doc)
        out = {}
        for p in prj:
            try:
                self.setKey(out, p, self.getKey(doc, p))
            except KeyError:
                continue
        return out

    def getOne(self, aid, prj=[]):
        doc = self.content.get(aid, None)
        if doc is None:
            return RetVal(False, None, None)
        else:
            if len(prj) > 0:
                doc = self.project(doc, prj)
            return RetVal(True, None, doc)

    def getMulti(self, aids, prj=[]):
        docs = {aid: self.content[aid] for aid in aids if aid in self.content}
        if len(prj) > 0:
            for doc in docs:
                docs[doc] = self.project(docs[doc], prj)

        return RetVal(True, None, docs)

    def getAll(self, prj=[]):
        docs = copy.deepcopy(self.content)
        if len(prj) > 0:
            for doc in docs:
                docs[doc] = self.project(docs[doc], prj)

        return RetVal(True, None, docs)

class DatabaseMongo:
    def __init__(self, name: tuple):
        import pymongo
        self.name_db, self.name_col = name
        client = pymongo.MongoClient()
        self.db = client[self.name_db][self.name_col]
        self.db.ensure_index([('objID', 1)])

    def reset(self):
        self.db.drop()
        return RetVal(True, None, None)

    def count(self):
        return RetVal(True, None, self.db.count())

    def put(self, ops: dict):
        ret = {}
        for aid, op in ops.items():
            exists, data = op['exists'], op['data']
            data = copy.deepcopy(data)
            data['objID'] = aid
            if exists:
                r = self.db.update({'objID': aid}, data, upsert=False)
                if r['updatedExisting']:
                    ret[aid] = True
                else:
                    ret[aid] = False
            else:
                r = self.db.update_one(
                    {'objID': aid},
                    {'$setOnInsert': data},
                    upsert=True
                )
                if r.upserted_id is None:
                    ret[aid] = False
                else:
                    ret[aid] = True
        return RetVal(True, None, ret)

    def mod(self, ops):
        ret = {}
        for aid, op_tmp in ops.items():
            query = {'.'.join(key): {'$exists': yes} for key, yes in op_tmp['exists'].items()}
            query['objID'] = aid

            # Update operations.
            op = {
                '$inc': {'.'.join(key): val for key, val in op_tmp['inc'].items()},
                '$set': {'.'.join(key): val for key, val in op_tmp['set'].items()},
                '$unset': {'.'.join(key): True for key in op_tmp['unset']},
            }
            # Prune update operations.
            op = {k: v for k, v in op.items() if len(v) > 0}

            # If no updates are necessary then skip this object.
            if len(op) == 0:
                continue

            # Issue the database query.
            r = self.db.update_one(query, op, upsert=False)

            # fixme: this used to be acknowledged. That was wrong yet it passed
            # all unit tests. Devise a test to identify this problem.
#            ret[aid] = r.acknowledged
            ret[aid] = (r.matched_count == 1)
        return RetVal(True, None, ret)

    def _removeAID(self, docs):
        docs = {doc['objID']: doc for doc in docs}
        for aid, doc in docs.items():
            del doc['objID']
        return docs

    def remove(self, aids: (tuple, list)):
        ret = self.db.delete_many({'objID': {'$in': aids}})
        return RetVal(True, None, ret.deleted_count)

    def allKeys(self):
        keys = self.db.distinct('objID')
        return RetVal(True, None, keys)

    def getOne(self, aid, prj=[]):
        prj = {'.'.join(_): True for _ in prj}
        prj['_id'] = False

        doc = self.db.find_one({'objID': aid}, prj)
        try:
            del doc['objID']
        except (KeyError, TypeError):
            pass

        if doc is None:
            return RetVal(False, None, None)
        else:
            return RetVal(True, None, doc)

    def getMulti(self, aids, prj=[]):
        prj = {'.'.join(_): True for _ in prj}
        if len(prj) > 0:
            prj['objID'] = True
        prj['_id'] = False
        cursor = self.db.find({'objID': {'$in': aids}}, prj)
        docs = self._removeAID(cursor)
        return RetVal(True, None, docs)

    # fixme: don't use a list default argument
    def getAll(self, prj=[]):
        prj = {'.'.join(_): True for _ in prj}
        if len(prj) > 0:
            prj['objID'] = True
        prj['_id'] = False
        cursor = self.db.find({}, prj)
        docs = self._removeAID(cursor)
        return RetVal(True, None, docs)

# Connect to MongoDB and store the relevant collection handles in the
# 'dbHandles' variables to avoid hard coded collection names in Azrael.
client = config.getMongoClient()
dbName = 'azrael'
dbHandles = {
    'Commands': client[dbName]['Cmd'],
    'Templates': DatabaseMongo((dbName, 'template')),
    'ObjInstances': DatabaseMongo((dbName, 'objinstances')),
    'Counters': client[dbName]['Counters'],
    'Constraints': client[dbName]['Constraints'],
}
<|endoftext|>