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

# Usage of ramcloud.py script:
# Need to run within a salloc allocation.
#
# setup: Download and install RAMCloud into a shared directory.
# start: Start a RAMCloud cluster.
# stop: Stop the RAMCloud cluster.
#
# Assumes zookeeper has been deployed

# Contributors:
# Joao Carreira <joao@eecs.berkeley.edu> (2014)

import argparse
import getpass
import os
import urllib
import subprocess
import shutil
import re
import datetime
import time
import atexit
import json

run_timestamp = time.time()

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

version = '1.0'

user_dir = '/nscratch/' + getpass.getuser()
outdir = user_dir + '/ramcloud_bits'
workdir = os.getcwd() + '/work'
rundir = workdir + '/ramcloud-' +  datetime.datetime.fromtimestamp(run_timestamp). \
	strftime('%Y-%m-%d-%H-%M-%S')

network_if = 'eth0'

ramcloud_url = 'git://fiz.stanford.edu/git/ramcloud.git'

# paths used to set zookeeper dependencies
# XXX zookeeper version should be parameter
zookeeper_lib_path = user_dir +\
                     '/zookeeper_bits/zookeeper-3.4.6/src/c/.libs/libzookeeper_mt.a'
zookeeper_dir_path = user_dir + \
                     '/zookeeper_bits/zookeeper-3.4.6/src/c/include'

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

def make_dir(mydir):
	if not os.path.exists(mydir):
                print "Creating " + mydir
    		os.makedirs(mydir)
        else:
                print "Directory " + mydir + " already exists"

# Return the list of nodes in the current SLURM allocation
def get_slurm_nodelist():
	nodelist = subprocess.check_output( \
		['scontrol', 'show', 'hostname', os.environ['SLURM_NODELIST']], \
		universal_newlines=True)
	nodelist = nodelist.strip().split('\n')
	return nodelist

# Return list of the Ip addresses for the chosen network_if on all nodes in nodelist
def get_ip_addresses(nodelist):
	results = subprocess.check_output( \
		['srun', '--nodelist=' + ','.join(nodelist), 'bash', '../common/get_ip_address.sh', \
		network_if], universal_newlines=True)
	json_str = '[' + ','.join(results.splitlines()) + ']'
	raw_data = json.loads(json_str)
	ip_map = {}
	for entry in raw_data:
		ip_map[entry['host']] = entry['ip']
	return ip_map

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

def do_setup():
	print '> Setting up RAMCloud in directory ' + outdir
	print '>'

        ramcloud_dir = outdir + '/ramcloud'

	if not os.path.exists(ramcloud_dir):
                make_dir(ramcloud_dir)
                print '> Entering directory: ' + outdir
                print '> Cloning RAMCloud..'
                p = subprocess.Popen(['git', 'clone', ramcloud_url], cwd = outdir)
                p.wait()

                print '> Updating RAMCloud modules..'
                p = subprocess.Popen(['git', 'submodule', 'update', \
                        '--init', '--recursive'], cwd = ramcloud_dir)
                p.wait()

                print '> Compiling RAMCloud..'
                cmd_str = 'make -j 16 DEBUG=no ZOOKEEPER_LIB=' + \
                          zookeeper_lib_path + ' ZOOKEEPER_DIR=' + zookeeper_dir_path
                print cmd_str

                p =  subprocess.Popen(cmd_str, cwd = ramcloud_dir, shell=True)
                p.wait()

                print '> DONE'
	else:
                print '> RAMCloud was already setup. To perform setup' + \
                        'again please remove folder'
	print '>'
	print '> DONE'
	print '>'

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

ramcloud_instances = {}

def check_slurm_allocation():
        if (not 'SLURM_NODELIST' in os.environ) or (not os.environ['SLURM_NODELIST']):
            print '[ERROR] Need to run script within SLURM allocation'
            exit(1)

def shutdown_ramcloud_instances():
	print '> Shutting down RAMCloud instances'
	for c in spark_instances.values():
		c['process'].terminate()

	print '> Waiting for processes to terminate...'
	all_done = True
	while not all_done:
		all_done = True
		for c in spark_instances.values():
			if (c['process'].poll() == None):
				all_done = False
		time.sleep(0.01)
	print '> DONE'

def do_start():
	print '> Launching RAMCloud cluster...'
	print '>'
	print '> Output redirected to ' + rundir
	print '>'

	# Create symlink for latest run
	subprocess.call(['ln', '-s', '-f', '-T', rundir, workdir + '/latest'])

	nodelist = get_slurm_nodelist()
	master_node = nodelist[0]
	worker_nodes = nodelist[1:]

	print '> Master Node: ' + master_node
	print '> Worker Nodes: ' + ', '.join(worker_nodes)
	print '>'

	print '> Getting IP addresses for network interface ' + network_if + '...'
	ip_addresses = get_ip_addresses(nodelist)
	print '> IPs: ' + ', '.join(ip_addresses.values())
	print '>'

	# When exiting, make sure all children are terminated cleanly
	atexit.register(shutdown_spark_instances)

	print '>'
	print '> Waiting for Master to finish starting up...'
	while True:
		with open(myerrfile, 'r') as fout:
			outdata = fout.read()
		if re.search('Successfully started service \'sparkMaster\'', outdata) != None:
			break
		time.sleep(0.01)
	print '> MASTER IS UP!'
	print '>'

	# Step II: Launch Spark Workers
	print '> Launching Spark worker nodes'
	print '>'

	for node in worker_nodes:
		print '> Launching Spark worker on ' + node

		srun_cmd = ['srun', '--nodelist=' + node, '-N1']
		srun_cmd += ['bash', './bin/spark-class']
		srun_cmd += ['org.apache.spark.deploy.worker.Worker']
		srun_cmd += ['--ip', ip_addresses[node]]
		srun_cmd += ['spark://' + ip_addresses[master_node] + ':' + spark_master_port]

		#myenv = {'CASSANDRA_HOME': cassandra_home, 'CASSANDRA_CONF': myconfdir}
		myenv.update(os.environ)

		myrundir = rundir + '/worker-' + node
		make_dir(myrundir)
		myoutfile = myrundir + '/stdout'
		myerrfile = myrundir + '/stderr'

		fout = open(myoutfile, 'w')
		ferr = open(myerrfile, 'w')
		p = subprocess.Popen(srun_cmd, stdout=fout, stderr=ferr, env=myenv, cwd=spark_home)
		spark_instances[node] = {'process': p, 'out': myoutfile, 'err': myerrfile, 'type': 'worker'}

	# When exiting, make sure all children are terminated cleanly
	atexit.register(shutdown_spark_instances)

	print '>'
	print '> Waiting for all Workers to finish starting up...'
	unfinished_nodes = worker_nodes

	while unfinished_nodes:
		done_nodes = []
		for node in unfinished_nodes:
			with open(spark_instances[node]['err'], 'r') as fout:
				outdata = fout.read()
			if re.search("Successfully registered with master", outdata) != None:
				done_nodes.append(node)
		for node in done_nodes:
			unfinished_nodes.remove(node)
		time.sleep(0.01)

	print '> ALL WORKERS ARE UP!'

	# Write a JSON description of the Cassandra instance that can be used by others.
	print '> Writing instance description to ' + instance_json
	#cassandra_instance = { \
	#	'nodes' : ip_addresses.values(), \
	#	'cli-path' : cassandra_home + '/bin/cassandra-cli', \
	#}

	#json_str = json.dumps(cassandra_instance)
	#with open(instance_json, 'w') as fjson:
	#	fjson.write(json_str)

	print '>'
	print '> ALL NODES ARE UP! TERMINATE THIS PROCESS TO SHUT DOWN SPARK CLUSTER.'
	while True:
		time.sleep(0.5)

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

parser = argparse.ArgumentParser(description='Run script for RAMCloud on FireBox-0 cluster.')
parser.add_argument('action', nargs=1, help='the action to perform (setup|start|stop)')

args = parser.parse_args()

print '> ================================================================================'
print '> RAMCLOUD RUN SCRIPT FOR FIREBOX-0 CLUSTER (VERSION ' + str(version) + ')'
print '> ================================================================================'
print '>'

git_rev = subprocess.check_output(['git', 'rev-parse', 'HEAD'])
print '> GIT revision: ' + git_rev.replace('\n','')
print '>'

print '>'

print '> COMMAND = ' + str(args.action)

if args.action[0] == 'setup':
	do_setup()
elif args.action[0] == 'start':
        check_slurm_allocation()
	do_start()
else:
	print '[ERROR] Unknown action \'' + args.action[0] + '\''
<commit_msg>Added 'start' and 'stop' options to the RAMCloud script<commit_after>#!/usr/bin/env python

# Usage of ramcloud.py script:
# Need to run within a salloc allocation.
#
# setup: Download and install RAMCloud into a shared directory.
# start: Start a RAMCloud cluster.
# stop: Stop the RAMCloud cluster.
#
# Assumes zookeeper has been deployed

# Contributors:
# Joao Carreira <joao@eecs.berkeley.edu> (2014)

import argparse
import getpass
import os
import urllib
import subprocess
import shutil
import re
import datetime
import time
import atexit
import json

run_timestamp = time.time()

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

version = '1.0'

user_dir = '/nscratch/' + getpass.getuser()
outdir = user_dir + '/ramcloud_bits'
ramcloud_dir = outdir + '/ramcloud'
workdir = os.getcwd() + '/work'
rundir = workdir + '/ramcloud-' +  datetime.datetime.fromtimestamp(run_timestamp). \
	strftime('%Y-%m-%d-%H-%M-%S')

network_if = 'eth0'

ramcloud_url = 'git://fiz.stanford.edu/git/ramcloud.git'

# paths used to set zookeeper dependencies
# XXX zookeeper version should be parameter
zookeeper_lib_path = user_dir +\
                     '/zookeeper_bits/zookeeper-3.4.6/src/c/.libs/libzookeeper_mt.a'
zookeeper_dir_path = user_dir + \
                     '/zookeeper_bits/zookeeper-3.4.6/src/c/include'

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

def make_dir(mydir):
	if not os.path.exists(mydir):
                print "Creating " + mydir
    		os.makedirs(mydir)
        else:
                print "Directory " + mydir + " already exists"

# Return the list of nodes in the current SLURM allocation
def get_slurm_nodelist():
	nodelist = subprocess.check_output( \
		['scontrol', 'show', 'hostname', os.environ['SLURM_NODELIST']], \
		universal_newlines=True)
	nodelist = nodelist.strip().split('\n')
	return nodelist

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

