<commit_before>#! /bin/python3/

"""
Author: @pme1123
Created: Jan 17th, 2017

Frangi Segmentation - combines various functions into a single one for convenience
Frangi Image Loop - For series analysis across directories.

"""

# TODO: Reminder - don't forget to update these imports!
import os
# from PIL import Image
import pandas as pd
from pyroots.image_manipulation import img_split
from pyroots.geometry_filters import morphology_filter, hollow_filter, diameter_filter
from pyroots.noise_filters import color_filter
from pyroots.summarize import bin_by_diameter, summarize_geometry
from pyroots.skeletonization import skeleton_with_distance
from skimage import io, color, filters, morphology, img_as_ubyte
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool
from numpy import array, uint8
import importlib

def frangi_segmentation(image, colors, frangi_args, threshold_args,
                        color_args_1=None, color_args_2=None,
                        morphology_args=None, hollow_args=None, hole_filling=None,
                        diameter_args=None, diameter_bins=None, image_name="image"):
    """
    Possible approach to object detection using frangi filters. Selects colorbands for
    analysis, runs frangi filter, thresholds to identify candidate objects, then removes
    spurrious objects by characteristics. See frangi_approach.ipynb.

    Parameters
    ----------
    image : ndarray
        RGB image to analyze
    colors : dict
        Parameters for picking the colorspace
    frangi_args : dict
        Parameters to pass to `skimage.filters.frangi`
    threshold_args : dict
        Parameters to pass to `skimage.filters.threshold_adaptive`
    color_args_1 : dict
        Parameters to pass to `pyroots.color_filter`.
    color_args_2 : list
        Parameters to pass to `pyroots.color_filter`. Combines with color_args_1
        in an 'and' statement.
    morphology_args : dict
        Parameters to pass to `pyroots.morphology_filter`
    hollow_args : dict
        Parameters to pass to `pyroots.hollow_filter`
    hole_filling : dict
        'selem' : ndarray to pass to `morphology.binary_closing`
        'min_size' : int to pass to `morphology.remove_small_holes`
    diameter_bins : list
        To pass to `pyroots.bin_by_diameter`
    optimize : bool
        Print images automatically, to facilitate parameter tweaking
    image_name : str
        Identifier of image for summarizing

    Returns
    -------
    A dictionary containing:
        1. `"geometry"` summary `pandas.DataFrame`
        2. `"objects"` binary image
        3. `"length"` medial axis image
        4. `"diameter"` medial axis image

    """

    working_image = image.copy()

    # Pull band from colorspace
    color_conversion = 'rgb2' + colors['colorspace'].lower()
    try:
        working_image = getattr(color, color_conversion)(working_image)
    except:
        pass

    working_image = img_split(working_image)[colors['band']]

    # Frangi vessel enhancement
    working_image = filters.frangi(working_image, **frangi_args)
    working_image = 1-(working_image/working_image.max())  # rescale and invert --> dark objects for thresholding

    # Threshold to ID candidate objects
    working_image = filters.threshold_adaptive(working_image, **threshold_args)
    working_image = ~working_image  # new line for clarity; --> objects = True

    # Filter candidate objects by color
    try:
        working_image = color_filter(image, working_image, **color_args_1)  #colorspace, target_band, low, high, percent)

        try:
            working_image = color_filter(image, working_image, **color_args_2)  # nesting equates to an "and" statement.
        except:
            pass
    except:
        pass

    # Filter candidate objects by morphology
    try:
        working_image = morphology_filter(working_image, **morphology_args)  #### ADD PR
    except:
        pass

    # Filter candidate objects by hollowness
    try:
        working_image = hollow_filter(working_image, **hollow_args)          #### ADD PR
    except:
        pass

    # Close small gaps and holes in accepted objects
    try:
        working_image = morphology.binary_closing(working_image, selem=hole_filling['selem'])
        working_image = morphology.remove_small_holes(working_image, min_size=hole_filling['min_size'])
    except:
        pass

    # Skeletonize
    skel = skeleton_with_distance(working_image)

    # Diameter filter
    diam = diameter_filter(skel, **diameter_args)

    # Summarize
    if diameter_bins is None:
        summary_df = summarize_geometry(diam['geometry'], image_name)

    else:
        diam_out, summary_df = bin_by_diameter(diam['length'],
                                               diam['diameter'],
                                               diameter_bins,
                                               image_name)
        diam['diameter'] = diam_out

    out = {'geometry' : summary_df,
           'objects'  : diam['objects'],
           'length'   : diam['length'],
           'diameter' : diam['diameter']}

    return(out)

##################################################################################################################################################################################
####################                                                                                ###############################
####################                                                 IMAGE LOOP - IN SERIES                                                        ###############################
####################                                                                                ###############################
##################################################################################################################################################################################

def frangi_image_loop(base_directory, image_extension, params=None, out_dir="Pyroots Analyzed", table_out="Pyroots_Analyzed.txt", mask=None, save_images=False, threads=1, extra_imports=None):
    """
    Reference function to loop through images in a directory. As it is written, it returns
    a dataframe from `pyroots.frangi_segmentation` and also writes images showing the objects analyzed.
    Note that all parameters in `pyroots.frangi_segmentation` must be defined in the params
    file, even as `None` for this to work.

    Use the parallel version for lots of images, if you have the capability (i.e. supercomputer or 4+ cores)

    Parameters
    ----------
    base_directory : str
        Directory to the images or subdirectories containing the images.
    image_extension : str
        Extension of images to analyze
    params : str
        Path + filename for parameters file for `pyroots.frangi_segmentation`. If None (default), define the parameters
        in the workspace using appropriately named dictionaries.
    out_dir : str
        Name of directory to write output images
    table_out : str
        Name of the table to which you want to export results. Currently only supports tab-delimeted.
    mask : ndarray
        TODO: For limiting the analysis
    save_images : bool
        Do you want to save images of the objects?
    threads : int
        For multiprocessing
    extra_imports : list
        If raises error importing params, then write a list of lists as
            [[lib1, fun1, fun2, ...],
             [lib2, fun2, fun2, ...],
             [...]].
    """
    if extra_imports is not None:
        for k in extra_imports:
            for i in range(1, len(k+1)):
                try:
                    importlib.import_module(k[0], k[i])
                except:
                    raise("%fn could not be loaded from %pkg" % (k[i], k[0]))

    try:
        exec(open(params).read(), globals())
    except:
        raise ValueError("Couldn't load params file. Try checking it for words like 'array' or\n'uint8' that need to be loaded with numpy and load these\n functions using `extra_imports`.... Or edit source.")

    #Make an output directory for the analyzed images and data.
    if save_images is True:
        if not os.path.exists(base_directory + os.sep + out_dir):
            os.mkdir(base_directory + os.sep + out_dir)

    #Make a placeholder DataFrame for the output data. Requires pandas.
    if diameter_bins is None:
        img_df = pd.DataFrame(columns=("ImageName", "Length", "NObjects", "MeanDiam"))
    else:
        img_df = pd.DataFrame(columns=("ImageName", "DiameterClass", "Length"))

    out = []  # for tracking progress
    #Begin looping
    for subdir, dirs, files in os.walk(base_directory):
        # What we'll do:
        def _core_fn(file_in):
            if file_in.endswith(image_extension) and not subdir.endswith(out_dir):

                path_in = subdir + os.sep + file_in  # what's the image called and where is it?
                file_name = subdir[len(base_directory):] + file_in

                #Create filein name
                filein = "DIR" + subdir[len(base_directory):] + os.sep + file_in
                #Import image
                img = io.imread(path_in)[:,:,0:3]  # load image

                #Run through the analysis. Requires pyroots.pyroots_analysis reference function.    ###############################################################
                objects_dict = frangi_segmentation(img, colors, frangi_args,                      ######   THIS IS WHERE YOU INSERT YOUR CUSTOM FUNCTION   ######
                                                      threshold_args, color_args_1,                    ###############################################################
                                                      color_args_2, morphology_args,
                                                      hollow_args, hole_filling,
                                                      diameter_args, diameter_bins,
                                                      image_name=file_name)
                #save images?
                if save_images is True:
                    #Convert boolean array to 8-bit, 1-band image. Requires PIL
    #                im = Image.fromarray((255 * objects_dict['objects']).astype('uint8')) #white roots, black background    #### MAY NEED TO UPDATE THIS.
                    path_out = base_directory + os.sep + out_dir + os.sep + subdir[len(base_directory):] + file_in
                    io.imsave(path_out, 255*objects_dict['objects'].astype('uint8'))

                print("Done: " + subdir[len(base_directory):] + file_in)
                df_out = objects_dict['geometry']
                return(df_out)

        # Init threads within each subdir (at file level)
        thread_pool = ThreadPool(threads)
        # Work on _core_fn
        out += thread_pool.map(_core_fn, files)
        # finish
        thread_pool.close()
        thread_pool.join()

    out = pd.concat([i for i in out])

    try:
        out.to_csv(base_directory + os.sep + out_dir + os.sep + table_out, sep = "\t", index = False)
    except:
        raise ValueError("Warning: Did not export output to a table.")
    #Done
    return(out)
<commit_msg>change image export to png to save space<commit_after>#! /bin/python3/

"""
Author: @pme1123
Created: Jan 17th, 2017

Frangi Segmentation - combines various functions into a single one for convenience
Frangi Image Loop - For series analysis across directories.

"""

# TODO: Reminder - don't forget to update these imports!
import os
# from PIL import Image
import pandas as pd
from pyroots.image_manipulation import img_split
from pyroots.geometry_filters import morphology_filter, hollow_filter, diameter_filter
from pyroots.noise_filters import color_filter
from pyroots.summarize import bin_by_diameter, summarize_geometry
from pyroots.skeletonization import skeleton_with_distance
from skimage import io, color, filters, morphology, img_as_ubyte
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool
from numpy import array, uint8
import importlib

