<commit_before>"""This script will call LAMMPS to calculate and visualize (with matplotlib) the Gay-Berne potential
for the side-by-side, end-to-end,tee and cross configurations for a uniaxial ellipsoid.

Note:
If the default intermolecular distances for a configuration are too far from it's minimum,
some tuning might be required in the move_molecule function

Usage:
#the script must be in python path or working directory
from gayberne import gb
gb = gb(width=1, length=3, sigma=1, cross_depth=3, end_depth=1, nu=1, mu=2)
lammps_df = gb.lammps_gb()
gb.visualize(lammps_df=lammps_df)

Requirement:
numpy
matplotlib
pandas
Python interface to LAMMPS (http://lammps.sandia.gov/doc/Section_python.html)
LAMMPS compiled with asphere and molecule package
system.data file (the defined parameters will be overwritten)

TODO:
The system.data file should not be required

"""
import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
from collections import OrderedDict
from lammps import lammps
import matplotlib.lines as mlines

class gb():
    def __init__(self, width, length, sigma, cross_depth, end_depth, nu, mu, cutoff=30):
        """ Definition of the parameters and initialization"""
        self.df = pd.DataFrame()
        # configuration with the color for the graph: red, green, blue, yellow
        self.colors = {"side": "r", "cross": "g", "end": "b", "tee": "y"}
        # parameters definition
        self.width = width
        self.length = length
        self.sigma = sigma
        self.cross_depth = cross_depth
        self.end_depth = end_depth
        self.nu = nu
        self.mu = mu
        self.cutoff = cutoff

    def lammps_gb(self):
        """ This function will call lammps to calculate the Gay-Berne potential for a uniaxial ellipsoid
        Returns:
        Pandas dataframe with the intermolecular distance and energy for the four configurations
        """
        # no output to screen or file
        lmp = lammps(cmdargs=["-sc", "none", "-log", "none"])
        lmp.command("units lj")
        lmp.command("boundary p p p")
        lmp.command("atom_style hybrid molecular ellipsoid")
        # define the gay berne pair_style
        lmp.command("pair_style gayberne 1.0 %f %f %f" %
                    (self.nu, self.mu, self.cutoff))
        lmp.command("read_data system.data")
        # define the ellipsoid's shape
        lmp.command("set type 1 shape %f %f %f" %
                    (self.width, self.width, self.length))
        # define initial orientation for the atoms
        lmp.command("set atom 1 quat 1 0 0 0")
        lmp.command("set atom 2 quat 1 0 0 0")
        # Define the gay-berne pair_coeff
        lmp.command("pair_coeff 1 1 1.0 %f %f %f %f %f %f %f" % (self.sigma, self.cross_depth, self.cross_depth,
                                                                 self.end_depth, self.cross_depth,
                                                                 self.cross_depth, self.end_depth))

        # function to move the first ellipsoid. direction: 0 is x, 1 is y and 2
        # is z.
        def move_molecule(configuration, direction):
            """ Function to translate the molecule in the desired direction"""
            i = 0
            pot_tot = []
            position = []
            i = 0
            while i < 150:
                xc = lmp.gather_atoms("x", 1, 3)
                lmp.command("run 0")
                pot = lmp.extract_compute("thermo_pe", 0, 0)
                pot_tot.append(pot)
                position.append(xc[direction])
                xc[direction] = xc[direction] + 0.01
                lmp.scatter_atoms("x", 1, 3, xc)
                i = i + 1
            position = pd.DataFrame(
                {str(configuration) + "_distance": position})
            pot_tot = pd.DataFrame(
                {str(configuration) + "_energy": pot_tot})
            self.df = pd.concat([self.df, position, pot_tot], axis=1)

        for configuration in self.colors:
            # reinitialze the position of the first atom
            lmp.command("set atom 1 x 0 y 0 z 0")
            if configuration == "side" or "end":
                # No rotatio for side and end configuration
                lmp.command("set atom 1 quat 1 0 0 0")
                if configuration == "side":
                    # move in x direction for side-by-side
                    direction = 0
                    xc = lmp.gather_atoms("x", 1, 3)
                    # Can be changed if the values are too far from the minimum
                    xc[direction] = xc[direction] + self.width / 1.5
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)
                if configuration == "end":
                    # move in the z direction for end-to-end
                    direction = 2
                    xc = lmp.gather_atoms("x", 1, 3)
                    # Can be changed if the values are too far from the minimum
                    xc[direction] = xc[direction] + self.length / 1.1
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)

            if configuration == "tee" or "cross":
                # The first ellipsoid is rotated by 90 degree in the x axis
                lmp.command("set atom 1 quat 1 0  0 90")
                if configuration == "cross":
                    # the first ellipsoid is translated in the x direction
                    direction = 0
                    xc = lmp.gather_atoms("x", 1, 3)
                    # Can be changed if the values are too far from the minimum
                    xc[direction] = xc[direction] + self.width / 1.5
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)
                if configuration == "tee":
                    # Translation in the y axis
                    direction = 1
                    xc = lmp.gather_atoms("x", 1, 3)
                    xc[direction] = xc[direction] + self.length - \
                        self.width  # Can be changed if the values are too far from the minimum
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)
        # close lammps
        return self.df
        lmp.close()

    def visualize(self, lammps_df=None):
        """ Function to visualize with matplotlib the gay-berne potential calculated from lammps"""
        if type(lammps_df).__name__ == "NoneType":
            lammps_df = self.lammps_gb()

        for configuration in self.colors:
            # Remove values with a energy superior to 5 for the graph
            df_graph = lammps_df.where(lammps_df[str(configuration) + "_energy"] < 5)
            plt.plot(df_graph[str(configuration) + "_distance"].dropna(),
            df_graph[str(configuration) + "_energy"].dropna(), label=str(configuration),
            color=self.colors[configuration], ms=5, markevery=0.02)
            plt.ylabel(r"Potentiel ($\epsilon_0$)")
            plt.xlabel(r"Distance ($\sigma_0$)")

        plt.legend()
        plt.tight_layout()

        plt.show()
<commit_msg>fixed for the correct key value<commit_after>#!/usr/bin/python2
"""This script will call LAMMPS to calculate and visualize (with matplotlib) the Gay-Berne potential
for the side-by-side (ll), end-to-end(z),tee(T) and cross(X) configurations for a uniaxial ellipsoid.

Note:
If the default intermolecular distances for a configuration are too far from it's minimum,
some tuning might be required in the move_molecule function

Usage:
#the script must be in python path or working directory
from gayberne import gb
gb = gb(width=1, length=3, sigma=1, X_depth=3, z_depth=1, nu=1, mu=2,x_range=150)
lammps_df = gb.lammps_gb()
gb.visualize(lammps_df=lammps_df)

Requirement:
numpy
matplotlib
pandas
Python interface to LAMMPS (http://lammps.sandia.gov/doc/Section_python.html)
LAMMPS compiled with asphere and molecule package
system.data file (the defined parameters will be overwritten)

TODO:
The system.data file should not be required
"""
import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
from collections import OrderedDict
from lammps import lammps
import matplotlib.lines as mlines

class gb():
    def __init__(self, width, length, sigma, X_depth, z_depth, nu, mu, cutoff=30,x_range=150):
        """ Definition of the parameters and initialization"""
        self.df = pd.DataFrame()
        # configuration with the color for the graph: red, green, blue, yellow
        self.colors = {"ll": "r", "X": "g", "z": "b", "T": "y"}
        # parameters definition
        self.width = width
        self.length = length
        self.sigma = sigma
        self.X_depth = X_depth
        self.z_depth = z_depth
        self.nu = nu
        self.mu = mu
        self.cutoff = cutoff
        # range for the graph
        self.x_range = x_range

    def lammps_gb(self):
        """ This function will call lammps to calculate the Gay-Berne potential for a uniaxial ellipsoid
        Returns:
        Pandas dataframe with the intermolecular distance and energy for the four configurations
        """
        # no output to screen or file
        lmp = lammps(cmdargs=["-sc", "none", "-log", "none"])
        lmp.command("units lj")
        lmp.command("boundary p p p")
        lmp.command("atom_style hybrid molecular ellipsoid")
        # define the gay berne pair_style
        lmp.command("pair_style gayberne 1.0 %f %f %f" %
                    (self.nu, self.mu, self.cutoff))
        lmp.command("read_data system.data")
        # define the ellipsoid's shape
        lmp.command("set type 1 shape %f %f %f" %
                    (self.width, self.width, self.length))
        # define initial orientation for the atoms
        lmp.command("set atom 1 quat 1 0 0 0")
        lmp.command("set atom 2 quat 1 0 0 0")
        # Define the gay-berne pair_coeff
        lmp.command("pair_coeff 1 1 1.0 %f %f %f %f %f %f %f" % (self.sigma, self.X_depth, self.X_depth,
                                                                 self.z_depth, self.X_depth,
                                                                 self.X_depth, self.z_depth))

        # function to move the first ellipsoid. direction: 0 is x, 1 is y and 2
        # is z.
        def move_molecule(configuration, direction):
            """ Function to translate the molecule in the desired direction"""
            i = 0
            pot_tot = []
            position = []
            i = 0
            while i < self.x_range:
                xc = lmp.gather_atoms("x", 1, 3)
                lmp.command("run 0")
                pot = lmp.extract_compute("thermo_pe", 0, 0)
                pot_tot.append(pot)
                position.append(xc[direction])
                xc[direction] = xc[direction] + 0.01
                lmp.scatter_atoms("x", 1, 3, xc)
                i = i + 1
            position = pd.DataFrame(
                {"MD_"+ str(configuration) + "_distance": position})
            pot_tot = pd.DataFrame(
                {"MD_" + str(configuration) + "_energy": pot_tot})
            self.df = pd.concat([self.df, position, pot_tot], axis=1)

        for configuration in self.colors:
            # reinitialze the position of the first atom
            lmp.command("set atom 1 x 0 y 0 z 0")
            if configuration == "ll" or "z":
                # No rotatio for ll and end configuration
                lmp.command("set atom 1 quat 1 0 0 0")
                if configuration == "ll":
                    # move in x direction for side-by-side
                    direction = 0
                    xc = lmp.gather_atoms("x", 1, 3)
                    # Can be changed if the values are too far from the minimum
                    xc[direction] = xc[direction] + self.width/1.5
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)
                if configuration == "z":
                    # move in the z direction for end-to-end
                    direction = 2
                    xc = lmp.gather_atoms("x", 1, 3)
                    # Can be changed if the values are too far from the minimum
                    xc[direction] = xc[direction] + self.length/1.1
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)

            if configuration == "T" or "X":
                # The first ellipsoid is rotated by 90 degree in the x axis
                lmp.command("set atom 1 quat 1 0  0 90")
                if configuration == "X":
                    # the first ellipsoid is translated in the x direction
                    direction = 0
                    xc = lmp.gather_atoms("x", 1, 3)
                    # Can be changed if the values are too far from the minimum
                    xc[direction] = xc[direction] + self.width/1.5
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)
                if configuration == "T":
                    # Translation in the y axis
                    direction = 1
                    xc = lmp.gather_atoms("x", 1, 3)
                    xc[direction] = xc[direction] + self.length - \
                        self.width  # Can be changed if the values are too far from the minimum
                    xc = lmp.scatter_atoms("x", 1, 3, xc)
                    move_molecule(configuration, direction)
        # close lammps
        print("LAMMPS done")
        return self.df
        lmp.close()

    def visualize(self, lammps_df=None):
        """ Function to visualize with matplotlib the gay-berne potential calculated from lammps"""
        if type(lammps_df).__name__ == "NoneType":
            lammps_df = self.lammps_gb()

        for configuration in self.colors:
            # Remove values with a energy superior to 5 for the graph
	    df_graph = lammps_df.where(lammps_df["MD_"+str(configuration) + "_energy"] < 5)
            plt.plot(df_graph["MD_"+str(configuration) + "_distance"].dropna(),
            df_graph["MD_"+str(configuration) + "_energy"].dropna(), label=str(configuration),
            color=self.colors[configuration], ms=5, markevery=0.02)
            plt.ylabel(r"Potentiel ($\epsilon_0$)")
            plt.xlabel(r"Distance ($\sigma_0$)")

        plt.legend()
        plt.tight_layout()

        plt.show()