def do_setup():
	print '> Setting up RAMCloud in directory ' + outdir
	print '>'

	if not os.path.exists(ramcloud_dir):
                make_dir(ramcloud_dir)
                print '> Entering directory: ' + outdir
                print '> Cloning RAMCloud..'
                p = subprocess.Popen(['git', 'clone', ramcloud_url], cwd = outdir)
                p.wait()

                print '> Updating RAMCloud modules..'
                p = subprocess.Popen(['git', 'submodule', 'update', \
                        '--init', '--recursive'], cwd = ramcloud_dir)
                p.wait()

                print '> Compiling RAMCloud..'
                cmd_str = 'make -j 16 DEBUG=no ZOOKEEPER_LIB=' + \
                          zookeeper_lib_path + ' ZOOKEEPER_DIR=' + zookeeper_dir_path
                print cmd_str

                p =  subprocess.Popen(cmd_str, cwd = ramcloud_dir, shell=True)
                p.wait()

                print '> DONE'
	else:
                print '> RAMCloud was already setup. To perform setup' + \
                        'again please remove folder'
	print '>'
	print '> DONE'
	print '>'

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

ramcloud_instances = {}

def check_slurm_allocation():
        if (not 'SLURM_NODELIST' in os.environ) or (not os.environ['SLURM_NODELIST']):
            print '[ERROR] Need to run script within SLURM allocation'
            exit(1)

def shutdown_ramcloud_instances():
	print '> Shutting down RAMCloud instances'
	for c in ramcloud_instances.values():
                print "Terminating node " + c['node_id']
		c['process'].terminate()

	print '> Waiting for processes to terminate...'
	all_done = True
	while not all_done:
		all_done = True
		for c in ramcloud_instances.values():
			if (c['process'].poll() == None):
				all_done = False
		time.sleep(0.01)
	print '> DONE'

def wait_for_coordinator(coordinator_log):
	print '>'
	print '> Waiting for Coordinator to finish starting up...'
	while True:
                try:
		        with open(coordinator_log, 'r') as fout:
		    	        outdata = fout.read()
		        if re.search('ZooKeeper connection opened', outdata) != None:
			        break
                except:
                        print 'Error opening Coordinator log file'
		time.sleep(0.1)
	print '> COORDINATOR IS UP!'
	print '>'

# note: ulimit -l unlimited must be added to
# /etc/init.d/slurm-lln
# to avoid ramcloud complaining about limitations on the locked memory
def get_ramcloud_coordinator_cmd(coordinator_node, coordinator_log):
        srun_cmd = ['srun', '--nodelist=' + coordinator_node, '-N1']
        srun_cmd += ['obj.master/coordinator']
        srun_cmd += ['-C', 'infrc:host='+coordinator_node+',port=11100']
        srun_cmd += ['-x', 'zk:f16:2181'] #XXX make Zookeeper node a parameter
        srun_cmd += ['--logFile', coordinator_log]
        return srun_cmd

def get_ramcloud_server_cmd(node, coordinator_node):
	srun_cmd = ['srun', '--nodelist=' + node, '-N1']
        srun_cmd += ['obj.master/server', '-L', 'infrc:host='+node+',port=11100']
        srun_cmd += ['zk:f16:2181'] #XXX make Zookeeper node a parameter
        srun_cmd += ['-C', 'infrc:host='+coordinator_node+',port=11100']
        srun_cmd += ['--totalMasterMemory', '16000']
        srun_cmd += ['-f', '/tmp/ramcloud_data']
        srun_cmd += ['--segmentFrames', '10000','-r','2']
        return srun_cmd

def wait_for_ramcloud_servers(worker_nodes):
	print '>'
	print '> Waiting for all Servers to finish starting up...'
	unfinished_nodes = worker_nodes

	while unfinished_nodes:
		done_nodes = []
		for node in unfinished_nodes:
			with open(ramcloud_instances[node]['err'], 'r') as fout:
				outdata = fout.read()
			if re.search("Opened session with coordinator", outdata) != None:
				done_nodes.append(node)
		for node in done_nodes:
			unfinished_nodes.remove(node)
		time.sleep(0.01)

	print '> ALL SERVERS ARE UP!'

def do_start():
	print '> Launching RAMCloud cluster...'
	print '>'
	print '> Output redirected to ' + rundir
	print '>'

	# Create symlink for latest run
	subprocess.call(['ln', '-s', '-f', '-T', rundir, workdir + '/latest'])

	nodelist = get_slurm_nodelist()
	coordinator_node = nodelist[0]
	worker_nodes = nodelist[1:]

	print '> Coordinator Node: ' + coordinator_node
	print '> Worker Nodes: ' + ', '.join(worker_nodes)
	print '>'

        # Step I: Launch RAMCloud coordinator
	print '> Launching RAMCloud coordinator on node: ' + coordinator_node

        myrundir = rundir + '/coordinator-' + coordinator_node
        make_dir(myrundir)
        myoutfile = myrundir + '/stdout'
        myerrfile = myrundir + '/stderr'
        coordinator_log = myrundir + '/coordinator_log'

        fout = open(myoutfile, 'w')
        ferr = open(myerrfile, 'w')

        srun_cmd = get_ramcloud_coordinator_cmd(coordinator_node, coordinator_log)

        print "Launching coordinator.."
        print ' '.join(srun_cmd)
        p = subprocess.Popen(srun_cmd, stdout=fout, stderr=ferr, cwd=ramcloud_dir)
        ramcloud_instances[coordinator_node] = {'process': p, 'out': myoutfile, \
                    'err': myerrfile, 'type': 'coordinator', 'node_id' : coordinator_node}

	# When exiting, make sure all children are terminated cleanly
	atexit.register(shutdown_ramcloud_instances)

        wait_for_coordinator(coordinator_log)

	# Step II: Launch RAMCloud servers
	print '> Launching RAMCloud server nodes'
	print '>'

	for node in worker_nodes:
		print '> Launching RAMCloud server on ' + node

                srun_cmd = get_ramcloud_server_cmd(node, coordinator_node)

		myrundir = rundir + '/server-' + node
		make_dir(myrundir)
		myoutfile = myrundir + '/stdout'
		myerrfile = myrundir + '/stderr'

		fout = open(myoutfile, 'w')
		ferr = open(myerrfile, 'w')
		p = subprocess.Popen(srun_cmd, stdout=fout, stderr=ferr, cwd=ramcloud_dir)
		ramcloud_instances[node] = {'process': p, 'out': myoutfile, \
                    'err': myerrfile, 'type': 'worker', 'node_id' : node}

        wait_for_ramcloud_servers(worker_nodes)

	print '>'
	print '> ALL NODES ARE UP! TERMINATE THIS PROCESS TO SHUT DOWN RAMCLOUD CLUSTER.'
	while True:
		time.sleep(0.5)

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

parser = argparse.ArgumentParser(description='Run script for RAMCloud on FireBox-0 cluster.')
parser.add_argument('action', nargs=1, help='the action to perform (setup|start|stop)')

args = parser.parse_args()

print '> ================================================================================'
print '> RAMCLOUD RUN SCRIPT FOR FIREBOX-0 CLUSTER (VERSION ' + str(version) + ')'
print '> ================================================================================'
print '>'

git_rev = subprocess.check_output(['git', 'rev-parse', 'HEAD'])
print '> GIT revision: ' + git_rev.replace('\n','')
print '>'

print '>'

print '> COMMAND = ' + str(args.action)

if args.action[0] == 'setup':
	do_setup()
elif args.action[0] == 'start':
        check_slurm_allocation()
	do_start()
elif args.action[0] == 'stop':
	shutdown_ramcloud_instances()
else:
	print '[ERROR] Unknown action \'' + args.action[0] + '\''
<|endoftext|>


<commit_before>import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import PolyCollection
from matplotlib.patches import Rectangle
from matplotlib import colors, cm

from scarce import geometry

def get_mesh_plot(fig, mesh, values=None, invert_y_axis=True):
    ax = fig.add_subplot(111)

    vertexIDs = mesh._orderedCellVertexIDs
    vertexCoords = mesh.vertexCoords
    xCoords = np.take(vertexCoords[0], vertexIDs)
    yCoords = np.take(vertexCoords[1], vertexIDs)

    polys = []

    for x, y in zip(xCoords.swapaxes(0, 1), yCoords.swapaxes(0, 1)):
        if hasattr(x, 'mask'):
            x = x.compressed()
        if hasattr(y, 'mask'):
            y = y.compressed()
        polys.append(zip(x, y))

    collection = PolyCollection(polys)
    collection.set_facecolors('None')
    if invert_y_axis:
        ax.invert_yaxis()
    ax.add_collection(collection)

    if values:
        pass
        #     rgba = cm.cmap(plt.colors.norm(Z))
#         rgba = cm.get_cmap('coolwarm')(colors.norm(Z))
#         collection.set_facecolors(rgba)
#         collection.set_edgecolors(rgba)

    ax.plot()

def plot_mesh(mesh, values=None, invert_y_axis=True):
    fig = plt.figure()
    get_mesh_plot(fig, mesh, values=values, invert_y_axis=invert_y_axis)
    fig.gca().set_aspect('equal')
    plt.show()

def plot_planar_sensor(width,
                       pitch,
                       thickness,
                       n_pixel,
                       V_backplane,
                       V_readout,
                       potential_function=None,
                       field_function=None,
                       depletion_function=None,
                       mesh=None,
                       title=None):
    fig = plt.figure()
    get_planar_sensor_plot(fig,
                           width,
                           pitch,
                           thickness,
                           n_pixel,
                           V_backplane,
                           V_readout,
                           potential_function=potential_function,
                           field_function=field_function,
                           depletion_function=depletion_function,
                           mesh=mesh,
                           title=title)
    plt.show()

def get_planar_sensor_plot(fig,
                           width,
                           pitch,
                           thickness,
                           n_pixel,
                           V_backplane,
                           V_readout,
                           potential_function=None,
                           field_function=None,
                           depletion_function=None,
                           mesh=None,
                           title=None):

    ax = fig.add_subplot(111)

    min_x, max_x = -width * float(n_pixel) / 2., width * float(n_pixel) / 2.
    min_y, max_y = 0., thickness

    # Define plot space
    x = np.linspace(min_x, max_x, 1000)
    y = np.linspace(min_y, max_y, 1000)