def frangi_segmentation(image, colors, frangi_args, threshold_args,
                        color_args_1=None, color_args_2=None,
                        morphology_args=None, hollow_args=None, hole_filling=None,
                        diameter_args=None, diameter_bins=None, image_name="image"):
    """
    Possible approach to object detection using frangi filters. Selects colorbands for
    analysis, runs frangi filter, thresholds to identify candidate objects, then removes
    spurrious objects by characteristics. See frangi_approach.ipynb.

    Parameters
    ----------
    image : ndarray
        RGB image to analyze
    colors : dict
        Parameters for picking the colorspace
    frangi_args : dict
        Parameters to pass to `skimage.filters.frangi`
    threshold_args : dict
        Parameters to pass to `skimage.filters.threshold_adaptive`
    color_args_1 : dict
        Parameters to pass to `pyroots.color_filter`.
    color_args_2 : list
        Parameters to pass to `pyroots.color_filter`. Combines with color_args_1
        in an 'and' statement.
    morphology_args : dict
        Parameters to pass to `pyroots.morphology_filter`
    hollow_args : dict
        Parameters to pass to `pyroots.hollow_filter`
    hole_filling : dict
        'selem' : ndarray to pass to `morphology.binary_closing`
        'min_size' : int to pass to `morphology.remove_small_holes`
    diameter_bins : list
        To pass to `pyroots.bin_by_diameter`
    optimize : bool
        Print images automatically, to facilitate parameter tweaking
    image_name : str
        Identifier of image for summarizing

    Returns
    -------
    A dictionary containing:
        1. `"geometry"` summary `pandas.DataFrame`
        2. `"objects"` binary image
        3. `"length"` medial axis image
        4. `"diameter"` medial axis image

    """

    working_image = image.copy()

    # Pull band from colorspace
    color_conversion = 'rgb2' + colors['colorspace'].lower()
    try:
        working_image = getattr(color, color_conversion)(working_image)
    except:
        pass

    working_image = img_split(working_image)[colors['band']]

    # Frangi vessel enhancement
    working_image = filters.frangi(working_image, **frangi_args)
    working_image = 1-(working_image/working_image.max())  # rescale and invert --> dark objects for thresholding

    # Threshold to ID candidate objects
    working_image = filters.threshold_adaptive(working_image, **threshold_args)
    working_image = ~working_image  # new line for clarity; --> objects = True

    # Filter candidate objects by color
    try:
        working_image = color_filter(image, working_image, **color_args_1)  #colorspace, target_band, low, high, percent)

        try:
            working_image = color_filter(image, working_image, **color_args_2)  # nesting equates to an "and" statement.
        except:
            pass
    except:
        pass

    # Filter candidate objects by morphology
    try:
        working_image = morphology_filter(working_image, **morphology_args)  #### ADD PR
    except:
        pass

    # Filter candidate objects by hollowness
    try:
        working_image = hollow_filter(working_image, **hollow_args)          #### ADD PR
    except:
        pass

    # Close small gaps and holes in accepted objects
    try:
        working_image = morphology.binary_closing(working_image, selem=hole_filling['selem'])
        working_image = morphology.remove_small_holes(working_image, min_size=hole_filling['min_size'])
    except:
        pass

    # Skeletonize
    skel = skeleton_with_distance(working_image)

    # Diameter filter
    diam = diameter_filter(skel, **diameter_args)

    # Summarize
    if diameter_bins is None:
        summary_df = summarize_geometry(diam['geometry'], image_name)

    else:
        diam_out, summary_df = bin_by_diameter(diam['length'],
                                               diam['diameter'],
                                               diameter_bins,
                                               image_name)
        diam['diameter'] = diam_out

    out = {'geometry' : summary_df,
           'objects'  : diam['objects'],
           'length'   : diam['length'],
           'diameter' : diam['diameter']}

    return(out)

##################################################################################################################################################################################
####################                                                                                ###############################
####################                                                 IMAGE LOOP - IN SERIES                                                        ###############################
####################                                                                                ###############################
##################################################################################################################################################################################

def frangi_image_loop(base_directory, image_extension, params=None, out_dir="Pyroots Analyzed", table_out="Pyroots_Analyzed.txt", mask=None, save_images=False, threads=1, extra_imports=None):
    """
    Reference function to loop through images in a directory. As it is written, it returns
    a dataframe from `pyroots.frangi_segmentation` and also writes images showing the objects analyzed.
    Note that all parameters in `pyroots.frangi_segmentation` must be defined in the params
    file, even as `None` for this to work.

    Use the parallel version for lots of images, if you have the capability (i.e. supercomputer or 4+ cores)

    Parameters
    ----------
    base_directory : str
        Directory to the images or subdirectories containing the images.
    image_extension : str
        Extension of images to analyze
    params : str
        Path + filename for parameters file for `pyroots.frangi_segmentation`. If None (default), define the parameters
        in the workspace using appropriately named dictionaries.
    out_dir : str
        Name of directory to write output images
    table_out : str
        Name of the table to which you want to export results. Currently only supports tab-delimeted.
    mask : ndarray
        TODO: For limiting the analysis
    save_images : bool
        Do you want to save images of the objects?
    threads : int
        For multiprocessing
    extra_imports : list
        If raises error importing params, then write a list of lists as
            [[lib1, fun1, fun2, ...],
             [lib2, fun2, fun2, ...],
             [...]].
    """
    if extra_imports is not None:
        for k in extra_imports:
            for i in range(1, len(k+1)):
                try:
                    importlib.import_module(k[0], k[i])
                except:
                    raise("%fn could not be loaded from %pkg" % (k[i], k[0]))

    try:
        exec(open(params).read(), globals())
    except:
        raise ValueError("Couldn't load params file. Try checking it for words like 'array' or\n'uint8' that need to be loaded with numpy and load these\n functions using `extra_imports`.... Or edit source.")

    #Make an output directory for the analyzed images and data.
    if save_images is True:
        if not os.path.exists(base_directory + os.sep + out_dir):
            os.mkdir(base_directory + os.sep + out_dir)

    #Make a placeholder DataFrame for the output data. Requires pandas.
    if diameter_bins is None:
        img_df = pd.DataFrame(columns=("ImageName", "Length", "NObjects", "MeanDiam"))
    else:
        img_df = pd.DataFrame(columns=("ImageName", "DiameterClass", "Length"))

    out = []  # for tracking progress
    #Begin looping
    for subdir, dirs, files in os.walk(base_directory):
        # What we'll do:
        def _core_fn(file_in):
            if file_in.endswith(image_extension) and not subdir.endswith(out_dir):

                path_in = subdir + os.sep + file_in  # what's the image called and where is it?
                file_name = subdir[len(base_directory):] + file_in

                #Create filein name
                filein = "DIR" + subdir[len(base_directory):] + os.sep + file_in
                #Import image
                img = io.imread(path_in)[:,:,0:3]  # load image

                #Run through the analysis. Requires pyroots.pyroots_analysis reference function.    ###############################################################
                objects_dict = frangi_segmentation(img, colors, frangi_args,                      ######   THIS IS WHERE YOU INSERT YOUR CUSTOM FUNCTION   ######
                                                      threshold_args, color_args_1,                    ###############################################################
                                                      color_args_2, morphology_args,
                                                      hollow_args, hole_filling,
                                                      diameter_args, diameter_bins,
                                                      image_name=file_name)
                #save images?
                if save_images is True:
                    #Convert boolean array to 8-bit, 1-band image. Requires PIL
    #                im = Image.fromarray((255 * objects_dict['objects']).astype('uint8')) #white roots, black background    #### MAY NEED TO UPDATE THIS.
                    path_out = base_directory + os.sep + out_dir + os.sep + subdir[len(base_directory):] + os.path.splitext(file_in)[0] + '.png'
                    io.imsave(path_out, 255*objects_dict['objects'].astype('uint8'))

                print("Done: " + subdir[len(base_directory):] + file_in)
                df_out = objects_dict['geometry']
                return(df_out)

        # Init threads within each subdir (at file level)
        thread_pool = ThreadPool(threads)
        # Work on _core_fn
        out += thread_pool.map(_core_fn, files)
        # finish
        thread_pool.close()
        thread_pool.join()

    out = pd.concat([i for i in out])

    try:
        out.to_csv(base_directory + os.sep + out_dir + os.sep + table_out, sep = "\t", index = False)
    except:
        raise ValueError("Warning: Did not export output to a table.")
    #Done
    return(out)
<|endoftext|>


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

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

import numpy

from pysc2.agents import base_agent
from pysc2.lib import actions
from pysc2.lib import features

_PLAYER_SELF = features.PlayerRelative.SELF
_PLAYER_NEUTRAL = features.PlayerRelative.NEUTRAL  # beacon/minerals
_PLAYER_ENEMY = features.PlayerRelative.ENEMY

FUNCTIONS = actions.FUNCTIONS

def _xy_locs(mask):
  """Mask should be a set of bools from comparison with a feature layer."""
  y, x = mask.nonzero()
  return zip(x, y)

class MoveToBeacon(base_agent.BaseAgent):
  """An agent specifically for solving the MoveToBeacon map."""

  def step(self, obs):
    super(MoveToBeacon, self).step(obs)
    if FUNCTIONS.Move_screen.id in obs.observation.available_actions:
      player_relative = obs.observation.feature_screen.player_relative
      beacon = _xy_locs(player_relative == _PLAYER_NEUTRAL)
      if not beacon:
        return FUNCTIONS.no_op()
      beacon_center = numpy.mean(beacon, axis=0).round()
      return FUNCTIONS.Move_screen("now", beacon_center)
    else:
      return FUNCTIONS.select_army("select")