#gb = gb(width=1, length=3, sigma=1, X_depth=3, z_depth=1, nu=1, mu=2,x_range=150)
#lammps_df = gb.lammps_gb()
#gb.visualize(lammps_df=lammps_df)
<|endoftext|>


<commit_before>from __future__ import unicode_literals, division, absolute_import
import logging

from flexget import plugin
from flexget.event import event

log = logging.getLogger('parsing')
PARSER_TYPES = ['movie', 'series']

# Mapping of parser type to (mapping of parser name to plugin instance)
parsers = {}
# Mapping from parser type to the name of the default/selected parser for that type
default_parsers = {}
selected_parsers = {}

# We need to wait until manager startup to access other plugin instances, to make sure they have all been loaded
@event('manager.startup')
def init_parsers(manager):
    """Prepare our list of parsing plugins and default parsers."""
    for parser_type in PARSER_TYPES:
        parsers[parser_type] = {}
        for p in plugin.get_plugins(group=parser_type + '_parser'):
            parsers[parser_type][p.name.replace('parser_', '')] = p.instance
        # Select default parsers based on priority
        func_name = 'parse_' + parser_type
        default_parsers[parser_type] = max(parsers[parser_type].iteritems(),
                                          key=lambda p: getattr(getattr(p[1], func_name), 'priority', 0))[0]

class PluginParsing(object):
    """Provides parsing framework"""

    @property
    def schema(self):
        # Create a schema allowing only our registered parsers to be used under the key of each parser type
        properties = {}
        for parser_type in PARSER_TYPES:
            parser_names = [p.name.replace('parser_', '') for p in plugin.get_plugins(group=parser_type + '_parser')]
            properties[parser_type] = {'type': 'string', 'enum': parser_names}
        s = {
            'type': 'object',
            'properties': properties,
            'additionalProperties': False
        }
        return s

    def on_task_start(self, task, config):
        # Set up user selected parsers from config for this task run
        if config:
            selected_parsers.update(config)

    def on_task_end(self, task, config):
        # Restore default parsers for next task run
        selected_parsers.clear()

    on_task_abort = on_task_end

    def parse_series(self, data, name=None, **kwargs):
        """
        Use the selected series parser to parse series information from `data`

        :param data: The raw string to parse information from.
        :param name: The series name to parse data for. If not supplied, parser will attempt to guess series name
            automatically from `data`.

        :returns: An object containing the parsed information. The `valid` attribute will be set depending on success.
        """
        parser = parsers['series'][selected_parsers.get('series', default_parsers.get('series'))]
        return parser.parse_series(data, name=name, **kwargs)

    def parse_movie(self, data, **kwargs):
        """
        Use the selected movie parser to parse movie information from `data`

        :param data: The raw string to parse information from

        :returns: An object containing the parsed information. The `valid` attribute will be set depending on success.
        """
        parser = parsers['movie'][selected_parsers.get('movie') or default_parsers['movie']]
        return parser.parse_movie(data, **kwargs)

@event('plugin.register')
def register_plugin():
    plugin.register(PluginParsing, 'parsing', api_ver=2)
<commit_msg>Add some debug info when picking default parsers.<commit_after>from __future__ import unicode_literals, division, absolute_import
import logging

from flexget import plugin
from flexget.event import event

log = logging.getLogger('parsing')
PARSER_TYPES = ['movie', 'series']

# Mapping of parser type to (mapping of parser name to plugin instance)
parsers = {}
# Mapping from parser type to the name of the default/selected parser for that type
default_parsers = {}
selected_parsers = {}

# We need to wait until manager startup to access other plugin instances, to make sure they have all been loaded
@event('manager.startup')
def init_parsers(manager):
    """Prepare our list of parsing plugins and default parsers."""
    for parser_type in PARSER_TYPES:
        parsers[parser_type] = {}
        for p in plugin.get_plugins(group=parser_type + '_parser'):
            parsers[parser_type][p.name.replace('parser_', '')] = p.instance
        # Select default parsers based on priority
        func_name = 'parse_' + parser_type
        default_parsers[parser_type] = max(parsers[parser_type].iteritems(),
                                          key=lambda p: getattr(getattr(p[1], func_name), 'priority', 0))[0]
        log.debug('setting default %s parser to %s. (options: %s)' %
                  (parser_type, default_parsers[parser_type], parsers[parser_type]))

class PluginParsing(object):
    """Provides parsing framework"""

    @property
    def schema(self):
        # Create a schema allowing only our registered parsers to be used under the key of each parser type
        properties = {}
        for parser_type in PARSER_TYPES:
            parser_names = [p.name.replace('parser_', '') for p in plugin.get_plugins(group=parser_type + '_parser')]
            properties[parser_type] = {'type': 'string', 'enum': parser_names}
        s = {
            'type': 'object',
            'properties': properties,
            'additionalProperties': False
        }
        return s

    def on_task_start(self, task, config):
        # Set up user selected parsers from config for this task run
        if config:
            selected_parsers.update(config)

    def on_task_end(self, task, config):
        # Restore default parsers for next task run
        selected_parsers.clear()

    on_task_abort = on_task_end

    def parse_series(self, data, name=None, **kwargs):
        """
        Use the selected series parser to parse series information from `data`

        :param data: The raw string to parse information from.
        :param name: The series name to parse data for. If not supplied, parser will attempt to guess series name
            automatically from `data`.

        :returns: An object containing the parsed information. The `valid` attribute will be set depending on success.
        """
        parser = parsers['series'][selected_parsers.get('series', default_parsers.get('series'))]
        return parser.parse_series(data, name=name, **kwargs)

    def parse_movie(self, data, **kwargs):
        """
        Use the selected movie parser to parse movie information from `data`

        :param data: The raw string to parse information from

        :returns: An object containing the parsed information. The `valid` attribute will be set depending on success.
        """
        parser = parsers['movie'][selected_parsers.get('movie') or default_parsers['movie']]
        return parser.parse_movie(data, **kwargs)

@event('plugin.register')
def register_plugin():
    plugin.register(PluginParsing, 'parsing', api_ver=2)
<|endoftext|>


<commit_before>from thundercat.response import process_response
from thundercat import addhost, discover
from thundercat.ssh import Transport
from importlib import import_module
import config
import sys
import os

def main(argv=sys.argv):
    if len(argv) < 2:
        print "thundercat - client for interfacing with nappingcat servers"
        print "usage:"
        print "\tthundercat addremote <nickname> <git@someurl.com>"
        print "\tthundercat help <nickname> <command>"
        print "\tthundercat <nickname> -- list commands"
        print "\tthundercat <nickname> <command> <args> -- send a command to the host"
        return
    subcommand = argv[1]
    if subcommand == 'addremote':
        addhost.add_host(argv[1:4])
        discover.discover_command(['discover', argv[2]])
    elif subcommand == 'help':
        host_info = config.get_host_info(argv[2])
        print host_info['endpoints'][argv[3]]['help_text'].replace('{EXECUTABLE}', os.path.basename(argv[0]))
    else:
        nickname = argv[1]
        host_info = config.get_host_info(nickname)
        if host_info:
            endpoint = host_info['endpoints'].get(argv[2], None) if len(argv) > 2 else None
            if endpoint:
                simple_regex = endpoint['regex']
                for possibility, arg_names in simple_regex:
                    response_dict = {}
                    available_args = argv[3:]
                    available_args.reverse()
                    for arg_name in arg_names:
                        parameter_type = endpoint['params'].get(arg_name, None)
                        if isinstance(parameter_type, list):
                            available_args.reverse()
                            response_dict[arg_name] = parameter_type[1].join(available_args)
                            available_args = []
                        else:
                            try:
                                response_dict[arg_name] = available_args.pop()
                            except IndexError:
                                pass
                    try:
                        result = possibility % response_dict
                        transport = Transport(host_info['url'])
                        response = transport.request(result)
                        try:
                            module = import_module('thundercat.plugins.%s' % argv[2])
                            target = getattr(module, 'main', lambda x: x)
                            print target(nickname, host_info, response)
                        except ImportError:
                            print response
                    except (TypeError, KeyError) as e:
                        pass
            else:
                print process_response({
                    'content': {
                        'message':'Bad endpoint %s\nAvailable endpoints:\n\t%s\n' % (argv[2] if len(argv) > 2 else "<none provided>", '\n\t'.join(host_info['endpoints'].keys())),
                    },
                    'status_code':404,
                })
        else:
            print process_response({
                'content':{
                    'message':'Bad host %s\nAvailable hosts:\n\t%s\n' % (argv[1], '\n\t'.join(config.get_available_host_list())),
                },
                'status_code':404,
            })

if __name__ == '__main__':
    main(sys.argv)