#     y = np.square(np.linspace(np.sqrt(min_y), np.sqrt(max_y), 1000)) TODO: make this work with streamplot

    # Create x,y plot grid
    xx, yy = np.meshgrid(x, y, sparse=True)

    def get_depletion_mask(depletion_function, x, y):
        ''' Returns true for all points outside of the depletion zone.
            If x, y in a sparse array representation are handled correctly.
        '''

        if x.shape != y.shape:  # Shapes do not fit
            if x.shape != y.T.shape:   # Sparse meshgrid assumption
                raise RuntimeError('The point representation in x,y in neither a grid nor a sparse grid.')
            x_dense, y_dense = np.meshgrid(x[0, :], y[:, 0], sparse=False)
        else:
            x_dense, y_dense = x, y

        mask = np.zeros_like(x_dense, dtype=np.bool)
        mask[y_dense > depletion_function(x_dense)] = True

        return mask

    if depletion_function:
        ax.plot(x, depletion_function(x), '-', color='black', linewidth=4)
        ax.plot(x, depletion_function(x), '--', color='blue', linewidth=4, label='Depletion')

    if potential_function:
        # Plot Potential
        phi = potential_function(xx, yy)

        if depletion_function:  # Mask potential in not depleted region, otherwise contour plot goes crazy
            phi = np.ma.masked_array(phi, mask=get_depletion_mask(depletion_function, xx, yy))

        # BUG in matplotlib: aspect to be set to equal, otherwise contour plot wrong aspect ratio
        # http://stackoverflow.com/questions/28857673/wrong-aspect-ratio-for-contour-plot-with-python-matplotlib
        ax.set_aspect('equal')
        ax.contour(x, y, phi, 10, colors='black')
        cmesh = ax.pcolormesh(x, y, phi, cmap=cm.get_cmap('Blues'), vmin=V_backplane, vmax=V_readout)
        cax = fig.add_axes([ax.get_position().xmax, 0.1, 0.05, ax.get_position().ymax - ax.get_position().ymin])
        fig.colorbar(cmesh, cax=cax, orientation='vertical')

    # Plot E-Field
    if field_function:
        E_x, E_y = field_function(xx, yy)
        if depletion_function:  # Do not plot field in not depleted region, since it is 0
            E_x = np.ma.masked_array(E_x, mask=get_depletion_mask(depletion_function, xx, yy))
            E_y = np.ma.masked_array(E_y, mask=get_depletion_mask(depletion_function, xx, yy))

        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')
    elif potential_function:
        E_y, E_x = np.gradient(-phi, np.gradient(y), np.gradient(x))
        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')

    if mesh:
        get_mesh_plot(fig, mesh, invert_y_axis=False)

    # Plot backside
    ax.add_patch(Rectangle((min_x, max_y), (max_x - min_x), max_y, color="grey", linewidth=2))

    # Plot pixel(s)
    for pixel in range(n_pixel):
        pixel_position = width * (pixel + 1. / 2.) - width * n_pixel / 2.
        ax.add_patch(Rectangle((pixel_position - pitch / 2, ax.get_ylim()[0]), pitch, 0, color="darkred", linewidth=5))
        ax.plot([pixel_position - width / 2, pixel_position - width / 2], [min_y, max_y], '--', color='black', linewidth=4)
    ax.plot([pixel_position + width / 2, pixel_position + width / 2], [min_y, max_y], '--', color='black', linewidth=4)  # Last pixel end

    ax.set_ylim((- 0.02 * (ax.get_ylim()[1] - ax.get_ylim()[0]), 1.02 * max_y))
    ax.set_xlabel('Position x/y [um]', fontsize=22)
    ax.set_ylabel('Position z [um]', fontsize=22)
    if title:
        ax.set_title(title, fontsize=22)
    ax.invert_yaxis()

def plot_3D_sensor(width_x, width_y,
                   radius, nD,
                   n_pixel_x, n_pixel_y,
                   V_bias, V_readout,
                   potential_function=None,
                   field_function=None,
                   mesh=None,
                   title=None):
    fig = plt.figure()
    get_3D_sensor_plot(fig,
                       width_x, width_y,
                       radius, nD,
                       n_pixel_x, n_pixel_y,
                       V_bias, V_readout,
                       potential_function=potential_function,
                       field_function=field_function,
                       mesh=mesh,
                       title=title)
    plt.show()

def get_3D_sensor_plot(fig,
                       width_x, width_y,
                       radius, nD,
                       n_pixel_x, n_pixel_y,
                       V_bias, V_readout,
                       potential_function=None,
                       field_function=None,
                       mesh=None,
                       title=None):

    ax = fig.add_subplot(111)

    desc = geometry.SensorDescription3D(width_x, width_y, n_pixel_x, n_pixel_y, radius, nD)

    min_x, max_x, min_y, max_y = desc.get_array_corners()

    # Define plot space
    x = np.linspace(min_x, max_x, 1000)
    y = np.linspace(min_y, max_y, 1000)

    # Create x,y plot grid
    xx, yy = np.meshgrid(x, y, sparse=True)

    # BUG in matplotlib: aspect to be set to equal, otherwise contour plot wrong aspect ratio
    # http://stackoverflow.com/questions/28857673/wrong-aspect-ratio-for-contour-plot-with-python-matplotlib
    ax.set_aspect('equal')

    if potential_function:
        # Plot Potential
        phi = potential_function(xx, yy)
        ax.contour(x, y, phi, 10, colors='black')
        cmesh = ax.pcolormesh(x, y, phi, cmap=cm.get_cmap('Blues'), vmin=V_bias, vmax=V_readout)
        cax = fig.add_axes([ax.get_position().xmax, 0.1, 0.05, ax.get_position().ymax - ax.get_position().ymin])
        fig.colorbar(cmesh, cax=cax, orientation='vertical')

    # Plot E-Field
    if field_function:
        E_x, E_y = field_function(xx, yy)
        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')
    elif potential_function:
        E_y, E_x = np.gradient(-phi, np.gradient(y), np.gradient(x))
        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')

    if mesh:
        get_mesh_plot(fig, mesh, invert_y_axis=False)

    # Plot pixel bondaries in x
    for pos_x in desc.get_pixel_x_offsets():
        ax.plot([pos_x - width_x / 2., pos_x - width_x / 2.], [min_y, max_y], '--', color='black', linewidth=4)
    ax.plot([pos_x + width_x / 2., pos_x + width_x / 2.], [min_y, max_y], '--', color='black', linewidth=4)  # Last pixel end

    # Plot pixel bondaries in y
    for pos_y in desc.get_pixel_y_offsets():
        ax.plot([min_x, max_x], [pos_y - width_y / 2., pos_y - width_y / 2.], '--', color='black', linewidth=4)
    ax.plot([min_x, max_x], [pos_y + width_y / 2., pos_y + width_y / 2.], '--', color='black', linewidth=4)  # Last pixel end

    # Plot readout pillars
    for pos_x, pos_y in desc.get_ro_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkred", linewidth=0))

    # Plot full bias pillars
    for pos_x, pos_y in desc.get_center_bias_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkblue", linewidth=0))

    # Plot side bias pillars
    for pos_x, pos_y in desc.get_side_bias_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkblue", linewidth=0))

    # Plot edge bias pillars
    for pos_x, pos_y in desc.get_edge_bias_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkblue", linewidth=0))

    ax.set_xlim((1.05 * min_x, 1.05 * max_x))
    ax.set_ylim((1.05 * min_y, 1.05 * max_y))
    ax.set_xlabel('Position x [um]', fontsize=22)
    ax.set_ylabel('Position y [um]', fontsize=22)
    if title:
        ax.set_title(title, fontsize=22)

if __name__ == '__main__':
    import meshio as mio
    import fipy
    width_x, width_y = 250., 50.
    n_pixel_x, n_pixel_y = 3, 3
    radius, nD = 6., 2
    resolution = 10

    mesh = geometry.mesh_3D_sensor(width_x=width_x,
                                   width_y=width_y,
                                   n_pixel_x=n_pixel_x,
                                   n_pixel_y=n_pixel_y,
                                   radius=radius,
                                   nD=nD,
                                   resolution=resolution)

    plot_3D_sensor(width_x=width_x,
                   width_y=width_y,
                   radius=radius,
                   nD=nD,
                   n_pixel_x=n_pixel_x,
                   n_pixel_y=n_pixel_y,
                   V_bias=0,
                   V_readout=1,
                   potential_function=None,
                   field_function=None,
                   mesh=mesh,
                   title=None)

#     from scarce import fields
#
#     thickness = 200
#     width = 50
#     pitch = 50
#     n_pixel = 5
#     thickness = 200
#     resolution = 50
#     V_backplane, V_readout = -1, 0
#
#     def potential_function(x, y):
#         return fields.get_weighting_potential_analytic(x, y, D=thickness, S=width, is_planar=True)
#
#     def field_function(x, y):
#         return fields.get_weighting_field_analytic(x, y, D=thickness, S=width, is_planar=True)
#
#     plot_planar_sensor(width=width,
#                        pitch=width,
#                        thickness=thickness,
#                        n_pixel=n_pixel,
#                        V_backplane=0,
#                        V_readout=1,
#                        potential_function=potential_function,
#                        field_function=field_function,
#                        mesh=None)
<commit_msg>ENH: correct 2d potential position (shift bin by binsize / 2) ENH: mask potential in columns for contour plot<commit_after>import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import PolyCollection
from matplotlib.patches import Rectangle
from matplotlib import colors, cm

from scarce import geometry

def get_mesh_plot(fig, mesh, values=None, invert_y_axis=True):
    ax = fig.add_subplot(111)

    vertexIDs = mesh._orderedCellVertexIDs
    vertexCoords = mesh.vertexCoords
    xCoords = np.take(vertexCoords[0], vertexIDs)
    yCoords = np.take(vertexCoords[1], vertexIDs)

    polys = []

    for x, y in zip(xCoords.swapaxes(0, 1), yCoords.swapaxes(0, 1)):
        if hasattr(x, 'mask'):
            x = x.compressed()
        if hasattr(y, 'mask'):
            y = y.compressed()
        polys.append(zip(x, y))

    collection = PolyCollection(polys)
    collection.set_facecolors('None')
    if invert_y_axis:
        ax.invert_yaxis()
    ax.add_collection(collection)

    if values:
        pass
        #     rgba = cm.cmap(plt.colors.norm(Z))
#         rgba = cm.get_cmap('coolwarm')(colors.norm(Z))
#         collection.set_facecolors(rgba)
#         collection.set_edgecolors(rgba)

    ax.plot()

def plot_mesh(mesh, values=None, invert_y_axis=True):
    fig = plt.figure()
    get_mesh_plot(fig, mesh, values=values, invert_y_axis=invert_y_axis)
    fig.gca().set_aspect('equal')
    plt.show()

def plot_planar_sensor(width,
                       pitch,
                       thickness,
                       n_pixel,
                       V_backplane,
                       V_readout,
                       potential_function=None,
                       field_function=None,
                       depletion_function=None,
                       mesh=None,
                       title=None):
    fig = plt.figure()
    get_planar_sensor_plot(fig,
                           width,
                           pitch,
                           thickness,
                           n_pixel,
                           V_backplane,
                           V_readout,
                           potential_function=potential_function,
                           field_function=field_function,
                           depletion_function=depletion_function,
                           mesh=mesh,
                           title=title)
    plt.show()