class CollectMineralShards(base_agent.BaseAgent):
  """An agent specifically for solving the CollectMineralShards map."""

  def step(self, obs):
    super(CollectMineralShards, self).step(obs)
    if FUNCTIONS.Move_screen.id in obs.observation.available_actions:
      player_relative = obs.observation.feature_screen.player_relative
      minerals = _xy_locs(player_relative == _PLAYER_NEUTRAL)
      if not minerals:
        return FUNCTIONS.no_op()
      marines = _xy_locs(player_relative == _PLAYER_SELF)
      marine_xy = numpy.mean(marines, axis=0).round()  # Average location.
      distances = numpy.linalg.norm(numpy.array(minerals) - marine_xy, axis=1)
      closest_mineral_xy = minerals[numpy.argmin(distances)]
      return FUNCTIONS.Move_screen("now", closest_mineral_xy)
    else:
      return FUNCTIONS.select_army("select")

class CollectMineralShardsFeatureUnits(base_agent.BaseAgent):
  """An agent for solving the CollectMineralShards map with feature units.

  Controls the two marines independently:
  - select marine
  - move to nearest mineral shard that wasn't the previous target
  - swap marine and repeat
  """

  def setup(self, obs_spec, action_spec):
    super(CollectMineralShardsFeatureUnits, self).setup(obs_spec, action_spec)
    if "feature_units" not in obs_spec:
      raise Exception("This agent requires the feature_units observation.")

  def reset(self):
    super(CollectMineralShardsFeatureUnits, self).reset()
    self._current_marine = 0
    self._previous_mineral_xy = [-1, -1]

  def step(self, obs):
    super(CollectMineralShardsFeatureUnits, self).step(obs)
    marines = [unit for unit in obs.observation.feature_units
               if unit.alliance == _PLAYER_SELF]
    if not marines:
      return FUNCTIONS.no_op()
    marine_unit = marines[self._current_marine]
    marine_xy = [marine_unit.x, marine_unit.y]

    if not marine_unit.is_selected:
      # Nothing selected or the wrong marine is selected.
      return FUNCTIONS.select_point("select", marine_xy)

    if FUNCTIONS.Move_screen.id in obs.observation.available_actions:
      # Find and move to the nearest mineral.
      minerals = [[unit.x, unit.y] for unit in obs.observation.feature_units
                  if unit.alliance == _PLAYER_NEUTRAL]

      if self._previous_mineral_xy in minerals:
        # Don't go for the same mineral shard as other marine.
        minerals.remove(self._previous_mineral_xy)

      if minerals:
        # Find the closest.
        distances = numpy.linalg.norm(
            numpy.array(minerals) - numpy.array(marine_xy), axis=1)
        closest_mineral_xy = minerals[numpy.argmin(distances)]

        # Swap to the other marine.
        self._current_marine = 1 - self._current_marine
        self._previous_mineral_xy = closest_mineral_xy
        return FUNCTIONS.Move_screen("now", closest_mineral_xy)

    return FUNCTIONS.no_op()

class DefeatRoaches(base_agent.BaseAgent):
  """An agent specifically for solving the DefeatRoaches map."""

  def step(self, obs):
    super(DefeatRoaches, self).step(obs)
    if FUNCTIONS.Attack_screen.id in obs.observation.available_actions:
      player_relative = obs.observation.feature_screen.player_relative
      roaches = _xy_locs(player_relative == _PLAYER_ENEMY)
      if not roaches:
        return FUNCTIONS.no_op()

      # Find the roach with max y coord.
      target = roaches[numpy.argmax(numpy.array(roaches)[:, 1])]
      return FUNCTIONS.Attack_screen("now", target)

    if FUNCTIONS.select_army.id in obs.observation.available_actions:
      return FUNCTIONS.select_army("select")

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

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

import numpy

from pysc2.agents import base_agent
from pysc2.lib import actions
from pysc2.lib import features

_PLAYER_SELF = features.PlayerRelative.SELF
_PLAYER_NEUTRAL = features.PlayerRelative.NEUTRAL  # beacon/minerals
_PLAYER_ENEMY = features.PlayerRelative.ENEMY

FUNCTIONS = actions.FUNCTIONS

def _xy_locs(mask):
  """Mask should be a set of bools from comparison with a feature layer."""
  y, x = mask.nonzero()
  return list(zip(x, y))

class MoveToBeacon(base_agent.BaseAgent):
  """An agent specifically for solving the MoveToBeacon map."""

  def step(self, obs):
    super(MoveToBeacon, self).step(obs)
    if FUNCTIONS.Move_screen.id in obs.observation.available_actions:
      player_relative = obs.observation.feature_screen.player_relative
      beacon = _xy_locs(player_relative == _PLAYER_NEUTRAL)
      if not beacon:
        return FUNCTIONS.no_op()
      beacon_center = numpy.mean(beacon, axis=0).round()
      return FUNCTIONS.Move_screen("now", beacon_center)
    else:
      return FUNCTIONS.select_army("select")

class CollectMineralShards(base_agent.BaseAgent):
  """An agent specifically for solving the CollectMineralShards map."""

  def step(self, obs):
    super(CollectMineralShards, self).step(obs)
    if FUNCTIONS.Move_screen.id in obs.observation.available_actions:
      player_relative = obs.observation.feature_screen.player_relative
      minerals = _xy_locs(player_relative == _PLAYER_NEUTRAL)
      if not minerals:
        return FUNCTIONS.no_op()
      marines = _xy_locs(player_relative == _PLAYER_SELF)
      marine_xy = numpy.mean(marines, axis=0).round()  # Average location.
      distances = numpy.linalg.norm(numpy.array(minerals) - marine_xy, axis=1)
      closest_mineral_xy = minerals[numpy.argmin(distances)]
      return FUNCTIONS.Move_screen("now", closest_mineral_xy)
    else:
      return FUNCTIONS.select_army("select")

class CollectMineralShardsFeatureUnits(base_agent.BaseAgent):
  """An agent for solving the CollectMineralShards map with feature units.

  Controls the two marines independently:
  - select marine
  - move to nearest mineral shard that wasn't the previous target
  - swap marine and repeat
  """

  def setup(self, obs_spec, action_spec):
    super(CollectMineralShardsFeatureUnits, self).setup(obs_spec, action_spec)
    if "feature_units" not in obs_spec:
      raise Exception("This agent requires the feature_units observation.")

  def reset(self):
    super(CollectMineralShardsFeatureUnits, self).reset()
    self._current_marine = 0
    self._previous_mineral_xy = [-1, -1]

  def step(self, obs):
    super(CollectMineralShardsFeatureUnits, self).step(obs)
    marines = [unit for unit in obs.observation.feature_units
               if unit.alliance == _PLAYER_SELF]
    if not marines:
      return FUNCTIONS.no_op()
    marine_unit = marines[self._current_marine]
    marine_xy = [marine_unit.x, marine_unit.y]

    if not marine_unit.is_selected:
      # Nothing selected or the wrong marine is selected.
      return FUNCTIONS.select_point("select", marine_xy)

    if FUNCTIONS.Move_screen.id in obs.observation.available_actions:
      # Find and move to the nearest mineral.
      minerals = [[unit.x, unit.y] for unit in obs.observation.feature_units
                  if unit.alliance == _PLAYER_NEUTRAL]

      if self._previous_mineral_xy in minerals:
        # Don't go for the same mineral shard as other marine.
        minerals.remove(self._previous_mineral_xy)

      if minerals:
        # Find the closest.
        distances = numpy.linalg.norm(
            numpy.array(minerals) - numpy.array(marine_xy), axis=1)
        closest_mineral_xy = minerals[numpy.argmin(distances)]

        # Swap to the other marine.
        self._current_marine = 1 - self._current_marine
        self._previous_mineral_xy = closest_mineral_xy
        return FUNCTIONS.Move_screen("now", closest_mineral_xy)

    return FUNCTIONS.no_op()

class DefeatRoaches(base_agent.BaseAgent):
  """An agent specifically for solving the DefeatRoaches map."""

  def step(self, obs):
    super(DefeatRoaches, self).step(obs)
    if FUNCTIONS.Attack_screen.id in obs.observation.available_actions:
      player_relative = obs.observation.feature_screen.player_relative
      roaches = _xy_locs(player_relative == _PLAYER_ENEMY)
      if not roaches:
        return FUNCTIONS.no_op()

      # Find the roach with max y coord.
      target = roaches[numpy.argmax(numpy.array(roaches)[:, 1])]
      return FUNCTIONS.Attack_screen("now", target)

    if FUNCTIONS.select_army.id in obs.observation.available_actions:
      return FUNCTIONS.select_army("select")

    return FUNCTIONS.no_op()
<|endoftext|>


<commit_before>#!/usr/bin/env python
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
import numpy
import os
import sys
import platform
import traceback
win_cuda_dir = None

def find_cuda():
    if 'CUDAHOME' in os.environ:
        home = os.environ['CUDAHOME']
        nvcc = os.path.join(home, 'bin', 'nvcc')
    else:
        nvcc = None
        for dir in os.environ['PATH'].split(os.pathsep):
            binpath = os.path.join(dir, 'nvcc')
            if os.path.exists(binpath):
                nvcc = os.path.abspath(binpath)
        if nvcc is None:
            raise EnvironmentError('The nvcc binary could not be located in '
                                   'your $PATH. Either add it to your path, or'
                                   'set $CUDAHOME')
        home = os.path.dirname(os.path.dirname(nvcc))
    cudaconfig = {'home': home, 'nvcc': nvcc,
                  'include': os.path.join(home, 'include')}
    for k, v in cudaconfig.items():
        if not os.path.exists(v):
            raise EnvironmentError('The CUDA %s path could not be located in '
                                   '%s' % (k, v))
    libdir = os.path.join(home, 'lib')
    arch = int(platform.architecture()[0][0:2])
    if sys.platform.startswith('win'):
        os.path.join(libdir, "x"+str(arch))
    if os.path.exists(os.path.join(home, libdir + "64")):
        cudaconfig['lib'] = libdir + "64"
    elif os.path.exists(os.path.join(home, libdir)):
        cudaconfig['lib'] = libdir
    else:
        raise EnvironmentError('The CUDA libraries could not be located')
    return cudaconfig