<commit_msg>added improvements to sending files over the stream<commit_after>from thundercat.response import process_response
from thundercat import addhost, discover
from thundercat.ssh import Transport
from importlib import import_module
import config
import sys
import os
import StringIO
def main(argv=sys.argv):
    if len(argv) < 2:
        print "thundercat - client for interfacing with nappingcat servers"
        print "usage:"
        print "\tthundercat addremote <nickname> <git@someurl.com>"
        print "\tthundercat help <nickname> <command>"
        print "\tthundercat <nickname> -- list commands"
        print "\tthundercat <nickname> <command> <args> -- send a command to the host"
        return
    subcommand = argv[1]
    if subcommand == 'addremote':
        addhost.add_host(argv[1:4])
        discover.discover_command(['discover', argv[2]])
    elif subcommand == 'help':
        host_info = config.get_host_info(argv[2])
        print host_info['endpoints'][argv[3]]['help_text'].replace('{EXECUTABLE}', os.path.basename(argv[0]))
    else:
        nickname = argv[1]
        host_info = config.get_host_info(nickname)
        if host_info:
            endpoint = host_info['endpoints'].get(argv[2], None) if len(argv) > 2 else None
            if endpoint:
                simple_regex = endpoint['regex']
                for possibility, arg_names in simple_regex:
                    response_dict = {}
                    available_args = argv[3:]
                    available_args.reverse()
                    for arg_name in arg_names:
                        parameter_type = endpoint['params'].pop(arg_name, None)
                        if isinstance(parameter_type, list):
                            available_args.reverse()
                            response_dict[arg_name] = parameter_type[1].join(available_args)
                            available_args = []
                        else:
                            try:
                                response_dict[arg_name] = available_args.pop()
                            except IndexError:
                                pass
                    file = None
                    for key, type in endpoint['params'].iteritems():
                        if type == 'stdin':
                            with open(available_args.pop(), 'r') as input:
                                file = StringIO.StringIO(input.read())
                            break
                    try:
                        result = possibility % response_dict
                        transport = Transport(host_info['url'])
                        response = transport.request(result, file)
                        try:
                            module = import_module('thundercat.plugins.%s' % argv[2])
                            target = getattr(module, 'main', lambda x: x)
                            print target(nickname, host_info, response)
                        except ImportError:
                            print response
                        return
                    except (TypeError, KeyError) as e:
                        pass
            else:
                print process_response({
                    'content': {
                        'message':'Bad endpoint %s\nAvailable endpoints:\n\t%s\n' % (argv[2] if len(argv) > 2 else "<none provided>", '\n\t'.join(host_info['endpoints'].keys())),
                    },
                    'status_code':404,
                })
        else:
            print process_response({
                'content':{
                    'message':'Bad host %s\nAvailable hosts:\n\t%s\n' % (argv[1], '\n\t'.join(config.get_available_host_list())),
                },
                'status_code':404,
            })

if __name__ == '__main__':
    main(sys.argv)
<|endoftext|>


<commit_before>import os
import sys
import glob
import datetime

import numpy

import arcpy

def progress_print(msg, verbose=False, asMessage=False):
    if verbose:
        if asMessage:
            arcpy.AddMessage(msg)
        else:
            print(msg)

class EasyMapDoc(object):
    def __init__(self, *args, **kwargs):
        self.mapdoc = arcpy.mapping.MapDocument(*args, **kwargs)

    @property
    def layers(self):
        return arcpy.mapping.ListLayers(self.mapdoc)

    @property
    def dataframes(self):
        return arcpy.mapping.ListDataFrames(self.mapdoc)

    def findLayerByName(self, name):
        for lyr in self.layers:
            if not lyr.isGroupLayer and lyr.name == name:
                return lyr

    def add_layer(self, layer, df=None, position='top'):
        # if no dataframe is provided, select the first
        if df is None:
            df = self.dataframes[0]

        # check that the position is valid
        valid_positions = ['auto_arrange', 'bottom', 'top']
        if position.lower() not in valid_positions:
            raise ValueError('Position: %s is not in %s' % (position.lower, valid_positions))

        # layer can be a path to a file. if so, convert to a Layer object
        if isinstance(layer, basestring):
            layer = arcpy.mapping.Layer(layer)
        elif not isinstance(layer, arcpy.mapping.Layer):
            raise ValueError("``layer`` be an arcpy Layer or a path to a file")

        # add the layer to the map
        arcpy.mapping.AddLayer(df, layer, position.upper())

        # return the layer
        return layer

class Extension(object):
    """ Context manager to facilitate the use of ArcGIS extensions

    Inside the context manager, the extension will be checked out. Once
    the interpreter leaves the code block by any means (e.g., sucessful
    execution, raised exception) the extension will be checked back in.

    Example
    -------
    >>> import tidegates, arcpy
    >>> with tidegates.utils.Extension("spatial"):
    ...     arcpy.sa.Hillshade("C:/data/dem.tif")

    """

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

    def __enter__(self):
        if arcpy.CheckExtension(self.name) == "Available":
            status = arcpy.CheckOutExtension(self.name)
            return status
        else:
            raise RuntimeError("%s license isn't available" % self.name)

    def __exit__(self, *args):
        arcpy.CheckOutExtension(self.name)

class OverwriteState(object):
    """ Context manager to temporarily set the ``overwriteOutput``
    environment variable.

    Inside the context manager, the ``arcpy.env.overwriteOutput`` will
    be set to the given value. Once the interpreter leaves the code
    block by any means (e.g., sucessful execution, raised exception),
    ``arcpy.env.overwriteOutput`` will reset to its original value.

    Example
    -------
    >>> import tidegates
    >>> with tidegates.utils.OverwriteState(False):
    ...     # some operation that should fail if output already exists

    """

    def __init__(self, overwrite):
        self.orig_state = arcpy.env.overwriteOutput
        self.new_state = bool(overwrite)

    def __enter__(self, *args, **kwargs):
        arcpy.env.overwriteOutput = self.new_state

    def __exit__(self, *args, **kwargs):
        arcpy.env.overwriteOutput = self.orig_state

class WorkSpace(object):
    """ Context manager to temporarily set the ``workspace``
    environment variable.

    Inside the context manager, the ``arcpy.env.workspace`` will
    be set to the given value. Once the interpreter leaves the code
    block by any means (e.g., sucessful execution, raised exception),
    ``arcpy.env.workspace`` will reset to its original value.

    Example
    -------
    >>> import tidegates
    >>> with tidegates.utils.OverwriteState(False):
    ...     # some operation that should fail if output already exists

    """

    def __init__(self, path):
        self.orig_workspace = arcpy.env.workspace
        self.new_workspace = path

    def __enter__(self, *args, **kwargs):
        arcpy.env.workspace = self.new_workspace

    def __exit__(self, *args, **kwargs):
        arcpy.env.workspace = self.orig_workspace

def result_to_raster(result):
    """ Gets the actual raster from an arcpy.Result object """
    return arcpy.Raster(result.getOutput(0))

def result_to_layer(result):
    """ Gets the actual layer from an arcpy.Result object """
    return arcpy.mapping.Layer(result.getOutput(0))

def rasters_to_arrays(*rasters, **kwargs):
    """ Converts an arbitrary number of rasters to numpy arrays"""
    verbose = kwargs.pop("verbose", False)
    asMessage = kwargs.pop("asMessage", False)
    squeeze = kwargs.pop("squeeze", False)

    arrays = []
    for n, r in enumerate(rasters):
        msg = 'Processing raster {} of {}: {}'.format(n, len(rasters), r)
        progress_print(msg, verbose=verbose, asMessage=asMessage)
        arrays.append(arcpy.RasterToNumPyArray(r, nodata_to_value=-999))

    if squeeze and len(arrays) == 1:
        arrays = arrays[0]

    return arrays

def array_to_raster(array, template):
    """ Create an arcpy.Raster from a numpy.ndarray based on a template.

    Parameters
    ----------
    array : numpy.ndarray
        The array of values to be coverted to a raster.
    template : arcpy.Raster
        The raster whose, extent, position, and cell size will be
        applied to ``array``.

    Returns
    -------
    newraster : arcpy.Raster

    """

    newraster = arcpy.NumPyArrayToRaster(
        in_array=array,
        lower_left_corner=template.extent.lowerLeft,
        x_cell_size=template.meanCellWidth,
        y_cell_size=template.meanCellHeight,
        value_to_nodata=0
    )

    return newraster

def load_data(datapath, datatype, greedyRasters=True):
    """ Prepare a DEM or other raster for masking floods

    Parameters
    ----------
    datapath : str, arcpy.Raster, or arcpy.mapping.Layer
        The (filepath to the) data you want to load.
    datatype : str
        The type of data you are trying to load. Must be either
        "shape" (for polygons) or "raster" (for rasters).
    greedyRasters : bool (default = True)
        Currently, arcpy lets you load raster data as a "Raster" or as a
        "Layer". When ``greedyRasters`` is True, rasters loaded as type
        "Layer" will be forced to type "Raster".

    Returns
    -------
    data : arcpy.Raster or arcpy.mapping.Layer
        The data loaded as an arcpy type.

    """

    dtype_lookup = {
        'raster': arcpy.Raster,
        'grid': arcpy.Raster,
        'shape': arcpy.mapping.Layer,
        'layer': arcpy.mapping.Layer,
    }

    try:
        objtype = dtype_lookup[datatype.lower()]
    except KeyError:
        raise ValueError("Datatype {} not supported. Must be raster or layer".format(datatype))

    if isinstance(datapath, objtype):
        data = datapath
    else:
        try:
            data = objtype(datapath)
        except:
            raise ValueError("could not load {} as a {}".format(datapath, objtype))

    if greedyRasters and isinstance(data, arcpy.mapping.Layer) and data.isRasterLayer:
        data = arcpy.Raster(datapath)

    return data

def process_polygons(polygons, tidegate_column, cellsize=4):
    """ Prepare tidegates' areas of influence polygons for flooding
    by converting to a raster.

    Parameters
    ----------
    polygons : str or arcpy.mapping.Layer
        The (filepath to the) zones that will be flooded. If a string,
        a Layer will be created.
    tidegate_column : str
        Name of the column in the ``polygons`` layer that associates
        each geomstry with a tidegate.
    cellsize : int
        Desired cell dimension of the output raster. Default is 4 m.

    Returns
    -------
    zones : arcpy.Raster
        The zones of influence as a raster dataset
    result : arcpy.Result
        The weird, crpyric results object that so many (but not all)
        ESRI arcpy function return.

    """

    _zones = load_data(polygons, 'shape')

    with OverwriteState(True), Extension("spatial"):
        result = arcpy.conversion.PolygonToRaster(
            in_features=_zones,
            value_field=tidegate_column,
            cellsize=cellsize,
        )

    # result object isn't actually the raster
    # need to read in the raster as an object
    # the file path embedded in the result
    zones = result_to_raster(result)

    return zones, result

def clip_dem_to_zones(dem, zones):
    """ Limits the extent of the topographic data (``dem``) to that of
    the zones of influence  so that we can easily use array
    representations of the rasters.

    Parameters
    ----------
    dem : arcpy.Raster
        Digital elevation model of the area of interest.
    zones : arcpy.Raster
        The raster whose cell values represent the zones of influence
        of each tidegate.

    Returns
    -------
    dem_clipped : arcpy.Raster
        The zones of influence as a raster dataset
    result : arcpy.Result
        The weird, crpyric results object that so many (but not all)
        ESRI arcpy function return.

    """

    _dem = load_data(dem, "raster")
    _zones = load_data(zones, "raster")

    with OverwriteState(True) as state:
        result = arcpy.management.Clip(
            in_raster=_dem,
            in_template_dataset=_zones,
            out_raster="_temp_clipped_dem",
            maintain_clipping_extent="MAINTAIN_EXTENT",
            clipping_geometry="NONE",
        )

    dem_clipped = result_to_raster(result)

    return dem_clipped, result