def get_planar_sensor_plot(fig,
                           width,
                           pitch,
                           thickness,
                           n_pixel,
                           V_backplane,
                           V_readout,
                           potential_function=None,
                           field_function=None,
                           depletion_function=None,
                           mesh=None,
                           title=None):

    ax = fig.add_subplot(111)

    min_x, max_x = -width * float(n_pixel) / 2., width * float(n_pixel) / 2.
    min_y, max_y = 0., thickness

    # Define plot space
    x = np.linspace(min_x, max_x, 1000)
    y = np.linspace(min_y, max_y, 1000)
#     y = np.square(np.linspace(np.sqrt(min_y), np.sqrt(max_y), 1000)) TODO: make this work with streamplot

    # Create x,y plot grid
    xx, yy = np.meshgrid(x, y, sparse=True)

    def get_depletion_mask(depletion_function, x, y):
        ''' Returns true for all points outside of the depletion zone.
            If x, y in a sparse array representation are handled correctly.
        '''

        if x.shape != y.shape:  # Shapes do not fit
            if x.shape != y.T.shape:   # Sparse meshgrid assumption
                raise RuntimeError('The point representation in x,y in neither a grid nor a sparse grid.')
            x_dense, y_dense = np.meshgrid(x[0, :], y[:, 0], sparse=False)
        else:
            x_dense, y_dense = x, y

        mask = np.zeros_like(x_dense, dtype=np.bool)
        mask[y_dense > depletion_function(x_dense)] = True

        return mask

    if depletion_function:
        ax.plot(x, depletion_function(x), '-', color='black', linewidth=4)
        ax.plot(x, depletion_function(x), '--', color='blue', linewidth=4, label='Depletion')

    if potential_function:
        # Plot Potential
        phi = potential_function(xx, yy)

        if depletion_function:  # Mask potential in not depleted region, otherwise contour plot goes crazy
            phi = np.ma.masked_array(phi, mask=get_depletion_mask(depletion_function, xx, yy))

        # BUG in matplotlib: aspect to be set to equal, otherwise contour plot wrong aspect ratio
        # http://stackoverflow.com/questions/28857673/wrong-aspect-ratio-for-contour-plot-with-python-matplotlib
        ax.set_aspect('equal')
        ax.contour(x, y, phi, 10, colors='black')
        cmesh = ax.pcolormesh(x - np.diff(x)[0] / 2., y - np.diff(y)[0] / 2., phi, cmap=cm.get_cmap('Blues'), vmin=V_backplane, vmax=V_readout)
        cax = fig.add_axes([ax.get_position().xmax, 0.1, 0.05, ax.get_position().ymax - ax.get_position().ymin])
        fig.colorbar(cmesh, cax=cax, orientation='vertical')

    # Plot E-Field
    if field_function:
        E_x, E_y = field_function(xx, yy)
        if depletion_function:  # Do not plot field in not depleted region, since it is 0
            E_x = np.ma.masked_array(E_x, mask=get_depletion_mask(depletion_function, xx, yy))
            E_y = np.ma.masked_array(E_y, mask=get_depletion_mask(depletion_function, xx, yy))

        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')
    elif potential_function:
        E_y, E_x = np.gradient(-phi, np.gradient(y), np.gradient(x))
        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')

    if mesh:
        get_mesh_plot(fig, mesh, invert_y_axis=False)

    # Plot backside
    ax.add_patch(Rectangle((min_x, max_y), (max_x - min_x), max_y, color="grey", linewidth=2))

    # Plot pixel(s)
    for pixel in range(n_pixel):
        pixel_position = width * (pixel + 1. / 2.) - width * n_pixel / 2.
        ax.add_patch(Rectangle((pixel_position - pitch / 2, ax.get_ylim()[0]), pitch, 0, color="darkred", linewidth=5))
        ax.plot([pixel_position - width / 2, pixel_position - width / 2], [min_y, max_y], '--', color='black', linewidth=4)
    ax.plot([pixel_position + width / 2, pixel_position + width / 2], [min_y, max_y], '--', color='black', linewidth=4)  # Last pixel end

    ax.set_ylim((- 0.02 * (ax.get_ylim()[1] - ax.get_ylim()[0]), 1.02 * max_y))
    ax.set_xlabel('Position x/y [um]', fontsize=22)
    ax.set_ylabel('Position z [um]', fontsize=22)
    if title:
        ax.set_title(title, fontsize=22)
    ax.invert_yaxis()

def plot_3D_sensor(width_x, width_y,
                   radius, nD,
                   n_pixel_x, n_pixel_y,
                   V_bias, V_readout,
                   potential_function=None,
                   field_function=None,
                   mesh=None,
                   title=None):
    fig = plt.figure()
    get_3D_sensor_plot(fig,
                       width_x, width_y,
                       radius, nD,
                       n_pixel_x, n_pixel_y,
                       V_bias, V_readout,
                       potential_function=potential_function,
                       field_function=field_function,
                       mesh=mesh,
                       title=title)
    plt.show()

def get_3D_sensor_plot(fig,
                       width_x, width_y,
                       radius, nD,
                       n_pixel_x, n_pixel_y,
                       V_bias, V_readout,
                       potential_function=None,
                       field_function=None,
                       mesh=None,
                       title=None):

    ax = fig.add_subplot(111)

    desc = geometry.SensorDescription3D(width_x, width_y, n_pixel_x, n_pixel_y, radius, nD)

    min_x, max_x, min_y, max_y = desc.get_array_corners()

    # Define plot space
    x = np.linspace(min_x, max_x, max(width_x * n_pixel_x, width_y * n_pixel_y))  # >= 0.5 um resolution
    y = np.linspace(min_y, max_y, max(width_x * n_pixel_x, width_y * n_pixel_y))  # >= 0.5 um resolution

    # Create x,y plot grid
    xx, yy = np.meshgrid(x, y, sparse=True)

    # BUG in matplotlib: aspect to be set to equal, otherwise contour plot wrong aspect ratio
    # http://stackoverflow.com/questions/28857673/wrong-aspect-ratio-for-contour-plot-with-python-matplotlib
    ax.set_aspect('equal')

    def get_column_mask(x, y):
        ''' Returns true for points inside a column.
            X, y as a sparse array representation are handled correctly.
        '''

        if x.shape != y.shape:  # Shapes do not fit
            if x.shape != y.T.shape:   # Sparse meshgrid assumption
                raise RuntimeError('The point representation in x,y in neither a grid nor a sparse grid.')
            x_dense, y_dense = np.meshgrid(x[0, :], y[:, 0], sparse=False)
        else:
            x_dense, y_dense = x, y
        # Reduce radius to prevent aliasing at round column edges
        return desc.position_in_column(x_dense, y_dense)

    if potential_function:
        # Plot Potential
        phi = potential_function(xx, yy)

        # Mask potential in columns, otherwise contour plot goes crazy
        phi_masked = np.ma.masked_array(phi, mask=get_column_mask(xx, yy))

        ax.contour(x, y, phi_masked, 10, colors='black')
        cmesh = ax.pcolormesh(x - np.diff(x)[0] / 2., y - np.diff(y)[0] / 2., phi, cmap=cm.get_cmap('Blues'), vmin=V_bias, vmax=V_readout)
        cax = fig.add_axes([ax.get_position().xmax, 0.1, 0.05, ax.get_position().ymax - ax.get_position().ymin])
        fig.colorbar(cmesh, cax=cax, orientation='vertical')

    # Plot E-Field
    if field_function:
        E_x, E_y = field_function(xx, yy)
        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')
    elif potential_function:
        E_y, E_x = np.gradient(-phi, np.gradient(y), np.gradient(x))
        ax.streamplot(x, y, E_x, E_y, density=1.0, color='gray', arrowstyle='-')

    if mesh:
        get_mesh_plot(fig, mesh, invert_y_axis=False)

    # Plot pixel bondaries in x
    for pos_x in desc.get_pixel_x_offsets():
        ax.plot([pos_x - width_x / 2., pos_x - width_x / 2.], [min_y, max_y], '--', color='black', linewidth=4)
    ax.plot([pos_x + width_x / 2., pos_x + width_x / 2.], [min_y, max_y], '--', color='black', linewidth=4)  # Last pixel end

    # Plot pixel bondaries in y
    for pos_y in desc.get_pixel_y_offsets():
        ax.plot([min_x, max_x], [pos_y - width_y / 2., pos_y - width_y / 2.], '--', color='black', linewidth=4)
    ax.plot([min_x, max_x], [pos_y + width_y / 2., pos_y + width_y / 2.], '--', color='black', linewidth=4)  # Last pixel end

    # Plot readout pillars
    for pos_x, pos_y in desc.get_ro_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkred", linewidth=0))

    # Plot full bias pillars
    for pos_x, pos_y in desc.get_center_bias_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkblue", linewidth=0))

    # Plot side bias pillars
    for pos_x, pos_y in desc.get_side_bias_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkblue", linewidth=0))

    # Plot edge bias pillars
    for pos_x, pos_y in desc.get_edge_bias_col_offsets():
        ax.add_patch(plt.Circle((pos_x, pos_y), radius, color="darkblue", linewidth=0))

    ax.set_xlim((1.05 * min_x, 1.05 * max_x))
    ax.set_ylim((1.05 * min_y, 1.05 * max_y))
    ax.set_xlabel('Position x [um]', fontsize=22)
    ax.set_ylabel('Position y [um]', fontsize=22)
    if title:
        ax.set_title(title, fontsize=22)

if __name__ == '__main__':
    import meshio as mio
    import fipy
    width_x, width_y = 250., 50.
    n_pixel_x, n_pixel_y = 3, 3
    radius, nD = 6., 2
    resolution = 10

    mesh = geometry.mesh_3D_sensor(width_x=width_x,
                                   width_y=width_y,
                                   n_pixel_x=n_pixel_x,
                                   n_pixel_y=n_pixel_y,
                                   radius=radius,
                                   nD=nD,
                                   resolution=resolution)

    plot_3D_sensor(width_x=width_x,
                   width_y=width_y,
                   radius=radius,
                   nD=nD,
                   n_pixel_x=n_pixel_x,
                   n_pixel_y=n_pixel_y,
                   V_bias=0,
                   V_readout=1,
                   potential_function=None,
                   field_function=None,
                   mesh=mesh,
                   title=None)