try:
    CUDA = find_cuda()
except EnvironmentError:
    CUDA = None
    print("Proceeding without CUDA")

try:
    numpy_include = numpy.get_include()
except AttributeError:
    numpy_include = numpy.get_numpy_include()

def customize_compiler_for_nvcc(self):
    '''This is a verbatim copy of the NVCC compiler extension from
    https://github.com/rmcgibbo/npcuda-example
    '''
    self.src_extensions.append('.cu')
    default_compiler_so = self.compiler_so
    super = self._compile

    def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
        if os.path.splitext(src)[1] == '.cu':
            self.set_executable('compiler_so', CUDA['nvcc'])
            postargs = extra_postargs['nvcc']
        else:
            postargs = extra_postargs['cc']

        super(obj, src, ext, cc_args, postargs, pp_opts)
        self.compiler_so = default_compiler_so
    self._compile = _compile

# run the customize_compiler
class custom_build_ext(build_ext):
    def build_extensions(self):
        customize_compiler_for_nvcc(self.compiler)
        build_ext.build_extensions(self)

cmdclass = {}
if sys.platform.startswith('win') and win_cuda_dir is not None:
    if win_cuda_dir == "":
        if 'CUDA_PATH' in os.environ:
            win_cuda_dir = os.environ['CUDA_PATH']
    elif os.path.exists(win_cuda_dir):
        pass
    else:
        win_cuda_dir = None
    if win_cuda_dir:
        arch = int(platform.architecture()[0][0:2])
        somoclu_module = Extension('_somoclu_wrap',
                                   sources=['somoclu/somoclu_wrap.cxx'],
                                   extra_objects=[
                                            'somoclu/src/denseCpuKernels.obj',
                                            'somoclu/src/sparseCpuKernels.obj',
                                            'somoclu/src/training.obj',
                                            'somoclu/src/mapDistanceFunctions.obj',
                                            'somoclu/src/uMatrix.obj',
                                            'somoclu/src/denseGpuKernels.cu.obj'],
                                   define_macros=[('CUDA', None)],
                                   library_dirs=[win_cuda_dir+"/lib/x"+str(arch)],
                                   libraries=['cudart', 'cublas'],
                                   include_dirs=[numpy_include])
else:
    sources_files = ['somoclu/src/denseCpuKernels.cpp',
                     'somoclu/src/sparseCpuKernels.cpp',
                     'somoclu/src/mapDistanceFunctions.cpp',
                     'somoclu/src/training.cpp',
                     'somoclu/src/uMatrix.cpp',
                     'somoclu/somoclu_wrap.cxx']
    if sys.platform.startswith('win'):
        extra_compile_args = ['-openmp']
        cmdclass = {}
        somoclu_module = Extension('_somoclu_wrap',
                               sources=sources_files,
                               include_dirs=[numpy_include, 'src'],
                               extra_compile_args=extra_compile_args,
                               )
    else:
        extra_compile_args = ['-fopenmp']
        if 'CC' in os.environ and 'clang-omp' in os.environ['CC']:
            openmp = 'iomp5'
        else:
            openmp = 'gomp'
        cmdclass = {'build_ext': custom_build_ext}
        somoclu_module = Extension('_somoclu_wrap',
                                   sources=sources_files,
                                   include_dirs=[numpy_include, 'src'],
                                   extra_compile_args={'cc': extra_compile_args},
                                   libraries=[openmp],
                                   )
    if CUDA is not None:
        somoclu_module.sources.append('somoclu/src/denseGpuKernels.cu')
        somoclu_module.define_macros = [('CUDA', None)]
        somoclu_module.include_dirs.append(CUDA['include'])
        somoclu_module.library_dirs = [CUDA['lib']]
        somoclu_module.libraries += ['cudart', 'cublas']
        somoclu_module.runtime_library_dirs = [CUDA['lib']]
        somoclu_module.extra_compile_args['nvcc']=['-use_fast_math',
                                                   '--ptxas-options=-v', '-c',
                                                   '--compiler-options','-fPIC ' +
                                                   extra_compile_args[0]]

try:
    setup(name='somoclu',
          version='1.7.4',
          license='GPL3',
          author="Peter Wittek, Shi Chao Gao",
          author_email="",
          maintainer="shichaogao",
          maintainer_email="xgdgsc@gmail.com",
          url="https://somoclu.readthedocs.io/",
          platforms=["unix", "windows"],
          description="Massively parallel implementation of self-organizing maps",
          ext_modules=[somoclu_module],
          packages=["somoclu"],
          install_requires=['numpy', 'matplotlib'],
          classifiers=[
            'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
            'Operating System :: OS Independent',
            'Development Status :: 5 - Production/Stable',
            'Topic :: Scientific/Engineering :: Artificial Intelligence',
            'Topic :: Scientific/Engineering :: Visualization',
            'Intended Audience :: Science/Research',
            'Programming Language :: Python',
            'Programming Language :: C++'
          ],
          cmdclass=cmdclass,
          test_suite="tests"
          )
except:
    traceback.print_exc()
    setup(name='somoclu',
          version='1.7.4',
          license='GPL3',
          author="Peter Wittek, Shi Chao Gao",
          author_email="",
          maintainer="shichaogao",
          maintainer_email="xgdgsc@gmail.com",
          url="https://somoclu.readthedocs.io/",
          platforms=["unix", "windows"],
          description="Massively parallel implementation of self-organizing maps",
          packages=["somoclu"],
          classifiers=[
            'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
            'Operating System :: OS Independent',
            'Development Status :: 5 - Production/Stable',
            'Topic :: Scientific/Engineering :: Artificial Intelligence',
            'Topic :: Scientific/Engineering :: Visualization',
            'Intended Audience :: Science/Research',
            'Programming Language :: Python',
            'Programming Language :: C++'
          ],
          install_requires=['numpy', 'matplotlib'],
          test_suite="tests"
          )
<commit_msg>add dep<commit_after>#!/usr/bin/env python
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
import numpy
import os
import sys
import platform
import traceback
win_cuda_dir = None

def find_cuda():
    if 'CUDAHOME' in os.environ:
        home = os.environ['CUDAHOME']
        nvcc = os.path.join(home, 'bin', 'nvcc')
    else:
        nvcc = None
        for dir in os.environ['PATH'].split(os.pathsep):
            binpath = os.path.join(dir, 'nvcc')
            if os.path.exists(binpath):
                nvcc = os.path.abspath(binpath)
        if nvcc is None:
            raise EnvironmentError('The nvcc binary could not be located in '
                                   'your $PATH. Either add it to your path, or'
                                   'set $CUDAHOME')
        home = os.path.dirname(os.path.dirname(nvcc))
    cudaconfig = {'home': home, 'nvcc': nvcc,
                  'include': os.path.join(home, 'include')}
    for k, v in cudaconfig.items():
        if not os.path.exists(v):
            raise EnvironmentError('The CUDA %s path could not be located in '
                                   '%s' % (k, v))
    libdir = os.path.join(home, 'lib')
    arch = int(platform.architecture()[0][0:2])
    if sys.platform.startswith('win'):
        os.path.join(libdir, "x"+str(arch))
    if os.path.exists(os.path.join(home, libdir + "64")):
        cudaconfig['lib'] = libdir + "64"
    elif os.path.exists(os.path.join(home, libdir)):
        cudaconfig['lib'] = libdir
    else:
        raise EnvironmentError('The CUDA libraries could not be located')
    return cudaconfig

try:
    CUDA = find_cuda()
except EnvironmentError:
    CUDA = None
    print("Proceeding without CUDA")

try:
    numpy_include = numpy.get_include()
except AttributeError:
    numpy_include = numpy.get_numpy_include()

def customize_compiler_for_nvcc(self):
    '''This is a verbatim copy of the NVCC compiler extension from
    https://github.com/rmcgibbo/npcuda-example
    '''
    self.src_extensions.append('.cu')
    default_compiler_so = self.compiler_so
    super = self._compile

    def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
        if os.path.splitext(src)[1] == '.cu':
            self.set_executable('compiler_so', CUDA['nvcc'])
            postargs = extra_postargs['nvcc']
        else:
            postargs = extra_postargs['cc']

        super(obj, src, ext, cc_args, postargs, pp_opts)
        self.compiler_so = default_compiler_so
    self._compile = _compile

# run the customize_compiler
class custom_build_ext(build_ext):
    def build_extensions(self):
        customize_compiler_for_nvcc(self.compiler)
        build_ext.build_extensions(self)

cmdclass = {}
if sys.platform.startswith('win') and win_cuda_dir is not None:
    if win_cuda_dir == "":
        if 'CUDA_PATH' in os.environ:
            win_cuda_dir = os.environ['CUDA_PATH']
    elif os.path.exists(win_cuda_dir):
        pass
    else:
        win_cuda_dir = None
    if win_cuda_dir:
        arch = int(platform.architecture()[0][0:2])
        somoclu_module = Extension('_somoclu_wrap',
                                   sources=['somoclu/somoclu_wrap.cxx'],
                                   extra_objects=[
                                            'somoclu/src/denseCpuKernels.obj',
                                            'somoclu/src/sparseCpuKernels.obj',
                                            'somoclu/src/training.obj',
                                            'somoclu/src/mapDistanceFunctions.obj',
                                            'somoclu/src/uMatrix.obj',
                                            'somoclu/src/denseGpuKernels.cu.obj'],
                                   define_macros=[('CUDA', None)],
                                   library_dirs=[win_cuda_dir+"/lib/x"+str(arch)],
                                   libraries=['cudart', 'cublas'],
                                   include_dirs=[numpy_include])