def mask_array_with_flood(zones_array, topo_array, elevation):
    """ Mask out non-flooded portions of rasters.

    Parameters
    ----------
    zones_array : numpy.ndarray
        Array of zone IDs from each zone of influence.
    topo_array : numpy.ndarray
        Digital elevation model (as an array) of the areas.
    elevation : float
        The flood elevation *above& which everything will be masked.

    Returns
    -------
    flooded_array : numpy.ndarray
        Array of zone IDs only where there is flooding.

    """

    # compute mask of non-zoned areas of topo
    nonzone_mask = zones_array <= 0

    invalid_mask = numpy.ma.masked_invalid(topo_array).mask
    topo_array[invalid_mask] = -999

    # mask out zoned areas above the flood elevation
    unflooded_mask = topo_array > elevation

    # apply the mask to the zone array
    final_mask = nonzone_mask | unflooded_mask
    flooded_array = zones_array.copy()
    flooded_array[final_mask] = 0

    return flooded_array
<commit_msg>fix docstring and rename flood function<commit_after>import os
import sys
import glob
import datetime

import numpy

import arcpy

def progress_print(msg, verbose=False, asMessage=False):
    if verbose:
        if asMessage:
            arcpy.AddMessage(msg)
        else:
            print(msg)

class EasyMapDoc(object):
    def __init__(self, *args, **kwargs):
        self.mapdoc = arcpy.mapping.MapDocument(*args, **kwargs)

    @property
    def layers(self):
        return arcpy.mapping.ListLayers(self.mapdoc)

    @property
    def dataframes(self):
        return arcpy.mapping.ListDataFrames(self.mapdoc)

    def findLayerByName(self, name):
        for lyr in self.layers:
            if not lyr.isGroupLayer and lyr.name == name:
                return lyr

    def add_layer(self, layer, df=None, position='top'):
        # if no dataframe is provided, select the first
        if df is None:
            df = self.dataframes[0]

        # check that the position is valid
        valid_positions = ['auto_arrange', 'bottom', 'top']
        if position.lower() not in valid_positions:
            raise ValueError('Position: %s is not in %s' % (position.lower, valid_positions))

        # layer can be a path to a file. if so, convert to a Layer object
        if isinstance(layer, basestring):
            layer = arcpy.mapping.Layer(layer)
        elif not isinstance(layer, arcpy.mapping.Layer):
            raise ValueError("``layer`` be an arcpy Layer or a path to a file")

        # add the layer to the map
        arcpy.mapping.AddLayer(df, layer, position.upper())

        # return the layer
        return layer

class Extension(object):
    """ Context manager to facilitate the use of ArcGIS extensions

    Inside the context manager, the extension will be checked out. Once
    the interpreter leaves the code block by any means (e.g., sucessful
    execution, raised exception) the extension will be checked back in.

    Example
    -------
    >>> import tidegates, arcpy
    >>> with tidegates.utils.Extension("spatial"):
    ...     arcpy.sa.Hillshade("C:/data/dem.tif")

    """

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

    def __enter__(self):
        if arcpy.CheckExtension(self.name) == "Available":
            status = arcpy.CheckOutExtension(self.name)
            return status
        else:
            raise RuntimeError("%s license isn't available" % self.name)

    def __exit__(self, *args):
        arcpy.CheckOutExtension(self.name)

class OverwriteState(object):
    """ Context manager to temporarily set the ``overwriteOutput``
    environment variable.

    Inside the context manager, the ``arcpy.env.overwriteOutput`` will
    be set to the given value. Once the interpreter leaves the code
    block by any means (e.g., sucessful execution, raised exception),
    ``arcpy.env.overwriteOutput`` will reset to its original value.

    Example
    -------
    >>> import tidegates
    >>> with tidegates.utils.OverwriteState(False):
    ...     # some operation that should fail if output already exists

    """

    def __init__(self, overwrite):
        self.orig_state = arcpy.env.overwriteOutput
        self.new_state = bool(overwrite)

    def __enter__(self, *args, **kwargs):
        arcpy.env.overwriteOutput = self.new_state

    def __exit__(self, *args, **kwargs):
        arcpy.env.overwriteOutput = self.orig_state

class WorkSpace(object):
    """ Context manager to temporarily set the ``workspace``
    environment variable.

    Inside the context manager, the ``arcpy.env.workspace`` will
    be set to the given value. Once the interpreter leaves the code
    block by any means (e.g., sucessful execution, raised exception),
    ``arcpy.env.workspace`` will reset to its original value.

    Example
    -------
    >>> import tidegates
    >>> with tidegates.utils.OverwriteState(False):
    ...     # some operation that should fail if output already exists

    """

    def __init__(self, path):
        self.orig_workspace = arcpy.env.workspace
        self.new_workspace = path

    def __enter__(self, *args, **kwargs):
        arcpy.env.workspace = self.new_workspace

    def __exit__(self, *args, **kwargs):
        arcpy.env.workspace = self.orig_workspace

def result_to_raster(result):
    """ Gets the actual raster from an arcpy.Result object """
    return arcpy.Raster(result.getOutput(0))

def result_to_layer(result):
    """ Gets the actual layer from an arcpy.Result object """
    return arcpy.mapping.Layer(result.getOutput(0))

def rasters_to_arrays(*rasters, **kwargs):
    """ Converts an arbitrary number of rasters to numpy arrays"""
    verbose = kwargs.pop("verbose", False)
    asMessage = kwargs.pop("asMessage", False)
    squeeze = kwargs.pop("squeeze", False)

    arrays = []
    for n, r in enumerate(rasters):
        msg = 'Processing raster {} of {}: {}'.format(n, len(rasters), r)
        progress_print(msg, verbose=verbose, asMessage=asMessage)
        arrays.append(arcpy.RasterToNumPyArray(r, nodata_to_value=-999))

    if squeeze and len(arrays) == 1:
        arrays = arrays[0]

    return arrays

def array_to_raster(array, template):
    """ Create an arcpy.Raster from a numpy.ndarray based on a template.

    Parameters
    ----------
    array : numpy.ndarray
        The array of values to be coverted to a raster.
    template : arcpy.Raster
        The raster whose, extent, position, and cell size will be
        applied to ``array``.

    Returns
    -------
    newraster : arcpy.Raster

    """

    newraster = arcpy.NumPyArrayToRaster(
        in_array=array,
        lower_left_corner=template.extent.lowerLeft,
        x_cell_size=template.meanCellWidth,
        y_cell_size=template.meanCellHeight,
        value_to_nodata=0
    )

    return newraster

def load_data(datapath, datatype, greedyRasters=True):
    """ Prepare a DEM or other raster for masking floods

    Parameters
    ----------
    datapath : str, arcpy.Raster, or arcpy.mapping.Layer
        The (filepath to the) data you want to load.
    datatype : str
        The type of data you are trying to load. Must be either
        "shape" (for polygons) or "raster" (for rasters).
    greedyRasters : bool (default = True)
        Currently, arcpy lets you load raster data as a "Raster" or as a
        "Layer". When ``greedyRasters`` is True, rasters loaded as type
        "Layer" will be forced to type "Raster".

    Returns
    -------
    data : arcpy.Raster or arcpy.mapping.Layer
        The data loaded as an arcpy type.

    """

    dtype_lookup = {
        'raster': arcpy.Raster,
        'grid': arcpy.Raster,
        'shape': arcpy.mapping.Layer,
        'layer': arcpy.mapping.Layer,
    }

    try:
        objtype = dtype_lookup[datatype.lower()]
    except KeyError:
        raise ValueError("Datatype {} not supported. Must be raster or layer".format(datatype))

    # if the input is already a Raster or Layer, just return it
    if isinstance(datapath, objtype):
        data = datapath
    # otherwise, load it as the datatype
    else:
        try:
            data = objtype(datapath)
        except:
            raise ValueError("could not load {} as a {}".format(datapath, objtype))

    if greedyRasters and isinstance(data, arcpy.mapping.Layer) and data.isRasterLayer:
        data = arcpy.Raster(datapath)

    return data

def process_polygons(polygons, tidegate_column, cellsize=4):
    """ Prepare tidegates' areas of influence polygons for flooding
    by converting to a raster.

    Parameters
    ----------
    polygons : str or arcpy.mapping.Layer
        The (filepath to the) zones that will be flooded. If a string,
        a Layer will be created.
    tidegate_column : str
        Name of the column in the ``polygons`` layer that associates
        each geomstry with a tidegate.
    cellsize : int
        Desired cell dimension of the output raster. Default is 4 m.

    Returns
    -------
    zones : arcpy.Raster
        The zones of influence as a raster dataset
    result : arcpy.Result
        The weird, crpyric results object that so many (but not all)
        ESRI arcpy function return.

    """

    _zones = load_data(polygons, 'shape')

    with OverwriteState(True), Extension("spatial"):
        result = arcpy.conversion.PolygonToRaster(
            in_features=_zones,
            value_field=tidegate_column,
            cellsize=cellsize,
        )

    # result object isn't actually the raster
    # need to read in the raster as an object
    # the file path embedded in the result
    zones = result_to_raster(result)

    return zones, result

def clip_dem_to_zones(dem, zones):
    """ Limits the extent of the topographic data (``dem``) to that of
    the zones of influence  so that we can easily use array
    representations of the rasters.

    Parameters
    ----------
    dem : arcpy.Raster
        Digital elevation model of the area of interest.
    zones : arcpy.Raster
        The raster whose cell values represent the zones of influence
        of each tidegate.

    Returns
    -------
    dem_clipped : arcpy.Raster
        The zones of influence as a raster dataset
    result : arcpy.Result
        The weird, cryptic results object that so many (but not all)
        ESRI arcpy function return.

    """

    _dem = load_data(dem, "raster")
    _zones = load_data(zones, "raster")

    with OverwriteState(True) as state:
        result = arcpy.management.Clip(
            in_raster=_dem,
            in_template_dataset=_zones,
            out_raster="_temp_clipped",
            maintain_clipping_extent="MAINTAIN_EXTENT",
            clipping_geometry="NONE",
        )

    dem_clipped = result_to_raster(result)

    return dem_clipped, result

def flood_zones(zones_array, topo_array, elevation):
    """ Mask out non-flooded portions of rasters.

    Parameters
    ----------
    zones_array : numpy.ndarray
        Array of zone IDs from each zone of influence.
    topo_array : numpy.ndarray
        Digital elevation model (as an array) of the areas.
    elevation : float
        The flood elevation *above& which everything will be masked.

    Returns
    -------
    flooded_array : numpy.ndarray
        Array of zone IDs only where there is flooding.

    """

    # compute mask of non-zoned areas of topo
    nonzone_mask = zones_array <= 0

    invalid_mask = numpy.ma.masked_invalid(topo_array).mask
    topo_array[invalid_mask] = -999

    # mask out zoned areas above the flood elevation
    unflooded_mask = topo_array > elevation

    # apply the mask to the zone array
    final_mask = nonzone_mask | unflooded_mask
    flooded_array = zones_array.copy()
    flooded_array[final_mask] = 0

    return flooded_array
<|endoftext|>


<commit_before>import os

import pygame
from intervaltree import Interval

from behaviours import Collide