#     from scarce import fields
#
#     thickness = 200
#     width = 50
#     pitch = 50
#     n_pixel = 5
#     thickness = 200
#     resolution = 50
#     V_backplane, V_readout = -1, 0
#
#     def potential_function(x, y):
#         return fields.get_weighting_potential_analytic(x, y, D=thickness, S=width, is_planar=True)
#
#     def field_function(x, y):
#         return fields.get_weighting_field_analytic(x, y, D=thickness, S=width, is_planar=True)
#
#     plot_planar_sensor(width=width,
#                        pitch=width,
#                        thickness=thickness,
#                        n_pixel=n_pixel,
#                        V_backplane=0,
#                        V_readout=1,
#                        potential_function=potential_function,
#                        field_function=field_function,
#                        mesh=None)
<|endoftext|>


<commit_before>from feature_format import featureFormat, targetFeatureSplit
import pickle
from sklearn.ensemble import RandomForestClassifier
from sklearn.cross_validation import KFold
from sklearn.metrics import precision_score, recall_score, f1_score
import numpy as np

# loading the enron data dictionary
with open("final_project_dataset.pkl", "r") as data_file:
    data_dict = pickle.load(data_file)

# removing 'TOTAL' outlier
del data_dict['TOTAL']

# function to train and predict POIs based on the features provided as input
def nbclassifier(feat_list):
    # creating list of labels and list of numpy arrays containing the features
    data = featureFormat(data_dict, feat_list, sort_keys = True)
    labels, features = targetFeatureSplit(data)

    # Fitting and testing Gaussian Naive Bayes Classifier
    accuracy = []
    precision = []
    recall = []
    f1 = []

    kf = KFold(len(data), n_folds=4, random_state=1)

    for train_indices, test_indices in kf:
        features_train = [features[i] for i in train_indices]
        features_test = [features[j] for j in test_indices]
        labels_train = [labels[i] for i in train_indices]
        labels_test = [labels[j] for j in test_indices]

        clf = RandomForestClassifier(n_estimators=10, random_state=10)
        clf.fit(features_train, labels_train)
        pred = clf.predict(features_test)

        accuracy.append(clf.score(features_test, labels_test))
        precision.append(precision_score(labels_test, pred))
        recall.append(recall_score(labels_test, pred))
        f1.append(f1_score(labels_test, pred))

    print "\n\nNB Classifier Result with features - "
    for f in feat_list[1:]:
        print f
    print "\naccuracy:", np.mean(accuracy)
    print "precision:", np.mean(precision)
    print "recall:", np.mean(recall)
    print "f1 score:", np.mean(f1)

# selecting only 2 features - total_stock_value and bonus for now
# total_stock_value - data available for all POIs and second best feature
# bonus - data available for 16 out of 18 POIs and third best feature
nbclassifier(['poi',
              'total_stock_value',
              'bonus'])

# changing features list to replace total_stock_value with
# exercised_stock_option
nbclassifier(['poi',
              'exercised_stock_options',
              'bonus'])
'''
Provided one of the best results with the highest precision and f1 scores
'''

# Adding saalary to the features list
nbclassifier(['poi',
              'exercised_stock_options',
              'bonus',
              'salary'])

# Removing everything except exercised stock options
nbclassifier(['poi',
              'exercised_stock_options'])

# Replacing exercised stock options with bonus
nbclassifier(['poi',
              'bonus'])
'''
Provided one of the best results with the precision and recall above 0.3
'''

# Replacing exercised stock options with bonus
nbclassifier(['poi',
              'salary'])<commit_msg>feat: Create new feature for RandomForestClassifier<commit_after>from feature_format import featureFormat, targetFeatureSplit
import pickle
from sklearn.ensemble import RandomForestClassifier
from sklearn.cross_validation import KFold
from sklearn.metrics import precision_score, recall_score, f1_score
import numpy as np

# loading the enron data dictionary
with open("final_project_dataset.pkl", "r") as data_file:
    data_dict = pickle.load(data_file)

# removing 'TOTAL' outlier
del data_dict['TOTAL']

# Adding new feature = ratio of from this person to POI and from messages
for name in data_dict:
    if data_dict[name]["from_this_person_to_poi"] != "NaN" and\
    data_dict[name]["from_messages"] != "NaN" and\
    data_dict[name]["from_messages"] != 0.0:
        data_dict[name]["toPOI_fromMsgs"] = \
        data_dict[name]["from_this_person_to_poi"] * 1.0\
        /data_dict[name]["from_messages"]

# function to train and predict POIs based on the features provided as input
def nbclassifier(feat_list):
    # creating list of labels and list of numpy arrays containing the features
    data = featureFormat(data_dict, feat_list, sort_keys = True)
    labels, features = targetFeatureSplit(data)

    # Fitting and testing Gaussian Naive Bayes Classifier
    accuracy = []
    precision = []
    recall = []
    f1 = []

    kf = KFold(len(data), n_folds=4, random_state=1)

    for train_indices, test_indices in kf:
        features_train = [features[i] for i in train_indices]
        features_test = [features[j] for j in test_indices]
        labels_train = [labels[i] for i in train_indices]
        labels_test = [labels[j] for j in test_indices]

        clf = RandomForestClassifier(n_estimators=10, random_state=10)
        clf.fit(features_train, labels_train)
        pred = clf.predict(features_test)

        accuracy.append(clf.score(features_test, labels_test))
        precision.append(precision_score(labels_test, pred))
        recall.append(recall_score(labels_test, pred))
        f1.append(f1_score(labels_test, pred))

    print "\n\nNB Classifier Result with features - "
    for f in feat_list[1:]:
        print f
    print "\naccuracy:", np.mean(accuracy)
    print "precision:", np.mean(precision)
    print "recall:", np.mean(recall)
    print "f1 score:", np.mean(f1)

# selecting only 2 features - total_stock_value and bonus for now
# total_stock_value - data available for all POIs and second best feature
# bonus - data available for 16 out of 18 POIs and third best feature
nbclassifier(['poi',
              'total_stock_value',
              'bonus'])

# changing features list to replace total_stock_value with
# exercised_stock_option
nbclassifier(['poi',
              'exercised_stock_options',
              'bonus'])
'''
Provided one of the best results with the highest precision and f1 scores
'''

# Adding saalary to the features list
nbclassifier(['poi',
              'exercised_stock_options',
              'bonus',
              'salary'])

# Removing everything except exercised stock options
nbclassifier(['poi',
              'exercised_stock_options'])

# Replacing exercised stock options with bonus
nbclassifier(['poi',
              'bonus'])
'''
Provided one of the best results with the precision and recall above 0.3
'''

# Replacing exercised stock options with bonus
nbclassifier(['poi',
              'salary'])<|endoftext|>


<commit_before>import unittest
import os
import shutil
import sys
import time
import subprocess
import json
import xml.etree.ElementTree as ET
import igorVar
import igorSetup
import igorCA

DEBUG_TEST=False
if DEBUG_TEST:
    igorVar.VERBOSE=True

FIXTURES=os.path.join(os.path.abspath(os.path.dirname(__file__)), 'fixtures')