else:
    sources_files = ['somoclu/src/denseCpuKernels.cpp',
                     'somoclu/src/sparseCpuKernels.cpp',
                     'somoclu/src/mapDistanceFunctions.cpp',
                     'somoclu/src/training.cpp',
                     'somoclu/src/uMatrix.cpp',
                     'somoclu/somoclu_wrap.cxx']
    if sys.platform.startswith('win'):
        extra_compile_args = ['-openmp']
        cmdclass = {}
        somoclu_module = Extension('_somoclu_wrap',
                               sources=sources_files,
                               include_dirs=[numpy_include, 'src'],
                               extra_compile_args=extra_compile_args,
                               )
    else:
        extra_compile_args = ['-fopenmp']
        if 'CC' in os.environ and 'clang-omp' in os.environ['CC']:
            openmp = 'iomp5'
        else:
            openmp = 'gomp'
        cmdclass = {'build_ext': custom_build_ext}
        somoclu_module = Extension('_somoclu_wrap',
                                   sources=sources_files,
                                   include_dirs=[numpy_include, 'src'],
                                   extra_compile_args={'cc': extra_compile_args},
                                   libraries=[openmp],
                                   )
    if CUDA is not None:
        somoclu_module.sources.append('somoclu/src/denseGpuKernels.cu')
        somoclu_module.define_macros = [('CUDA', None)]
        somoclu_module.include_dirs.append(CUDA['include'])
        somoclu_module.library_dirs = [CUDA['lib']]
        somoclu_module.libraries += ['cudart', 'cublas']
        somoclu_module.runtime_library_dirs = [CUDA['lib']]
        somoclu_module.extra_compile_args['nvcc']=['-use_fast_math',
                                                   '--ptxas-options=-v', '-c',
                                                   '--compiler-options','-fPIC ' +
                                                   extra_compile_args[0]]

try:
    setup(name='somoclu',
          version='1.7.4',
          license='GPL3',
          author="Peter Wittek, Shi Chao Gao",
          author_email="",
          maintainer="shichaogao",
          maintainer_email="xgdgsc@gmail.com",
          url="https://somoclu.readthedocs.io/",
          platforms=["unix", "windows"],
          description="Massively parallel implementation of self-organizing maps",
          ext_modules=[somoclu_module],
          packages=["somoclu"],
          install_requires=['numpy', 'matplotlib', 'scipy'],
          classifiers=[
            'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
            'Operating System :: OS Independent',
            'Development Status :: 5 - Production/Stable',
            'Topic :: Scientific/Engineering :: Artificial Intelligence',
            'Topic :: Scientific/Engineering :: Visualization',
            'Intended Audience :: Science/Research',
            'Programming Language :: Python',
            'Programming Language :: C++'
          ],
          cmdclass=cmdclass,
          test_suite="tests"
          )
except:
    traceback.print_exc()
    setup(name='somoclu',
          version='1.7.4',
          license='GPL3',
          author="Peter Wittek, Shi Chao Gao",
          author_email="",
          maintainer="shichaogao",
          maintainer_email="xgdgsc@gmail.com",
          url="https://somoclu.readthedocs.io/",
          platforms=["unix", "windows"],
          description="Massively parallel implementation of self-organizing maps",
          packages=["somoclu"],
          classifiers=[
            'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
            'Operating System :: OS Independent',
            'Development Status :: 5 - Production/Stable',
            'Topic :: Scientific/Engineering :: Artificial Intelligence',
            'Topic :: Scientific/Engineering :: Visualization',
            'Intended Audience :: Science/Research',
            'Programming Language :: Python',
            'Programming Language :: C++'
          ],
          install_requires=['numpy', 'matplotlib', 'scipy'],
          test_suite="tests"
          )
<|endoftext|>


<commit_before>from pyshelf.bucket_update.container import Container
import logging
import os.path

def update_search_index(bucket_config):
    """
        Updates the search index for a specific bucket.

        Args:
            bucket_config(schemas/search-bucket-update-config.json)
    """
    container = _create_container(bucket_config)
    container.search_updater.run()

def _create_container(bucket_config):
    filename = os.path.join(bucket_config["bulkUpdateLogDirectory"], "{0}.log".format(bucket_config["name"]))
    logging.basicConfig(
        format="%(asctime)s:%(name)s:%(levelname)s:%(message)s",
        filename=filename,
        level=bucket_config["logLevel"],
    )
    logger = logging.getLogger(bucket_config["name"])
    container = Container(bucket_config, logger)
    return container
<commit_msg>specifically setting up elasticsearch instead<commit_after>from pyshelf.bucket_update.container import Container
import logging
import os.path

def update_search_index(bucket_config):
    """
        Updates the search index for a specific bucket.

        Args:
            bucket_config(schemas/search-bucket-update-config.json)
    """
    container = _create_container(bucket_config)
    container.search_updater.run()

def _create_container(bucket_config):
    filename = os.path.join(bucket_config["bulkUpdateLogDirectory"], "{0}.log".format(bucket_config["name"]))
    log_level = bucket_config["logLevel"]
    _configure_logger("elasticsearch", filename, log_level)
    logger = _configure_logger(bucket_config["name"], filename, log_level)
    container = Container(bucket_config, logger)
    return container