class Tile:
    """A simple tile in the gameworld"""

    def __init__(self, x, y, name):
        self.name = name
        self.x = x
        self.y = y
        self.width = 1
        self.height = 1
        self.interval_x = Interval(self.x, self.x + 1, self)
        self.interval_y = Interval(self.y, self.y + 1, self)
        Collide.tx.add(self.interval_x)
        Collide.ty.add(self.interval_y)

        path = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../" + name.lower() + "/" + name + ".png")
        try:
            self.sprite = pygame.image.load(path)
        except:
            print("Could not load sprite for " + name)

    def draw(self, screen):
        screen.blit(self.sprite, (self.x, self.y))

    def get_top(self):
        return self.y

    def get_vertical_center(self):
        return self.y + self.height / 2

    def get_bottom(self):
        return self.y + self.height

    def get_left(self):
        return self.x

    def get_horizontal_center(self):
        return self.x + self.width / 2

    def get_right(self):
        return self.x + self.width

    def __del__(self):
        Collide.tx.discard(self.interval_x)
        Collide.ty.discard(self.interval_y)
<commit_msg>Add caching for tile sprites<commit_after>import os

import pygame
from intervaltree import Interval

from behaviours import Collide

sprite_map = {}

class Tile:
    """A simple tile in the gameworld"""

    def __init__(self, x, y, name):
        self.name = name
        self.x = x
        self.y = y
        self.width = 1
        self.height = 1
        self.interval_x = Interval(self.x, self.x + 1, self)
        self.interval_y = Interval(self.y, self.y + 1, self)
        Collide.tx.add(self.interval_x)
        Collide.ty.add(self.interval_y)
        if name in sprite_map:
            self.sprite = sprite_map[name]
        else:
            path = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../" + name.lower() + "/" + name + ".png")
            try:
                self.sprite = pygame.image.load(path)
                sprite_map[name] = self.sprite
            except:
                print("Could not load sprite for " + name)

    def draw(self, screen):
        screen.blit(self.sprite, (self.x, self.y))

    def get_top(self):
        return self.y

    def get_vertical_center(self):
        return self.y + self.height / 2

    def get_bottom(self):
        return self.y + self.height

    def get_left(self):
        return self.x

    def get_horizontal_center(self):
        return self.x + self.width / 2

    def get_right(self):
        return self.x + self.width

    def __del__(self):
        Collide.tx.discard(self.interval_x)
        Collide.ty.discard(self.interval_y)
<|endoftext|>


<commit_before>import argparse
from timeflow.helpers import write_to_log_file

def log(args):
    write_to_log_file(args.message)

def edit(args):
    print("Parsing edit command")
    if args.editor:
        print("Parsing edit's --editor option")

def stats(args):
    print("Parsing stats command")
    if args.today:
        print("Parsing stats's --today option")
    elif args.yesterday:
        print("Parsing stats's --yesterday option")
    elif args.day:
        print("Parsing stats's --day option")
    elif args.week:
        print("Parsing stats's --week option")
    elif args.last_week:
        print("Parsing stats's --last-week option")
    elif args.month:
        print("Parsing stats's --month option")
    elif args.last_month:
        print("Parsing stats's --last-month option")
    elif args._from:
        print("Parsing stats's --from option")
    elif args.to:
        print("Parsing stats's --to option")
    else:
        print("Parsing stats's --today option")

def set_log_parser(subparser):
    log_parser = subparser.add_parser("log", help="Create timelog message")
    log_parser.add_argument("message", help="message that will be logged")
    # call log() function, when processing log command
    log_parser.set_defaults(func=log)

def set_edit_parser(subparser):
    edit_parser = subparser.add_parser("edit", help="Edit timelog file")
    edit_parser.add_argument("-e", "--editor", help="Explicitly set editor")
    # call edit() function, when processing edit command
    edit_parser.set_defaults(func=edit)

def set_stats_parser(subparser):
    stats_parser = subparser.add_parser(
        "stats",
        help="Show how much time was spent working or slacking"
    )

    stats_parser.add_argument("--today",
                              action="store_true",
                              help="Show today's work times (default)")
    stats_parser.add_argument("-y", "--yesterday",
                              action="store_true",
                              help="Show yesterday's work times")
    stats_parser.add_argument("-d", "--day",
                              help="Show specific day's work times")

    stats_parser.add_argument("--week",
                              help="Show specific week's work times")
    stats_parser.add_argument("--last-week",
                              action="store_true",
                              help="Show last week's work times")

    stats_parser.add_argument("--month",
                              help="Show specific month's work times")
    stats_parser.add_argument("--last-month",
                              action="store_true",
                              help="Show last month's work times")

    stats_parser.add_argument("-f", "--from",
                              help="Show work times from specific date",
                              dest="_from")
    stats_parser.add_argument("-t", "--to",
                              help="Show work times from to specific date")

    # call stats() function, when processing stats command
    stats_parser.set_defaults(func=stats)

def parse_args(args):
    parser = argparse.ArgumentParser()

    subparser = parser.add_subparsers(help="sub-command help")
    set_log_parser(subparser)
    set_edit_parser(subparser)
    set_stats_parser(subparser)

    return parser.parse_args(args)
<commit_msg>Create edit option functionality<commit_after>import argparse
import os
import subprocess

from timeflow.helpers import write_to_log_file, LOG_FILE

def log(args):
    write_to_log_file(args.message)

def edit(args):
    if args.editor:
        subprocess.call([args.editor, LOG_FILE])
    else:
        subprocess.call(['echo', 'Trying to open $EDITOR'])
        if os.environ.get('EDITOR'):
            subprocess.call([os.environ.get('EDITOR'), LOG_FILE])
        else:
            subprocess.call([
                "echo",
                "Set your default editor in EDITOR environment variable or \n"
                "call edit command with -e option and pass your editor, e.g.:\n"
                "timeflow edit -e vim",
            ])

def stats(args):
    print("Parsing stats command")
    if args.today:
        print("Parsing stats's --today option")
    elif args.yesterday:
        print("Parsing stats's --yesterday option")
    elif args.day:
        print("Parsing stats's --day option")
    elif args.week:
        print("Parsing stats's --week option")
    elif args.last_week:
        print("Parsing stats's --last-week option")
    elif args.month:
        print("Parsing stats's --month option")
    elif args.last_month:
        print("Parsing stats's --last-month option")
    elif args._from:
        print("Parsing stats's --from option")
    elif args.to:
        print("Parsing stats's --to option")
    else:
        print("Parsing stats's --today option")

def set_log_parser(subparser):
    log_parser = subparser.add_parser("log", help="Create timelog message")
    log_parser.add_argument("message", help="message that will be logged")
    # call log() function, when processing log command
    log_parser.set_defaults(func=log)

def set_edit_parser(subparser):
    edit_parser = subparser.add_parser("edit", help="Edit timelog file")
    edit_parser.add_argument("-e", "--editor", help="Explicitly set editor")
    # call edit() function, when processing edit command
    edit_parser.set_defaults(func=edit)

def set_stats_parser(subparser):
    stats_parser = subparser.add_parser(
        "stats",
        help="Show how much time was spent working or slacking"
    )

    stats_parser.add_argument("--today",
                              action="store_true",
                              help="Show today's work times (default)")
    stats_parser.add_argument("-y", "--yesterday",
                              action="store_true",
                              help="Show yesterday's work times")
    stats_parser.add_argument("-d", "--day",
                              help="Show specific day's work times")

    stats_parser.add_argument("--week",
                              help="Show specific week's work times")
    stats_parser.add_argument("--last-week",
                              action="store_true",
                              help="Show last week's work times")

    stats_parser.add_argument("--month",
                              help="Show specific month's work times")
    stats_parser.add_argument("--last-month",
                              action="store_true",
                              help="Show last month's work times")

    stats_parser.add_argument("-f", "--from",
                              help="Show work times from specific date",
                              dest="_from")
    stats_parser.add_argument("-t", "--to",
                              help="Show work times from to specific date")

    # call stats() function, when processing stats command
    stats_parser.set_defaults(func=stats)

def parse_args(args):
    parser = argparse.ArgumentParser()

    subparser = parser.add_subparsers(help="sub-command help")
    set_log_parser(subparser)
    set_edit_parser(subparser)
    set_stats_parser(subparser)

    return parser.parse_args(args)
<|endoftext|>


<commit_before># Copyright (c) 2008 Joost Cassee
# Licensed under the terms of the MIT License (see LICENSE.txt)

"""
This TinyMCE widget was copied and extended from this code by John D'Agostino:
http://code.djangoproject.com/wiki/CustomWidgetsTinyMCE
"""

from django import forms
from django.conf import settings
from django.contrib.admin import widgets as admin_widgets
from django.core.urlresolvers import reverse
from django.forms.widgets import flatatt
try:
    from django.utils.encoding import smart_unicode
except ImportError:
    from django.forms.util import smart_unicode
from django.utils.html import escape
from django.utils import simplejson
from django.utils.datastructures import SortedDict
from django.utils.safestring import mark_safe
from django.utils.translation import get_language, ugettext as _
import tinymce.settings

class TinyMCE(forms.Textarea):
    """
    TinyMCE widget. Set settings.TINYMCE_JS_URL to set the location of the
    javascript file. Default is "MEDIA_URL + 'js/tiny_mce/tiny_mce.js'".
    You can customize the configuration with the mce_attrs argument to the
    constructor.

    In addition to the standard configuration you can set the
    'content_language' parameter. It takes the value of the 'language'
    parameter by default.

    In addition to the default settings from settings.TINYMCE_DEFAULT_CONFIG,
    this widget sets the 'language', 'directionality' and
    'spellchecker_languages' parameters by default. The first is derived from
    the current Django language, the others from the 'content_language'
    parameter.
    """

    def __init__(self, content_language=None, attrs=None, mce_attrs=None):
        super(TinyMCE, self).__init__(attrs)
        if mce_attrs is None:
            mce_attrs = {}
        self.mce_attrs = mce_attrs
        if content_language is None:
            content_language = mce_attrs.get('language', None)
        self.content_language = content_language

    def render(self, name, value, attrs=None):
        if value is None: value = ''
        value = smart_unicode(value)
        final_attrs = self.build_attrs(attrs)
        final_attrs['name'] = name
        assert 'id' in final_attrs, "TinyMCE widget attributes must contain 'id'"

        mce_config = tinymce.settings.DEFAULT_CONFIG.copy()
        mce_config.update(get_language_config(self.content_language))
        if tinymce.settings.USE_FILEBROWSER:
            mce_config['file_browser_callback'] = "djangoFileBrowser"
        mce_config.update(self.mce_attrs)
        if not 'mode' in mce_config:
            mce_config['mode'] = 'exact'
        if mce_config['mode'] == 'exact':
            mce_config['elements'] = final_attrs['id']
        mce_config['strict_loading_mode'] = 1

        # Fix for js functions
        js_functions = {}
        for k in ('paste_preprocess','paste_postprocess'):
            if k in mce_config:
               js_functions[k] = mce_config[k]
               del mce_config[k]
        mce_json = simplejson.dumps(mce_config)

        pos = final_attrs['id'].find('__prefix__')
        if pos != -1:
            mce_json = mce_json.replace(u'"%s"' % final_attrs['id'], u'elements')

        for k in js_functions:
            index = mce_json.rfind('}')
            mce_json = mce_json[:index]+', '+k+':'+js_functions[k].strip()+mce_json[index:]

        html = [u'<textarea%s>%s</textarea>' % (flatatt(final_attrs), escape(value))]
        if tinymce.settings.USE_COMPRESSOR:
            compressor_config = {
                'plugins': mce_config.get('plugins', ''),
                'themes': mce_config.get('theme', 'advanced'),
                'languages': mce_config.get('language', ''),
                'diskcache': True,
                'debug': False,
            }
            compressor_json = simplejson.dumps(compressor_config)

        if pos != -1:
            html.append(u'''<script type="text/javascript">
setTimeout(function () {
    var id = '%s';

    if (typeof(window._tinymce_inited) == 'undefined') {
        window._tinymce_inited = [];
    }

    if (typeof(window._tinymce_inited[id]) == 'undefined') {
        window._tinymce_inited[id] = true;
    } else {
        var elements = id.replace(/__prefix__/, parseInt(document.getElementById('%sTOTAL_FORMS').value) - 1);
        console.log(elements);
        if (document.getElementById(elements)) {
            tinymce.init(%s);
        }
    }
}, 0);
</script>''' % (final_attrs['id'], final_attrs['id'][0:pos], mce_json))
        else:
            html.append(u'<script type="text/javascript">tinyMCE.init(%s)</script>' % mce_json)

        return mark_safe(u'\n'.join(html))

    def _media(self):
        if tinymce.settings.USE_COMPRESSOR:
            js = [reverse('tinymce-compressor')]
        else:
            js = [tinymce.settings.JS_URL]
        if tinymce.settings.USE_FILEBROWSER:
            js.append(reverse('tinymce-filebrowser'))
        return forms.Media(js=js)
    media = property(_media)