class IgorTest(unittest.TestCase):
    igorDir = os.path.join(FIXTURES, 'testIgor')
    igorLogFile = os.path.join(FIXTURES, 'testIgor.log')
    igorPort = 19333
    igorProtocol = "http"
    igorVarArgs = {}

    @classmethod
    def setUpClass(cls):
        super(IgorTest, cls).setUpClass()

        if DEBUG_TEST: print 'IgorTest: Delete old database and logfile'
        shutil.rmtree(cls.igorDir, True)
        shutil.rmtree(cls.igorLogFile, True)

        cls.igorUrl = "%s://localhost:%d/data/" % (cls.igorProtocol, cls.igorPort)

        if DEBUG_TEST: print 'IgorTest: Setup database'
        setup = igorSetup.IgorSetup(database=cls.igorDir)
        ok = setup.cmd_initialize()
        assert ok
        setup.postprocess(run=True)

        if cls.igorProtocol == 'https':
            if DEBUG_TEST: print 'IgorTest: setup self-signed signature'
            ok = setup.cmd_certificateSelfsigned('/C=NL/O=igor/CN=localhost', 'localhost', '127.0.0.1', '::1')
            assert ok
            setup.postprocess(run=True)
            cls.igorVarArgs['certificate'] = os.path.join(cls.igorDir, 'igor.crt')
            cls.igorVarArgs['noverify'] = True

        if DEBUG_TEST: print 'IgorTest: Check database consistency'
        subprocess.check_call([sys.executable, "-m", "igor", "--check", "--database", cls.igorDir, "--port", str(cls.igorPort)], stdout=open(cls.igorLogFile, 'a'), stderr=subprocess.STDOUT)

        if DEBUG_TEST: print 'IgorTest: Start server'
        cls.igorProcess = subprocess.Popen([sys.executable, "-u", "-m", "igor", "--database", cls.igorDir, "--port", str(cls.igorPort)], stdout=open(cls.igorLogFile, 'a'), stderr=subprocess.STDOUT)
        time.sleep(2)

    @classmethod
    def tearDownClass(cls):
        # Gracefully stop server
        if DEBUG_TEST: print 'IgorTest: Request server to stop'
        try:
            p = igorVar.IgorServer(cls.igorUrl, **cls.igorVarArgs)
            result = p.get('/internal/stop')
        except:
            if DEBUG_TEST: print 'IgorTest: Ignoring exception during stop request'
        time.sleep(2)

        sts = cls.igorProcess.poll()
        if sts is None:
            if DEBUG_TEST: print 'IgorTest: Terminate server'
            cls.igorProcess.terminate()
            time.sleep(2)
        sts = cls.igorProcess.wait()
        assert sts != None

        super(IgorTest, cls).tearDownClass()
        pass

    def _igorVar(self):
        return igorVar.IgorServer(self.igorUrl, **self.igorVarArgs)

    def test01_get_static(self):
        p = self._igorVar()
        result = p.get('/')
        self.assertTrue(result)
        self.assertEqual(result[0], "<")

    def notyet_test01_get_static_nonexistent(self):
        p = self._igorVar()
        result = p.get('/nonexistent.html')
        self.assertTrue(result)
        self.assertEqual(result[0], "<")

    def test11_get_xml(self):
        p = self._igorVar()
        result = p.get('environment/systemHealth', format='application/xml')
        self.assertTrue(result)
        root = ET.fromstring(result)
        self.assertEqual(root.tag, "systemHealth")

    def test12_get_text(self):
        p = self._igorVar()
        result = p.get('environment/systemHealth', format='text/plain')
        self.assertTrue(result)

    def test13_get_json(self):
        p = self._igorVar()
        result = p.get('environment/systemHealth', format='application/json')
        self.assertTrue(result)
        root = json.loads(result)
        self.assertIsInstance(root, dict)
        self.assertEqual(root.keys(), ["systemHealth"])

    def test21_put_xml(self):
        p = self._igorVar()
        data = '<test21>21</test21>'
        p.put('sandbox/test21', data, datatype='application/xml')
        result = p.get('sandbox/test21', format='application/xml')
        self.assertEqual(data.strip(), result.strip())
        result2 = p.get('sandbox/test21', format='text/plain')
        self.assertEqual('21', result2.strip())
        result3 = p.get('sandbox/test21', format='application/json')
        result3dict = json.loads(result3)
        self.assertEqual({"test21" : 21}, result3dict)

    def test22_put_text(self):
        p = self._igorVar()
        data = 'twenty two'
        p.put('sandbox/test22', data, datatype='text/plain')
        result = p.get('sandbox/test22', format='text/plain')
        self.assertEqual(data.strip(), result.strip())
        result2 = p.get('sandbox/test22', format='application/xml')
        self.assertEqual('<test22>twenty two</test22>', result2.strip())
        result3 = p.get('sandbox/test22', format='application/json')
        result3dict = json.loads(result3)
        self.assertEqual({'test22':'twenty two'}, result3dict)

    def test23_put_json(self):
        p = self._igorVar()
        data = json.dumps({"test23" : 23})
        p.put('sandbox/test23', data, datatype='application/json')
        result = p.get('sandbox/test23', format='application/json')
        resultDict = json.loads(result)
        self.assertEqual({"test23" : 23}, resultDict)
        result2 = p.get('sandbox/test23', format='application/xml')
        self.assertEqual("<test23>23</test23>", result2.strip())

    def test24_put_multi(self):
        p = self._igorVar()
        data = '<test24>24</test24>'
        p.put('sandbox/test24', data, datatype='application/xml')
        result = p.get('sandbox/test24', format='application/xml')
        self.assertEqual(data.strip(), result.strip())
        data = '<test24>twentyfour</test24>'
        p.put('sandbox/test24', data, datatype='application/xml')
        result = p.get('sandbox/test24', format='application/xml')
        self.assertEqual(data.strip(), result.strip())

    def test31_post_text(self):
        p = self._igorVar()
        p.put('sandbox/test31', '', datatype='text/plain')
        p.post('sandbox/test31/item', 'thirty', datatype='text/plain')
        p.post('sandbox/test31/item', 'one', datatype='text/plain')
        result = p.get('sandbox/test31/item', format='text/plain')
        self.assertEqual('thirtyone', result.translate(None, ' \n'))

        self.assertRaises(igorVar.IgorError, p.get, 'sandbox/test31/item', format='application/xml')

        result2 = p.get('sandbox/test31/item', format='application/xml', variant='multi')
        self.assertIn('thirty', result2)
        self.assertIn('one', result2)

        result3 = p.get('sandbox/test31/item', format='application/json', variant='multi')
        result3list = json.loads(result3)
        self.assertEqual(len(result3list), 2)
        self.assertIsInstance(result3list, list)
        self.assertEqual(result3list[0]['item'], 'thirty')
        self.assertEqual(result3list[1]['item'], 'one')

    def test41_action(self):
        p = self._igorVar()
        content = {'test41':{'src':'', 'sink':''}}
        action = {'action':dict(name='test41action', url='/data/sandbox/test41/sink', xpath='/data/sandbox/test41/src', method='PUT', data='copy-{.}-copy')}
        p.put('sandbox/test41', json.dumps(content), datatype='application/json')
        p.post('actions/action', json.dumps(action), datatype='application/json')
        p.put('sandbox/test41/src', 'forty-one', datatype='text/plain')

        time.sleep(1)

        result = p.get('sandbox/test41', format='application/json')
        resultDict = json.loads(result)
        wantedContent = {'test41':{'src':'forty-one', 'sink':'copy-forty-one-copy'}}
        self.assertEqual(resultDict, wantedContent)

    def test42_action_post(self):
        p = self._igorVar()
        content = {'test42':{'src':''}}
        action = {'action':dict(name='test42action', url='/data/sandbox/test42/sink', xpath='/data/sandbox/test42/src', method='POST', data='{.}')}
        p.put('sandbox/test42', json.dumps(content), datatype='application/json')
        p.post('actions/action', json.dumps(action), datatype='application/json')
        time.sleep(1)
        p.put('sandbox/test42/src', '42a', datatype='text/plain')
        time.sleep(1)
        p.put('sandbox/test42/src', '42b', datatype='text/plain')
        time.sleep(1)
        p.put('sandbox/test42/src', '42c', datatype='text/plain')

        time.sleep(1)

        result = p.get('sandbox/test42', format='application/json')
        resultDict = json.loads(result)
        wantedContent = {'test42':{'src':'42c', 'sink':['42a','42b','42c']}}
        self.assertEqual(resultDict, wantedContent)

    def test43_action_indirect(self):
        p = self._igorVar()
        content = {'test43':{'src':'', 'sink':''}}
        action1 = {'action':dict(name='test43first', url='/action/test43second', xpath='/data/sandbox/test43/src')}
        action2 = {'action':dict(name='test43second', url='/data/sandbox/test43/sink', method='PUT', data='copy-{/data/sandbox/test43/src}-copy')}
        p.put('sandbox/test43', json.dumps(content), datatype='application/json')
        p.post('actions/action', json.dumps(action1), datatype='application/json')
        p.post('actions/action', json.dumps(action2), datatype='application/json')
        p.put('sandbox/test43/src', 'forty-three', datatype='text/plain')

        time.sleep(1)

        result = p.get('sandbox/test43', format='application/json')
        resultDict = json.loads(result)
        wantedContent = {'test43':{'src':'forty-three', 'sink':'copy-forty-three-copy'}}
        self.assertEqual(resultDict, wantedContent)

class IgorTestHttps(IgorTest):
    igorDir = os.path.join(FIXTURES, 'testIgorHttps')
    igorLogFile = os.path.join(FIXTURES, 'testIgorHttps.log')
    igorPort = 29333
    igorProtocol = "https"

if __name__ == '__main__':
    unittest.main()

<commit_msg>Added test for external actions (actually igor talking to itself over http/s)<commit_after>import unittest
import os
import shutil
import sys
import time
import subprocess
import json
import socket
import xml.etree.ElementTree as ET
import igorVar
import igorSetup
import igorCA

DEBUG_TEST=False
if DEBUG_TEST:
    igorVar.VERBOSE=True

FIXTURES=os.path.join(os.path.abspath(os.path.dirname(__file__)), 'fixtures')