def _configure_logger(name, filename, log_level):
    logger = logging.getLogger(name)
    handler = logging.FileHandler(filename)
    formatter = logging.Formatter(fmt="%(asctime)s:%(name)s:%(levelname)s:%(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(log_level)
    return logger
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from .base import Exporter

from ..subtitle import HumanTime

class SubRipExporter(Exporter):
  """Exported for SubRip format."""
  FORMAT = 'SubRip'

  def _init(self, encoding = 'utf-8', line_ending = b'\n'):
    self._encoding = encoding
    self._line_ending = line_ending

  @staticmethod
  def _convert_time(time):
    output = []

    if isinstance(time, (float, int)):
      time = HumanTime.from_seconds(time)
    elif not isinstance(time, HumanTime):
      raise TypeError("Expecting time")

    output.append('{:02d}'.format(time.hours))
    output.append('{:02d}'.format(time.minutes))
    seconds = int(time.seconds)
    miliseconds = int((time.seconds - seconds) * 1000)
    output.append('{:02d},{:03d}'.format(seconds, miliseconds))

    return ':'.join(output)

  def _export_metadata(self, metadata):
    # No subtitle wide metadata, just reset counter
    self._unit = 0
    return b''

  def _export_unit(self, unit):
    output = []

    if self._unit:
      # An empty line at the beginning
      output.append(b'')
    self._unit += 1

    # Sequence
    output.append(str(self._unit).encode(self._encoding))
    # Timing
    # TODO 3D positions
    output.append("{} --> {}".format(self._convert_time(unit.start),
                                     self._convert_time(unit.end)).encode(self._encoding))
    # Text
    output.append(self._line_ending.join([i.encode(self._encoding) for i in unit.lines]))

    # End of line
    output.append(b'')

    # All done
    return self._line_ending.join(output)

  def _export_end(self, metadata):
    # No specific footer also
    return b''
<commit_msg>Ignore encoding errors. It is probably a result from a unicode non-printable character anyway.<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from .base import Exporter

from ..subtitle import HumanTime

class SubRipExporter(Exporter):
  """Exported for SubRip format."""
  FORMAT = 'SubRip'

  def _init(self, encoding = 'utf-8', line_ending = b'\n'):
    self._encoding = encoding
    self._line_ending = line_ending

  @staticmethod
  def _convert_time(time):
    output = []

    if isinstance(time, (float, int)):
      time = HumanTime.from_seconds(time)
    elif not isinstance(time, HumanTime):
      raise TypeError("Expecting time")

    output.append('{:02d}'.format(time.hours))
    output.append('{:02d}'.format(time.minutes))
    seconds = int(time.seconds)
    miliseconds = int((time.seconds - seconds) * 1000)
    output.append('{:02d},{:03d}'.format(seconds, miliseconds))

    return ':'.join(output)

  def _export_metadata(self, metadata):
    # No subtitle wide metadata, just reset counter
    self._unit = 0
    return b''

  def _export_unit(self, unit):
    output = []

    if self._unit:
      # An empty line at the beginning
      output.append(b'')
    self._unit += 1

    # Sequence
    output.append(str(self._unit).encode(self._encoding))
    # Timing
    # TODO 3D positions
    output.append("{} --> {}".format(self._convert_time(unit.start),
                                     self._convert_time(unit.end)).encode(self._encoding))
    # Text
    output.append(self._line_ending.join([i.encode(self._encoding, 'ignore') for i in unit.lines]))

    # End of line
    output.append(b'')

    # All done
    return self._line_ending.join(output)

  def _export_end(self, metadata):
    # No specific footer also
    return b''
<|endoftext|>


<commit_before>from os.path import isdir, join
from glob import iglob
from xml.etree.ElementTree import parse
from datetime import datetime, timedelta

# syspol reference
file_ext = '.tsk'

DEFAULT_TIMEDELTA = timedelta(weeks=1)
DEFAULT_DATETIME_FMT = '%Y-%m-%d %H:%M:%S'

def get_categories():
    return ()

def get_category_efforts(categories=(), *, paths=None):
    if paths is not None:
        for path in paths:
            if isdir(path):
                for tsk in iglob(join(path, '*'+file_ext)):
                    return _tsk_file_get_category_efforts(categories, tsk)

def _tsk_file_get_category_efforts(categories, tskfp):
    effort_time = timedelta()
    doc = parse(tskfp)
    categories = tuple(categories)
    for tskcategory in doc.iterfind('category'):
        subject = tskcategory.get('subject')
        if subject not in categories:
            continue
        tasks = tuple(tskcategory.get('categorizables').split())
        for task in tasks:
            #print(task)
            #print(get_task_effort(task, tskfp))
            effort_time += get_task_effort(task, tskfp)
            pass
        yield (subject, effort_time)

def get_task_effort(tskid, tskfp, start=None, end=None):
    effort_time = timedelta()
    if start is None:
        start = datetime.now() - DEFAULT_TIMEDELTA

    doc = parse(tskfp)
    task = None
    for tsktask in doc.iterfind('task'):
        if tsktask.get('id') == tskid:
            task = tsktask
            break
    if task is None:
        return effort_time

    #print('Task `{}`'.format(task.get('subject')))

    for effort in task.iterfind('effort'):
        try:
            effort_start = datetime.strptime(effort.get('start'),
                    DEFAULT_DATETIME_FMT)
            effort_end = datetime.strptime(effort.get('stop'),
                    DEFAULT_DATETIME_FMT)
        except TypeError:
            print('Effort id `{}` has no stop attribute'.format(effort.get('id')))
            continue
        if effort_start.date() < start.date():
            continue
        if end is not None and effort_end.date() > end.date():
            continue
        effort_time += effort_end - effort_start

    return effort_time
<commit_msg>`pysyspol.taskcoach.get_task_effort` is now recursive, using XPath<commit_after>from os.path import isdir, join
from glob import iglob
from xml.etree.ElementTree import parse
from datetime import datetime, timedelta

# syspol reference
file_ext = '.tsk'

DEFAULT_TIMEDELTA = timedelta(weeks=1)
DEFAULT_DATETIME_FMT = '%Y-%m-%d %H:%M:%S'

def get_categories():
    return ()

def get_category_efforts(categories=(), *, paths=None):
    if paths is not None:
        for path in paths:
            if isdir(path):
                for tsk in iglob(join(path, '*'+file_ext)):
                    return _tsk_file_get_category_efforts(categories, tsk)

def _tsk_file_get_category_efforts(categories, tskfp):
    doc = parse(tskfp)
    categories = tuple(categories)
    for tskcategory in doc.iterfind('category'):
        effort_time = timedelta()
        subject = tskcategory.get('subject')
        if subject not in categories:
            continue
        tasks = tuple(tskcategory.get('categorizables').split())
        for task in tasks:
            effort_time += get_task_effort(task, tskfp)
        yield (subject, effort_time)

def get_task_effort(tskid, tskfp, start=None, end=None):
    effort_time = timedelta()
    if start is None:
        start = datetime.now() - DEFAULT_TIMEDELTA

    doc = parse(tskfp)
    task = doc.find(".//task[@id='{}']".format(tskid))
    if task is None:
        return effort_time

    for effort in task.iterfind('effort'):
        try:
            effort_start = datetime.strptime(effort.get('start'),
                    DEFAULT_DATETIME_FMT)
            effort_end = datetime.strptime(effort.get('stop'),
                    DEFAULT_DATETIME_FMT)
        except TypeError:
            print('Effort id `{}` has no stop attribute'.format(
                effort.get('id')), file=sys.stderr)
            continue
        if effort_start < start:
            continue
        if end is not None and effort_end > end:
            continue
        effort_time += effort_end - effort_start

    return effort_time
<|endoftext|>


<commit_before><commit_msg>scrapepinterest wasn't using password<commit_after><|endoftext|>


<commit_before>from sqlalchemy.dialects.postgresql import UUID
from podhub.website import db
from . import Base

class User(Base):
    profile = db.relationship('Profile', backref='user', lazy='dynamic',
                              uselist=False)

class Profile(Base):
    username = db.Column(db.String(80), index=True, unique=True)
    email = db.Column(db.String(120), index=True, unique=True)
    display_name = db.Column(db.String(80), index=True, unique=True)
    first_name = db.Column(db.String(80), index=True)
    last_name = db.Column(db.String(80), index=True)
    user_id = db.Column(UUID, db.ForeignKey('user.id'))
    author_id = db.Column(UUID, db.ForeignKey('author.id'))

    def __init__(self, username, email, first_name=None, last_name=None,
                 display_name=None, user_id=None, author_id=None):
        self.username = username
        self.email = email
        if not display_name:
            self.display_name = username
        else:
            self.display_name = display_name
        self.first_name = first_name
        self.last_name = last_name

        self.user_id = user_id
        self.author_id = author_id

class Author(Base):
    name = db.Column(db.String(80), index=True)
    Profile = db.relationship('Profile', backref='author', lazy='dynamic',
                              uselist=False)

    def __init__(self, name):
        self.name = name
<commit_msg>add Publisher class<commit_after>from sqlalchemy.dialects.postgresql import UUID
from podhub.website import db
from . import Base

class User(Base):
    profile = db.relationship('Profile', backref='user', lazy='dynamic',
                              uselist=False)

class Profile(Base):
    username = db.Column(db.String(80), index=True, unique=True)
    email = db.Column(db.String(120), index=True, unique=True)
    display_name = db.Column(db.String(80), index=True, unique=True)
    first_name = db.Column(db.String(80), index=True)
    last_name = db.Column(db.String(80), index=True)
    user_id = db.Column(UUID, db.ForeignKey('user.id'))
    author_id = db.Column(UUID, db.ForeignKey('author.id'))

    def __init__(self, username, email, first_name=None, last_name=None,
                 display_name=None, user_id=None, author_id=None):
        self.username = username
        self.email = email
        if not display_name:
            self.display_name = username
        else:
            self.display_name = display_name
        self.first_name = first_name
        self.last_name = last_name

        self.user_id = user_id
        self.author_id = author_id

class Author(Base):
    name = db.Column(db.String(80), index=True)
    profile = db.relationship('Profile', backref='author', lazy='dynamic',
                              uselist=False)

    def __init__(self, name):
        self.name = name

class Publisher(Base):
    name = db.Column(db.String(80), index=True)
    profile = db.relationship('Profile', backref='publisher', lazy='dynamic',
                              uselist=False)

    def __init__(self, name):
        self.name = name
<|endoftext|>


<commit_before><commit_msg>Hardcoded Microscope control<commit_after><|endoftext|>


<commit_before>import httplib2
import json
import os
import sys

from functools import wraps
from codemangler import app
from codemangler.models.user import User, UserModel
from config import MongoConfig
from flask import request, render_template, url_for, redirect, session
from oauth2client.client import OAuth2WebServerFlow

# TODO: move values to config
FLOW = OAuth2WebServerFlow(client_id=os.environ['UTEACH_OAUTH2_CLIENT_ID'],
                           client_secret=os.environ['UTEACH_OAUTH2_CLIENT_SECRET'],
                           auth_uri='https://umairidris.auth0.com/authorize',
                           token_uri='https://umairidris.auth0.com/oauth/token',
                           scope='openid name email nickname',
                           redirect_uri='http://localhost:8000/oauth2callback')
AUTH_URL = FLOW.step1_get_authorize_url() + '&connection=Username-Password-Authentication'

def login_required(f):
    """ (function) -> function

    Wrap views for users so that only
    logged in users can view those pages
    """

    @wraps(f)
    def wrap(*args, **kwargs):
        if 'logged_in' in session:
            return f(*args, **kwargs)
        else:
            return redirect(AUTH_URL)

    return wrap

@app.route('/login')
def get_login():
    return redirect(AUTH_URL)

@app.route('/oauth2callback')
def oauth2_callback():
    error = request.args.get('error')
    if error:
        return '{error}: {error_description}'.format(error=error, error_description=request.args.get('error_description')), 401

    code = request.args.get('code')
    if not code:
        return 'No code given', 401

    credentials = FLOW.step2_exchange(code)
    http = httplib2.Http()
    http = credentials.authorize(http)
    response, content = http.request('https://umairidris.auth0.com/userinfo')

    if response.status != 200:
        return response.reason, response.status

    user_dict = json.loads(content.decode("utf-8"))

    # username is email for now
    user = UserModel.get({'username': user_dict['email']})

    if not user: # new user
        user = User(user_dict['email'], user_dict['nickname'], user_dict['email'])
        user = UserModel.create(user)

        if not user:
            return 'Failed to create user', 500

    session['username'] = user.username
    session['logged_in'] = True

    if user.user_type == 'admin':
        session['admin'] = True
        return redirect(url_for('get_admin'))

    return redirect('/')

@app.route('/logout')
@login_required
def logout():
    """ () -> rendered_template

    Returns the rendered template of login.html,
    after the user makes a POST request to 'logout'
    """
    session.pop('username', None)
    session.pop('logged_in', None)
    session.pop('admin', None)
    return redirect('/')
<commit_msg>Add keys check<commit_after>import httplib2
import json
import os
import sys

from functools import wraps
from codemangler import app
from codemangler.models.user import User, UserModel
from config import MongoConfig
from flask import request, render_template, url_for, redirect, session
from oauth2client.client import OAuth2WebServerFlow

# TODO: move values to config
FLOW = OAuth2WebServerFlow(client_id=os.environ['UTEACH_OAUTH2_CLIENT_ID'],
                           client_secret=os.environ['UTEACH_OAUTH2_CLIENT_SECRET'],
                           auth_uri='https://umairidris.auth0.com/authorize',
                           token_uri='https://umairidris.auth0.com/oauth/token',
                           scope='openid name email nickname',
                           redirect_uri='http://localhost:8000/oauth2callback')
AUTH_URL = FLOW.step1_get_authorize_url() + '&connection=Username-Password-Authentication'
REQUIRED_RESPONSE_KEYS = ['email', 'name', 'nickname']

def login_required(f):
    """ (function) -> function

    Wrap views for users so that only
    logged in users can view those pages
    """

    @wraps(f)
    def wrap(*args, **kwargs):
        if 'logged_in' in session:
            return f(*args, **kwargs)
        else:
            return redirect(AUTH_URL)

    return wrap

@app.route('/login')
def get_login():
    return redirect(AUTH_URL)

@app.route('/oauth2callback')
def oauth2_callback():
    error = request.args.get('error')
    if error:
        return '{error}: {error_description}'.format(error=error, error_description=request.args.get('error_description')), 401

    code = request.args.get('code')
    if not code:
        return 'No code given', 401

    credentials = FLOW.step2_exchange(code)
    http = httplib2.Http()
    http = credentials.authorize(http)
    response, content = http.request('https://umairidris.auth0.com/userinfo')

    if response.status != 200:
        return response.reason, response.status

    user_dict = json.loads(content.decode("utf-8"))

    for key in REQUIRED_RESPONSE_KEYS:
        if key not in user_dict:
            return 'Missing {key} in response'.format(key=key), 401

    user = UserModel.get({'email': user_dict['email']})

    if not user: # new user
        # username is email for now
        user = User(user_dict['email'], user_dict['nickname'], user_dict['email'])
        user = UserModel.create(user)

        if not user:
            return 'Failed to create user', 500

    session['username'] = user.username
    session['logged_in'] = True

    if user.user_type == 'admin':
        session['admin'] = True
        return redirect(url_for('get_admin'))

    return redirect('/')

@app.route('/logout')
@login_required
def logout():
    """ () -> rendered_template

    Returns the rendered template of login.html,
    after the user makes a POST request to 'logout'
    """
    session.pop('username', None)
    session.pop('logged_in', None)
    session.pop('admin', None)
    return redirect('/')
<|endoftext|>


<commit_before>from datetime import datetime

from django.conf import settings
from django.shortcuts import render, redirect
from django.contrib import messages
from django.core.urlresolvers import reverse
from django.contrib.auth.decorators import login_required

from vumi.persist.redis_manager import RedisManager

from go.base.utils import make_read_only_form, conversation_or_404
from go.vumitools.exceptions import ConversationSendError
from go.conversation.forms import ConversationForm, ConversationGroupForm
from go.apps.surveys import forms as sforms

from vxpolls.content import forms
from vxpolls.manager import PollManager

redis = RedisManager.from_config(settings.VXPOLLS_REDIS_CONFIG)

def link_poll_to_conversation(poll_name, poll_id, conversation):
    metadata = conversation.get_metadata(default={})
    vxpolls_metadata = metadata.setdefault('vxpolls', {})
    polls = vxpolls_metadata.setdefault('polls', {})
    polls.update({
        poll_name: poll_id,
    })
    conversation.set_metadata(metadata)
    conversation.save()

def unlink_poll_from_conversation(poll_name, conversation):
    metadata = conversation.get_metadata(default={})
    vxpolls_metadata = metadata.setdefault('vxpolls', {})
    polls = vxpolls_metadata.setdefault('polls', {})
    del polls[poll_name]
    conversation.set_metadata(metadata)
    conversation.save()

def get_polls_for_conversation(conversation):
    metadata = conversation.get_metadata(default={})
    vxpolls_metadata = metadata.setdefault('vxpolls', {})
    return vxpolls_metadata.get('polls', {})

def generate_poll_id(conversation, suffix):
    return 'poll-%s_%s' % (conversation.key, suffix)

def get_poll_config(poll_id):
    pm = PollManager(redis, settings.VXPOLLS_PREFIX)
    config = pm.get_config(poll_id)
    config.update({
        'poll_id': poll_id,
    })
    return pm, config

@login_required
def new(request):
    if request.POST:
        form = ConversationForm(request.user_api, request.POST)
        if form.is_valid():
            conversation_data = {}
            copy_keys = [
                'subject',
                'message',
                'delivery_class',
            ]

            for key in copy_keys:
                conversation_data[key] = form.cleaned_data[key]

            tag_info = form.cleaned_data['delivery_tag_pool'].partition(':')
            conversation_data['delivery_tag_pool'] = tag_info[0]
            if tag_info[2]:
                conversation_data['delivery_tag'] = tag_info[2]

            start_date = form.cleaned_data['start_date'] or datetime.utcnow()
            start_time = (form.cleaned_data['start_time'] or
                            datetime.utcnow().time())
            conversation_data['start_timestamp'] = datetime(
                start_date.year, start_date.month, start_date.day,
                start_time.hour, start_time.minute, start_time.second,
                start_time.microsecond)

            conversation = request.user_api.new_conversation(
                u'multi_survey', **conversation_data)
            messages.add_message(request, messages.INFO,
                'Survey Created')
            return redirect(reverse('multi_survey:surveys',
                kwargs={'conversation_key': conversation.key}))

    else:
        form = ConversationForm(request.user_api, initial={
            'start_date': datetime.utcnow().date(),
            'start_time': datetime.utcnow().time().replace(second=0,
                                                            microsecond=0),
        })
    return render(request, 'multi_surveys/new.html', {
        'form': form,
    })

@login_required
def surveys(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    polls = get_polls_for_conversation(conversation)

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))

    return render(request, 'multi_surveys/surveys.html', {
        'survey_form': survey_form,
        'conversation': conversation,
        'polls': polls,
    })