class AdminTinyMCE(admin_widgets.AdminTextareaWidget, TinyMCE):
    pass

def get_language_config(content_language=None):
    language = get_language()[:2]
    if content_language:
        content_language = content_language[:2]
    else:
        content_language = language

    config = {}
    config['language'] = language

    lang_names = SortedDict()
    for lang, name in settings.LANGUAGES:
        if lang[:2] not in lang_names: lang_names[lang[:2]] = []
        lang_names[lang[:2]].append(_(name))
    sp_langs = []
    for lang, names in lang_names.items():
        if lang == content_language:
            default = '+'
        else:
            default = ''
        sp_langs.append(u'%s%s=%s' % (default, ' / '.join(names), lang))

    config['spellchecker_languages'] = ','.join(sp_langs)

    if content_language in settings.LANGUAGES_BIDI:
        config['directionality'] = 'rtl'
    else:
        config['directionality'] = 'ltr'

    if tinymce.settings.USE_SPELLCHECKER:
        config['spellchecker_rpc_url'] = reverse('tinymce.views.spell_check')

    return config
<commit_msg>smart_unicode method has disapeared in Django1.5, now is called smart_text (python3 support). Try to import the last one and maintaing the old imports for backward compatibility. I'm sure that there is a better/cleaner way of doing this<commit_after># Copyright (c) 2008 Joost Cassee
# Licensed under the terms of the MIT License (see LICENSE.txt)

"""
This TinyMCE widget was copied and extended from this code by John D'Agostino:
http://code.djangoproject.com/wiki/CustomWidgetsTinyMCE
"""

from django import forms
from django.conf import settings
from django.contrib.admin import widgets as admin_widgets
from django.core.urlresolvers import reverse
from django.forms.widgets import flatatt
try:
    from django.utils.encoding import smart_text as smart_unicode
except ImportError:
    try:
        from django.utils.encoding import smart_unicode
    except ImportError:
        from django.forms.util import smart_unicode
from django.utils.html import escape
from django.utils import simplejson
from django.utils.datastructures import SortedDict
from django.utils.safestring import mark_safe
from django.utils.translation import get_language, ugettext as _
import tinymce.settings

class TinyMCE(forms.Textarea):
    """
    TinyMCE widget. Set settings.TINYMCE_JS_URL to set the location of the
    javascript file. Default is "MEDIA_URL + 'js/tiny_mce/tiny_mce.js'".
    You can customize the configuration with the mce_attrs argument to the
    constructor.

    In addition to the standard configuration you can set the
    'content_language' parameter. It takes the value of the 'language'
    parameter by default.

    In addition to the default settings from settings.TINYMCE_DEFAULT_CONFIG,
    this widget sets the 'language', 'directionality' and
    'spellchecker_languages' parameters by default. The first is derived from
    the current Django language, the others from the 'content_language'
    parameter.
    """

    def __init__(self, content_language=None, attrs=None, mce_attrs=None):
        super(TinyMCE, self).__init__(attrs)
        if mce_attrs is None:
            mce_attrs = {}
        self.mce_attrs = mce_attrs
        if content_language is None:
            content_language = mce_attrs.get('language', None)
        self.content_language = content_language

    def render(self, name, value, attrs=None):
        if value is None: value = ''
        value = smart_unicode(value)
        final_attrs = self.build_attrs(attrs)
        final_attrs['name'] = name
        assert 'id' in final_attrs, "TinyMCE widget attributes must contain 'id'"

        mce_config = tinymce.settings.DEFAULT_CONFIG.copy()
        mce_config.update(get_language_config(self.content_language))
        if tinymce.settings.USE_FILEBROWSER:
            mce_config['file_browser_callback'] = "djangoFileBrowser"
        mce_config.update(self.mce_attrs)
        if not 'mode' in mce_config:
            mce_config['mode'] = 'exact'
        if mce_config['mode'] == 'exact':
            mce_config['elements'] = final_attrs['id']
        mce_config['strict_loading_mode'] = 1

        # Fix for js functions
        js_functions = {}
        for k in ('paste_preprocess','paste_postprocess'):
            if k in mce_config:
               js_functions[k] = mce_config[k]
               del mce_config[k]
        mce_json = simplejson.dumps(mce_config)

        pos = final_attrs['id'].find('__prefix__')
        if pos != -1:
            mce_json = mce_json.replace(u'"%s"' % final_attrs['id'], u'elements')

        for k in js_functions:
            index = mce_json.rfind('}')
            mce_json = mce_json[:index]+', '+k+':'+js_functions[k].strip()+mce_json[index:]

        html = [u'<textarea%s>%s</textarea>' % (flatatt(final_attrs), escape(value))]
        if tinymce.settings.USE_COMPRESSOR:
            compressor_config = {
                'plugins': mce_config.get('plugins', ''),
                'themes': mce_config.get('theme', 'advanced'),
                'languages': mce_config.get('language', ''),
                'diskcache': True,
                'debug': False,
            }
            compressor_json = simplejson.dumps(compressor_config)

        if pos != -1:
            html.append(u'''<script type="text/javascript">
setTimeout(function () {
    var id = '%s';

    if (typeof(window._tinymce_inited) == 'undefined') {
        window._tinymce_inited = [];
    }

    if (typeof(window._tinymce_inited[id]) == 'undefined') {
        window._tinymce_inited[id] = true;
    } else {
        var elements = id.replace(/__prefix__/, parseInt(document.getElementById('%sTOTAL_FORMS').value) - 1);
        console.log(elements);
        if (document.getElementById(elements)) {
            tinymce.init(%s);
        }
    }
}, 0);
</script>''' % (final_attrs['id'], final_attrs['id'][0:pos], mce_json))
        else:
            html.append(u'<script type="text/javascript">tinyMCE.init(%s)</script>' % mce_json)

        return mark_safe(u'\n'.join(html))

    def _media(self):
        if tinymce.settings.USE_COMPRESSOR:
            js = [reverse('tinymce-compressor')]
        else:
            js = [tinymce.settings.JS_URL]
        if tinymce.settings.USE_FILEBROWSER:
            js.append(reverse('tinymce-filebrowser'))
        return forms.Media(js=js)
    media = property(_media)

class AdminTinyMCE(admin_widgets.AdminTextareaWidget, TinyMCE):
    pass

def get_language_config(content_language=None):
    language = get_language()[:2]
    if content_language:
        content_language = content_language[:2]
    else:
        content_language = language

    config = {}
    config['language'] = language

    lang_names = SortedDict()
    for lang, name in settings.LANGUAGES:
        if lang[:2] not in lang_names: lang_names[lang[:2]] = []
        lang_names[lang[:2]].append(_(name))
    sp_langs = []
    for lang, names in lang_names.items():
        if lang == content_language:
            default = '+'
        else:
            default = ''
        sp_langs.append(u'%s%s=%s' % (default, ' / '.join(names), lang))

    config['spellchecker_languages'] = ','.join(sp_langs)

    if content_language in settings.LANGUAGES_BIDI:
        config['directionality'] = 'rtl'
    else:
        config['directionality'] = 'ltr'

    if tinymce.settings.USE_SPELLCHECKER:
        config['spellchecker_rpc_url'] = reverse('tinymce.views.spell_check')

    return config
<|endoftext|>


<commit_before>import wpilib

def step(value, min):
    if abs(value) < min:
        value = 0
    return value

def step_range(value, min, max, default):
    if not (min <= value <= max):
        value = default
    return value

