<commit_before>#!/usr/bin/env python2
# -*- coding: utf-8 -*-

####
# Author: Edvard Holst
# Project Site: http://github.com/Zyg0te/pyand
# Version: 0.9.1.3
###

try:
    import sys
    import subprocess
    import re
    from os import popen3 as pipe
except ImportError as e:
    print "[!] Required module missing. %s" % e.args[0]
    sys.exit(-1)

class ADB(object):

    __adb_path = None
    __output = None
    __error = None
    __devices = None
    __target = None

    # reboot modes
    REBOOT_NORMAL = 0
    REBOOT_RECOVERY = 1
    REBOOT_BOOTLOADER = 2

    # default TCP/IP port
    DEFAULT_TCP_PORT = 5555
    # default TCP/IP host
    DEFAULT_TCP_HOST = "localhost"

    def __init__(self, adb_path="adb"):
        #By default we assume adb is in $PATH
        self.__adb_path = adb_path
        if not self.check_path():
            self.__error = "[!] adb path not valid"

    def __clean__(self):
        self.__output = None
        self.__error = None

    def __read_output__(self, fd):
        ret = ''
        while 1:
            line = fd.readline()
            if not line:
                break
            ret += line

        if len(ret) == 0:
            ret = None

        return ret

    def __build_command__(self, cmd):
        """
        Build command parameters
        """
        if self.__devices is not None and len(self.__devices) > 1 and self.__target is None:
            self.__error = "[!] Must set target device first"
            return None

        if type(cmd) is tuple:
            a = list(cmd)
        elif type(cmd) is list:
            a = cmd
        else:
            #All arguments must be single list items
            a = cmd.split(" ")

        a.insert(0, self.__adb_path)
        if self.__target is not None:
            # add target device arguments to the command
            a.insert(1, '-s')
            a.insert(2, self.__target)

        return a

    def run_cmd(self, cmd):
        """
        Run a command against the adb tool ($ adb <cmd>)
        """
        self.__clean__()

        if self.__adb_path is None:
            self.__error = "[!] ADB path not set"
            return False

        try:
            args = self.__build_command__(cmd)
            if args is None:
                return
            #Print out args for debug purposes
            #print 'args>', args
            cmdp = subprocess.Popen(args, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
            self.__output, self.__error = cmdp.communicate()
            retcode = cmdp.wait()
            return self.__output.rstrip('\n')
        except OSError, e:
            self.__error = str(e)

        return

    def get_version(self):
        """
        Returns ADB tool version
        adb version
        """
        ret = self.run_cmd("version")
        try:
            pattern = re.compile(r"version\s(.+)")
            version = pattern.findall(ret)[0]
        except:
            version = None
        return version

    def check_path(self):
        """
        Verify if adb path is valid
        """

        if self.get_version() is None:
            print "[-] adb executable not found"
            return False
        return True

    def set_adb_path(self,adb_path):
        """
        Set the ADB tool path
        """
        self.__adb_path = adb_path
        self.check_path()

    def get_adb_path(self):
        """
        Returns the ADB tool path
        """
        return self.__adb_path

    def start_server(self):
        """
        Starts the ADB server
        adb start-server
        """
        self.run_cmd('start-server')
        return self.__output

    def kill_server(self):
        """
        Kills the ADB server
        adb kill-server
        """
        self.run_cmd('kill-server')

    def restart_server(self):
        """
        Restarts the ADB server
        """
        self.kill_server()
        return self.start_server()

    def restore_file(self,file_name):
        """
        Restore device contents from the <file> backup archive
        adb restore <file>
        """
        self.run_cmd('restore %s' % file_name)
        return self.__output

    def wait_for_device(self):
        """
        Block operations until device is online
        adb wait-for-device
        """
        self.run_cmd('wait-for-device')
        return self.__output

    def get_help(self):
        """
        Returns ADB help
        adb help
        """
        self.run_cmd('help')
        return self.__output

    def get_devices(self):
        """
        Return a dictionary of connected devices along with an incremented Id.
        adb devices
        """
        error = 0
        #Clear existing list of devices
        self.__devices = None
        self.run_cmd("devices")
        device_dict =  {}
        if self.__error is not None:
            return None
        try:
            n = 0
            output_list = self.__output.split("\n")
            for line in output_list:
                pattern = re.compile(r"([^\s]+)\s+device$")
                device = pattern.findall(line)
                if device:
                    device_dict[n] = device[0]
                    n += 1
        except:
            self.__devices = None
            error = 1
        self.__devices = device_dict
        return self.__devices

    def set_target_by_name(self, device):
        """
        Specify the device name to target
        example: set_target_device('emulator-5554')
        """
        if device is None or not device in self.__devices.values():

            self.__error = 'Must get device list first'
            print "[!] Device not found in device list"
            return False
        self.__target = device
        return "[+] Target device set: %s" % self.get_target_device()

    def set_target_by_id(self, device):
        """
        Specify the device ID to target.
        The ID should be one from the device list.
        """
        if device is None or not device in self.__devices:
            self.__error = 'Must get device list first'
            print "[!] Device not found in device list"
            return False
        self.__target = self.__devices[device]
        return "[+] Target device set: %s" % self.get_target_device()

    def get_target_device(self):
        """
        Returns the selected device to work with
        """
        if self.__target == None:
            print "[*] No device target set"

        return self.__target

    def get_state(self):
        """
        Get ADB state. Returns either offline | offline | device
        adb get-state
        """
        return self.run_cmd('get-state')

    def get_model(self):
        """
        Get Model name from taget device
        """
        self.run_cmd("devices -l")
        device_model = ""
        if self.__error is not None:
            return self.__error
        try:
            for line in self.__output.split("\n"):
                if line.startswith(self.__target):
                    pattern = r"model:(.+)\sdevice"
                    pat = re.compile(pattern)
                    device_model = pat.findall(line)
                    device_model = re.sub("[\[\]\'\{\}<>]", '', str(device_model))
        except Exception as e:
            return "[-] Error: %s" %e.args[0]

        return device_model

    def get_serialno(self):
        """
        Get serialno from target device
        adb get-serialno
        """
        return self.run_cmd('get-serialno')

    def reboot_device(self,mode=0):
        """
        Reboot the target device
        Specify mode to reboot normally, recovery or bootloader
        adb reboot <normally (0)/recovery (1) /bootloader (2)>
        """
        if not mode in (self.REBOOT_NORMAL, self.REBOOT_RECOVERY,self.REBOOT_BOOTLOADER):
            self.__error = "mode must be REBOOT_NORMAL/REBOOT_RECOVERY/REBOOT_BOOTLOADER"
            return self.__output

        cmd_str = "reboot"
        if mode == self.REBOOT_RECOVERY:
            cmd_str += " recovery"
        elif mode == self.REBOOT_BOOTLOADER:
            cmd_str += " bootloader"

        return self.run_cmd(cmd_str)

    def set_adb_root(self, mode):
        """
        restarts the adbd daemon with root permissions
        adb root
        """
        return self.run_cmd('root')

    def set_system_rw(self):
        """
        Mounts /system as rw
        adb remount
        """
        self.run_cmd("remount")
        return self.__output

    def get_remote_file(self,remote,local):
        """
        Pulls a remote file
        adb pull remote local
        """
        self.run_cmd('pull \"%s\" \"%s\"' % (remote,local) )
        if "bytes in" in self.__error:
            self.__output = self.__error
            self.__error = None
        return self.__output

    def push_local_file(self,local,remote):
        """
        Push a local file
        adb push local remote
        """
        self.run_cmd('push \"%s\" \"%s\"' % (local,remote) )
        return self.__output

    def shell_command(self,cmd):
        """
        Executes a shell command
        adb shell <cmd>
        """
        self.run_cmd('shell %s' % cmd)
        return self.__output

    def listen_usb(self):
        """
        Restarts the adbd daemon listening on USB
        adb usb
        """
        self.run_cmd("usb")
        return self.__output

    def listen_tcp(self,port=DEFAULT_TCP_PORT):
        """
        Restarts the adbd daemon listening on the specified port
        adb tcpip <port>
        """
        self.run_cmd("tcpip %s" % port)
        return self.__output

    def get_bugreport(self):
        """
        Return all information from the device that should be included in a bug report
        adb bugreport
        """
        self.run_cmd("bugreport")
        return self.__output

    def get_jdwp(self):
        """
        List PIDs of processes hosting a JDWP transport
        adb jdwp
        """
        return self.run_cmd("jdwp")

    def get_logcat(self,lcfilter=""):
        """
        View device log
        adb logcat <filter>
        """
        self.run_cmd("logcat %s" % lcfilter)
        return self.__output

    def run_emulator(self,cmd=""):
        """
        Run emulator console command
        """
        self.run_cmd("emu %s" % cmd)
        return self.__output

    def connect_remote (self,host=DEFAULT_TCP_HOST,port=DEFAULT_TCP_PORT):
        """
        Connect to a device via TCP/IP
        adb connect host:port
        """
        self.run_cmd("connect %s:%s" % ( host , port ) )
        return self.__output

    def disconnect_remote (self , host=DEFAULT_TCP_HOST , port=DEFAULT_TCP_PORT):
        """
        Disconnect from a TCP/IP device
        adb disconnect host:port
        """
        self.run_cmd("disconnect %s:%s" % ( host , port ) )
        return self.__output

    def ppp_over_usb(self,tty=None,params=""):
        """
        Run PPP over USB
        adb ppp <tty> <params>
        """
        if tty is None:
            return self.__output

        cmd = "ppp %s" % tty
        if params != "":
            cmd += " %s" % params

        self.run_cmd(cmd)
        return self.__output

    def sync_directory(self,directory=""):
        """
        Copy host->device only if changed (-l means list but don't copy)
        adb sync <dir>
        """
        self.run_cmd("sync %s" % directory )
        return self.__output

    def forward_socket(self,local=None,remote=None):
        """
        Forward socket connections
        adb forward <local> <remote>
        """
        if local is None or remote is None:
            return self.__output
        self.run_cmd("forward %s %s" % (local,remote) )
        return self.__output

    def uninstall(self,package=None,keepdata=False):
        """
        Remove this app package from the device
        adb uninstall [-k] package
        """
        if package is None:
            return self.__output
        cmd = "uninstall %s" % (package if keepdata is True else "-k %s" % package )
        self.run_cmd(cmd)
        return self.__output

    def install(self,pkgapp=None,fwdlock=False,reinstall=False,sdcard=False):
        """
        Push this package file to the device and install it
        adb install [-l] [-r] [-s] <file>
        -l -> forward-lock the app
        -r -> reinstall the app, keeping its data
        -s -> install on sdcard instead of internal storage
        """

        if pkgapp is None:
            return self.__output

        cmd = "install"
        if fwdlock is True:
            cmd += " -l "
        if reinstall is True:
            cmd += " -r "
        if sdcard is True:
            cmd += " -s "

        self.run_cmd("%s %s" % (cmd , pkgapp) )
        return self.__output

    def find_binary(self,name=None):
        """
        Look for a binary file on the device
        """

        self.shell_command("which %s" % name)

        if self.__output is None: # not found
            self.__error = "'%s' was not found" % name
        elif self.__output.strip() == "which: not found": # which binary not available
            self.__output = None
            self.__error = "which binary not found"
        else:
            self.__output = self.__output.strip()

        return self.__output
<commit_msg>Added wakup adb command.<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-

####
# Author: Edvard Holst
# Project Site: http://github.com/Zyg0te/pyand
# Version: 0.9.1.3
###

try:
    import sys
    import subprocess
    import re
    from os import popen3 as pipe
except ImportError as e:
    print "[!] Required module missing. %s" % e.args[0]
    sys.exit(-1)

class ADB(object):

    __adb_path = None
    __output = None
    __error = None
    __devices = None
    __target = None

    # reboot modes
    REBOOT_NORMAL = 0
    REBOOT_RECOVERY = 1
    REBOOT_BOOTLOADER = 2

    # default TCP/IP port
    DEFAULT_TCP_PORT = 5555
    # default TCP/IP host
    DEFAULT_TCP_HOST = "localhost"

    def __init__(self, adb_path="adb"):
        #By default we assume adb is in $PATH
        self.__adb_path = adb_path
        if not self.check_path():
            self.__error = "[!] adb path not valid"

    def __clean__(self):
        self.__output = None
        self.__error = None

    def __read_output__(self, fd):
        ret = ''
        while 1:
            line = fd.readline()
            if not line:
                break
            ret += line

        if len(ret) == 0:
            ret = None

        return ret

    def __build_command__(self, cmd):
        """
        Build command parameters
        """
        if self.__devices is not None and len(self.__devices) > 1 and self.__target is None:
            self.__error = "[!] Must set target device first"
            return None

        if type(cmd) is tuple:
            a = list(cmd)
        elif type(cmd) is list:
            a = cmd
        else:
            #All arguments must be single list items
            a = cmd.split(" ")

        a.insert(0, self.__adb_path)
        if self.__target is not None:
            # add target device arguments to the command
            a.insert(1, '-s')
            a.insert(2, self.__target)

        return a

    def run_cmd(self, cmd):
        """
        Run a command against the adb tool ($ adb <cmd>)
        """
        self.__clean__()

        if self.__adb_path is None:
            self.__error = "[!] ADB path not set"
            return False

        try:
            args = self.__build_command__(cmd)
            if args is None:
                return
            #Print out args for debug purposes
            #print 'args>', args
            cmdp = subprocess.Popen(args, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
            self.__output, self.__error = cmdp.communicate()
            retcode = cmdp.wait()
            return self.__output.rstrip('\n')
        except OSError, e:
            self.__error = str(e)

        return

    def get_version(self):
        """
        Returns ADB tool version
        adb version
        """
        ret = self.run_cmd("version")
        try:
            pattern = re.compile(r"version\s(.+)")
            version = pattern.findall(ret)[0]
        except:
            version = None
        return version

    def check_path(self):
        """
        Verify if adb path is valid
        """

        if self.get_version() is None:
            print "[-] adb executable not found"
            return False
        return True

    def set_adb_path(self,adb_path):
        """
        Set the ADB tool path
        """
        self.__adb_path = adb_path
        self.check_path()

    def get_adb_path(self):
        """
        Returns the ADB tool path
        """
        return self.__adb_path

    def start_server(self):
        """
        Starts the ADB server
        adb start-server
        """
        self.run_cmd('start-server')
        return self.__output

    def kill_server(self):
        """
        Kills the ADB server
        adb kill-server
        """
        self.run_cmd('kill-server')

    def restart_server(self):
        """
        Restarts the ADB server
        """
        self.kill_server()
        return self.start_server()

    def restore_file(self,file_name):
        """
        Restore device contents from the <file> backup archive
        adb restore <file>
        """
        self.run_cmd('restore %s' % file_name)
        return self.__output

    def wait_for_device(self):
        """
        Block operations until device is online
        adb wait-for-device
        """
        self.run_cmd('wait-for-device')
        return self.__output

    def get_help(self):
        """
        Returns ADB help
        adb help
        """
        self.run_cmd('help')
        return self.__output

    def get_devices(self):
        """
        Return a dictionary of connected devices along with an incremented Id.
        adb devices
        """
        error = 0
        #Clear existing list of devices
        self.__devices = None
        self.run_cmd("devices")
        device_dict =  {}
        if self.__error is not None:
            return None
        try:
            n = 0
            output_list = self.__output.split("\n")
            for line in output_list:
                pattern = re.compile(r"([^\s]+)\s+device$")
                device = pattern.findall(line)
                if device:
                    device_dict[n] = device[0]
                    n += 1
        except:
            self.__devices = None
            error = 1
        self.__devices = device_dict
        return self.__devices

    def set_target_by_name(self, device):
        """
        Specify the device name to target
        example: set_target_device('emulator-5554')
        """
        if device is None or not device in self.__devices.values():

            self.__error = 'Must get device list first'
            print "[!] Device not found in device list"
            return False
        self.__target = device
        return "[+] Target device set: %s" % self.get_target_device()

    def set_target_by_id(self, device):
        """
        Specify the device ID to target.
        The ID should be one from the device list.
        """
        if device is None or not device in self.__devices:
            self.__error = 'Must get device list first'
            print "[!] Device not found in device list"
            return False
        self.__target = self.__devices[device]
        return "[+] Target device set: %s" % self.get_target_device()

    def get_target_device(self):
        """
        Returns the selected device to work with
        """
        if self.__target == None:
            print "[*] No device target set"

        return self.__target

    def get_state(self):
        """
        Get ADB state. Returns either offline | offline | device
        adb get-state
        """
        return self.run_cmd('get-state')

    def get_model(self):
        """
        Get Model name from taget device
        """
        self.run_cmd("devices -l")
        device_model = ""
        if self.__error is not None:
            return self.__error
        try:
            for line in self.__output.split("\n"):
                if line.startswith(self.__target):
                    pattern = r"model:(.+)\sdevice"
                    pat = re.compile(pattern)
                    device_model = pat.findall(line)
                    device_model = re.sub("[\[\]\'\{\}<>]", '', str(device_model))
        except Exception as e:
            return "[-] Error: %s" %e.args[0]

        return device_model

    def get_serialno(self):
        """
        Get serialno from target device
        adb get-serialno
        """
        return self.run_cmd('get-serialno')

    def reboot_device(self,mode=0):
        """
        Reboot the target device
        Specify mode to reboot normally, recovery or bootloader
        adb reboot <normally (0)/recovery (1) /bootloader (2)>
        """
        if not mode in (self.REBOOT_NORMAL, self.REBOOT_RECOVERY,self.REBOOT_BOOTLOADER):
            self.__error = "mode must be REBOOT_NORMAL/REBOOT_RECOVERY/REBOOT_BOOTLOADER"
            return self.__output

        cmd_str = "reboot"
        if mode == self.REBOOT_RECOVERY:
            cmd_str += " recovery"
        elif mode == self.REBOOT_BOOTLOADER:
            cmd_str += " bootloader"

        return self.run_cmd(cmd_str)

    def set_adb_root(self, mode):
        """
        restarts the adbd daemon with root permissions
        adb root
        """
        return self.run_cmd('root')

    def set_system_rw(self):
        """
        Mounts /system as rw
        adb remount
        """
        self.run_cmd("remount")
        return self.__output

    def get_remote_file(self,remote,local):
        """
        Pulls a remote file
        adb pull remote local
        """
        self.run_cmd('pull \"%s\" \"%s\"' % (remote,local) )
        if "bytes in" in self.__error:
            self.__output = self.__error
            self.__error = None
        return self.__output

    def push_local_file(self,local,remote):
        """
        Push a local file
        adb push local remote
        """
        self.run_cmd('push \"%s\" \"%s\"' % (local,remote) )
        return self.__output

    def shell_command(self,cmd):
        """
        Executes a shell command
        adb shell <cmd>
        """
        self.run_cmd('shell %s' % cmd)
        return self.__output

    def listen_usb(self):
        """
        Restarts the adbd daemon listening on USB
        adb usb
        """
        self.run_cmd("usb")
        return self.__output

    def listen_tcp(self,port=DEFAULT_TCP_PORT):
        """
        Restarts the adbd daemon listening on the specified port
        adb tcpip <port>
        """
        self.run_cmd("tcpip %s" % port)
        return self.__output

    def get_bugreport(self):
        """
        Return all information from the device that should be included in a bug report
        adb bugreport
        """
        self.run_cmd("bugreport")
        return self.__output

    def get_jdwp(self):
        """
        List PIDs of processes hosting a JDWP transport
        adb jdwp
        """
        return self.run_cmd("jdwp")

    def get_logcat(self,lcfilter=""):
        """
        View device log
        adb logcat <filter>
        """
        self.run_cmd("logcat %s" % lcfilter)
        return self.__output

    def run_emulator(self,cmd=""):
        """
        Run emulator console command
        """
        self.run_cmd("emu %s" % cmd)
        return self.__output

    def connect_remote (self,host=DEFAULT_TCP_HOST,port=DEFAULT_TCP_PORT):
        """
        Connect to a device via TCP/IP
        adb connect host:port
        """
        self.run_cmd("connect %s:%s" % ( host , port ) )
        return self.__output

    def disconnect_remote (self , host=DEFAULT_TCP_HOST , port=DEFAULT_TCP_PORT):
        """
        Disconnect from a TCP/IP device
        adb disconnect host:port
        """
        self.run_cmd("disconnect %s:%s" % ( host , port ) )
        return self.__output

    def ppp_over_usb(self,tty=None,params=""):
        """
        Run PPP over USB
        adb ppp <tty> <params>
        """
        if tty is None:
            return self.__output

        cmd = "ppp %s" % tty
        if params != "":
            cmd += " %s" % params

        self.run_cmd(cmd)
        return self.__output

    def sync_directory(self,directory=""):
        """
        Copy host->device only if changed (-l means list but don't copy)
        adb sync <dir>
        """
        self.run_cmd("sync %s" % directory )
        return self.__output

    def forward_socket(self,local=None,remote=None):
        """
        Forward socket connections
        adb forward <local> <remote>
        """
        if local is None or remote is None:
            return self.__output
        self.run_cmd("forward %s %s" % (local,remote) )
        return self.__output

    def uninstall(self,package=None,keepdata=False):
        """
        Remove this app package from the device
        adb uninstall [-k] package
        """
        if package is None:
            return self.__output
        cmd = "uninstall %s" % (package if keepdata is True else "-k %s" % package )
        self.run_cmd(cmd)
        return self.__output

    def install(self,pkgapp=None,fwdlock=False,reinstall=False,sdcard=False):
        """
        Push this package file to the device and install it
        adb install [-l] [-r] [-s] <file>
        -l -> forward-lock the app
        -r -> reinstall the app, keeping its data
        -s -> install on sdcard instead of internal storage
        """

        if pkgapp is None:
            return self.__output

        cmd = "install"
        if fwdlock is True:
            cmd += " -l "
        if reinstall is True:
            cmd += " -r "
        if sdcard is True:
            cmd += " -s "

        self.run_cmd("%s %s" % (cmd , pkgapp) )
        return self.__output

    def find_binary(self,name=None):
        """
        Look for a binary file on the device
        """

        self.shell_command("which %s" % name)

        if self.__output is None: # not found
            self.__error = "'%s' was not found" % name
        elif self.__output.strip() == "which: not found": # which binary not available
            self.__output = None
            self.__error = "which binary not found"
        else:
            self.__output = self.__output.strip()

        return self.__output

    def wake_device(self):
        return self.run_cmd('shell input keyevent 26')

<|endoftext|>


<commit_before># -*- coding: iso-8859-15 -*-
# =================================================================
#
# $Id$
#
# Authors: Tom Kralidis <tomkralidis@hotmail.com>
#
# Copyright (c) 2011 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

import os
import logging

MSG_FORMAT = '%(asctime)s] [%(levelname)s] file=%(pathname)s \
line=%(lineno)s module=%(module)s function=%(funcName)s %(message)s'

TIME_FORMAT = '%a, %d %b %Y %H:%M:%S'

LOGLEVELS = {
    'CRITICAL': logging.CRITICAL,
    'ERROR': logging.ERROR,
    'WARNING': logging.WARNING,
    'INFO': logging.INFO,
    'DEBUG': logging.DEBUG,
    'NOTSET': logging.NOTSET,
}

class Log(logging.Logger):
    ''' Logging facility   '''
    def __init__(self, config=None):
        ''' Initialize logging facility '''
        if config is None:
            return None

        logfile = None
        loglevel = 'NOTSET'

        if config.has_option('server', 'loglevel'):
            loglevel = config.get('server', 'loglevel')

            if loglevel not in LOGLEVELS.keys():
                raise RuntimeError, \
                ('Invalid server configuration (server.loglevel).')

            if not config.has_option('server', 'logfile'):
                raise RuntimeError\
                ('Invalid server configuration (server.loglevel set,\
                  but server.logfile missing).')

        if config.has_option('server', 'logfile'):
            if not config.has_option('server', 'loglevel'):
                raise RuntimeError, \
                ('Invalid server configuration (server.logfile set,\
                  but server.loglevel missing).')

            logfile = config.get('server', 'logfile')

        if loglevel != 'NOTSET' and logfile is None:
            raise RuntimeError, \
            ('Invalid server configuration \
            (server.loglevel set, but server.logfile is not).')

        logging.Logger.__init__(self, 'pycsw', LOGLEVELS[loglevel])

        if logfile:
            try:
                filehandler = logging.FileHandler(logfile)
                filehandler.setLevel(LOGLEVELS[loglevel])
                filehandler.setFormatter(logging.Formatter(MSG_FORMAT,
                    TIME_FORMAT))
                self.addHandler(filehandler)
            except Exception, err:
                raise RuntimeError, \
                ('Invalid server configuration: server.logfile access denied.\
                Make sure filepath exists and is writable. %s', str(err))
        else:
            logging.basicConfig(stream=open(os.devnull, 'w'))

        self.info('Logging initialized (level: %s).' % loglevel)

        if loglevel == 'DEBUG':  #turn on CGI debugging
            self.info('CGI debugging enabled.')
            import cgitb
            cgitb.enable()
<commit_msg>fix nullhandler safeguaring<commit_after># -*- coding: iso-8859-15 -*-
# =================================================================
#
# $Id$
#
# Authors: Tom Kralidis <tomkralidis@hotmail.com>
#
# Copyright (c) 2011 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

import logging

MSG_FORMAT = '%(asctime)s] [%(levelname)s] file=%(pathname)s \
line=%(lineno)s module=%(module)s function=%(funcName)s %(message)s'

TIME_FORMAT = '%a, %d %b %Y %H:%M:%S'

LOGLEVELS = {
    'CRITICAL': logging.CRITICAL,
    'ERROR': logging.ERROR,
    'WARNING': logging.WARNING,
    'INFO': logging.INFO,
    'DEBUG': logging.DEBUG,
    'NOTSET': logging.NOTSET,
}

class NullHandlerLocal(logging.Handler):
    ''' safeguarded nullhandler for Python 2.6/2.7  '''
    def emit(self, record):
        pass

class Log(logging.Logger):
    ''' Logging facility   '''
    def __init__(self, config=None):
        ''' Initialize logging facility '''
        if config is None:
            return None

        logfile = None
        loglevel = 'NOTSET'

        if config.has_option('server', 'loglevel'):
            loglevel = config.get('server', 'loglevel')

            if loglevel not in LOGLEVELS.keys():
                raise RuntimeError, \
                ('Invalid server configuration (server.loglevel).')

            if not config.has_option('server', 'logfile'):
                raise RuntimeError\
                ('Invalid server configuration (server.loglevel set,\
                  but server.logfile missing).')

        if config.has_option('server', 'logfile'):
            if not config.has_option('server', 'loglevel'):
                raise RuntimeError, \
                ('Invalid server configuration (server.logfile set,\
                  but server.loglevel missing).')

            logfile = config.get('server', 'logfile')

        if loglevel != 'NOTSET' and logfile is None:
            raise RuntimeError, \
            ('Invalid server configuration \
            (server.loglevel set, but server.logfile is not).')

        logging.Logger.__init__(self, 'pycsw', LOGLEVELS[loglevel])

        if logfile:
            try:
                filehandler = logging.FileHandler(logfile)
                filehandler.setLevel(LOGLEVELS[loglevel])
                filehandler.setFormatter(logging.Formatter(MSG_FORMAT,
                    TIME_FORMAT))
                self.addHandler(filehandler)
            except Exception, err:
                raise RuntimeError, \
                ('Invalid server configuration: server.logfile access denied.\
                Make sure filepath exists and is writable. %s', str(err))
        else:
            self.addHandler(NullHandlerLocal())

        self.info('Logging initialized (level: %s).' % loglevel)

        if loglevel == 'DEBUG':  #turn on CGI debugging
            self.info('CGI debugging enabled.')
            import cgitb
            cgitb.enable()
<|endoftext|>


<commit_before>#!/usr/local/sickbeard/env/bin/python
import os
import subprocess
import cgi
import configobj
from pyextdirect.configuration import create_configuration, expose, LOAD, SUBMIT
from pyextdirect.router import Router
from pyextdirect.api import create_api

sickbeard_sss = '/var/packages/sickbeard/scripts/start-stop-status'
sickbeard_path = '/usr/local/sickbeard/var/config.ini'
autoprocesstv_path = '/usr/local/sickbeard/share/SickBeard/autoProcessTV/autoProcessTV.cfg'
sabnzbd_path = '/usr/local/sabnzbd/var/config.ini'
nzbget_path = ''  #TODO

Base = create_configuration()

class SickBeardCfg(Base):
    def __init__(self):
        '''Load available configuration files'''
        self.configs = {'sickbeard': configobj.ConfigObj(sickbeard_path),
                        'autoprocesstv': configobj.ConfigObj(autoprocesstv_path)}
        if os.path.exists(sabnzbd_path):
            self.configs['sabnzbd'] = configobj.ConfigObj(sabnzbd_path)
        if os.path.exists(nzbget_path):
            self.configs['nzbget'] = configobj.ConfigObj(nzbget_path)

    @expose(kind=LOAD)
    def load(self):
        '''Read configuration files to determine current state'''
        autoprocesstv = False
        configure_for = self.configs['sickbeard']['General']['nzb_method']
        if 'sabnzbd' in self.configs and self.configs['sabnzbd']['misc']['script_dir'] == os.path.dirname(autoprocesstv_path):
            autoprocesstv = True
        if configure_for not in ['sabnzbd', 'nzbget']:
            configure_for = 'nothing'
        return {'configure_for': configure_for, 'autoprocesstv': autoprocesstv}

    @expose(kind=SUBMIT)
    def save(self, configure_for=None, autoprocesstv=None):
        '''Change configuration files as requested'''
        if configure_for == 'sabnzbd' and 'sabnzbd' in self.configs:
            devnull = open(os.devnull, 'w')
            subprocess.call([sickbeard_sss, 'stop'], stdout=devnull, stderr=devnull)
            self.configs['sickbeard']['General']['nzb_method'] = 'sabnzbd'
            self.configs['sickbeard']['SABnzbd']['sab_username'] = self.configs['sabnzbd']['misc']['username']
            self.configs['sickbeard']['SABnzbd']['sab_password'] = self.configs['sabnzbd']['misc']['password']
            self.configs['sickbeard']['SABnzbd']['sab_apikey'] = self.configs['sabnzbd']['misc']['api_key']
            self.configs['sickbeard']['SABnzbd']['sab_host'] = 'http://localhost:' + self.configs['sabnzbd']['misc']['port'] + '/'
            self.configs['sickbeard'].write()
            self.configs['autoprocesstv']['SickBeard']['username'] = self.configs['sickbeard']['General']['web_username']
            self.configs['autoprocesstv']['SickBeard']['password'] = self.configs['sickbeard']['General']['web_password']
            self.configs['autoprocesstv'].write()
            subprocess.call([sickbeard_sss, 'start'], stdout=devnull, stderr=devnull)
        elif configure_for == 'nzbget' and 'nzbget' in self.configs:
            #TODO
            pass
        if autoprocesstv and 'sabnzbd' in self.configs:
            self.configs['sabnzbd']['misc']['script_dir'] = os.path.dirname(autoprocesstv_path)
            self.configs['sabnzbd'].write()

    @expose
    def available_configs(self):
        '''List available configurations'''
        #TODO: Use this to disable radio butons/checkboxes
        return self.configs.keys()

if __name__ == '__main__':
    print 'Content-type: application/json'
    print
    router = Router(Base)
    fs = cgi.FieldStorage()
    request = fs.value
    if isinstance(request, list):
        request = dict((mfs.name, mfs.value) for mfs in request)
    print router.route(request)

<commit_msg>Fix permission issue during postprocessing in sickbeard SPK<commit_after>#!/usr/local/sickbeard/env/bin/python
import os
import subprocess
import pwd
import cgi
import configobj
from pyextdirect.configuration import create_configuration, expose, LOAD, SUBMIT
from pyextdirect.router import Router
from pyextdirect.api import create_api

sickbeard_sss = '/var/packages/sickbeard/scripts/start-stop-status'
sickbeard_path = '/usr/local/sickbeard/var/config.ini'
autoprocesstv_path = '/usr/local/sickbeard/share/SickBeard/autoProcessTV/autoProcessTV.cfg'
sabnzbd_path = '/usr/local/sabnzbd/var/config.ini'
nzbget_path = ''  #TODO

Base = create_configuration()

class SickBeardCfg(Base):
    def __init__(self):
        '''Load available configuration files'''
        self.configs = {'sickbeard': configobj.ConfigObj(sickbeard_path),
                        'autoprocesstv': configobj.ConfigObj(autoprocesstv_path)}
        if os.path.exists(sabnzbd_path):
            self.configs['sabnzbd'] = configobj.ConfigObj(sabnzbd_path)
        if os.path.exists(nzbget_path):
            self.configs['nzbget'] = configobj.ConfigObj(nzbget_path)

    @expose(kind=LOAD)
    def load(self):
        '''Read configuration files to determine current state'''
        autoprocesstv = False
        configure_for = self.configs['sickbeard']['General']['nzb_method']
        if 'sabnzbd' in self.configs and self.configs['sabnzbd']['misc']['script_dir'] == os.path.dirname(autoprocesstv_path):
            autoprocesstv = True
        if configure_for not in ['sabnzbd', 'nzbget']:
            configure_for = 'nothing'
        return {'configure_for': configure_for, 'autoprocesstv': autoprocesstv}

    @expose(kind=SUBMIT)
    def save(self, configure_for=None, autoprocesstv=None):
        '''Change configuration files as requested'''
        if configure_for == 'sabnzbd' and 'sabnzbd' in self.configs:
            devnull = open(os.devnull, 'w')
            subprocess.call([sickbeard_sss, 'stop'], stdout=devnull, stderr=devnull)
            self.configs['sickbeard']['General']['nzb_method'] = 'sabnzbd'
            self.configs['sickbeard']['SABnzbd']['sab_username'] = self.configs['sabnzbd']['misc']['username']
            self.configs['sickbeard']['SABnzbd']['sab_password'] = self.configs['sabnzbd']['misc']['password']
            self.configs['sickbeard']['SABnzbd']['sab_apikey'] = self.configs['sabnzbd']['misc']['api_key']
            self.configs['sickbeard']['SABnzbd']['sab_host'] = 'http://localhost:' + self.configs['sabnzbd']['misc']['port'] + '/'
            self.configs['sickbeard'].write()
            self.configs['autoprocesstv']['SickBeard']['username'] = self.configs['sickbeard']['General']['web_username']
            self.configs['autoprocesstv']['SickBeard']['password'] = self.configs['sickbeard']['General']['web_password']
            self.configs['autoprocesstv'].write()
            # Change ownership of autoProcessTV so SABnzbd doesn't fail opening it during postprocessing
            os.chown(autoprocesstv_path, pwd.getpwnam('sabnzbd')[2], -1)
            subprocess.call([sickbeard_sss, 'start'], stdout=devnull, stderr=devnull)
        elif configure_for == 'nzbget' and 'nzbget' in self.configs:
            #TODO
            pass
        if autoprocesstv and 'sabnzbd' in self.configs:
            self.configs['sabnzbd']['misc']['script_dir'] = os.path.dirname(autoprocesstv_path)
            self.configs['sabnzbd'].write()

    @expose
    def available_configs(self):
        '''List available configurations'''
        #TODO: Use this to disable radio butons/checkboxes
        return self.configs.keys()

if __name__ == '__main__':
    print 'Content-type: application/json'
    print
    router = Router(Base)
    fs = cgi.FieldStorage()
    request = fs.value
    if isinstance(request, list):
        request = dict((mfs.name, mfs.value) for mfs in request)
    print router.route(request)

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

#http://www.beyondlogic.org/usbnutshell/usb1.shtml

from IPython import embed
import com
import cmd
import hpstr
import protocol
from time import sleep
import sys

class HP49( object ):

  def __init__( self, autoconnect=False ):
    if autoconnect == True:
      self.waitforhp()
      self.connect()
      com.reset()
      try:
        cmd.version()
        return
      except:
        print "Please enable XMODEM server mode (right-shift-release arrow-right)...",
        sys.stdout.flush()
        sys.exit(5) # TODO: fix detection and remove hackish workaround
      while True:
        try:
          protocol.cancel()
          com.flush()
          com.dev.reset()
          cmd.version()
          break
        except:
          sleep( 1 )
      print "OK"
      sys.stdout.flush()

  def connected( self ):
    return com.dev != None

  def find( self ):
    hps = com.find()
    if len( hps ) == 0:
      print "No HP49-compatible USB devices connected."
      sys.stdout.flush()
    else:
      print "Number of HP49-compatible USB devices: %d" % len( hps )
      sys.stdout.flush()
    return len( hps )

  def waitforhp( self ):
    if len( com.find() ) == 0:
      print "Please connect HP49-compatible USB device...",
      sys.stdout.flush()
      while len( com.find() ) == 0:
        sleep( 1 )
      print "OK"

  def connect( self, id=0 ):
    com.connect( id )
    sleep( 1 )

  def cmd( self, cmd, args=None ):
    protocol.cmd( cmd, args=args )

  def waitack( self ):
    print hex( com.read( 1 )[0] )

  def read( self, len=0 ):
    if len == 0:
      data = protocol.readpacket()
    else:
      data = com.read( len )
    print hpstr.tohexstr( data )
    print hpstr.torepr( data )
    return data

  def sendack( self ):
    protocol.sendack()

  def sendnack( self ):
    protocol.sendnack()

  def write( self, data ):
    return com.write( data )

  def xeq( self, args ):
    protocol.cmd( "E", args=args )
    return protocol.waitack()

  def download( self ):
    return protocol.download()

  def ls( self ):
    ls = cmd.ls()
    for f in ls:
      print '{0:016b}  {1:7}  {2:7} '.format(
        f[3], cmd.objtype( f[1] ), f[2] ), hpstr.hptoutf( f[0] )

  def tohp( self, s ):
    if isinstance( s, unicode ):
      return hpstr.utftohp( s )
    elif isinstance( s, str ):
      return hpstr.utftohp( s.decode( "utf-8" ) )
    else:
      print "Unsupported encoding:", s
      return s

  def cd( self, remoteobj ):
    remoteobj = self.tohp( remoteobj )
    protocol.cmd( "E", remoteobj )
    return protocol.waitack()

  def rm( self, remoteobj ):
    remoteobj = self.tohp( remoteobj )
    protocol.cmd( "E", "'"+remoteobj+"' PURGE" )
    return protocol.waitack()

  def get( self, remoteobj, hexdump=False ):
    remoteobj = self.tohp( remoteobj )
    data = cmd.get( remoteobj )
    if hexdump == False:
      return data
    else:
      print hpstr.tohexstr( data )
      return

  def put( self, remoteobj, data ):
    remoteobj = self.tohp( remoteobj )
    return cmd.put( remoteobj, hpstr.toarr( data ) )

  def version( self ):
    print cmd.version()

  def meminfo( self ):
    print cmd.meminfo(), "bytes free."

  def info( self ):
    self.version()
    self.meminfo()

<commit_msg>return raw info instead of printing it<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

#http://www.beyondlogic.org/usbnutshell/usb1.shtml

from IPython import embed
import com
import cmd
import hpstr
import protocol
from time import sleep
import sys

class HP49( object ):

  def __init__( self, autoconnect=False ):
    if autoconnect == True:
      self.waitforhp()
      self.connect()
      com.reset()
      try:
        cmd.version()
        return
      except:
        print "Please enable XMODEM server mode (right-shift-release arrow-right)...",
        sys.stdout.flush()
        sys.exit(5) # TODO: fix detection and remove hackish workaround
      while True:
        try:
          protocol.cancel()
          com.flush()
          com.dev.reset()
          cmd.version()
          break
        except:
          sleep( 1 )
      print "OK"
      sys.stdout.flush()

  def connected( self ):
    return com.dev != None

  def find( self ):
    hps = com.find()
    if len( hps ) == 0:
      print "No HP49-compatible USB devices connected."
      sys.stdout.flush()
    else:
      print "Number of HP49-compatible USB devices: %d" % len( hps )
      sys.stdout.flush()
    return len( hps )

  def waitforhp( self ):
    if len( com.find() ) == 0:
      print "Please connect HP49-compatible USB device...",
      sys.stdout.flush()
      while len( com.find() ) == 0:
        sleep( 1 )
      print "OK"

  def connect( self, id=0 ):
    com.connect( id )
    sleep( 1 )

  def cmd( self, cmd, args=None ):
    protocol.cmd( cmd, args=args )

  def waitack( self ):
    print hex( com.read( 1 )[0] )

  def read( self, len=0 ):
    if len == 0:
      data = protocol.readpacket()
    else:
      data = com.read( len )
    print hpstr.tohexstr( data )
    print hpstr.torepr( data )
    return data

  def sendack( self ):
    protocol.sendack()

  def sendnack( self ):
    protocol.sendnack()

  def write( self, data ):
    return com.write( data )

  def xeq( self, args ):
    protocol.cmd( "E", args=args )
    return protocol.waitack()

  def download( self ):
    return protocol.download()

  def ls( self ):
    ls = cmd.ls()
    for f in ls:
      print '{0:016b}  {1:7}  {2:7} '.format(
        f[3], cmd.objtype( f[1] ), f[2] ), hpstr.hptoutf( f[0] )

  def tohp( self, s ):
    if isinstance( s, unicode ):
      return hpstr.utftohp( s )
    elif isinstance( s, str ):
      return hpstr.utftohp( s.decode( "utf-8" ) )
    else:
      print "Unsupported encoding:", s
      return s

  def cd( self, remoteobj ):
    remoteobj = self.tohp( remoteobj )
    protocol.cmd( "E", remoteobj )
    return protocol.waitack()

  def rm( self, remoteobj ):
    remoteobj = self.tohp( remoteobj )
    protocol.cmd( "E", "'"+remoteobj+"' PURGE" )
    return protocol.waitack()

  def get( self, remoteobj, hexdump=False ):
    remoteobj = self.tohp( remoteobj )
    data = cmd.get( remoteobj )
    if hexdump == False:
      return data
    else:
      print hpstr.tohexstr( data )
      return

  def put( self, remoteobj, data ):
    remoteobj = self.tohp( remoteobj )
    return cmd.put( remoteobj, hpstr.toarr( data ) )

  def version( self ):
    return cmd.version()

  def meminfo( self ):
    return cmd.meminfo()

  def info( self ):
    self.version()
    self.meminfo()

<|endoftext|>


<commit_before>import inspect
import sys
import socket
import string
import re

import threads

class Bot(object):
  def __init__(self, host, **kwargs):
    '''
    Initializes a new pyrc.Bot.
    '''
    nick = "PyrcBot" if self.__class__ == Bot else self.__class__.__name__
    self.config = dict(kwargs)
    self.config.setdefault('host', host)
    self.config.setdefault('port', 6667)
    self.config.setdefault('nick', nick)
    self.config.setdefault('names', [self.config['nick']])
    self.config.setdefault('ident', nick.lower())
    self.config.setdefault('realname', "A Pyrc Bot")
    self.config.setdefault('channels', [])
    self.config.setdefault('password', None)

    self._inbuffer = ""
    self._commands = []
    self._threads = []
    self.socket = None
    self.initialized = False

    self.addhooks()

  def message(self, recipient, s):
    "High level interface to sending an IRC message."
    self.cmd("PRIVMSG %s :%s" % (recipient, s))

  def connect(self):
    '''
    Connects to the IRC server with the options defined in `config`
    '''
    self._connect()

    try:
      self.listen()
    except (KeyboardInterrupt, SystemExit):
      pass
    finally:
      self.close()

  def listen(self):
    """
    Constantly listens to the input from the server. Since the messages come
    in pieces, we wait until we receive 1 or more full lines to start parsing.

    A new line is defined as ending in \r\n in the RFC, but some servers
    separate by \n. This script takes care of both.
    """
    while True:
      self._inbuffer = self._inbuffer + self.socket.recv(1024)
      # Some IRC servers disregard the RFC and split lines by \n rather than \r\n.

      temp = self._inbuffer.split("\n")
      self._inbuffer = temp.pop()

      for line in temp:
        # Strip \r from \r\n for RFC-compliant IRC servers.
        line = line.rstrip('\r')
        self.parseline(line)
        print line

  def parseline(self, line):
    if line.startswith("PING"):
      length = len("PING ")
      host = line[length:]
      self.cmd("PONG %s" % host)
    elif re.match(r"^:\S+ PRIVMSG", line):
      msg_regex = re.compile(r"^:(\S+)!\S+ PRIVMSG (\S+) :(.*)")
      nick, channel, message = re.match(msg_regex, line).groups()
      self.receivemessage(channel, nick, message)
    elif re.match(r"^:\S+ INVITE %s" % self.config['nick'], line):
      msg_regex = re.compile(r"^:(\S+)!\S+ INVITE \S+ :?(.*)")
      inviter, channel = re.match(msg_regex, line).groups()
      self.cmd("JOIN %s" % channel)
    elif re.match(r"^\S+ MODE %s :\+\w*i" % self.config['nick'], line)\
        and self.config['password']:
      # Autoidentify if a password is provided
      self.cmd("PRIVMSG NickServ :identify %s" % self.config['password'])
    elif re.match(r"^:\S+ MODE", line) and not self.initialized:
      self.initialized = True
      if self.config['channels']:
        self.cmd("JOIN %s" % ' '.join(self.config['channels']))
      # TODO: This doesn't ensure that threads run at the right time, e.g.
      # after the bot has joined every channel it needs to.
      for thread in self._threads:
        thread.start()

  def addhooks(self):
    for func in self.__class__.__dict__.values():
      if callable(func) and hasattr(func, '_type'):
        if func._type == 'COMMAND':
          self._commands.append(func)
        elif func._type == "REPEAT":
          thread = threads.JobThread(func, self)
          self._threads.append(thread)
        else:
          raise "This is not a type I've ever heard of."

  def receivemessage(self, channel, nick, message):
    self.parsecommand(channel, message)

  def parsecommand(self, channel, message):
    # sort names so names that are substrings work
    names = sorted(self.config['names'], key=len, reverse=True)

    name_used = None
    for name in names:
      name_regex_str = r'^(%s)[,:]?\s+' % re.escape(name)
      name_regex = re.compile(name_regex_str, re.IGNORECASE)
      if name_regex.match(message):
        name_used = name_regex.match(message).group(1)
        break

    if not name_used:
      return

    _,_,message = message.partition(name_used)
    command = re.match(r'^[,:]?\s+(.*)', message).group(1)
    for command_func in self._commands:
      if command_func._matcher.search(command):
        command_func(self, channel)

  def cmd(self, raw_line):
    self.socket.send(raw_line + "\r\n")

  def _connect(self):
    "Connects a socket to the server using options defined in `config`."
    self.socket = socket.socket()
    self.socket.connect((self.config['host'], self.config['port']))
    self.cmd("NICK %s" % self.config['nick'])
    self.cmd("USER %s %s bla :%s" %
        (self.config['ident'], self.config['host'], self.config['realname']))

  def close(self):
    for thread in self._threads:
      thread.shutdown()
    self.socket.shutdown(socket.SHUT_RDWR)
    self.socket.close()
<commit_msg>@hooks.command() regex matches are passed along to the decorated function.<commit_after>import inspect
import sys
import socket
import string
import re

import threads

class Bot(object):
  def __init__(self, host, **kwargs):
    '''
    Initializes a new pyrc.Bot.
    '''
    nick = "PyrcBot" if self.__class__ == Bot else self.__class__.__name__
    self.config = dict(kwargs)
    self.config.setdefault('host', host)
    self.config.setdefault('port', 6667)
    self.config.setdefault('nick', nick)
    self.config.setdefault('names', [self.config['nick']])
    self.config.setdefault('ident', nick.lower())
    self.config.setdefault('realname', "A Pyrc Bot")
    self.config.setdefault('channels', [])
    self.config.setdefault('password', None)

    self._inbuffer = ""
    self._commands = []
    self._threads = []
    self.socket = None
    self.initialized = False

    self.addhooks()

  def message(self, recipient, s):
    "High level interface to sending an IRC message."
    self.cmd("PRIVMSG %s :%s" % (recipient, s))

  def connect(self):
    '''
    Connects to the IRC server with the options defined in `config`
    '''
    self._connect()

    try:
      self.listen()
    except (KeyboardInterrupt, SystemExit):
      pass
    finally:
      self.close()

  def listen(self):
    """
    Constantly listens to the input from the server. Since the messages come
    in pieces, we wait until we receive 1 or more full lines to start parsing.

    A new line is defined as ending in \r\n in the RFC, but some servers
    separate by \n. This script takes care of both.
    """
    while True:
      self._inbuffer = self._inbuffer + self.socket.recv(1024)
      # Some IRC servers disregard the RFC and split lines by \n rather than \r\n.

      temp = self._inbuffer.split("\n")
      self._inbuffer = temp.pop()

      for line in temp:
        # Strip \r from \r\n for RFC-compliant IRC servers.
        line = line.rstrip('\r')
        self.parseline(line)
        print line

  def parseline(self, line):
    if line.startswith("PING"):
      length = len("PING ")
      host = line[length:]
      self.cmd("PONG %s" % host)
    elif re.match(r"^:\S+ PRIVMSG", line):
      msg_regex = re.compile(r"^:(\S+)!\S+ PRIVMSG (\S+) :(.*)")
      nick, channel, message = re.match(msg_regex, line).groups()
      self.receivemessage(channel, nick, message)
    elif re.match(r"^:\S+ INVITE %s" % self.config['nick'], line):
      msg_regex = re.compile(r"^:(\S+)!\S+ INVITE \S+ :?(.*)")
      inviter, channel = re.match(msg_regex, line).groups()
      self.cmd("JOIN %s" % channel)
    elif re.match(r"^\S+ MODE %s :\+\w*i" % self.config['nick'], line)\
        and self.config['password']:
      # Autoidentify if a password is provided
      self.cmd("PRIVMSG NickServ :identify %s" % self.config['password'])
    elif re.match(r"^:\S+ MODE", line) and not self.initialized:
      self.initialized = True
      if self.config['channels']:
        self.cmd("JOIN %s" % ' '.join(self.config['channels']))
      # TODO: This doesn't ensure that threads run at the right time, e.g.
      # after the bot has joined every channel it needs to.
      for thread in self._threads:
        thread.start()

  def addhooks(self):
    for func in self.__class__.__dict__.values():
      if callable(func) and hasattr(func, '_type'):
        if func._type == 'COMMAND':
          self._commands.append(func)
        elif func._type == "REPEAT":
          thread = threads.JobThread(func, self)
          self._threads.append(thread)
        else:
          raise "This is not a type I've ever heard of."

  def receivemessage(self, channel, nick, message):
    self.parsecommand(channel, message)

  def parsecommand(self, channel, message):
    # sort names so names that are substrings work
    names = sorted(self.config['names'], key=len, reverse=True)

    name_used = None
    for name in names:
      name_regex_str = r'^(%s)[,:]?\s+' % re.escape(name)
      name_regex = re.compile(name_regex_str, re.IGNORECASE)
      if name_regex.match(message):
        name_used = name_regex.match(message).group(1)
        break

    if not name_used:
      return

    _,_,message = message.partition(name_used)
    command = re.match(r'^[,:]?\s+(.*)', message).group(1)
    for command_func in self._commands:
      match = command_func._matcher.search(command)
      if match:
        command_func(self, channel, *match.groups(), **match.groupdict())

  def cmd(self, raw_line):
    self.socket.send(raw_line + "\r\n")

  def _connect(self):
    "Connects a socket to the server using options defined in `config`."
    self.socket = socket.socket()
    self.socket.connect((self.config['host'], self.config['port']))
    self.cmd("NICK %s" % self.config['nick'])
    self.cmd("USER %s %s bla :%s" %
        (self.config['ident'], self.config['host'], self.config['realname']))

  def close(self):
    for thread in self._threads:
      thread.shutdown()
    self.socket.shutdown(socket.SHUT_RDWR)
    self.socket.close()
<|endoftext|>


<commit_before>import os
import sys
from path import path

def show_not_dir(p):
	while not p.isdir():
		prev = p
		p = p.parent
		if p == '/': break
	raise ValueError('%s is not a directory' % prev)

def find_dir(dirr,sub_dir=None):
	if dirr == '-':
		try: dirr = os.environ['OLDPWD']
		except KeyError: dirr = '~'
	whither = path(dirr)
	if whither.isfile():
		whither = whither.parent
	if not whither.isdir():
		show_not_dir(whither)
	if sub_dir:
		possibles = [ x for x in whither.walkdirs() if sub_dir in x.name ]
		if len(possibles) == 1:
			whither = possibles[0]
		elif not possibles:
			pass
		else:
			raise NotImplementedError('Too many possiblities:\n\t%s' % '\n\t'.join(possibles) )
	return whither

def parse_command_line(args):
	if len(args) < 1 or len(args) > 2:
		raise ValueError('Usage: kd directory [sub-dir]')
	if len(args) == 1:
		args += [ None ]
	return args

def chdir(whither):
	if whither.isfile():
		whither = whither.parent
	if not whither.isdir():
		show_not_dir(whither)
	oldpwd = os.environ['PWD']
	os.chdir(whither)
	os.environ['OLDPWD'] = oldpwd
	os.environ['PWD'] = whither

def main():
	try:
		args = parse_command_line(sys.argv[1:])
		whither = find_dir(args[0], args[1])
		print str(whither)
		return 0
	except Exception, e:
		print os.getcwd()
		print >> sys.stderr, str(e)
		return 1

if __name__ == '__main__':
	sys.exit(main())
<commit_msg>allow going nowhere<commit_after>import os
import sys
from path import path

def show_not_dir(p):
	while not p.isdir():
		prev = p
		p = p.parent
		if p == '/': break
	raise ValueError('%s is not a directory' % prev)

def find_dir(dirr,sub_dir=None):
	if dirr == '':
		return path('~').abspath()
	if dirr == '-':
		try: dirr = os.environ['OLDPWD']
		except KeyError: dirr = '~'
	whither = path(dirr)
	if whither.isfile():
		whither = whither.parent
	if not whither.isdir():
		show_not_dir(whither)
	if sub_dir:
		possibles = [ x for x in whither.walkdirs() if sub_dir in x.name ]
		if len(possibles) == 1:
			whither = possibles[0]
		elif not possibles:
			pass
		else:
			raise NotImplementedError('Too many possiblities:\n\t%s' % '\n\t'.join(possibles) )
	return whither

def parse_command_line(args):
	if len(args) > 2:
		raise ValueError('Usage: kd directory [sub-dir]')
	elif len(args) < 1:
		args = '', None
	elif len(args) == 1:
		args += [ None ]
	else:
		raise ValueError('I cannot count')
	return args

def chdir(whither):
	if whither.isfile():
		whither = whither.parent
	if not whither.isdir():
		show_not_dir(whither)
	oldpwd = os.environ['PWD']
	os.chdir(whither)
	os.environ['OLDPWD'] = oldpwd
	os.environ['PWD'] = whither

def main():
	try:
		args = parse_command_line(sys.argv[1:])
		whither = find_dir(args[0], args[1])
		print str(whither)
		return 0
	except Exception, e:
		print os.getcwd()
		print >> sys.stderr, str(e)
		return 1

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


<commit_before>from django.db import models
from django.conf import settings
from django_markdown.models import MarkdownField

from taggit.managers import TaggableManager
from annoying.fields import AutoOneToOneField

class UserQAProfile(models.Model):
    user = AutoOneToOneField(settings.AUTH_USER_MODEL, primary_key=True)
    points = models.IntegerField(default=0)
    # The additional attributes we wish to include.
    website = models.URLField(blank=True)
    picture = models.ImageField(upload_to='qa/static/profile_images',
                                blank=True)

    def __str__(self):  # pragma: no cover
        return self.user.username

class Question(models.Model):
    title = models.CharField(max_length=200, blank=False)
    description = MarkdownField()
    pub_date = models.DateTimeField('date published', auto_now_add=True)
    tags = TaggableManager()
    reward = models.IntegerField(default=0)
    views = models.IntegerField(default=0)
    user = models.ForeignKey(settings.AUTH_USER_MODEL)
    closed = models.BooleanField(default=False)
    positive_votes = models.IntegerField(default=0)
    negative_votes = models.IntegerField(default=0)
    total_points = models.IntegerField(default=0)

    def save(self, *args, **kwargs):
        self.total_points = self.positive_votes - self.negative_votes
        super(Question, self).save(*args, **kwargs)

    def __str__(self):
        return self.title

class Answer(models.Model):
    question = models.ForeignKey(Question)
    answer_text = MarkdownField()
    pub_date = models.DateTimeField('date published', auto_now_add=True)
    updated = models.DateTimeField('date updated', auto_now=True)
    user = models.ForeignKey(settings.AUTH_USER_MODEL)
    answer = models.BooleanField(default=False)
    positive_votes = models.IntegerField(default=0)
    negative_votes = models.IntegerField(default=0)
    total_points = models.IntegerField(default=0)

    def save(self, *args, **kwargs):
        self.total_points = self.positive_votes - self.negative_votes
        super(Answer, self).save(*args, **kwargs)

    def __str__(self):  # pragma: no cover
        return self.answer_text

class VoteParent(models.Model):
    user = models.ForeignKey(settings.AUTH_USER_MODEL)
    value = models.BooleanField(default=True)

    class Meta:
        abstract = True

class AnswerVote(VoteParent):
    answer = models.ForeignKey(Answer)

    class Meta:
        unique_together = (('user', 'answer'),)

class QuestionVote(VoteParent):
    question = models.ForeignKey(Question)

    class Meta:
        unique_together = (('user', 'question'),)

class BaseComment(models.Model):
    pub_date = models.DateTimeField('date published', auto_now_add=True)
    user = models.ForeignKey(settings.AUTH_USER_MODEL)

    class Meta:
        abstract = True

    def __str__(self):  # pragma: no cover
        return self.comment_text

class AnswerComment(BaseComment):
    comment_text = MarkdownField()
    answer = models.ForeignKey(Answer)

class QuestionComment(BaseComment):
    comment_text = models.CharField(max_length=250)
    question = models.ForeignKey(Question)
<commit_msg>Deleting the field picture to take away unnecesary components on the model, this way is leave to the main project to implement the whole user profile<commit_after>from django.db import models
from django.conf import settings
from django_markdown.models import MarkdownField

from taggit.managers import TaggableManager
from annoying.fields import AutoOneToOneField

class UserQAProfile(models.Model):
    user = AutoOneToOneField(settings.AUTH_USER_MODEL, primary_key=True)
    points = models.IntegerField(default=0)
    # The additional attributes we wish to include.
    website = models.URLField(blank=True)

    def __str__(self):  # pragma: no cover
        return self.user.username

class Question(models.Model):
    title = models.CharField(max_length=200, blank=False)
    description = MarkdownField()
    pub_date = models.DateTimeField('date published', auto_now_add=True)
    tags = TaggableManager()
    reward = models.IntegerField(default=0)
    views = models.IntegerField(default=0)
    user = models.ForeignKey(settings.AUTH_USER_MODEL)
    closed = models.BooleanField(default=False)
    positive_votes = models.IntegerField(default=0)
    negative_votes = models.IntegerField(default=0)
    total_points = models.IntegerField(default=0)

    def save(self, *args, **kwargs):
        self.total_points = self.positive_votes - self.negative_votes
        super(Question, self).save(*args, **kwargs)

    def __str__(self):
        return self.title

class Answer(models.Model):
    question = models.ForeignKey(Question)
    answer_text = MarkdownField()
    pub_date = models.DateTimeField('date published', auto_now_add=True)
    updated = models.DateTimeField('date updated', auto_now=True)
    user = models.ForeignKey(settings.AUTH_USER_MODEL)
    answer = models.BooleanField(default=False)
    positive_votes = models.IntegerField(default=0)
    negative_votes = models.IntegerField(default=0)
    total_points = models.IntegerField(default=0)

    def save(self, *args, **kwargs):
        self.total_points = self.positive_votes - self.negative_votes
        super(Answer, self).save(*args, **kwargs)

    def __str__(self):  # pragma: no cover
        return self.answer_text

class VoteParent(models.Model):
    user = models.ForeignKey(settings.AUTH_USER_MODEL)
    value = models.BooleanField(default=True)

    class Meta:
        abstract = True

class AnswerVote(VoteParent):
    answer = models.ForeignKey(Answer)

    class Meta:
        unique_together = (('user', 'answer'),)

class QuestionVote(VoteParent):
    question = models.ForeignKey(Question)

    class Meta:
        unique_together = (('user', 'question'),)

class BaseComment(models.Model):
    pub_date = models.DateTimeField('date published', auto_now_add=True)
    user = models.ForeignKey(settings.AUTH_USER_MODEL)

    class Meta:
        abstract = True

    def __str__(self):  # pragma: no cover
        return self.comment_text

class AnswerComment(BaseComment):
    comment_text = MarkdownField()
    answer = models.ForeignKey(Answer)

class QuestionComment(BaseComment):
    comment_text = models.CharField(max_length=250)
    question = models.ForeignKey(Question)
<|endoftext|>


<commit_before>import torch
from torch.autograd import Variable
import qelos as q
import scipy as sp
from scipy import optimize as spopt

class GANTrainer(object):
    def __init__(self,  mode="DRAGAN",      # WGAN, WGAN-GP, DRAGAN, DRAGAN-G, DRAGAN-LG
                        modeD="critic",  # disc or critic
                        one_sided=False,
                        penalty_weight=None,
                        perturb_both=False,
                        perturb_symmetric=False,
                        perturb_scale=1.,
                        clamp_weights_rng=(-0.01, 0.01),
                        optimizerD=None,
                        optimizerG=None,
                        logger=None,
                        data_iter=None,
                        valid_data_iter=None,
                        validinter=1):
        if one_sided:
            self.clip_fn = lambda x: x.clamp(min=0)
        else:
            self.clip_fn = lambda x: x
        self.optimizerD = optimizerD
        self.optimizerG = optimizerG
        self.penalty_weight = penalty_weight
        self.perturb_both = perturb_both
        self.perturb_symmetric = perturb_symmetric
        self.perturb_scale = perturb_scale
        self.clamp_weights_rng = clamp_weights_rng
        self.noise_dim = 2
        self.mode = mode
        self.modeD = modeD
        self.logger = logger
        self._dragan_lg_var = 0.6
        self.data_iter = data_iter
        self.valid_data_iter = valid_data_iter
        self.validinter = validinter

    def perturb(self, x):
        if self.mode == "DRAGAN":
            return self.perturb_dragan(x)
        elif self.mode == "DRAGAN-G":
            return self.perturb_dragan_g(x)
        elif self.mode == "DRAGAN-LG":
            return self.perturb_dragan_lg(x)

    def perturb_dragan(self, x):  # Variable (batsize, dim)
        perturbation = x.data.new(x.size())
        if not self.perturb_symmetric:
            perturbation.uniform_(0, to=1)
        else:
            perturbation.uniform_(-1, to=1)
        perturbation *= self.perturb_scale * x.data.std()
        interp_alpha = perturbation.new(perturbation.size(0), 1)
        interp_alpha.uniform_(0, to=1)
        perturbation = interp_alpha * perturbation
        perturbed = x.data + perturbation
        return Variable(perturbed, requires_grad=True)

    def perturb_dragan_g(self, x):
        perturbation = x.data.new(x.size())
        perturbation.normal_()
        perturbation *= 0.25 * self.perturb_scale * x.data.std()
        ret = x.data + perturbation
        return Variable(ret, requires_grad=True)

    def perturb_dragan_lg(self, x):
        perturbation = x.data.new(x.size())
        perturbation.log_normal_(0, self._dragan_lg_var)
        interp_alpha = perturbation.new(perturbation.size())
        interp_alpha.uniform_(-1, to=1)
        interp_alpha.sign_()
        perturbation *= interp_alpha * 0.25 * self.perturb_scale * x.data.std()
        ret = x.data + perturbation
        return Variable(ret, requires_grad=True)

    def train(self, netD, netG, niter=0, niterD=10, batsizeG=100,
              data_gen=None, valid_data_gen=None, cuda=False):
        data_gen = data_gen if data_gen is not None else self.data_iter
        valid_data_gen = valid_data_gen if valid_data_gen is not None else self.valid_data_iter
        if cuda:
            netD.cuda()
            netG.cuda()

        valid_EMD = 0.
        valid_fake2real = 0.    # mean of distances from each fake to closest real
        valid_real2fake = 0.    # mean of distances from each real to closest fake

        vnoise = q.var(torch.zeros(batsizeG, self.noise_dim)).cuda(cuda).v

        for _iter in range(niter):
            ##########  Update D net ##########
            lip_loss = None
            for p in netD.parameters():
                p.requires_grad = True

            if _iter < 25 or _iter % 500 == 0:
                _niterD = 100
            else:
                _niterD = niterD
            for j in range(_niterD):
                netD.zero_grad()
                if self.mode == "WGAN":
                    for p in netD.parameters():
                        p.data.clamp_(*self.clamp_weights_rng)
                data = next(data_gen)
                real = q.var(data).cuda(cuda).v
                num_examples = data.size(0)
                vnoise = Variable(real.data.new(num_examples, self.noise_dim))
                vnoise.data.normal_(0, 1)
                fake = netG(vnoise)
                scoreD_real_vec = netD(real)        # (batsize,)
                #scoreD_real = scoreD_real_vec.mean()
                # scoreD_real.backward(one, retain_graph=(lc> 0))
                scoreD_fake_vec = netD(fake)
                #scoreD_fake = scoreD_fake_vec.mean()
                # scoreD_fake.backward(mone, retain_graph=(lc > 0))
                if self.mode == "WGAN":
                    errD = scoreD_fake_vec.mean() - scoreD_real_vec.mean()
                elif self.mode == "WGAN-GP":
                    errD = scoreD_fake_vec.mean() - scoreD_real_vec.mean()
                    interp_alpha = real.data.new(num_examples, 1)
                    interp_alpha.uniform_(0, 1)
                    interp_points = interp_alpha * real.data + (1 - interp_alpha) * fake.data
                    grad_points = interp_points.clone()
                    interp_points = Variable(interp_points, requires_grad=True)
                    errD_interp_vec = netD(interp_points)
                    errD_gradient, = torch.autograd.grad(errD_interp_vec.sum(),
                                                         interp_points,
                                                         create_graph=True)
                    lip_grad_norm = errD_gradient.norm(2, dim=1)
                    lip_loss = self.penalty_weight * (self.clip_fn(lip_grad_norm - 1) ** 2).mean()
                    errD = errD + lip_loss
                elif self.mode == "DRAGAN" or self.mode == "DRAGAN-G" or self.mode == "DRAGAN-LG":
                    if self.modeD == "disc":
                        errD = - torch.log(scoreD_real_vec).mean() - torch.log(1 - scoreD_fake_vec).mean()
                    elif self.modeD == "critic":
                        errD = scoreD_fake_vec.mean() - scoreD_real_vec.mean()

                    real_perturbed = self.perturb(real)
                    grad_points = real_perturbed.data
                    errD_real_perturbed_vec = netD(real_perturbed)
                    errD_gradient, = torch.autograd.grad(
                        errD_real_perturbed_vec.sum(), real_perturbed, create_graph=True)
                    lip_grad_norm = errD_gradient.norm(2)
                    lip_loss = self.penalty_weight * (self.clip_fn(lip_grad_norm - 1) ** 2).mean()
                    if self.perturb_both:
                        fake_perturbed = self.perturb(fake)
                        errD_fake_perturbed_vec = netD(fake_perturbed)
                        errD_fake_gradient, = torch.autograd.grad(
                            errD_fake_perturbed_vec.sum(), fake_perturbed, create_graph=True)
                        lip_fake_grad_norm = errD_fake_gradient.norm(2)
                        lip_fake_loss = self.penalty_weight * (self.clip_fn(lip_fake_grad_norm - 1) ** 2).mean()
                        lip_loss = lip_loss * 0.5 + lip_fake_loss * 0.5
                    errD = errD + lip_loss
                    #errD = lip_loss
                errD.backward()
                self.optimizerD.step()

            ##########  Update G net ##########
            for p in netD.parameters():
                p.requires_grad = False
            netG.zero_grad()

            vnoise.data.normal_(0, 1)
            fake = netG(vnoise)
            errG = netD(fake)
            if self.modeD == "critic":
                errG = -errG.mean()
            elif self.modeD == "disc":
                errG = torch.log(1.0 - errG).mean()
            errG.backward()
            self.optimizerG.step()

            ######### Validate G net ##########
            if valid_data_gen is not None:
                if _iter % self.validinter == 0 or self.validinter == 1:
                    validdata = next(valid_data_gen)
                    validreal = q.var(validdata, volatile=True).cuda(cuda).v
                    vnoise = q.var(validreal.data.new(validreal.size(0), self.noise_dim),
                                   volatile=True).v
                    vnoise.data.normal_(0, 1)
                    validfake = netG(vnoise)
                    # compute distance matrix
                    distmat = -q.LNormDistance(2)(
                        validreal.unsqueeze(0),
                        validfake.unsqueeze(0)).squeeze(0)
                    npdistmat = distmat.cpu().data.numpy()
                    ass_x, ass_y = spopt.linear_sum_assignment(npdistmat)
                    valid_EMD = npdistmat[ass_x, ass_y].mean()

                    #real2fake and fake2real
                    valid_fake2real, _ = torch.min(distmat, 0)
                    valid_real2fake, _ = torch.min(distmat, 1)
                    valid_fake2real = valid_fake2real.mean()
                    valid_real2fake = valid_real2fake.mean()

            if self.logger is not None:
                self.logger.log(_iter=_iter, niter=niter,
                                real=real.data, fake=fake.data, grad_points=grad_points,
                                errD=errD.data[0], errG=errG.data[0],
                                scoreD_real=scoreD_real_vec.mean().data[0],
                                scoreD_fake=scoreD_fake_vec.mean().data[0],
                                lip_loss=lip_loss.data[0] if lip_loss is not None else 0.,
                                valid_EMD=valid_EMD, valid_fake2real=valid_fake2real.data[0], valid_real2fake=valid_real2fake.data[0],
                                when="after_G")

<commit_msg>* added validation for EM distance based on euclidean distances in gan<commit_after>import torch
from torch.autograd import Variable
import qelos as q
import scipy as sp
from scipy import optimize as spopt

class GANTrainer(object):
    def __init__(self,  mode="DRAGAN",      # WGAN, WGAN-GP, DRAGAN, DRAGAN-G, DRAGAN-LG
                        modeD="critic",  # disc or critic
                        one_sided=False,
                        penalty_weight=None,
                        perturb_both=False,
                        perturb_symmetric=False,
                        perturb_scale=1.,
                        clamp_weights_rng=(-0.01, 0.01),
                        optimizerD=None,
                        optimizerG=None,
                        logger=None,
                        data_iter=None,
                        valid_data_iter=None,
                        validinter=1):
        if one_sided:
            self.clip_fn = lambda x: x.clamp(min=0)
        else:
            self.clip_fn = lambda x: x
        self.optimizerD = optimizerD
        self.optimizerG = optimizerG
        self.penalty_weight = penalty_weight
        self.perturb_both = perturb_both
        self.perturb_symmetric = perturb_symmetric
        self.perturb_scale = perturb_scale
        self.clamp_weights_rng = clamp_weights_rng
        self.noise_dim = 2
        self.mode = mode
        self.modeD = modeD
        self.logger = logger
        self._dragan_lg_var = 0.6
        self.data_iter = data_iter
        self.valid_data_iter = valid_data_iter
        self.validinter = validinter

    def perturb(self, x):
        if self.mode == "DRAGAN":
            return self.perturb_dragan(x)
        elif self.mode == "DRAGAN-G":
            return self.perturb_dragan_g(x)
        elif self.mode == "DRAGAN-LG":
            return self.perturb_dragan_lg(x)

    def perturb_dragan(self, x):  # Variable (batsize, dim)
        perturbation = x.data.new(x.size())
        if not self.perturb_symmetric:
            perturbation.uniform_(0, to=1)
        else:
            perturbation.uniform_(-1, to=1)
        perturbation *= self.perturb_scale * x.data.std()
        interp_alpha = perturbation.new(perturbation.size(0), 1)
        interp_alpha.uniform_(0, to=1)
        perturbation = interp_alpha * perturbation
        perturbed = x.data + perturbation
        return Variable(perturbed, requires_grad=True)

    def perturb_dragan_g(self, x):
        perturbation = x.data.new(x.size())
        perturbation.normal_()
        perturbation *= 0.25 * self.perturb_scale * x.data.std()
        ret = x.data + perturbation
        return Variable(ret, requires_grad=True)

    def perturb_dragan_lg(self, x):
        perturbation = x.data.new(x.size())
        perturbation.log_normal_(0, self._dragan_lg_var)
        interp_alpha = perturbation.new(perturbation.size())
        interp_alpha.uniform_(-1, to=1)
        interp_alpha.sign_()
        perturbation *= interp_alpha * 0.25 * self.perturb_scale * x.data.std()
        ret = x.data + perturbation
        return Variable(ret, requires_grad=True)

    def train(self, netD, netG, niter=0, niterD=10, batsizeG=100,
              data_gen=None, valid_data_gen=None, cuda=False):
        data_gen = data_gen if data_gen is not None else self.data_iter
        valid_data_gen = valid_data_gen if valid_data_gen is not None else self.valid_data_iter
        if cuda:
            netD.cuda()
            netG.cuda()

        valid_EMD = 0.
        valid_fake2real = 0.    # mean of distances from each fake to closest real
        valid_real2fake = 0.    # mean of distances from each real to closest fake
        valid_fakeandreal = 0.  # harmonic mean of above

        vnoise = q.var(torch.zeros(batsizeG, self.noise_dim)).cuda(cuda).v

        for _iter in range(niter):
            ##########  Update D net ##########
            lip_loss = None
            for p in netD.parameters():
                p.requires_grad = True

            if _iter < 25 or _iter % 500 == 0:
                _niterD = 100
            else:
                _niterD = niterD
            for j in range(_niterD):
                netD.zero_grad()
                if self.mode == "WGAN":
                    for p in netD.parameters():
                        p.data.clamp_(*self.clamp_weights_rng)
                data = next(data_gen)
                real = q.var(data).cuda(cuda).v
                num_examples = data.size(0)
                vnoise = Variable(real.data.new(num_examples, self.noise_dim))
                vnoise.data.normal_(0, 1)
                fake = netG(vnoise)
                scoreD_real_vec = netD(real)        # (batsize,)
                #scoreD_real = scoreD_real_vec.mean()
                # scoreD_real.backward(one, retain_graph=(lc> 0))
                scoreD_fake_vec = netD(fake)
                #scoreD_fake = scoreD_fake_vec.mean()
                # scoreD_fake.backward(mone, retain_graph=(lc > 0))
                if self.mode == "WGAN":
                    errD = scoreD_fake_vec.mean() - scoreD_real_vec.mean()
                elif self.mode == "WGAN-GP":
                    errD = scoreD_fake_vec.mean() - scoreD_real_vec.mean()
                    interp_alpha = real.data.new(num_examples, 1)
                    interp_alpha.uniform_(0, 1)
                    interp_points = interp_alpha * real.data + (1 - interp_alpha) * fake.data
                    grad_points = interp_points.clone()
                    interp_points = Variable(interp_points, requires_grad=True)
                    errD_interp_vec = netD(interp_points)
                    errD_gradient, = torch.autograd.grad(errD_interp_vec.sum(),
                                                         interp_points,
                                                         create_graph=True)
                    lip_grad_norm = errD_gradient.norm(2, dim=1)
                    lip_loss = self.penalty_weight * (self.clip_fn(lip_grad_norm - 1) ** 2).mean()
                    errD = errD + lip_loss
                elif self.mode == "DRAGAN" or self.mode == "DRAGAN-G" or self.mode == "DRAGAN-LG":
                    if self.modeD == "disc":
                        errD = - torch.log(scoreD_real_vec).mean() - torch.log(1 - scoreD_fake_vec).mean()
                    elif self.modeD == "critic":
                        errD = scoreD_fake_vec.mean() - scoreD_real_vec.mean()

                    real_perturbed = self.perturb(real)
                    grad_points = real_perturbed.data
                    errD_real_perturbed_vec = netD(real_perturbed)
                    errD_gradient, = torch.autograd.grad(
                        errD_real_perturbed_vec.sum(), real_perturbed, create_graph=True)
                    lip_grad_norm = errD_gradient.norm(2)
                    lip_loss = self.penalty_weight * (self.clip_fn(lip_grad_norm - 1) ** 2).mean()
                    if self.perturb_both:
                        fake_perturbed = self.perturb(fake)
                        errD_fake_perturbed_vec = netD(fake_perturbed)
                        errD_fake_gradient, = torch.autograd.grad(
                            errD_fake_perturbed_vec.sum(), fake_perturbed, create_graph=True)
                        lip_fake_grad_norm = errD_fake_gradient.norm(2)
                        lip_fake_loss = self.penalty_weight * (self.clip_fn(lip_fake_grad_norm - 1) ** 2).mean()
                        lip_loss = lip_loss * 0.5 + lip_fake_loss * 0.5
                    errD = errD + lip_loss
                    #errD = lip_loss
                errD.backward()
                self.optimizerD.step()

            ##########  Update G net ##########
            for p in netD.parameters():
                p.requires_grad = False
            netG.zero_grad()

            vnoise.data.normal_(0, 1)
            fake = netG(vnoise)
            errG = netD(fake)
            if self.modeD == "critic":
                errG = -errG.mean()
            elif self.modeD == "disc":
                errG = torch.log(1.0 - errG).mean()
            errG.backward()
            self.optimizerG.step()

            ######### Validate G net ##########
            if valid_data_gen is not None:
                if _iter % self.validinter == 0 or self.validinter == 1:
                    validdata = next(valid_data_gen)
                    validreal = q.var(validdata, volatile=True).cuda(cuda).v
                    vnoise = q.var(validreal.data.new(validreal.size(0), self.noise_dim),
                                   volatile=True).v
                    vnoise.data.normal_(0, 1)
                    validfake = netG(vnoise)
                    # compute distance matrix
                    distmat = -q.LNormDistance(2)(
                        validreal.unsqueeze(0),
                        validfake.unsqueeze(0)).squeeze(0)
                    npdistmat = distmat.cpu().data.numpy()
                    # ass_x, ass_y = spopt.linear_sum_assignment(npdistmat)
                    # valid_EMD = npdistmat[ass_x, ass_y].mean()

                    #real2fake and fake2real
                    valid_fake2real, _ = torch.min(distmat, 0)
                    valid_real2fake, _ = torch.min(distmat, 1)
                    valid_fake2real = valid_fake2real.mean()
                    valid_real2fake = valid_real2fake.mean()
                    valid_fakeandreal = 2 * valid_fake2real * valid_real2fake / (valid_fake2real + valid_real2fake).clamp(min=1e-6)

            if self.logger is not None:
                self.logger.log(_iter=_iter, niter=niter,
                                real=real.data, fake=fake.data, grad_points=grad_points,
                                errD=errD.data[0], errG=errG.data[0],
                                scoreD_real=scoreD_real_vec.mean().data[0],
                                scoreD_fake=scoreD_fake_vec.mean().data[0],
                                lip_loss=lip_loss.data[0] if lip_loss is not None else 0.,
                                valid_EMD=valid_EMD, valid_fake2real=valid_fake2real.data[0],
                                valid_real2fake=valid_real2fake.data[0], valid_fakeandreal=valid_fakeandreal.data[0],
                                when="after_G")

<|endoftext|>


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

from numpy.random import choice as numpy_choice
from random import choice

# class RandomisedPopulation(object):
#     """ Class to deal with stratified block randomisation """
#     def __init__(self, variables, population):
#         super(RandomisedPopulation, self).__init__()
#         self.variables = variables
#         self.population = population

class Participant(object):
    """Class for new Participant, member of population"""
    def __init__(self, participant_id):
        super(Participant, self).__init__()
        self.participant_id = participant_id

def generate_population():
    """Generate a set of participants to define a population.
    Need to generate these automatically"""

    population = []

    return population

def define_stratified_groups():

    """
    Create groups so population is stratified by randomisation variables
    """

    groups = {}

    # Tuples are hashable in python
    # http://stackoverflow.com/questions/2388302/best-way-to-define-multidimensional-dictionaries-in-python
    variable_set_1 = ['option1', 'option2']
    variable_set_2 = ['optionA', 'optionB']
    variables = [variable_set_1, variable_set_2]

    # Possible to loop through any number of variables?
    # Initialise groups
    for var_1 in variable_set_1:
        for var_2 in variable_set_2:
            groups[(var_1, var_2)] = []

    return groups

def population_to_groups(population, groups):
    """Assign each participant in a population to a stratified group.
    Currently doing this by hardcoding. Need to feed in the
    stratification variables for dynamic allocation."""

    for participant in population:
        if participant.var_1 == 'A':
            if participant.var_2 == 'X':
                groups[('A', 'X')].append(
                    participant.participant_id)
            else:
                groups[('A', 'Y')].append(
                    participant.participant_id)
        else:
            if participant.var_2 == 'X':
                groups[('B', 'X')].append(
                    participant.participant_id)
            else:
                groups[('B', 'Y')].append(
                    participant.participant_id)

    return groups

def define_block():

    # Todo: possible option -> hardcode block size?
    block_size = 8

    # Todo: possible option -> define other arms?
    arms = ['intervention', 'control']
    arm_assignment = {'intervention': 0, 'control': 0}
    arm_list = []

    # Todo: possible option -> include seed for reproducing results
    # seed(74832.9191663033)
    # Need to pass the state of the random generator back to main
    # so that it isn't reset for each loop, which would be pointless

    for x in xrange(0, block_size):

        if (arm_assignment['intervention'] == block_size / 2):
            arm_assignment['control'] += 1
            arm_list.append('control')

        elif (arm_assignment[arms[1]] == block_size / 2):
            arm_assignment[arms[0]] += 1
            arm_list.append('intervention')

        else:
            arm_choice = choice(arms)
            arm_assignment[arm_choice] += 1
            arm_list.append(arm_choice)

    return arm_list

def population_to_blocks(population):

    block_size = 8

    for features, participants in population.items():
        print "\nFeatures: ", features
        print "Participants:", participants

        # Slice notation
        # http://stackoverflow.com/questions/509211/explain-pythons-slice-notation
        blocks = [
                participants[x:x+block_size]
                for x in xrange(0, len(participants), block_size)
                ]
        print "\nBlocks:"
        for block in blocks:

            # Generate new block assignments
            block_assignment = define_block()

            for i, participant in enumerate(block):
                print participant, " - ", block_assignment[i]

    print "\n"

def main():

    # Generate population
    population = generate_population()

    # Define stratified groups
    groups = define_stratified_groups()

    # Population to groups
    population_in_groups = population_to_groups(population, groups)

    # Population to blocks
    population_to_blocks(population_in_groups)

    # Define blocks
    block = define_block()
    print "\nFinished\n"

if __name__ == '__main__':
    main()
<commit_msg>Added class for Population<commit_after>#!/usr/bin/env python

from numpy.random import choice as numpy_choice
from random import choice

class Population(object):
    """ Class for a population (of participants) """
    def __init__(self, variables, population):
        super(Population, self).__init__()
        self.variables = variables
        self.population = population

class Participant(object):
    """Class for new Participant, member of population"""
    def __init__(self, participant_id):
        super(Participant, self).__init__()
        self.participant_id = participant_id

def generate_population():
    """Generate a set of participants to define a population.
    Need to generate these automatically"""

    population = []

    return population

def define_stratified_groups():

    """
    Create groups so population is stratified by randomisation variables
    """

    groups = {}

    # Tuples are hashable in python
    # http://stackoverflow.com/questions/2388302/best-way-to-define-multidimensional-dictionaries-in-python
    variable_set_1 = ['option1', 'option2']
    variable_set_2 = ['optionA', 'optionB']
    variables = [variable_set_1, variable_set_2]

    # Possible to loop through any number of variables?
    # Initialise groups
    for var_1 in variable_set_1:
        for var_2 in variable_set_2:
            groups[(var_1, var_2)] = []

    return groups

def population_to_groups(population, groups):
    """Assign each participant in a population to a stratified group.
    Currently doing this by hardcoding. Need to feed in the
    stratification variables for dynamic allocation."""

    for participant in population:
        if participant.var_1 == 'A':
            if participant.var_2 == 'X':
                groups[('A', 'X')].append(
                    participant.participant_id)
            else:
                groups[('A', 'Y')].append(
                    participant.participant_id)
        else:
            if participant.var_2 == 'X':
                groups[('B', 'X')].append(
                    participant.participant_id)
            else:
                groups[('B', 'Y')].append(
                    participant.participant_id)

    return groups

def define_block():

    # Todo: possible option -> hardcode block size?
    block_size = 8

    # Todo: possible option -> define other arms?
    arms = ['intervention', 'control']
    arm_assignment = {'intervention': 0, 'control': 0}
    arm_list = []

    # Todo: possible option -> include seed for reproducing results
    # seed(74832.9191663033)
    # Need to pass the state of the random generator back to main
    # so that it isn't reset for each loop, which would be pointless

    for x in xrange(0, block_size):

        if (arm_assignment['intervention'] == block_size / 2):
            arm_assignment['control'] += 1
            arm_list.append('control')

        elif (arm_assignment[arms[1]] == block_size / 2):
            arm_assignment[arms[0]] += 1
            arm_list.append('intervention')

        else:
            arm_choice = choice(arms)
            arm_assignment[arm_choice] += 1
            arm_list.append(arm_choice)

    return arm_list

def population_to_blocks(population):

    block_size = 8

    for features, participants in population.items():
        print "\nFeatures: ", features
        print "Participants:", participants

        # Slice notation
        # http://stackoverflow.com/questions/509211/explain-pythons-slice-notation
        blocks = [
                participants[x:x+block_size]
                for x in xrange(0, len(participants), block_size)
                ]
        print "\nBlocks:"
        for block in blocks:

            # Generate new block assignments
            block_assignment = define_block()

            for i, participant in enumerate(block):
                print participant, " - ", block_assignment[i]

    print "\n"

def main():

    # Generate population
    population = generate_population()

    # Define stratified groups
    groups = define_stratified_groups()

    # Population to groups
    population_in_groups = population_to_groups(population, groups)

    # Population to blocks
    population_to_blocks(population_in_groups)

    # Define blocks
    block = define_block()
    print "\nFinished\n"

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