@login_required
def new_survey(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    initial_config = {
        'poll_name': '',
    }
    if request.method == 'POST':
        pm = PollManager(redis, settings.VXPOLLS_PREFIX)
        post_data = request.POST.copy()
        form = forms.make_form(data=post_data, initial=initial_config)
        form.fields['poll_name'].required = True
        if form.is_valid():
            poll_name = form.cleaned_data['poll_name']
            poll_id = generate_poll_id(conversation, poll_name)
            link_poll_to_conversation(poll_name, poll_id, conversation)
            pm.set(poll_id, form.export())
            if request.POST.get('_save_contents'):
                return redirect(reverse('multi_survey:survey', kwargs={
                    'conversation_key': conversation.key,
                    'poll_name': poll_name,
                }))
            else:
                return redirect(reverse('multi_survey:surveys', kwargs={
                    'conversation_key': conversation.key,
                }))
    else:
        form = forms.make_form(data=initial_config, initial=initial_config)
        form.fields['poll_name'].required = True

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))
    return render(request, 'multi_surveys/contents.html', {
        'form': form,
        'survey_form': survey_form,
    })

@login_required
def survey(request, conversation_key, poll_name):
    conversation = conversation_or_404(request.user_api, conversation_key)
    poll_id = generate_poll_id(conversation, poll_name)
    pm, config = get_poll_config(poll_id)
    questions_data = config.get('questions', [])
    completed_response_data = config.get('survey_completed_responses', [])

    if request.method == 'POST':
        if request.POST.get('_delete_survey'):
            unlink_poll_from_conversation(poll_name, conversation)
            return redirect(reverse('multi_survey:surveys', kwargs={
                'conversation_key': conversation.key,
            }))

        post_data = request.POST.copy()
        post_data.update({
            'poll_id': poll_id,
        })
        form = forms.make_form(data=post_data, initial=config)

        questions_formset = forms.make_form_set(data=post_data)
        completed_response_formset = forms.make_completed_response_form_set(
            data=post_data)
        #poll_form = forms.SurveyPollForm(data=post_data)
        if form.is_valid():
            pm.set(poll_id, form.export())
            link_poll_to_conversation(poll_name, poll_id, conversation)
            if request.POST.get('_save_contents'):
                for k, v in post_data.items():
                    print k, v
                return redirect(reverse('multi_survey:survey', kwargs={
                    'conversation_key': conversation.key,
                    'poll_name': poll_name,
                }))
            else:
                return redirect(reverse('multi_survey:surveys', kwargs={
                    'conversation_key': conversation.key,
                }))
    else:
        form = forms.make_form(data=config, initial=config)
        questions_formset = forms.make_form_set(initial=questions_data)
        completed_response_formset = forms.make_completed_response_form_set(
            initial=completed_response_data)

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))
    return render(request, 'multi_surveys/contents.html', {
        'form': form,
        'questions_formset': questions_formset,
        'completed_response_formset': completed_response_formset,
        'survey_form': survey_form,
    })

@login_required
def people(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    groups = request.user_api.list_groups()

    poll_id = "poll-%s" % (conversation.key,)
    pm, config = get_poll_config(poll_id)

    if request.method == 'POST':
        if conversation.is_client_initiated():
            try:
                conversation.start()
            except ConversationSendError as error:
                if str(error) == 'No spare messaging tags.':
                    error = 'You have maxed out your available ' \
                            '%(delivery_class)s addresses. ' \
                            'End one or more running %(delivery_class)s ' \
                            'conversations to free one up.' % {
                                'delivery_class': conversation.delivery_class,
                            }
                messages.add_message(request, messages.ERROR, str(error))
                return redirect(reverse('multi_survey:people', kwargs={
                    'conversation_key': conversation.key}))

            addresses = [tag[1] for tag in conversation.get_tags()]
            messages.add_message(request, messages.INFO,
                'Survey started on %s' % (', '.join(addresses),))
            return redirect(reverse('multi_survey:show', kwargs={
                'conversation_key': conversation.key}))
        else:
            group_form = ConversationGroupForm(
                request.POST, groups=request.user_api.list_groups())

            if group_form.is_valid():
                for group in group_form.cleaned_data['groups']:
                    conversation.groups.add_key(group)
                conversation.save()
                messages.add_message(request, messages.INFO,
                    'The selected groups have been added to the survey')
                return redirect(reverse('multi_survey:start', kwargs={
                                    'conversation_key': conversation.key}))

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))
    content_form = forms.make_form(data=config, initial=config, extra=0)
    read_only_content_form = make_read_only_form(content_form)
    return render(request, 'multi_surveys/people.html', {
        'conversation': conversation,
        'survey_form': survey_form,
        'content_form': read_only_content_form,
        'groups': groups,
        'polls': get_polls_for_conversation(conversation),
    })