#Robot object definition
class Robot(wpilib.IterativeRobot):
    #Initialize all of the sensors and controllers on the robot
    def robotInit(self):
        #Initialize the Joysticks
        self.left_joystick = wpilib.Joystick(0)
        self.right_joystick = wpilib.Joystick(1)

        #Initialize the drive motors
        self.left_motor = wpilib.Talon(0)
        self.right_motor = wpilib.Talon(1)

        #Initialize the drive system
        self.robotdrive = wpilib.RobotDrive(self.left_motor, self.right_motor)

        #Invert the motors so that they drive the right way
        self.robotdrive.setInvertedMotor(
            wpilib.RobotDrive.MotorType.kRearLeft, False)
        self.robotdrive.setInvertedMotor(
            wpilib.RobotDrive.MotorType.kRearRight, False)

        #Initialize the winch motor
        self.winch_motor = wpilib.Talon(2)

        #Initialize the arm motor
        self.arm_motor = wpilib.Talon(3)

        #Initialize the accelerometer
        self.accel = wpilib.BuiltInAccelerometer()
        self.a_x_sum = 0.0
        self.a_x_count = 0
        self.a_y_sum = 0.0
        self.a_y_count = 0

        self.last_winch_signal = 0

        #Initialize the gyro
        self.gyro = wpilib.Gyro(0)

        #Initialize the winch encoder
        self.winch_encoder = wpilib.Encoder(1,2)
        self.winch_setpoint = -self.winch_encoder.get()

        #Initialize the compressor
        self.compressor = wpilib.Compressor(0)

        #Initialize the pneumatic solenoids
        self.solenoid1 = wpilib.Solenoid(1)
        self.solenoid2 = wpilib.Solenoid(2)

        #Initialize the ultrasonic sensors
        self.left_ultrasonic_sensor = wpilib.AnalogInput(1)
        self.right_ultrasonic_sensor = wpilib.AnalogInput(2)

        #Initialize the optical sensors
        self.left_optical_sensor = wpilib.DigitalInput(3)
        self.right_optical_sensor = wpilib.DigitalInput(4)

        #Initialize the limit switches
        print ("ofoof")
        self.left_limit_switch = wpilib.DigitalInput(6)
        self.right_limit_switch = wpilib.DigitalInput(5)

        #Initialize the compressor watchdog
        self.dog = wpilib.MotorSafety()
        self.dog.setSafetyEnabled(False)
        self.dog.setExpiration(1.75)

        #Set run variables
        self.auto_mode = "tote"
        self.claw_state = True
        self.x_pressed_last = False

    #Autonomous Mode
    def autonomousInit(self):
        self.auto_state = "start"
        self.positioned_count = 0

    def autonomousPeriodic(self):
        self.dog.feed()
        if self.auto_mode == "container":
            self.autoContainerPeriodic()
        elif self.auto_mode == "tote":
            self.autoTotePeriodic()

    #Autonomous mode for picking up recycling containers
    #Note: run variable "auto_mode" should be set to "container"
    def autoContainerPeriodic(self):
        pass

    def autoTotePeriodic2(self):
        while not (self.left_claw_whisker() and self.right_claw_whisker()):
            if not self.left_claw_whisker() and not self.right_claw_whisker():
                pass
            yield
    #Autonomous mode for picking up totes
    #Note: run variable "auto_mode" should be set to "tote"
    def autoTotePeriodic(self):
        if self.auto_state == "start":
            right_dist = step_range(self.left_ultrasonic_sensor.getValue(), 50, 200, 200)
            left_dist = step_range(self.right_ultrasonic_sensor.getValue(), 50, 200, 200)
            if right_dist < 70 and left_dist < 70:
                self.auto_state = "positioned"
            elif right_dist >= 70 and left_dist < 70:
                val1 = (0.5 * abs(right_dist - 70) / 200.0)
                if abs(val1) < 0.2:
                    val1 = 0.2
                val0 = 0.0
                print (val0, val1)
                self.left_motor.set(val0)
                self.right_motor.set(val1)
            elif right_dist < 70 and left_dist >= 70:
                val0 = 0.0
                val1 = (0.5 * abs(left_dist - 70) / 200.0)
                if abs(val1) < 0.2:
                    val1 = 0.2
                self.left_motor.set(val0)
                self.right_motor.set(val1)
            else:
                val0 = (-0.5 * abs(right_dist - 70) / 200.0)
                if abs(val0) < 0.2:
                    val0 = -0.2
                val1 = (0.5 * abs(left_dist - 70) / 200.0)
                if abs(val1) < 0.2:
                    val1 = 0.2
                self.left_motor.set(val0)
                self.right_motor.set(val1)
        elif self.auto_mode == "positioned":
            self.positioned_count += 1
            self.winch_set(0.5);
            if self.positioned_count > 40:
                self.claw_up()

    #Teleop Mode
    def teleopInit(self):
        self.compressor.start()

    def teleopPeriodic(self):
        #If left trigger pulled, run brake algorithm, otherwise use joystick values to drive
        if self.left_joystick.getRawButton(1):
            rotation_values = self.brake_rotation()
            linear_values = self.brake_linear()
            left_wheel = rotation_values[0] + linear_values[0]
            right_wheel = rotation_values[1] + linear_values[1]
        else:
            left_wheel, right_wheel = self.drive_values()

        #Feed joystick values into drive system
        self.robotdrive.tankDrive(left_wheel, right_wheel)

        #Reset winch encoder value to 0 if right button 7 is pressed
        if self.right_joystick.getRawButton(7):
            self.winch_encoder.reset()

        #Feed winch controller raw values from the joystick
        #Right joystick button 3 raises winch, button 2 lowers winch
        winch_signal = self.right_joystick.getRawButton(3) + -self.right_joystick.getRawButton(2)
        self.winch_set(winch_signal)

        #Feed arm controller raw values from the joystick
        #Left joystick button 3 goes forward, 2 goes backward
        arm_signal = self.left_joystick.getRawButton(3) + -self.left_joystick.getRawButton(2)
        self.arm_motor.set(0.3 * arm_signal)

        #Handle piston in and out
        #Right joystick trigger button toggles claw in or out
        if self.right_joystick.getRawButton(1):
            self.x_pressed_last = True
        elif self.x_pressed_last:
            self.x_pressed_last = False
            self.claw_state = not self.claw_state
        self.set_claw()

        #Reset winch encoder value to 0 if right button 7 is pressed
        if self.right_joystick.getRawButton(7):
            self.winch_encoder.reset()

        #If the right joystick slider is down, go to test mode
        #print('slider: ' +  self.right_joystick.getRawAxis(2))
        if self.right_joystick.getRawAxis(2) > .5:
            self.testMode()

    #Required method for robotpy
    def testPeriodic():
        pass

    #Test Mode
    def testMode(self):
        #Calculate x and y distance travelled using accelerometer
        a_x = self.accel.getX()
        self.a_x_sum += a_x
        self.a_x_count += 1
        a_y = self.accel.getY()
        self.a_y_sum += a_y
        self.a_y_count += 1

        #Prints acceleration values when right button 8 is pressed
        if self.right_joystick.getRawButton(8):
            print ('x sum: ', self.a_x_sum, ' x count: ', self.a_x_count)
            print ('y sum: ', self.a_y_sum, ' y count: ', self.a_y_count)

        #Resets accelerometer counts when right button 9 is pressed
        if self.right_joystick.getRawButton(9):
            self.a_x_sum = 0.0
            self.a_x_count = 0
            self.a_y_sum = 0.0
            self.a_y_count = 0

        #Prints left ultrasonic sensor values when left button 6 is pressed
        if self.left_joystick.getRawButton(6):
            print ("left_ultrasonic_sensor: ", self.left_ultrasonic_sensor.getValue())
            print ("left limit: ", self.left_claw_whisker())
            print ("right limit: ", self.right_claw_whisker())

        #Prints right ultrasonic sensor values when left button 7 is pressed
        if self.left_joystick.getRawButton(7):
           print ("right_ultrasonic_sensor: ", self.right_ultrasonic_sensor.getValue())

        #Prints left optical sensor values when left button 11 is pressed
        if self.left_joystick.getRawButton(11):
           print ("left_optical_sensor: ", self.right_optical_sensor.get())

        #Prints right optical sensor values when left button 10 is pressed
        if self.left_joystick.getRawButton(10):
           print ("right_optical_sensor: ", self.right_optical_sensor.get())

        #Print current winch encoder value if right button 6 is pressed
        if self.right_joystick.getRawButton(6):
            revs = -self.winch_encoder.get()
            print ('revs: ', revs)
        #Print current gyro value if left button 8 is pressed
        if self.left_joystick.getRawButton(8):
            angle = self.gyro.getAngle()
            print ('angle: ', angle)

        #Reset gyro to 0 if left button 9 is pressed
        if self.left_joystick.getRawButton(9):
            self.gyro.reset()

    #Disabled Mode
    def disabledPeriodic(self):
        self.compressor.stop()

    #Helper Functions

    #Set "left" and "right" variables to the left and right
    #joystick outputs provided they are more than "threshold"
    def drive_values(self):
        joystick_threshold = 0.2
        left = step(
            self.left_joystick.getRawAxis(1),
            joystick_threshold,
        )
        right = step(
            self.right_joystick.getRawAxis(1),
            -joystick_threshold,
        )

        #Fuzzy match where if the left and right joysticks are moved about the same,
        #    then it moves the tankDrive the average of the two values
        if abs(left - right) < .1 :
            left = right = (left + right) / 2.0

        return (left, right)

    def brake_rotation(self):
        gyro_rate = self.gyro.getRate()
        wheel_rotation = gyro_rate * .1
        return (wheel_rotation, -wheel_rotation)

    def brake_linear(self):
        accel_y = self.accel.getY()
        wheel_motion = accel_y * .1
        return (-wheel_motion, -wheel_motion)

    def set_claw(self):
        self.solenoid1.set(not self.claw_state)
        self.solenoid2.set(self.claw_state)

    def claw_up(self):
        self.claw_state = False

    def claw_down(self):
        self.claw_state = True

    def left_claw_whisker(self):
        return self.left_limit_switch.get()

    def right_claw_whisker(self):
        return self.right_limit_switch.get()

    def winch_set(self, winch_signal):
        """
        Set winch controller safely by taking max and min encoder values into account
        (unless you're pressing the override button - right joystick, button 6)

        winch_signal=0 -> maintain winch position
        winch_signal>0 -> winch up?
        winch_signal<0 -> winch down?
        """
        revs = -self.winch_encoder.get()
        if self.last_winch_signal != 0 and winch_signal == 0:
            self.winch_setpoint = revs
        self.last_winch_signal = winch_signal
        if winch_signal == 0:
            val = 0.1 - 0.01 * (revs - self.winch_setpoint)
            self.winch_motor.set(val)
        else:
            if not (self.right_joystick.getRawButton(6)):
                if winch_signal > 0.1 and revs >= 1170:
                    winch_signal = 0
                if winch_signal < -0.1 and revs <= 8:
                    winch_signal = 0
            val = 0.5 * winch_signal
            self.winch_motor.set(val)

if __name__ == "__main__":
    wpilib.run(Robot)
<commit_msg>container autonomous mode<commit_after>import wpilib

def step(value, min):
    if abs(value) < min:
        value = 0
    return value

def step_range(value, min, max, default):
    if not (min <= value <= max):
        value = default
    return value