class IgorTest(unittest.TestCase):
    igorDir = os.path.join(FIXTURES, 'testIgor')
    igorLogFile = os.path.join(FIXTURES, 'testIgor.log')
    igorPort = 19333
    igorProtocol = "http"
    igorVarArgs = {}

    @classmethod
    def setUpClass(cls):
        super(IgorTest, cls).setUpClass()

        if DEBUG_TEST: print 'IgorTest: Delete old database and logfile'
        shutil.rmtree(cls.igorDir, True)
        if os.path.exists(cls.igorLogFile):
            os.unlink(cls.igorLogFile)

        cls.igorUrl = "%s://localhost:%d/data/" % (cls.igorProtocol, cls.igorPort)

        if DEBUG_TEST: print 'IgorTest: Setup database'
        setup = igorSetup.IgorSetup(database=cls.igorDir)
        ok = setup.cmd_initialize()
        assert ok
        setup.postprocess(run=True)

        if cls.igorProtocol == 'https':
            if DEBUG_TEST: print 'IgorTest: setup self-signed signature'
            ok = setup.cmd_certificateSelfsigned('/C=NL/O=igor/CN=localhost', 'localhost', socket.gethostname(), '127.0.0.1', '::1')
            assert ok
            setup.postprocess(run=True)
            certFile = os.path.join(cls.igorDir, 'igor.crt')
            cls.igorVarArgs['certificate'] = certFile
            cls.igorVarArgs['noverify'] = True
            os.putenv('SSL_CERT_FILE', certFile)
            os.putenv('IGOR_TEST_NO_SSL_VERIFY', '1')

        if DEBUG_TEST: print 'IgorTest: Check database consistency'
        subprocess.check_call([sys.executable, "-m", "igor", "--check", "--database", cls.igorDir, "--port", str(cls.igorPort)], stdout=open(cls.igorLogFile, 'a'), stderr=subprocess.STDOUT)

        if DEBUG_TEST: print 'IgorTest: Start server'
        cls.igorProcess = subprocess.Popen([sys.executable, "-u", "-m", "igor", "--database", cls.igorDir, "--port", str(cls.igorPort)], stdout=open(cls.igorLogFile, 'a'), stderr=subprocess.STDOUT)
        time.sleep(2)

    @classmethod
    def tearDownClass(cls):
        # Gracefully stop server
        if DEBUG_TEST: print 'IgorTest: Request server to stop'
        try:
            p = igorVar.IgorServer(cls.igorUrl, **cls.igorVarArgs)
            result = p.get('/internal/stop')
        except:
            if DEBUG_TEST: print 'IgorTest: Ignoring exception during stop request'
        time.sleep(2)

        sts = cls.igorProcess.poll()
        if sts is None:
            if DEBUG_TEST: print 'IgorTest: Terminate server'
            cls.igorProcess.terminate()
            time.sleep(2)
        sts = cls.igorProcess.wait()
        assert sts != None

        super(IgorTest, cls).tearDownClass()
        pass

    def _igorVar(self):
        return igorVar.IgorServer(self.igorUrl, **self.igorVarArgs)

    def test01_get_static(self):
        p = self._igorVar()
        result = p.get('/')
        self.assertTrue(result)
        self.assertEqual(result[0], "<")

    def notyet_test01_get_static_nonexistent(self):
        p = self._igorVar()
        result = p.get('/nonexistent.html')
        self.assertTrue(result)
        self.assertEqual(result[0], "<")

    def test11_get_xml(self):
        p = self._igorVar()
        result = p.get('environment/systemHealth', format='application/xml')
        self.assertTrue(result)
        root = ET.fromstring(result)
        self.assertEqual(root.tag, "systemHealth")

    def test12_get_text(self):
        p = self._igorVar()
        result = p.get('environment/systemHealth', format='text/plain')
        self.assertTrue(result)

    def test13_get_json(self):
        p = self._igorVar()
        result = p.get('environment/systemHealth', format='application/json')
        self.assertTrue(result)
        root = json.loads(result)
        self.assertIsInstance(root, dict)
        self.assertEqual(root.keys(), ["systemHealth"])

    def test21_put_xml(self):
        p = self._igorVar()
        data = '<test21>21</test21>'
        p.put('sandbox/test21', data, datatype='application/xml')
        result = p.get('sandbox/test21', format='application/xml')
        self.assertEqual(data.strip(), result.strip())
        result2 = p.get('sandbox/test21', format='text/plain')
        self.assertEqual('21', result2.strip())
        result3 = p.get('sandbox/test21', format='application/json')
        result3dict = json.loads(result3)
        self.assertEqual({"test21" : 21}, result3dict)

    def test22_put_text(self):
        p = self._igorVar()
        data = 'twenty two'
        p.put('sandbox/test22', data, datatype='text/plain')
        result = p.get('sandbox/test22', format='text/plain')
        self.assertEqual(data.strip(), result.strip())
        result2 = p.get('sandbox/test22', format='application/xml')
        self.assertEqual('<test22>twenty two</test22>', result2.strip())
        result3 = p.get('sandbox/test22', format='application/json')
        result3dict = json.loads(result3)
        self.assertEqual({'test22':'twenty two'}, result3dict)

    def test23_put_json(self):
        p = self._igorVar()
        data = json.dumps({"test23" : 23})
        p.put('sandbox/test23', data, datatype='application/json')
        result = p.get('sandbox/test23', format='application/json')
        resultDict = json.loads(result)
        self.assertEqual({"test23" : 23}, resultDict)
        result2 = p.get('sandbox/test23', format='application/xml')
        self.assertEqual("<test23>23</test23>", result2.strip())

    def test24_put_multi(self):
        p = self._igorVar()
        data = '<test24>24</test24>'
        p.put('sandbox/test24', data, datatype='application/xml')
        result = p.get('sandbox/test24', format='application/xml')
        self.assertEqual(data.strip(), result.strip())
        data = '<test24>twentyfour</test24>'
        p.put('sandbox/test24', data, datatype='application/xml')
        result = p.get('sandbox/test24', format='application/xml')
        self.assertEqual(data.strip(), result.strip())

    def test31_post_text(self):
        p = self._igorVar()
        p.put('sandbox/test31', '', datatype='text/plain')
        p.post('sandbox/test31/item', 'thirty', datatype='text/plain')
        p.post('sandbox/test31/item', 'one', datatype='text/plain')
        result = p.get('sandbox/test31/item', format='text/plain')
        self.assertEqual('thirtyone', result.translate(None, ' \n'))

        self.assertRaises(igorVar.IgorError, p.get, 'sandbox/test31/item', format='application/xml')

        result2 = p.get('sandbox/test31/item', format='application/xml', variant='multi')
        self.assertIn('thirty', result2)
        self.assertIn('one', result2)

        result3 = p.get('sandbox/test31/item', format='application/json', variant='multi')
        result3list = json.loads(result3)
        self.assertEqual(len(result3list), 2)
        self.assertIsInstance(result3list, list)
        self.assertEqual(result3list[0]['item'], 'thirty')
        self.assertEqual(result3list[1]['item'], 'one')

    def test41_action(self):
        p = self._igorVar()
        content = {'test41':{'src':'', 'sink':''}}
        action = {'action':dict(name='test41action', url='/data/sandbox/test41/sink', xpath='/data/sandbox/test41/src', method='PUT', data='copy-{.}-copy')}
        p.put('sandbox/test41', json.dumps(content), datatype='application/json')
        p.post('actions/action', json.dumps(action), datatype='application/json')
        p.put('sandbox/test41/src', 'forty-one', datatype='text/plain')

        time.sleep(1)

        result = p.get('sandbox/test41', format='application/json')
        resultDict = json.loads(result)
        wantedContent = {'test41':{'src':'forty-one', 'sink':'copy-forty-one-copy'}}
        self.assertEqual(resultDict, wantedContent)

    def test42_action_post(self):
        p = self._igorVar()
        content = {'test42':{'src':''}}
        action = {'action':dict(name='test42action', url='/data/sandbox/test42/sink', xpath='/data/sandbox/test42/src', method='POST', data='{.}')}
        p.put('sandbox/test42', json.dumps(content), datatype='application/json')
        p.post('actions/action', json.dumps(action), datatype='application/json')
        time.sleep(1)
        p.put('sandbox/test42/src', '42a', datatype='text/plain')
        time.sleep(1)
        p.put('sandbox/test42/src', '42b', datatype='text/plain')
        time.sleep(1)
        p.put('sandbox/test42/src', '42c', datatype='text/plain')

        time.sleep(1)

        result = p.get('sandbox/test42', format='application/json')
        resultDict = json.loads(result)
        wantedContent = {'test42':{'src':'42c', 'sink':['42a','42b','42c']}}
        self.assertEqual(resultDict, wantedContent)

    def test43_action_indirect(self):
        p = self._igorVar()
        content = {'test43':{'src':'', 'sink':''}}
        action1 = {'action':dict(name='test43first', url='/action/test43second', xpath='/data/sandbox/test43/src')}
        action2 = {'action':dict(name='test43second', url='/data/sandbox/test43/sink', method='PUT', data='copy-{/data/sandbox/test43/src}-copy')}
        p.put('sandbox/test43', json.dumps(content), datatype='application/json')
        p.post('actions/action', json.dumps(action1), datatype='application/json')
        p.post('actions/action', json.dumps(action2), datatype='application/json')
        p.put('sandbox/test43/src', 'forty-three', datatype='text/plain')

        time.sleep(1)

        result = p.get('sandbox/test43', format='application/json')
        resultDict = json.loads(result)
        wantedContent = {'test43':{'src':'forty-three', 'sink':'copy-forty-three-copy'}}
        self.assertEqual(resultDict, wantedContent)

    def test44_action_external(self):
        p = self._igorVar()
        content = {'test44':{'src':'', 'sink':''}}
        action1 = {'action':dict(name='test44first', url='{/data/services/igor/protocol}://{/data/services/igor/host}:{/data/services/igor/port}/action/test44second', xpath='/data/sandbox/test44/src')}
        action2 = {'action':dict(name='test44second', url='/data/sandbox/test44/sink', method='PUT', data='copy-{/data/sandbox/test44/src}-copy')}
        p.put('sandbox/test44', json.dumps(content), datatype='application/json')
        p.post('actions/action', json.dumps(action1), datatype='application/json')
        p.post('actions/action', json.dumps(action2), datatype='application/json')
        p.put('sandbox/test44/src', 'forty-four', datatype='text/plain')

        time.sleep(2)

        result = p.get('sandbox/test44', format='application/json')
        resultDict = json.loads(result)
        wantedContent = {'test44':{'src':'forty-four', 'sink':'copy-forty-four-copy'}}
        self.assertEqual(resultDict, wantedContent)

class IgorTestHttps(IgorTest):
    igorDir = os.path.join(FIXTURES, 'testIgorHttps')
    igorLogFile = os.path.join(FIXTURES, 'testIgorHttps.log')
    igorPort = 29333
    igorProtocol = "https"

if __name__ == '__main__':
    unittest.main()

<|endoftext|>


<commit_before>import time
import logging
import platform

import requests

from .leakybucket import LeakyBucket
from .__version__ import __version__

class TokenExpireException(Exception):
    pass

class NoResultsException(Exception):
    pass

# pylint: disable=too-many-instance-attributes,too-few-public-methods
class Torrent(object):
    '''
    brief
    {
        "filename":"Off.Piste.2016.iNTERNAL.BDRip.x264-LiBRARiANS",
        "category":"Movies/x264",
        "download":"magnet:..."
    }

    extended
    {
        "title":"Off.Piste.2016.iNTERNAL.BDRip.x264-LiBRARiANS",
        "category":"Movies/x264",
        "download":"magnet:...",
        "seeders":12,
        "leechers":6,
        "size":504519520,
        "pubdate":"2017-05-21 02:13:49 +0000",
        "episode_info":{
            "imdb":"tt4443856",
            "tvrage":null,
            "tvdb":null,
            "themoviedb":"430293"
        },
        "ranked":1,
        "info_page":"https://torrentapi.org/...."
    }
    '''
    def __init__(self, mapping):
        self._raw = mapping
        self.is_extended = 'title' in self._raw
        self.category = self._raw['category']
        self.download = self._raw['download']
        self.filename = self._raw.get('filename') or self._raw.get('title')
        self.size = self._raw.get('size')
        self.pubdate = self._raw.get('pubdate')
        self.page = self._raw.get('info_page')

    def __str__(self):
        return '%s(%s)' % (self.filename, self.category)

def json_hook(dct):
    error_code = dct.get('error_code')
    if error_code in [2, 4]:
        raise TokenExpireException('Token expired')
    if error_code == 20:
        raise NoResultsException('No results found')
    if 'download' in dct:
        return Torrent(dct)
    return dct

class _RarbgAPIv2(object):
    '''
    API reference
    https://torrentapi.org/apidocs_v2.txt
    '''
    ENDPOINT = 'http://torrentapi.org/pubapi_v2.php'
    APP_ID = 'rarbgapi'

    def __init__(self):
        super(_RarbgAPIv2, self).__init__()
        self._endpoint = self.ENDPOINT

    def _get_user_agent(self):
        return '{appid}/{version} ({uname}) python {pyver}'.format(
            appid=self.APP_ID,
            version=__version__,
            uname='; '.join(platform.uname()),
            pyver=platform.python_version())

    # pylint: disable=no-self-use
    def _requests(self, method, url, params=None):
        if not params:
            params = dict()
        params.update({
            'app_id': self.APP_ID
        })

        headers = {
            'user-agent': self._get_user_agent()
        }
        sess = requests.Session()
        req = requests.Request(method, url, params=params, headers=headers)
        preq = req.prepare()
        resp = sess.send(preq)
        resp.raise_for_status()
        return resp

    def _get_token(self):
        '''
        {"token":"xxxxx"}
        '''
        params = {
            'get_token': 'get_token'
        }
        return self._requests('GET', self._endpoint, params)

    def _query(self, mode, token=None, **kwargs):
        params = {
            'mode': mode,
            'token': token
        }
        for key, value in kwargs.items():
            if key not in ['search_string', 'sort', 'limit',
                           'category', 'format_', 'search_imdb',
                           'search_tvdb', 'search_themoviedb']:
                raise ValueError('unsupported parameter %s' % key)

            if key == 'format_':
                key = 'format'

            if value is None:
                continue

            params[key] = value

        return self._requests('GET', self._endpoint, params)