@login_required
def start(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    if request.method == 'POST':
        try:
            conversation.start()
        except ConversationSendError as error:
            messages.add_message(request, messages.ERROR, str(error))
            return redirect(reverse('multi_survey:start', kwargs={
                'conversation_key': conversation.key}))
        messages.add_message(request, messages.INFO, 'Survey started')
        return redirect(reverse('multi_survey:show', kwargs={
            'conversation_key': conversation.key}))
    return render(request, 'multi_surveys/start.html', {
        'conversation': conversation,
    })

@login_required
def end(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    if request.method == 'POST':
        conversation.end_conversation()
        messages.add_message(request, messages.INFO, 'Survey ended')
    return redirect(reverse('multi_survey:show', kwargs={
        'conversation_key': conversation.key}))

@login_required
def show(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    return render(request, 'multi_surveys/show.html', {
        'conversation': conversation,
    })
<commit_msg>just the new version<commit_after>from datetime import datetime

from django.conf import settings
from django.shortcuts import render, redirect
from django.contrib import messages
from django.core.urlresolvers import reverse
from django.contrib.auth.decorators import login_required

from vumi.persist.redis_manager import RedisManager

from go.base.utils import make_read_only_form, conversation_or_404
from go.vumitools.exceptions import ConversationSendError
from go.conversation.forms import ConversationForm, ConversationGroupForm
from go.apps.surveys import forms

from vxpolls.manager import PollManager

redis = RedisManager.from_config(settings.VXPOLLS_REDIS_CONFIG)

def link_poll_to_conversation(poll_name, poll_id, conversation):
    metadata = conversation.get_metadata(default={})
    vxpolls_metadata = metadata.setdefault('vxpolls', {})
    polls = vxpolls_metadata.setdefault('polls', {})
    polls.update({
        poll_name: poll_id,
    })
    conversation.set_metadata(metadata)
    conversation.save()

def unlink_poll_from_conversation(poll_name, conversation):
    metadata = conversation.get_metadata(default={})
    vxpolls_metadata = metadata.setdefault('vxpolls', {})
    polls = vxpolls_metadata.setdefault('polls', {})
    del polls[poll_name]
    conversation.set_metadata(metadata)
    conversation.save()

def get_polls_for_conversation(conversation):
    metadata = conversation.get_metadata(default={})
    vxpolls_metadata = metadata.setdefault('vxpolls', {})
    return vxpolls_metadata.get('polls', {})

def generate_poll_id(conversation, suffix):
    return 'poll-%s_%s' % (conversation.key, suffix)

def get_poll_config(poll_id):
    pm = PollManager(redis, settings.VXPOLLS_PREFIX)
    config = pm.get_config(poll_id)
    config.update({
        'poll_id': poll_id,
    })
    return pm, config

@login_required
def new(request):
    if request.POST:
        form = ConversationForm(request.user_api, request.POST)
        if form.is_valid():
            conversation_data = {}
            copy_keys = [
                'subject',
                'message',
                'delivery_class',
            ]

            for key in copy_keys:
                conversation_data[key] = form.cleaned_data[key]

            tag_info = form.cleaned_data['delivery_tag_pool'].partition(':')
            conversation_data['delivery_tag_pool'] = tag_info[0]
            if tag_info[2]:
                conversation_data['delivery_tag'] = tag_info[2]

            start_date = form.cleaned_data['start_date'] or datetime.utcnow()
            start_time = (form.cleaned_data['start_time'] or
                            datetime.utcnow().time())
            conversation_data['start_timestamp'] = datetime(
                start_date.year, start_date.month, start_date.day,
                start_time.hour, start_time.minute, start_time.second,
                start_time.microsecond)

            conversation = request.user_api.new_conversation(
                u'multi_survey', **conversation_data)
            messages.add_message(request, messages.INFO,
                'Survey Created')
            return redirect(reverse('multi_survey:surveys',
                kwargs={'conversation_key': conversation.key}))

    else:
        form = ConversationForm(request.user_api, initial={
            'start_date': datetime.utcnow().date(),
            'start_time': datetime.utcnow().time().replace(second=0,
                                                            microsecond=0),
        })
    return render(request, 'multi_surveys/new.html', {
        'form': form,
    })

@login_required
def surveys(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    polls = get_polls_for_conversation(conversation)

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))

    return render(request, 'multi_surveys/surveys.html', {
        'survey_form': survey_form,
        'conversation': conversation,
        'polls': polls,
    })

@login_required
def new_survey(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    initial_config = {
        'poll_name': '',
    }
    if request.method == 'POST':
        pm = PollManager(redis, settings.VXPOLLS_PREFIX)
        post_data = request.POST.copy()
        form = forms.make_form(data=post_data, initial=initial_config)
        form.fields['poll_name'].required = True
        if form.is_valid():
            poll_name = form.cleaned_data['poll_name']
            poll_id = generate_poll_id(conversation, poll_name)
            link_poll_to_conversation(poll_name, poll_id, conversation)
            pm.set(poll_id, form.export())
            if request.POST.get('_save_contents'):
                return redirect(reverse('multi_survey:survey', kwargs={
                    'conversation_key': conversation.key,
                    'poll_name': poll_name,
                }))
            else:
                return redirect(reverse('multi_survey:surveys', kwargs={
                    'conversation_key': conversation.key,
                }))
    else:
        form = forms.make_form(data=initial_config, initial=initial_config)
        form.fields['poll_name'].required = True

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))
    return render(request, 'multi_surveys/contents.html', {
        'form': form,
        'survey_form': survey_form,
    })

def _clear_empties(cleaned_data):
    """
    FIXME:  this is a work around because for some reason Django is seeing
            the new (empty) forms in the formsets as stuff that is to be
            stored when it really should be discarded.
    """
    return [cd for cd in cleaned_data if cd.get('copy')]

@login_required
def survey(request, conversation_key, poll_name):
    conversation = conversation_or_404(request.user_api, conversation_key)
    poll_id = generate_poll_id(conversation, poll_name)
    pm, config = get_poll_config(poll_id)
    questions_data = config.get('questions', [])
    completed_response_data = config.get('survey_completed_responses', [])

    if request.method == 'POST':
        if request.POST.get('_delete_survey'):
            unlink_poll_from_conversation(poll_name, conversation)
            return redirect(reverse('multi_survey:surveys', kwargs={
                'conversation_key': conversation.key,
            }))

        post_data = request.POST.copy()
        post_data.update({
            'poll_id': poll_id,
        })

        questions_formset = forms.make_form_set(data=post_data)
        completed_response_formset = forms.make_completed_response_form_set(
            data=post_data)
        poll_form = forms.SurveyPollForm(data=post_data)
        if poll_form.is_valid():
            data = poll_form.cleaned_data.copy()
            data.update({
                'questions': _clear_empties(questions_formset.cleaned_data),
                'survey_completed_responses': _clear_empties(
                    completed_response_formset.cleaned_data)
                })
            pm.set(poll_id, data)
            link_poll_to_conversation(poll_name, poll_id, conversation)
            if request.POST.get('_save_contents'):
                for k, v in post_data.items():
                    print k, v
                return redirect(reverse('multi_survey:survey', kwargs={
                    'conversation_key': conversation.key,
                    'poll_name': poll_name,
                }))
            else:
                return redirect(reverse('multi_survey:surveys', kwargs={
                    'conversation_key': conversation.key,
                }))
    else:
        poll_form = forms.SurveyPollForm(initial=config)
        questions_formset = forms.make_form_set(initial=questions_data)
        completed_response_formset = forms.make_completed_response_form_set(
            initial=completed_response_data)

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))

    return render(request, 'multi_surveys/contents.html', {
        'poll_form': poll_form,
        'questions_formset': questions_formset,
        'completed_response_formset': completed_response_formset,
        'survey_form': survey_form,
    })

@login_required
def people(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    groups = request.user_api.list_groups()

    poll_id = "poll-%s" % (conversation.key,)
    pm, config = get_poll_config(poll_id)

    if request.method == 'POST':
        if conversation.is_client_initiated():
            try:
                conversation.start()
            except ConversationSendError as error:
                if str(error) == 'No spare messaging tags.':
                    error = 'You have maxed out your available ' \
                            '%(delivery_class)s addresses. ' \
                            'End one or more running %(delivery_class)s ' \
                            'conversations to free one up.' % {
                                'delivery_class': conversation.delivery_class,
                            }
                messages.add_message(request, messages.ERROR, str(error))
                return redirect(reverse('multi_survey:people', kwargs={
                    'conversation_key': conversation.key}))

            addresses = [tag[1] for tag in conversation.get_tags()]
            messages.add_message(request, messages.INFO,
                'Survey started on %s' % (', '.join(addresses),))
            return redirect(reverse('multi_survey:show', kwargs={
                'conversation_key': conversation.key}))
        else:
            group_form = ConversationGroupForm(
                request.POST, groups=request.user_api.list_groups())

            if group_form.is_valid():
                for group in group_form.cleaned_data['groups']:
                    conversation.groups.add_key(group)
                conversation.save()
                messages.add_message(request, messages.INFO,
                    'The selected groups have been added to the survey')
                return redirect(reverse('multi_survey:start', kwargs={
                                    'conversation_key': conversation.key}))

    survey_form = make_read_only_form(ConversationForm(request.user_api,
        instance=conversation, initial={
            'start_date': conversation.start_timestamp.date(),
            'start_time': conversation.start_timestamp.time(),
        }))
    content_form = forms.make_form(data=config, initial=config, extra=0)
    read_only_content_form = make_read_only_form(content_form)
    return render(request, 'multi_surveys/people.html', {
        'conversation': conversation,
        'survey_form': survey_form,
        'content_form': read_only_content_form,
        'groups': groups,
        'polls': get_polls_for_conversation(conversation),
    })

@login_required
def start(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    if request.method == 'POST':
        try:
            conversation.start()
        except ConversationSendError as error:
            messages.add_message(request, messages.ERROR, str(error))
            return redirect(reverse('multi_survey:start', kwargs={
                'conversation_key': conversation.key}))
        messages.add_message(request, messages.INFO, 'Survey started')
        return redirect(reverse('multi_survey:show', kwargs={
            'conversation_key': conversation.key}))
    return render(request, 'multi_surveys/start.html', {
        'conversation': conversation,
    })

@login_required
def end(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    if request.method == 'POST':
        conversation.end_conversation()
        messages.add_message(request, messages.INFO, 'Survey ended')
    return redirect(reverse('multi_survey:show', kwargs={
        'conversation_key': conversation.key}))

@login_required
def show(request, conversation_key):
    conversation = conversation_or_404(request.user_api, conversation_key)
    return render(request, 'multi_surveys/show.html', {
        'conversation': conversation,
    })
<|endoftext|>