#Robot object definition
class Robot(wpilib.IterativeRobot):
    #Initialize all of the sensors and controllers on the robot
    def robotInit(self):
        #Initialize the Joysticks
        self.left_joystick = wpilib.Joystick(0)
        self.right_joystick = wpilib.Joystick(1)

        #Initialize the drive motors
        self.left_motor = wpilib.Talon(0)
        self.right_motor = wpilib.Talon(1)

        #Initialize the drive system
        self.robotdrive = wpilib.RobotDrive(self.left_motor, self.right_motor)

        #Invert the motors so that they drive the right way
        self.robotdrive.setInvertedMotor(
            wpilib.RobotDrive.MotorType.kRearLeft, False)
        self.robotdrive.setInvertedMotor(
            wpilib.RobotDrive.MotorType.kRearRight, False)

        #Initialize the winch motor
        self.winch_motor = wpilib.Talon(2)

        #Initialize the arm motor
        self.arm_motor = wpilib.Talon(3)

        #Initialize the accelerometer
        self.accel = wpilib.BuiltInAccelerometer()
        self.a_x_sum = 0.0
        self.a_x_count = 0
        self.a_y_sum = 0.0
        self.a_y_count = 0

        self.last_winch_signal = 0

        #Initialize the gyro
        self.gyro = wpilib.Gyro(0)

        #Initialize the winch encoder
        self.winch_encoder = wpilib.Encoder(1,2)
        self.winch_setpoint = -self.winch_encoder.get()

        #Initialize the compressor
        self.compressor = wpilib.Compressor(0)

        #Initialize the pneumatic solenoids
        self.solenoid1 = wpilib.Solenoid(1)
        self.solenoid2 = wpilib.Solenoid(2)

        #Initialize the ultrasonic sensors
        self.left_ultrasonic_sensor = wpilib.AnalogInput(1)
        self.right_ultrasonic_sensor = wpilib.AnalogInput(2)

        #Initialize the optical sensors
        self.left_optical_sensor = wpilib.DigitalInput(3)
        self.right_optical_sensor = wpilib.DigitalInput(4)

        #Initialize the limit switches
        print ("ofoof")
        self.left_limit_switch = wpilib.DigitalInput(6)
        self.right_limit_switch = wpilib.DigitalInput(5)

        #Initialize the compressor watchdog
        self.dog = wpilib.MotorSafety()
        self.dog.setSafetyEnabled(False)
        self.dog.setExpiration(1.75)

        #Set run variables
        self.auto_mode = "container"
        self.auto_counter = 0
        self.auto_state = 1

        self.claw_state = True
        self.x_pressed_last = False

    #Autonomous Mode
    def autonomousInit(self):
        self.auto_state = "start"
        self.positioned_count = 0

    def autonomousPeriodic(self):
        self.dog.feed()
        if self.auto_mode == "container":
            self.autoContainerPeriodic()
        elif self.auto_mode == "tote":
            self.autoTotePeriodic()

    #Autonomous mode for picking up recycling containers
    #Note: run variable "auto_mode" should be set to "container"
    def autoContainerPeriodic(self):
        # state 1: lift up to pick up container
        if self.auto_state == 1 and -self.winch_encoder.get() < 500:
            self.winch_motor.set(.5)
        elif self.auto_state == 1:
            self.auto_state = 2

        # state 2: fire right trigger - push out solenoid
        if self.auto_state == 2:
            self.claw_up()
            self.set_claw()
            self.auto_state = 3
			self.auto_counter = 0

        # state 3: drive backward over the bump
        if self.auto_state == 3 and self.auto_counter < 500: # 500 is an arbitrary test value, needs to be tested in IRL
            self.robotdrive.tankDrive(-1, -1)
            self.auto_counter += 1
        elif self.auto_state == 3:
            self.auto_state = 4

        # state 4: set container down
        if self.auto_state == 4 and -self.winch_encoder.get() > 10:
            self.winch_motor.set(-.5)
        elif self.auto_state == 4:
            self.auto_state = 5
			self.auto_counter = 0

        # state 5: back up
        if self.auto_state == 5 and self.auto_counter < 5:
            self.robotdrive.tankDrive(-1, -1)
            self.auto_counter += 1
        elif self.auto_state == 5:
            self.auto_state = 6

    def autoTotePeriodic2(self):
        while not (self.left_claw_whisker() and self.right_claw_whisker()):
            if not self.left_claw_whisker() and not self.right_claw_whisker():
                pass
            yield
    #Autonomous mode for picking up totes
    #Note: run variable "auto_mode" should be set to "tote"
    def autoTotePeriodic(self):
        if self.auto_state == "start":
            right_dist = step_range(self.left_ultrasonic_sensor.getValue(), 50, 200, 200)
            left_dist = step_range(self.right_ultrasonic_sensor.getValue(), 50, 200, 200)
            if right_dist < 70 and left_dist < 70:
                self.auto_state = "positioned"
            elif right_dist >= 70 and left_dist < 70:
                val1 = (0.5 * abs(right_dist - 70) / 200.0)
                if abs(val1) < 0.2:
                    val1 = 0.2
                val0 = 0.0
                print (val0, val1)
                self.left_motor.set(val0)
                self.right_motor.set(val1)
            elif right_dist < 70 and left_dist >= 70:
                val0 = 0.0
                val1 = (0.5 * abs(left_dist - 70) / 200.0)
                if abs(val1) < 0.2:
                    val1 = 0.2
                self.left_motor.set(val0)
                self.right_motor.set(val1)
            else:
                val0 = (-0.5 * abs(right_dist - 70) / 200.0)
                if abs(val0) < 0.2:
                    val0 = -0.2
                val1 = (0.5 * abs(left_dist - 70) / 200.0)
                if abs(val1) < 0.2:
                    val1 = 0.2
                self.left_motor.set(val0)
                self.right_motor.set(val1)
        elif self.auto_mode == "positioned":
            self.positioned_count += 1
            self.winch_set(0.5);
            if self.positioned_count > 40:
                self.claw_up()

    #Teleop Mode
    def teleopInit(self):
        self.compressor.start()

    def teleopPeriodic(self):
        #If left trigger pulled, run brake algorithm, otherwise use joystick values to drive
        if self.left_joystick.getRawButton(1):
            rotation_values = self.brake_rotation()
            linear_values = self.brake_linear()
            left_wheel = rotation_values[0] + linear_values[0]
            right_wheel = rotation_values[1] + linear_values[1]
        else:
            left_wheel, right_wheel = self.drive_values()

        #Feed joystick values into drive system
        self.robotdrive.tankDrive(left_wheel, right_wheel)

        #Reset winch encoder value to 0 if right button 7 is pressed
        if self.right_joystick.getRawButton(7):
            self.winch_encoder.reset()

        #Feed winch controller raw values from the joystick
        #Right joystick button 3 raises winch, button 2 lowers winch
        winch_signal = self.right_joystick.getRawButton(3) + -self.right_joystick.getRawButton(2)
        self.winch_set(winch_signal)

        #Feed arm controller raw values from the joystick
        #Left joystick button 3 goes forward, 2 goes backward
        arm_signal = self.left_joystick.getRawButton(3) + -self.left_joystick.getRawButton(2)
        self.arm_motor.set(0.3 * arm_signal)

        #Handle piston in and out
        #Right joystick trigger button toggles claw in or out
        if self.right_joystick.getRawButton(1):
            self.x_pressed_last = True
        elif self.x_pressed_last:
            self.x_pressed_last = False
            self.claw_state = not self.claw_state
        self.set_claw()

        #Reset winch encoder value to 0 if right button 7 is pressed
        if self.right_joystick.getRawButton(7):
            self.winch_encoder.reset()

        #If the right joystick slider is down, go to test mode
        #print('slider: ' +  self.right_joystick.getRawAxis(2))
        if self.right_joystick.getRawAxis(2) > .5:
            self.testMode()

    #Required method for robotpy
    def testPeriodic():
        pass

    #Test Mode
    def testMode(self):
        #Calculate x and y distance travelled using accelerometer
        a_x = self.accel.getX()
        self.a_x_sum += a_x
        self.a_x_count += 1
        a_y = self.accel.getY()
        self.a_y_sum += a_y
        self.a_y_count += 1

        #Prints acceleration values when right button 8 is pressed
        if self.right_joystick.getRawButton(8):
            print ('x sum: ', self.a_x_sum, ' x count: ', self.a_x_count)
            print ('y sum: ', self.a_y_sum, ' y count: ', self.a_y_count)

        #Resets accelerometer counts when right button 9 is pressed
        if self.right_joystick.getRawButton(9):
            self.a_x_sum = 0.0
            self.a_x_count = 0
            self.a_y_sum = 0.0
            self.a_y_count = 0

        #Prints left ultrasonic sensor values when left button 6 is pressed
        if self.left_joystick.getRawButton(6):
            print ("left_ultrasonic_sensor: ", self.left_ultrasonic_sensor.getValue())
            print ("left limit: ", self.left_claw_whisker())
            print ("right limit: ", self.right_claw_whisker())

        #Prints right ultrasonic sensor values when left button 7 is pressed
        if self.left_joystick.getRawButton(7):
           print ("right_ultrasonic_sensor: ", self.right_ultrasonic_sensor.getValue())

        #Prints left optical sensor values when left button 11 is pressed
        if self.left_joystick.getRawButton(11):
           print ("left_optical_sensor: ", self.right_optical_sensor.get())

        #Prints right optical sensor values when left button 10 is pressed
        if self.left_joystick.getRawButton(10):
           print ("right_optical_sensor: ", self.right_optical_sensor.get())

        #Print current winch encoder value if right button 6 is pressed
        if self.right_joystick.getRawButton(6):
            revs = -self.winch_encoder.get()
            print ('revs: ', revs)
        #Print current gyro value if left button 8 is pressed
        if self.left_joystick.getRawButton(8):
            angle = self.gyro.getAngle()
            print ('angle: ', angle)

        #Reset gyro to 0 if left button 9 is pressed
        if self.left_joystick.getRawButton(9):
            self.gyro.reset()

    #Disabled Mode
    def disabledPeriodic(self):
        self.compressor.stop()

    #Helper Functions

    #Set "left" and "right" variables to the left and right
    #joystick outputs provided they are more than "threshold"
    def drive_values(self):
        joystick_threshold = 0.2
        left = step(
            self.left_joystick.getRawAxis(1),
            joystick_threshold,
        )
        right = step(
            self.right_joystick.getRawAxis(1),
            -joystick_threshold,
        )

        #Fuzzy match where if the left and right joysticks are moved about the same,
        #    then it moves the tankDrive the average of the two values
        if abs(left - right) < .1 :
            left = right = (left + right) / 2.0

        return (left, right)

    def brake_rotation(self):
        gyro_rate = self.gyro.getRate()
        wheel_rotation = gyro_rate * .1
        return (wheel_rotation, -wheel_rotation)

    def brake_linear(self):
        accel_y = self.accel.getY()
        wheel_motion = accel_y * .1
        return (-wheel_motion, -wheel_motion)

    def set_claw(self):
        self.solenoid1.set(not self.claw_state)
        self.solenoid2.set(self.claw_state)

    def claw_up(self):
        self.claw_state = False

    def claw_down(self):
        self.claw_state = True

    def left_claw_whisker(self):
        return self.left_limit_switch.get()

    def right_claw_whisker(self):
        return self.right_limit_switch.get()

    def winch_set(self, winch_signal):
        """
        Set winch controller safely by taking max and min encoder values into account
        (unless you're pressing the override button - right joystick, button 6)

        winch_signal=0 -> maintain winch position
        winch_signal>0 -> winch up?
        winch_signal<0 -> winch down?
        """
        revs = -self.winch_encoder.get()
        if self.last_winch_signal != 0 and winch_signal == 0:
            self.winch_setpoint = revs
        self.last_winch_signal = winch_signal
        if winch_signal == 0:
            val = 0.1 - 0.01 * (revs - self.winch_setpoint)
            self.winch_motor.set(val)
        else:
            if not (self.right_joystick.getRawButton(6)):
                if winch_signal > 0.1 and revs >= 1170:
                    winch_signal = 0
                if winch_signal < -0.1 and revs <= 8:
                    winch_signal = 0
            val = 0.5 * winch_signal
            self.winch_motor.set(val)

if __name__ == "__main__":
    wpilib.run(Robot)
<|endoftext|>