def request(func):
    # pylint: disable=protected-access, too-many-branches
    def wrapper(self, *args, **kwargs):
        max_retries = retries = self._options['retries']
        while retries > 0:
            try:
                backoff = 2**(max_retries - retries)
                if not self._bucket.acquire(1, timeout=2):
                    raise ValueError('accquire token timeout')

                if not self._token:
                    raise TokenExpireException('Empty token')

                resp = func(self, token=self._token, *args, **kwargs)
                json_ = resp.json(object_hook=json_hook)
                error_code = json_.get('error_code')
                if error_code:
                    if error_code == 5:
                        # {
                        #     u'error_code': 5,
                        #     u'error': u'Too many requests per second.
                        #            Maximum requests allowed are 1req/2sec
                        #            Please try again later!'
                        # }
                        self._log.debug('Retry due to throttle')
                        continue

                    self._log.warn('error %s', json_)
                    raise ValueError('error')
                elif 'torrent_results' not in json_:
                    self._log.info('Bad response %s', json_)
                return json_['torrent_results']
            except ValueError:
                # bad arguments, not necessary to retry
                raise
            except NoResultsException:
                return []
            except TokenExpireException:
                resp = self._get_token()
                content = resp.json()
                self._token = content['token']
                self._log.debug('token=%s', self._token)
            except Exception as exp:  # pylint: disable=broad-except
                self._log.exception('Unexpected exception %s', exp)
                retries -= 1
                if not retries:
                    raise
            else:
                retries -= 1
            finally:
                time.sleep(backoff)

        assert 0, 'Unexpected response'
    return wrapper

class RarbgAPI(_RarbgAPIv2):

    CATEGORY_ADULT = 4
    CATEGORY_MOVIE_XVID = 14
    CATEGORY_MOVIE_XVID_720P = 48
    CATEGORY_MOVIE_H264 = 17
    CATEGORY_MOVIE_H264_1080P = 44
    CATEGORY_MOVIE_H264_720P = 45
    CATEGORY_MOVIE_H264_3D = 47
    CATEGORY_MOVIE_H264_4K = 50
    CATEGORY_MOVIE_H265_4K = 51
    CATEGORY_MOVIE_H265_4K_HDR = 52
    CATEGORY_MOVIE_FULL_BD = 42
    CATEGORY_MOVIE_BD_REMUX = 46
    CATEGORY_TV_EPISODES = 18
    CATEGORY_TV_EPISODES_HD = 41
    CATEGORY_TV_EPISODES_UHD = 49
    CATEGORY_MUSIC_MP3 = 23
    CATEGORY_MUSIC_FLAC = 25
    CATEGORY_GAMES_PC_ISO = 27
    CATEGORY_GAMES_PC_RIP = 28
    CATEGORY_GAMES_PS3 = 40
    CATEGORY_GAMES_PS4 = 53
    CATEGORY_GAMES_XBOX = 32
    CATEGORY_SOFTWARE = 33
    CATEGORY_EBOOK = 35

    def __init__(self, **options):
        super(RarbgAPI, self).__init__()
        self._token = None
        self._bucket = LeakyBucket(0.5)
        self._log = logging.getLogger(__name__)
        default_options = {
            'retries': 5,
        }
        if options:
            default_options.update(options)
        self._options = default_options

    @request
    # pylint: disable=too-many-arguments
    def list(self, search_string=None, sort=None, limit=None,
             category=None, format_=None, **kwargs):
        return self._query('list', search_string=search_string, sort=sort,
                           limit=limit, category=category, format_=format_,
                           **kwargs)

    @request
    # pylint: disable=too-many-arguments
    def search(self, search_string=None, sort=None, limit=None,
               category=None, format_=None, **kwargs):
        return self._query('search', search_string=search_string,
                           sort=sort, limit=limit, category=category,
                           format_=format_, **kwargs)
<commit_msg>rarbgapi: fix pylint<commit_after>import time
import logging
import platform

import requests

from .leakybucket import LeakyBucket
from .__version__ import __version__

class TokenExpireException(Exception):
    pass

class NoResultsException(Exception):
    pass

# pylint: disable=too-many-instance-attributes,too-few-public-methods
class Torrent(object):
    '''
    brief
    {
        "filename":"Off.Piste.2016.iNTERNAL.BDRip.x264-LiBRARiANS",
        "category":"Movies/x264",
        "download":"magnet:..."
    }

    extended
    {
        "title":"Off.Piste.2016.iNTERNAL.BDRip.x264-LiBRARiANS",
        "category":"Movies/x264",
        "download":"magnet:...",
        "seeders":12,
        "leechers":6,
        "size":504519520,
        "pubdate":"2017-05-21 02:13:49 +0000",
        "episode_info":{
            "imdb":"tt4443856",
            "tvrage":null,
            "tvdb":null,
            "themoviedb":"430293"
        },
        "ranked":1,
        "info_page":"https://torrentapi.org/...."
    }
    '''
    def __init__(self, mapping):
        self._raw = mapping
        self.is_extended = 'title' in self._raw
        self.category = self._raw['category']
        self.download = self._raw['download']
        self.filename = self._raw.get('filename') or self._raw.get('title')
        self.size = self._raw.get('size')
        self.pubdate = self._raw.get('pubdate')
        self.page = self._raw.get('info_page')

    def __str__(self):
        return '%s(%s)' % (self.filename, self.category)

def json_hook(dct):
    error_code = dct.get('error_code')
    if error_code in [2, 4]:
        raise TokenExpireException('Token expired')
    if error_code == 20:
        raise NoResultsException('No results found')
    if 'download' in dct:
        return Torrent(dct)
    return dct

class _RarbgAPIv2(object):
    '''
    API reference
    https://torrentapi.org/apidocs_v2.txt
    '''
    ENDPOINT = 'http://torrentapi.org/pubapi_v2.php'
    APP_ID = 'rarbgapi'

    def __init__(self):
        super(_RarbgAPIv2, self).__init__()
        self._endpoint = self.ENDPOINT

    def _get_user_agent(self):
        return '{appid}/{version} ({uname}) python {pyver}'.format(
            appid=self.APP_ID,
            version=__version__,
            uname='; '.join(platform.uname()),
            pyver=platform.python_version())

    # pylint: disable=no-self-use
    def _requests(self, method, url, params=None):
        if not params:
            params = dict()
        params.update({
            'app_id': self.APP_ID
        })

        headers = {
            'user-agent': self._get_user_agent()
        }
        sess = requests.Session()
        req = requests.Request(method, url, params=params, headers=headers)
        preq = req.prepare()
        resp = sess.send(preq)
        resp.raise_for_status()
        return resp

    def _get_token(self):
        '''
        {"token":"xxxxx"}
        '''
        params = {
            'get_token': 'get_token'
        }
        return self._requests('GET', self._endpoint, params)

    def _query(self, mode, token=None, **kwargs):
        params = {
            'mode': mode,
            'token': token
        }
        for key, value in kwargs.items():
            if key not in ['search_string', 'sort', 'limit',
                           'category', 'format_', 'search_imdb',
                           'search_tvdb', 'search_themoviedb']:
                raise ValueError('unsupported parameter %s' % key)

            if key == 'format_':
                key = 'format'

            if value is None:
                continue

            params[key] = value

        return self._requests('GET', self._endpoint, params)

def request(func):
    # pylint: disable=protected-access, too-many-branches
    def wrapper(self, *args, **kwargs):
        max_retries = retries = self._options['retries']
        while retries > 0:
            try:
                backoff = 2**(max_retries - retries)
                if not self._bucket.acquire(1, timeout=2):
                    raise ValueError('accquire token timeout')

                if not self._token:
                    raise TokenExpireException('Empty token')

                resp = func(self, token=self._token, *args, **kwargs)
                json_ = resp.json(object_hook=json_hook)
                error_code = json_.get('error_code')
                if error_code:  # pylint: disable=no-else-raise
                    if error_code == 5:
                        # {
                        #     u'error_code': 5,
                        #     u'error': u'Too many requests per second.
                        #            Maximum requests allowed are 1req/2sec
                        #            Please try again later!'
                        # }
                        self._log.debug('Retry due to throttle')
                        continue

                    self._log.warn('error %s', json_)
                    raise ValueError('error')
                elif 'torrent_results' not in json_:
                    self._log.info('Bad response %s', json_)
                return json_['torrent_results']
            except ValueError:
                # bad arguments, not necessary to retry
                raise
            except NoResultsException:
                return []
            except TokenExpireException:
                resp = self._get_token()
                content = resp.json()
                self._token = content['token']
                self._log.debug('token=%s', self._token)
            except Exception as exp:  # pylint: disable=broad-except
                self._log.exception('Unexpected exception %s', exp)
                retries -= 1
                if not retries:
                    raise
            else:
                retries -= 1
            finally:
                time.sleep(backoff)

        assert 0, 'Unexpected response'
    return wrapper

class RarbgAPI(_RarbgAPIv2):

    CATEGORY_ADULT = 4
    CATEGORY_MOVIE_XVID = 14
    CATEGORY_MOVIE_XVID_720P = 48
    CATEGORY_MOVIE_H264 = 17
    CATEGORY_MOVIE_H264_1080P = 44
    CATEGORY_MOVIE_H264_720P = 45
    CATEGORY_MOVIE_H264_3D = 47
    CATEGORY_MOVIE_H264_4K = 50
    CATEGORY_MOVIE_H265_4K = 51
    CATEGORY_MOVIE_H265_4K_HDR = 52
    CATEGORY_MOVIE_FULL_BD = 42
    CATEGORY_MOVIE_BD_REMUX = 46
    CATEGORY_TV_EPISODES = 18
    CATEGORY_TV_EPISODES_HD = 41
    CATEGORY_TV_EPISODES_UHD = 49
    CATEGORY_MUSIC_MP3 = 23
    CATEGORY_MUSIC_FLAC = 25
    CATEGORY_GAMES_PC_ISO = 27
    CATEGORY_GAMES_PC_RIP = 28
    CATEGORY_GAMES_PS3 = 40
    CATEGORY_GAMES_PS4 = 53
    CATEGORY_GAMES_XBOX = 32
    CATEGORY_SOFTWARE = 33
    CATEGORY_EBOOK = 35

    def __init__(self, **options):
        super(RarbgAPI, self).__init__()
        self._token = None
        self._bucket = LeakyBucket(0.5)
        self._log = logging.getLogger(__name__)
        default_options = {
            'retries': 5,
        }
        if options:
            default_options.update(options)
        self._options = default_options

    @request
    # pylint: disable=too-many-arguments
    def list(self, search_string=None, sort=None, limit=None,
             category=None, format_=None, **kwargs):
        return self._query('list', search_string=search_string, sort=sort,
                           limit=limit, category=category, format_=format_,
                           **kwargs)

    @request
    # pylint: disable=too-many-arguments
    def search(self, search_string=None, sort=None, limit=None,
               category=None, format_=None, **kwargs):
        return self._query('search', search_string=search_string,
                           sort=sort, limit=limit, category=category,
                           format_=format_, **kwargs)
<|endoftext|>