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

"""gen_pattern.py
Usage example:
python gen_pattern.py -o out.svg -r 11 -c 8 -T circles -s 20.0 -R 5.0 -u mm -w 216 -h 279
-o, --output - output file (default out.svg)
-r, --rows - pattern rows (default 11)
-c, --columns - pattern columns (default 8)
-T, --type - type of pattern, circles, acircles, checkerboard (default circles)
-s, --square_size - size of squares in pattern (default 20.0)
-R, --radius_rate - circles_radius = square_size/radius_rate (default 5.0)
-u, --units - mm, inches, px, m (default mm)
-w, --page_width - page width in units (default 216)
-h, --page_height - page height in units (default 279)
-a, --page_size - page size (default A4), supercedes -h -w arguments
-H, --help - show help
"""

from svgfig import *

import sys
import getopt

class PatternMaker:
  def __init__(self, cols,rows,output,units,square_size,radius_rate,page_width,page_height):
    self.cols = cols
    self.rows = rows
    self.output = output
    self.units = units
    self.square_size = square_size
    self.radius_rate = radius_rate
    self.width = page_width
    self.height = page_height
    self.g = SVG("g") # the svg group container

  def makeCirclesPattern(self):
    spacing = self.square_size
    r = spacing / self.radius_rate
    for x in range(1,self.cols+1):
      for y in range(1,self.rows+1):
        dot = SVG("circle", cx=x * spacing, cy=y * spacing, r=r, fill="black")
        self.g.append(dot)

  def makeACirclesPattern(self):
    spacing = self.square_size
    r = spacing / self.radius_rate
    for i in range(0,self.rows):
      for j in range(0,self.cols):
        dot = SVG("circle", cx= ((j*2 + i%2)*spacing) + spacing, cy=self.height - (i * spacing + spacing), r=r, fill="black")
        self.g.append(dot)

  def makeCheckerboardPattern(self):
    spacing = self.square_size
    xspacing = (self.width - self.cols * self.square_size) / 2.0
    yspacing = (self.height - self.rows * self.square_size) / 2.0
    for x in range(0,self.cols):
      for y in range(0,self.rows):
        if x%2 == y%2:
          square = SVG("rect", x=x * spacing + xspacing, y=y * spacing + yspacing, width=spacing, height=spacing, fill="black")
          self.g.append(square)

  def save(self):
    c = canvas(self.g,width="%d%s"%(self.width,self.units),height="%d%s"%(self.height,self.units),viewBox="0 0 %d %d"%(self.width,self.height))
    c.inkview(self.output)

def main():
    # parse command line options, TODO use argparse for better doc
    try:
        opts, args = getopt.getopt(sys.argv[1:], "Ho:c:r:T:u:s:R:w:h:a:", ["help","output=","columns=","rows=",
                                                                      "type=","units=","square_size=","radius_rate=",
                                                                      "page_width=","page_height=", "page_size="])
    except getopt.error, msg:
        print msg
        print "for help use --help"
        sys.exit(2)
    output = "out.svg"
    columns = 8
    rows = 11
    p_type = "circles"
    units = "mm"
    square_size = 20.0
    radius_rate = 5.0
    page_size = "A4"
    # page size dict (ISO standard, mm) for easy lookup. format - size: [width, height]
    page_sizes = {"A0": [840, 1188], "A1": [594, 840], "A2": [420, 594], "A3": [297, 420], "A4": [210, 297], "A5": [148, 210]}
    page_width = page_sizes[page_size.upper()][0]
    page_height = page_sizes[page_size.upper()][1]
    # process options
    for o, a in opts:
        if o in ("-H", "--help"):
            print __doc__
            sys.exit(0)
        elif o in ("-r", "--rows"):
            rows = int(a)
        elif o in ("-c", "--columns"):
            columns = int(a)
        elif o in ("-o", "--output"):
            output = a
        elif o in ("-T", "--type"):
            p_type = a
        elif o in ("-u", "--units"):
            units = a
        elif o in ("-s", "--square_size"):
            square_size = float(a)
        elif o in ("-R", "--radius_rate"):
            radius_rate = float(a)
        elif o in ("-w", "--page_width"):
            page_width = float(a)
        elif o in ("-h", "--page_height"):
            page_height = float(a)
        elif o in ("-a", "--page_size"):
            units = "mm"
            page_size = a.upper()
            page_width = page_sizes[page_size][0]
            page_height = page_sizes[page_size][1]
    pm = PatternMaker(columns,rows,output,units,square_size,radius_rate,page_width,page_height)
    #dict for easy lookup of pattern type
    mp = {"circles":pm.makeCirclesPattern,"acircles":pm.makeACirclesPattern,"checkerboard":pm.makeCheckerboardPattern}
    mp[p_type]()
    #this should save pattern to output
    pm.save()

if __name__ == "__main__":
    main()
<commit_msg>Fixed gen_pattern.py file permission back to 755<commit_after>#!/usr/bin/env python

"""gen_pattern.py
Usage example:
python gen_pattern.py -o out.svg -r 11 -c 8 -T circles -s 20.0 -R 5.0 -u mm -w 216 -h 279
-o, --output - output file (default out.svg)
-r, --rows - pattern rows (default 11)
-c, --columns - pattern columns (default 8)
-T, --type - type of pattern, circles, acircles, checkerboard (default circles)
-s, --square_size - size of squares in pattern (default 20.0)
-R, --radius_rate - circles_radius = square_size/radius_rate (default 5.0)
-u, --units - mm, inches, px, m (default mm)
-w, --page_width - page width in units (default 216)
-h, --page_height - page height in units (default 279)
-a, --page_size - page size (default A4), supercedes -h -w arguments
-H, --help - show help
"""

from svgfig import *

import sys
import getopt

class PatternMaker:
  def __init__(self, cols,rows,output,units,square_size,radius_rate,page_width,page_height):
    self.cols = cols
    self.rows = rows
    self.output = output
    self.units = units
    self.square_size = square_size
    self.radius_rate = radius_rate
    self.width = page_width
    self.height = page_height
    self.g = SVG("g") # the svg group container

  def makeCirclesPattern(self):
    spacing = self.square_size
    r = spacing / self.radius_rate
    for x in range(1,self.cols+1):
      for y in range(1,self.rows+1):
        dot = SVG("circle", cx=x * spacing, cy=y * spacing, r=r, fill="black")
        self.g.append(dot)

  def makeACirclesPattern(self):
    spacing = self.square_size
    r = spacing / self.radius_rate
    for i in range(0,self.rows):
      for j in range(0,self.cols):
        dot = SVG("circle", cx= ((j*2 + i%2)*spacing) + spacing, cy=self.height - (i * spacing + spacing), r=r, fill="black")
        self.g.append(dot)

  def makeCheckerboardPattern(self):
    spacing = self.square_size
    xspacing = (self.width - self.cols * self.square_size) / 2.0
    yspacing = (self.height - self.rows * self.square_size) / 2.0
    for x in range(0,self.cols):
      for y in range(0,self.rows):
        if x%2 == y%2:
          square = SVG("rect", x=x * spacing + xspacing, y=y * spacing + yspacing, width=spacing, height=spacing, fill="black")
          self.g.append(square)

  def save(self):
    c = canvas(self.g,width="%d%s"%(self.width,self.units),height="%d%s"%(self.height,self.units),viewBox="0 0 %d %d"%(self.width,self.height))
    c.inkview(self.output)

def main():
    # parse command line options, TODO use argparse for better doc
    try:
        opts, args = getopt.getopt(sys.argv[1:], "Ho:c:r:T:u:s:R:w:h:a:", ["help","output=","columns=","rows=",
                                                                      "type=","units=","square_size=","radius_rate=",
                                                                      "page_width=","page_height=", "page_size="])
    except getopt.error, msg:
        print msg
        print "for help use --help"
        sys.exit(2)
    output = "out.svg"
    columns = 8
    rows = 11
    p_type = "circles"
    units = "mm"
    square_size = 20.0
    radius_rate = 5.0
    page_size = "A4"
    # page size dict (ISO standard, mm) for easy lookup. format - size: [width, height]
    page_sizes = {"A0": [840, 1188], "A1": [594, 840], "A2": [420, 594], "A3": [297, 420], "A4": [210, 297], "A5": [148, 210]}
    page_width = page_sizes[page_size.upper()][0]
    page_height = page_sizes[page_size.upper()][1]
    # process options
    for o, a in opts:
        if o in ("-H", "--help"):
            print __doc__
            sys.exit(0)
        elif o in ("-r", "--rows"):
            rows = int(a)
        elif o in ("-c", "--columns"):
            columns = int(a)
        elif o in ("-o", "--output"):
            output = a
        elif o in ("-T", "--type"):
            p_type = a
        elif o in ("-u", "--units"):
            units = a
        elif o in ("-s", "--square_size"):
            square_size = float(a)
        elif o in ("-R", "--radius_rate"):
            radius_rate = float(a)
        elif o in ("-w", "--page_width"):
            page_width = float(a)
        elif o in ("-h", "--page_height"):
            page_height = float(a)
        elif o in ("-a", "--page_size"):
            units = "mm"
            page_size = a.upper()
            page_width = page_sizes[page_size][0]
            page_height = page_sizes[page_size][1]
    pm = PatternMaker(columns,rows,output,units,square_size,radius_rate,page_width,page_height)
    #dict for easy lookup of pattern type
    mp = {"circles":pm.makeCirclesPattern,"acircles":pm.makeACirclesPattern,"checkerboard":pm.makeCheckerboardPattern}
    mp[p_type]()
    #this should save pattern to output
    pm.save()

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


<commit_before># coding=utf-8
import json

import config
import requests

headers = {
    'app_version': '6.9.4',
    'platform': 'ios',
    "content-type": "application/json",
    "User-agent": "Tinder/7.5.3 (iPhone; iOS 10.3.2; Scale/2.00)",
}

def get_auth_token(fb_auth_token, fb_user_id):
    if "error" in fb_auth_token:
        return {"error": "could not retrieve fb_auth_token"}
    if "error" in fb_user_id:
        return {"error": "could not retrieve fb_user_id"}
    url = config.host + '/auth'
    req = requests.post(url,
                        headers=headers,
                        data=json.dumps(
                            {'facebook_token': fb_auth_token, 'facebook_id': fb_user_id})
                        )
    try:
        tinder_auth_token = req.json()["token"]
        headers.update({"X-Auth-Token": tinder_auth_token})
        print("You have been successfully authorized!")
        return tinder_auth_token
    except Exception as e:
        print(e)
        return {"error": "Something went wrong. Sorry, but we could not authorize you."}

def authverif():
    res = get_auth_token(config.fb_access_token, config.fb_user_id)
    if "error" in res:
        return False
    return True

def get_recommendations():
    '''
    Returns a list of users that you can swipe on
    '''
    try:
        r = requests.get('https://api.gotinder.com/user/recs', headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong with getting recomendations:", e)

def get_updates(last_activity_date=""):
    '''
    Returns all updates since the given activity date.
    The last activity date is defaulted at the beginning of time.
    Format for last_activity_date: "2017-07-09T10:28:13.392Z"
    '''
    try:
        url = config.host + '/updates'
        r = requests.post(url,
                          headers=headers,
                          data=json.dumps({"last_activity_date": last_activity_date}))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong with getting updates:", e)

def get_self():
    '''
    Returns your own profile data
    '''
    try:
        url = config.host + '/profile'
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get your data:", e)

def change_preferences(**kwargs):
    '''
    ex: change_preferences(age_filter_min=30, gender=0)
    kwargs: a dictionary - whose keys become separate keyword arguments and the values become values of these arguments
    age_filter_min: 18..46
    age_filter_max: 22..55
    age_filter_min <= age_filter_max - 4
    gender: 0 == seeking males, 1 == seeking females
    distance_filter: 1..100
    discoverable: true | false
    {"photo_optimizer_enabled":false}
    '''
    try:
        url = config.host + '/profile'
        r = requests.post(url, headers=headers, data=json.dumps(kwargs))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not change your preferences:", e)

def get_meta():
    '''
    Returns meta data on yourself. Including the following keys:
    ['globals', 'client_resources', 'versions', 'purchases',
    'status', 'groups', 'products', 'rating', 'tutorials',
    'travel', 'notifications', 'user']
    '''
    try:
        url = config.host + '/meta'
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get your metadata:", e)

def update_location(lat, lon):
    '''
    Updates your location to the given float inputs
    Note: Requires a passport / Tinder Plus
    '''
    try:
        url = config.host + '/passport/user/travel'
        r = requests.post(url, headers=headers, data=json.dumps({"lat": lat, "lon": lon}))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not update your location:", e)

def reset_real_location():
    try:
        url = config.host + '/passport/user/reset'
        r = requests.post(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not update your location:", e)

def get_recs_v2():
    '''
    This works more consistently then the normal get_recommendations becuase it seeems to check new location
    '''
    try:
        url = config.host + '/v2/recs/core?locale=en-US'
        r = requests.get(url, headers=headers)
        return r.json()
    except Exception as e:
        print('excepted')

def get_person(id):
    '''
    Gets a user's profile via their id
    '''
    try:
        url = config.host + '/user/%s' % id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get that person:", e)

def send_msg(match_id, msg):
    try:
        url = config.host + '/user/matches/%s' % match_id
        r = requests.post(url, headers=headers,
                          data=json.dumps({"message": msg}))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not send your message:", e)

def superlike(person_id):
    try:
        url = config.host + '/like/%s/super' % person_id
        r = requests.post(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not superlike:", e)

def like(person_id):
    try:
        url = config.host + '/like/%s' % person_id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not like:", e)

def dislike(person_id):
    try:
        url = config.host + '/pass/%s' % person_id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not dislike:", e)

def report(person_id, cause, explanation=''):
    '''
    There are three options for cause:
        0 : Other and requires an explanation
        1 : Feels like spam and no explanation
        4 : Inappropriate Photos and no explanation
    '''
    try:
        url = config.host + '/report/%s' % person_id
        r = requests.post(url, headers=headers, data={
                          "cause": cause, "text": explanation})
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not report:", e)

def match_info(match_id):
    try:
        url = config.host + '/matches/%s' % match_id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get your match info:", e)

# def see_friends():
#     try:
#         url = config.host + '/group/friends'
#         r = requests.get(url, headers=headers)
#         return r.json()['results']
#     except requests.exceptions.RequestException as e:
#         print("Something went wrong. Could not get your Facebook friends:", e)
<commit_msg>Update tinder_api.py<commit_after># coding=utf-8
import json

import config
import requests

headers = {
    'app_version': '6.9.4',
    'platform': 'ios',
    "content-type": "application/json",
    "User-agent": "Tinder/7.5.3 (iPhone; iOS 10.3.2; Scale/2.00)",
}

def get_auth_token(fb_auth_token, fb_user_id):
    if "error" in fb_auth_token:
        return {"error": "could not retrieve fb_auth_token"}
    if "error" in fb_user_id:
        return {"error": "could not retrieve fb_user_id"}
    url = config.host + '/auth'
    req = requests.post(url,
                        headers=headers,
                        data=json.dumps(
                            {'facebook_token': fb_auth_token, 'facebook_id': fb_user_id})
                        )
    try:
        tinder_auth_token = req.json()["token"]
        headers.update({"X-Auth-Token": tinder_auth_token})
        print("You have been successfully authorized!")
        return tinder_auth_token
    except Exception as e:
        print(e)
        return {"error": "Something went wrong. Sorry, but we could not authorize you."}

def authverif():
    res = get_auth_token(config.fb_access_token, config.fb_user_id)
    if "error" in res:
        return False
    return True

def get_recommendations():
    '''
    Returns a list of users that you can swipe on
    '''
    try:
        r = requests.get('https://api.gotinder.com/user/recs', headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong with getting recomendations:", e)

def get_updates(last_activity_date=""):
    '''
    Returns all updates since the given activity date.
    The last activity date is defaulted at the beginning of time.
    Format for last_activity_date: "2017-07-09T10:28:13.392Z"
    '''
    try:
        url = config.host + '/updates'
        r = requests.post(url,
                          headers=headers,
                          data=json.dumps({"last_activity_date": last_activity_date}))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong with getting updates:", e)

def get_self():
    '''
    Returns your own profile data
    '''
    try:
        url = config.host + '/profile'
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get your data:", e)

def change_preferences(**kwargs):
    '''
    ex: change_preferences(age_filter_min=30, gender=0)
    kwargs: a dictionary - whose keys become separate keyword arguments and the values become values of these arguments
    age_filter_min: 18..46
    age_filter_max: 22..55
    age_filter_min <= age_filter_max - 4
    gender: 0 == seeking males, 1 == seeking females
    distance_filter: 1..100
    discoverable: true | false
    {"photo_optimizer_enabled":false}
    '''
    try:
        url = config.host + '/profile'
        r = requests.post(url, headers=headers, data=json.dumps(kwargs))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not change your preferences:", e)

def get_meta():
    '''
    Returns meta data on yourself. Including the following keys:
    ['globals', 'client_resources', 'versions', 'purchases',
    'status', 'groups', 'products', 'rating', 'tutorials',
    'travel', 'notifications', 'user']
    '''
    try:
        url = config.host + '/meta'
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get your metadata:", e)

def update_location(lat, lon):
    '''
    Updates your location to the given float inputs
    Note: Requires a passport / Tinder Plus
    '''
    try:
        url = config.host + '/passport/user/travel'
        r = requests.post(url, headers=headers, data=json.dumps({"lat": lat, "lon": lon}))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not update your location:", e)

def reset_real_location():
    try:
        url = config.host + '/passport/user/reset'
        r = requests.post(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not update your location:", e)

def get_recs_v2():
    '''
    This works more consistently then the normal get_recommendations becuase it seeems to check new location
    '''
    try:
        url = config.host + '/v2/recs/core?locale=en-US'
        r = requests.get(url, headers=headers)
        return r.json()
    except Exception as e:
        print('excepted')

def set_webprofileusername(username):
    '''
    Sets the username for the webprofile: https://www.gotinder.com/@YOURUSERNAME
    '''
    try:
        url = config.host + '/profile/username'
        r = requests.put(url, headers=headers,
                         data=json.dumps({"username": username}))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not set webprofile username:", e)

def reset_webprofileusername(username):
    '''
    Resets the username for the webprofile
    '''
    try:
        url = config.host + '/profile/username'
        r = requests.delete(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not delete webprofile username:", e)

def get_person(id):
    '''
    Gets a user's profile via their id
    '''
    try:
        url = config.host + '/user/%s' % id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get that person:", e)

def send_msg(match_id, msg):
    try:
        url = config.host + '/user/matches/%s' % match_id
        r = requests.post(url, headers=headers,
                          data=json.dumps({"message": msg}))
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not send your message:", e)

def superlike(person_id):
    try:
        url = config.host + '/like/%s/super' % person_id
        r = requests.post(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not superlike:", e)

def like(person_id):
    try:
        url = config.host + '/like/%s' % person_id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not like:", e)

def dislike(person_id):
    try:
        url = config.host + '/pass/%s' % person_id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not dislike:", e)

def report(person_id, cause, explanation=''):
    '''
    There are three options for cause:
        0 : Other and requires an explanation
        1 : Feels like spam and no explanation
        4 : Inappropriate Photos and no explanation
    '''
    try:
        url = config.host + '/report/%s' % person_id
        r = requests.post(url, headers=headers, data={
                          "cause": cause, "text": explanation})
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not report:", e)

def match_info(match_id):
    try:
        url = config.host + '/matches/%s' % match_id
        r = requests.get(url, headers=headers)
        return r.json()
    except requests.exceptions.RequestException as e:
        print("Something went wrong. Could not get your match info:", e)

# def see_friends():
#     try:
#         url = config.host + '/group/friends'
#         r = requests.get(url, headers=headers)
#         return r.json()['results']
#     except requests.exceptions.RequestException as e:
#         print("Something went wrong. Could not get your Facebook friends:", e)
<|endoftext|>


<commit_before>import ads
from unidecode import unidecode
from itertools import combinations
import numpy as np

##### bibcodes.txt has the bibcodes of the publications, one per line
with open('bibcodes_opt.txt') as f:
	bibcodes = f.read().splitlines()

# PI last name. Does not need to be the full name to help with the woes
# of accented characters, but be careful that it does not match other co-authors
# in your papers
PI_name='Jord'
#############

def P(c):
    return min(5,3 + 2*c/12.)

num_pubs = len(bibcodes)
if (num_pubs < 11):
	raise ValueError("La lista de publicaciones debe tener al menos 10 articulos")

PI_name_length=len(PI_name)

l_array=[]
p_array=[]
N=10

for b in bibcodes:
	res = ads.SearchQuery(bibcode=b,fl=['year', 'author', 'citation_count', 'title'])
	for paper in res:
		print "Processing paper: "+ b
		print 'Title: "'+unidecode(paper.title[0])+'"'
		yr = int(paper.year)
		years = 2016 - yr
		if (yr < 2011):
			raise ValueError("No papers prior to 2011 can be used")
		if years == 0:
			years = 1
		p_i = P(paper.citation_count/years)
		if (yr == 2016) and (p_i < 4):
			p_i = 4.0
		p_array.append(p_i)
		author_list_length = len(paper.author)
		bis = min(author_list_length,5)
		l_i=0
		for i in range(bis):
			last_name = paper.author[i].split(',')[0]
			if (last_name[:PI_name_length] == PI_name):
				if (i <= 2):
					l_i = 1
				if (i > 2) and (i < 5):
					l_i = 0.5
		l_array.append(l_i)
		print "\tP_i={0:4.2f} l_i={1:2.1f}".format(p_i,l_i)

l_array = np.array(l_array)
p_array = np.array(p_array)

def score(l_array,p_array,I):
	liderazgo = l_array[I].sum()
	av_score  = p_array[I].mean()
	deltaL = 0
	if (liderazgo < 5):
		deltaL = 1.75*((1-0.2*liderazgo)**(1.25))
	return av_score - deltaL

best_score=-1e10
C = combinations(range(num_pubs),10)
Ci = np.array(list(C))
I_best = None
ni = len(Ci)

for ind in range(ni):
	II = Ci[ind]
	try_score = score(l_array,p_array,II)
	if (try_score > best_score):
		best_score = try_score
		I_best = II

print "\nBest set of papers:"
for i in I_best:
	print bibcodes[i]

print "\nBest score:", best_score<commit_msg>Improved optimizer.py<commit_after>from __future__ import print_function
import ads
from unidecode import unidecode
from itertools import combinations
import numpy as np
import copy

max_n = 25

##### bibcodes.txt has the bibcodes of the publications, one per line
with open('bibcodes_opt.txt') as f:
	bibcodes = f.read().splitlines()

# PI last name. Does not need to be the full name to help with the woes
# of accented characters, but be careful that it does not match other co-authors
# in your papers
PI_name='Jord'
#############

def P(c):
    return min(5,3 + 2*c/12.)

num_pubs = len(bibcodes)
if (num_pubs < 11):
	raise ValueError("La lista de publicaciones debe tener al menos 10 articulos")

PI_name_length=len(PI_name)

l_array=[]
p_array=[]
N=10

for b in bibcodes:
	res = ads.SearchQuery(bibcode=b,fl=['year', 'author', 'citation_count', 'title'])
	for paper in res:
		print("Processing paper: "+ b)
		print('Title: "'+unidecode(paper.title[0])+'"')
		yr = int(paper.year)
		years = 2016 - yr
		if (yr < 2011):
			raise ValueError("No papers prior to 2011 can be used")
		if years == 0:
			years = 1
		p_i = P(paper.citation_count/years)
		if (yr == 2016) and (p_i < 4):
			p_i = 4.0
		p_array.append(p_i)
		author_list_length = len(paper.author)
		bis = min(author_list_length,5)
		l_i=0
		for i in range(bis):
			last_name = paper.author[i].split(',')[0]
			if (last_name[:PI_name_length] == PI_name):
				if (i <= 2):
					l_i = 1
				if (i > 2) and (i < 5):
					l_i = 0.5
		l_array.append(l_i)
		print("\tP_i={0:4.2f} l_i={1:2.1f}".format(p_i,l_i))

l_array = np.array(l_array)
p_array = np.array(p_array)

# Now we select a subset of maximum 20 papers to make the calculation
# in a reasonable time
n_actual_papers = len(l_array)
if (n_actual_papers > max_n):
	# include all papers with leadership > 0
	print("Restricting your list to {0:d} papers".format(max_n))
	i_l = np.where(l_i > 0)
	tmp_array = copy.copy(p_array)
	tmp_array[i_l] = 5.1
	idx = np.argsort(tmp_array)
	l_array = l_array[idx][-max_n:]
	p_array = p_array[idx][-max_n:]
	#bibcodes = bibcodes[idx][max_n:]
	tmp=[]
	for k in range(max_n):
		tmp.append(bibcodes[idx[-max_n+k + n_actual_papers]])
	bibcodes=tmp
	num_pubs = max_n

def score(l_array,p_array,I):
	liderazgo = l_array[I].sum()
	av_score  = p_array[I].mean()
	deltaL = 0
	if (liderazgo < 5):
		deltaL = 1.75*((1-0.2*liderazgo)**(1.25))
	return av_score - deltaL

best_score=-1e10
C = combinations(range(num_pubs),10)
Ci = np.array(list(C))
I_best = None
ni = len(Ci)

for ind in range(ni):
	II = Ci[ind]
	try_score = score(l_array,p_array,II)
	if (try_score > best_score):
		best_score = try_score
		I_best = II

print("\nBest set of papers:")
for i in I_best:
	print(bibcodes[i])

print("\nBest score:", best_score)

<|endoftext|>


<commit_before><commit_msg>comentário<commit_after><|endoftext|>


<commit_before>
import flexceptions
import bz2

def init(config, compression=None):
    if 'compression' not in config:
        compression = 'bz2'
    if compression is not None:
        config['compression'] = compression

def compress(config, data):
    if config['compression'] == 'bz2':
        compressed = bz2.compress(data)
        if len(compressed) < data:
            return b'bz2' + bz2.compress(data)
        else:  # bz2'ed random data is larger than not
            return b'off' + data
    if config['compression'] == 'off':
        return b'off' + data
    raise flexceptions.CompressionError

def decompress(config, data):
    if data[:3] == b'bz2':
        return bz2.decompress(data[3:])
    if data[:3] == b'off':
        return data[3:]
    raise flexceptions.BlockCorruptionError
<commit_msg>Add zstd compression<commit_after>
import flexceptions
import bz2

COMPRESSORS = {b'off': lambda x: x}
DECOMPRESSORS = {b'off': lambda x: x}

BEST = b'off'

try:
    import bz2
    COMPRESSORS[b'bz2'] = lambda x: bz2.compress(x)
    DECOMPRESSORS[b'bz2'] = lambda x: bz2.decompress(x)
    BEST = b'bz2'
except ImportError:
    pass

try:
    import zstd
    COMPRESSORS[b'zst'] = lambda x: zst.compress(x)
    DECOMPRESSORS[b'zst'] = lambda x: zst.decompress(x)
    BEST = b'zst'
except ImportError:
    pass

def init(config, compression=None):
    if 'compression' not in config:
        compression = BEST
    if compression is not None:
        config['compression'] = compression

def compress(config, data):
    compression = bytes(config['compression'])
    compressed = COMPRESSORS[compression](data)
    if len(compressed) < len(data):
        return compression + compressed
    else:  # bz2'ed random data is larger than not
        return b'off' + COMPRESSORS[b'off'](data)

def decompress(config, data):
    return DECOMPRESSORS[data[:3]](data[3:])
<|endoftext|>


<commit_before># coding=utf-8

class Command:

    def __init__(self):
        self.twist = [0.0, 0.0, 0.0]
        self.wrench = [0.0, 0.0, 0.0]
        self.string = ''
        self.manual = True

def process_string(val, cmd):
    pass

def process_vec3(val, cmd):
    for i in range(len(val)):
        if val[i] is 'r':
            if cmd[i] != 0.0:
                cmd[i] = max(min(-cmd[i] / abs(cmd[i]), 1), -1)
        else:
            cmd[i] += val[i] / 20.0

        cmd[i] = max(min(cmd[i], 1), -1)

    return cmd

def process_twist(val, cmd):
    return process_vec3(val, cmd.twist)

def process_wrench(val, cmd):
    return process_vec3(val, cmd.wrench)

def process(msg, key_map, topic, cmd):

    # Gestion des erreurs : le message est vide, non reconnu ou pas du bon topic
    if len(msg.data) == 0 or not key_map.has_key(msg.data):
        print 'len(msg.data) == 0 : ', len(msg.data) == 0
        print 'key_map.has_key(msg.data) : ', key_map.has_key(msg.data)
        print 'unknown key : ', msg.data
        return None, None # unknown key
    elif key_map[msg.data]['topic'] is not topic \
            and key_map[msg.data]['topic'] is not 'state':
        return None, None # key that shouldn't be translated in this instance

    # Récupération des données
    val = key_map[msg.data]['value']
    topic = key_map[msg.data]['topic']
    print "received mgs.data : ", msg.data
    print 'vals : ', val

    # Gestion des entrées
    if topic is "twist":
        cmd = process_twist(val, cmd)
    elif topic is "wrench":
        cmd = process_wrench(val, cmd)
    elif topic is "state":
        cmd.manual = val

    # Envoi de la commande si en manuel
    if cmd.manual:
        return True, cmd
    return False, None
    print "auto mode :", cmd.manual
<commit_msg>Debug 4<commit_after># coding=utf-8

class Command:

    def __init__(self):
        self.twist = [0.0, 0.0, 0.0]
        self.wrench = [0.0, 0.0, 0.0]
        self.string = ''
        self.manual = True

def process_string(val, cmd):
    pass

def process_vec3(val, cmd):
    for i in range(len(val)):
        if val[i] is 'r':
            if cmd[i] != 0.0:
                cmd[i] = max(min(-cmd[i] / abs(cmd[i]), 1), -1)
        else:
            cmd[i] += val[i] / 20.0

        cmd[i] = max(min(cmd[i], 1), -1)

    return cmd

def process_twist(val, cmd):
    return process_vec3(val, cmd.twist)

def process_wrench(val, cmd):
    return process_vec3(val, cmd.wrench)

def process(msg, key_map, topic, cmd):

    # Gestion des erreurs : le message est vide, non reconnu ou pas du bon topic
    if len(msg.data) == 0 or not key_map.has_key(msg.data):
        print 'len(msg.data) == 0 : ', len(msg.data) == 0
        print 'key_map.has_key(msg.data) : ', key_map.has_key(msg.data)
        print 'unknown key : ', msg.data
        return None, None # unknown key
    elif key_map[msg.data]['topic'] is not topic \
            and key_map[msg.data]['topic'] is not 'state':
        return None, None # key that shouldn't be translated in this instance

    # Récupération des données
    val = key_map[msg.data]['value']
    topic = key_map[msg.data]['topic']
    print "received mgs.data : ", msg.data
    print 'vals : ', val

    # Gestion des entrées
    if topic is "twist":
        cmd = process_twist(val, cmd)
    elif topic is "wrench":
        cmd = process_wrench(val, cmd)
    elif topic is "state":
        print "Changing state command :", cmd.manual
        cmd.manual = val

    # Envoi de la commande si en manuel
    if cmd.manual:
        return True, cmd
    print "auto mode :", cmd.manual
    return False, None

<|endoftext|>


<commit_before>import boto
import moto
import unittest
import os
import uuid
import shutil

import pickle
from sframe.util import cloudpickle

import sframe as gl
from .. import _gl_pickle as gl_pickle
from ..util import _assert_sframe_equal as assert_sframe_equal

class GLPicklingTest(unittest.TestCase):

    def setUp(self):
        self.filename = str(uuid.uuid4())
        self.dir_mode = False

    def tearDown(self):
        if os.path.isdir(self.filename):
            shutil.rmtree(self.filename)
        elif os.path.exists(self.filename):
            os.remove(self.filename)

    def test_pickling_simple_types(self):

        obj_list = [
            1, "hello", 5.0,
            (1, 2), ("i", "love", "cricket"),
            [1, 2, "hello"], [1.3, (1,2), "foo"], ["bar", {"foo": "bar"}],
            {"cricket": "best-sport", "test": [1,2,3]}, {"foo": 1.3},
        ]
        for obj in obj_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert obj == obj_ret, "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_pickling_sarray_types(self):

        sarray_list = [
            gl.SArray([1,2,3]),
            gl.SArray([1.0,2.0,3.5]),
            gl.SArray(["foo", "bar"]),
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert list(obj) ==  list(obj_ret), \
                       "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_pickling_sframe_types(self):

        sarray_list = [
            gl.SFrame([1,2,3]),
            gl.SFrame([1.0,2.0,3.5]),
            gl.SFrame(["foo", "bar"]),
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj, obj_ret)

    def test_pickling_sgraph_types(self):

        sg_test_1 = gl.SGraph().add_vertices([
                                gl.Vertex(0, {'fluffy': 1}),
                                gl.Vertex(1, {'fluffy': 1, 'woof': 1}),
                                gl.Vertex(2, {})])

        sg_test_2 = gl.SGraph()
        sg_test_2 = sg_test_2.add_vertices([
                            gl.Vertex(x) for x in [0, 1, 2]])
        sg_test_2 = sg_test_2.add_edges([
                        gl.Edge(0, 1, attr={'relationship': 'dislikes'}),
                        gl.Edge(1, 2, attr={'relationship': 'likes'}),
                        gl.Edge(1, 0, attr={'relationship': 'likes'})])

        sarray_list = [
            sg_test_1,
            sg_test_2
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj.get_vertices(), obj_ret.get_vertices())
            assert_sframe_equal(obj.get_edges(), obj_ret.get_edges())

    def test_combination_gl_python_types(self):

        sg_test_1 = gl.SGraph().add_vertices([
                                gl.Vertex(1, {'fluffy': 1}),
                                gl.Vertex(2, {'fluffy': 1, 'woof': 1}),
                                gl.Vertex(3, {})])
        sarray_test_1 = gl.SArray([1,2,3])
        sframe_test_1 = gl.SFrame([1,2,3])

        obj_list = [
            [sg_test_1, sframe_test_1, sarray_test_1],
            {0:sg_test_1, 1:sframe_test_1, 2:sarray_test_1}
        ]

        for obj in obj_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj[0].get_vertices(), obj_ret[0].get_vertices())
            assert_sframe_equal(obj[0].get_edges(), obj_ret[0].get_edges())
            assert_sframe_equal(obj[1], obj_ret[1])
            assert list(obj[2]) ==  list(obj_ret[2])

    def test_pickle_compatibility(self):
        obj_list = [
            1, "hello", 5.0,
            (1, 2), ("i", "love", "cricket"),
            [1, 2, "hello"], [1.3, (1,2), "foo"], ["bar", {"foo": "bar"}],
            {"cricket": "best-sport", "test": [1,2,3]}, {"foo": 1.3},
        ]
        for obj in obj_list:
            file = open(self.filename, 'wb')
            pickler = pickle.Pickler(file)
            pickler.dump(obj)
            file.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert obj == obj_ret, \
                "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_cloud_pickle_compatibility(self):
        obj_list = [
            1, "hello", 5.0,
            (1, 2), ("i", "love", "cricket"),
            [1, 2, "hello"], [1.3, (1,2), "foo"], ["bar", {"foo": "bar"}],
            {"cricket": "best-sport", "test": [1,2,3]}, {"foo": 1.3},
        ]
        for obj in obj_list:
            file = open(self.filename, 'wb')
            pickler = cloudpickle.CloudPickler(file)
            pickler.dump(obj)
            file.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert obj == obj_ret, \
                "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_relative_path(self):
        # Arrange
        sf1 = gl.SFrame(range(10))
        relative_path = 'tmp/%s' % self.filename

        # Act
        pickler = gl_pickle.GLPickler(relative_path)
        pickler.dump(sf1)
        pickler.close()
        sf2 = gl_pickle.GLUnpickler(relative_path).load()

        # Assert
        assert_sframe_equal(sf1, sf2)

        # Clean up
        shutil.rmtree(relative_path)

    @moto.mock_s3
    def test_save_to_s3(self):
        # Arrange
        os.environ['GRAPHLAB_UNIT_TEST'] = 'PredictiveService'
        os.environ['AWS_ACCESS_KEY_ID'] = ''
        os.environ['AWS_SECRET_ACCESS_KEY'] = ''
        BUCKET_NAME = 'foo'
        boto.connect_s3().create_bucket(BUCKET_NAME)
        sf1 = gl.SFrame(range(10))
        S3_PATH = "s3://%s/foobar" % BUCKET_NAME

        # Act
        pickler = gl_pickle.GLPickler(S3_PATH)
        pickler.dump(sf1)
        pickler.close()
        sf2 = gl_pickle.GLUnpickler(S3_PATH).load()

        # Assert
        assert_sframe_equal(sf1, sf2)

        del os.environ['GRAPHLAB_UNIT_TEST']

    def _test_backward_compatibility(self):

        # Arrange
        file_name = 's3://gl-internal-datasets/models/1.3/gl-pickle.gl'
        obj = {'foo': gl.SFrame([1,2,3]),
               'bar': gl.SFrame(),
               'foo-bar': ['foo-and-bar', gl.SFrame([1])]}

        # Act
        unpickler = gl_pickle.GLUnpickler(file_name)
        obj_ret = unpickler.load()

        # Assert
        assert_sframe_equal(obj['foo'], obj_ret['foo'])
        assert_sframe_equal(obj['bar'], obj_ret['bar'])
        assert_sframe_equal(obj['foo-bar'][1], obj_ret['foo-bar'][1])
        self.assertEqual(obj['foo-bar'][0], obj_ret['foo-bar'][0])

    def test_save_over_previous(self):

        sarray_list = [
            gl.SFrame([1,2,3]),
            gl.SFrame([1.0,2.0,3.5]),
            gl.SFrame(["foo", "bar"]),
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()

            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj, obj_ret)

            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()

<commit_msg>Fix file_utile and gl_pickle test to not import sframe directory<commit_after>import boto
import moto
import unittest
import os
import uuid
import shutil

import pickle
from ..util import cloudpickle

from ..data_structures.sframe import SFrame
from ..data_structures.sarray import SArray
from ..data_structures.sgraph import SGraph, Vertex, Edge

from .. import _gl_pickle as gl_pickle
from ..util import _assert_sframe_equal as assert_sframe_equal

class GLPicklingTest(unittest.TestCase):

    def setUp(self):
        self.filename = str(uuid.uuid4())
        self.dir_mode = False

    def tearDown(self):
        if os.path.isdir(self.filename):
            shutil.rmtree(self.filename)
        elif os.path.exists(self.filename):
            os.remove(self.filename)

    def test_pickling_simple_types(self):

        obj_list = [
            1, "hello", 5.0,
            (1, 2), ("i", "love", "cricket"),
            [1, 2, "hello"], [1.3, (1,2), "foo"], ["bar", {"foo": "bar"}],
            {"cricket": "best-sport", "test": [1,2,3]}, {"foo": 1.3},
        ]
        for obj in obj_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert obj == obj_ret, "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_pickling_sarray_types(self):

        sarray_list = [
            SArray([1,2,3]),
            SArray([1.0,2.0,3.5]),
            SArray(["foo", "bar"]),
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert list(obj) ==  list(obj_ret), \
                       "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_pickling_sframe_types(self):

        sarray_list = [
            SFrame([1,2,3]),
            SFrame([1.0,2.0,3.5]),
            SFrame(["foo", "bar"]),
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj, obj_ret)

    def test_pickling_sgraph_types(self):

        sg_test_1 = SGraph().add_vertices([
                                Vertex(0, {'fluffy': 1}),
                                Vertex(1, {'fluffy': 1, 'woof': 1}),
                                Vertex(2, {})])

        sg_test_2 = SGraph()
        sg_test_2 = sg_test_2.add_vertices([
                            Vertex(x) for x in [0, 1, 2]])
        sg_test_2 = sg_test_2.add_edges([
                        Edge(0, 1, attr={'relationship': 'dislikes'}),
                        Edge(1, 2, attr={'relationship': 'likes'}),
                        Edge(1, 0, attr={'relationship': 'likes'})])

        sarray_list = [
            sg_test_1,
            sg_test_2
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj.get_vertices(), obj_ret.get_vertices())
            assert_sframe_equal(obj.get_edges(), obj_ret.get_edges())

    def test_combination_gl_python_types(self):

        sg_test_1 = SGraph().add_vertices([
                                Vertex(1, {'fluffy': 1}),
                                Vertex(2, {'fluffy': 1, 'woof': 1}),
                                Vertex(3, {})])
        sarray_test_1 = SArray([1,2,3])
        sframe_test_1 = SFrame([1,2,3])

        obj_list = [
            [sg_test_1, sframe_test_1, sarray_test_1],
            {0:sg_test_1, 1:sframe_test_1, 2:sarray_test_1}
        ]

        for obj in obj_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj[0].get_vertices(), obj_ret[0].get_vertices())
            assert_sframe_equal(obj[0].get_edges(), obj_ret[0].get_edges())
            assert_sframe_equal(obj[1], obj_ret[1])
            assert list(obj[2]) == list(obj_ret[2])

    def test_pickle_compatibility(self):
        obj_list = [
            1, "hello", 5.0,
            (1, 2), ("i", "love", "cricket"),
            [1, 2, "hello"], [1.3, (1,2), "foo"], ["bar", {"foo": "bar"}],
            {"cricket": "best-sport", "test": [1,2,3]}, {"foo": 1.3},
        ]
        for obj in obj_list:
            file = open(self.filename, 'wb')
            pickler = pickle.Pickler(file)
            pickler.dump(obj)
            file.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert obj == obj_ret, \
                "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_cloud_pickle_compatibility(self):
        obj_list = [
            1, "hello", 5.0,
            (1, 2), ("i", "love", "cricket"),
            [1, 2, "hello"], [1.3, (1,2), "foo"], ["bar", {"foo": "bar"}],
            {"cricket": "best-sport", "test": [1,2,3]}, {"foo": 1.3},
        ]
        for obj in obj_list:
            file = open(self.filename, 'wb')
            pickler = cloudpickle.CloudPickler(file)
            pickler.dump(obj)
            file.close()
            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert obj == obj_ret, \
                "Failed pickling in %s (Got back %s)" % (obj, obj_ret)

    def test_relative_path(self):
        # Arrange
        sf1 = SFrame(range(10))
        relative_path = 'tmp/%s' % self.filename

        # Act
        pickler = gl_pickle.GLPickler(relative_path)
        pickler.dump(sf1)
        pickler.close()
        sf2 = gl_pickle.GLUnpickler(relative_path).load()

        # Assert
        assert_sframe_equal(sf1, sf2)

        # Clean up
        shutil.rmtree(relative_path)

    @moto.mock_s3
    def test_save_to_s3(self):
        # Arrange
        os.environ['GRAPHLAB_UNIT_TEST'] = 'PredictiveService'
        os.environ['AWS_ACCESS_KEY_ID'] = ''
        os.environ['AWS_SECRET_ACCESS_KEY'] = ''
        BUCKET_NAME = 'foo'
        boto.connect_s3().create_bucket(BUCKET_NAME)
        sf1 = SFrame(range(10))
        S3_PATH = "s3://%s/foobar" % BUCKET_NAME

        # Act
        pickler = gl_pickle.GLPickler(S3_PATH)
        pickler.dump(sf1)
        pickler.close()
        sf2 = gl_pickle.GLUnpickler(S3_PATH).load()

        # Assert
        assert_sframe_equal(sf1, sf2)

        del os.environ['GRAPHLAB_UNIT_TEST']

    def _test_backward_compatibility(self):

        # Arrange
        file_name = 's3://gl-internal-datasets/models/1.3/gl-pickle.gl'
        obj = {'foo': SFrame([1,2,3]),
               'bar': SFrame(),
               'foo-bar': ['foo-and-bar', SFrame([1])]}

        # Act
        unpickler = gl_pickle.GLUnpickler(file_name)
        obj_ret = unpickler.load()

        # Assert
        assert_sframe_equal(obj['foo'], obj_ret['foo'])
        assert_sframe_equal(obj['bar'], obj_ret['bar'])
        assert_sframe_equal(obj['foo-bar'][1], obj_ret['foo-bar'][1])
        self.assertEqual(obj['foo-bar'][0], obj_ret['foo-bar'][0])

    def test_save_over_previous(self):

        sarray_list = [
            SFrame([1,2,3]),
            SFrame([1.0,2.0,3.5]),
            SFrame(["foo", "bar"]),
        ]
        for obj in sarray_list:
            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()

            obj_ret = gl_pickle.GLUnpickler(self.filename).load()
            assert_sframe_equal(obj, obj_ret)

            pickler = gl_pickle.GLPickler(self.filename)
            pickler.dump(obj)
            pickler.close()

<|endoftext|>


<commit_before>"""Helpers for writing templates (not for build rules directly)."""

__all__ = [
    'parse_common_args',
    'tapeout_files',
    'write_json_to',
    'render_template',
]

import functools
import json

from garage import scripts

def parse_common_args(template):
    """Parse template arguments by the convention."""

    # Don't use setdefault() in parsers since arguments may be
    # None-valued.

    def parse_root(kwargs, arg):
        root = kwargs.get(arg)
        kwargs[arg] = root or '//base:root'

    def parse_name(kwargs, arg):
        name = kwargs.get(arg)
        if not name:
            name = ''
        elif not name.endswith('/'):
            name += '/'
        kwargs[arg] = name

    parsers = []
    for arg, anno in template.__annotations__.items():
        if anno == 'root':
            parsers.append(functools.partial(parse_root, arg=arg))
        elif anno == 'name':
            parsers.append(functools.partial(parse_name, arg=arg))
        else:
            raise AssertionError('cannot parse: %s' % anno)

    @functools.wraps(template)
    def wrapper(*args, **kwargs):
        for parser in parsers:
            parser(kwargs)
        return template(*args, **kwargs)

    return wrapper

def tapeout_files(parameters, paths, excludes=()):
    with scripts.using_sudo():
        rootfs = parameters['//base:drydock/rootfs']
        scripts.rsync(paths, rootfs, relative=True, excludes=excludes)

def write_json_to(obj, path):
    if scripts.is_dry_run():
        return
    with scripts.ensure_path(path).open('w') as json_file:
        json_file.write(json.dumps(obj, indent=4, sort_keys=True))
        json_file.write('\n')

def render_template(
        parameters, *,
        template_path,
        template_dirs=(),
        template_vars=(),
        output_path):
    """Render a template file (with Mako)."""
    cmd = [
        scripts.ensure_file(
            parameters['//host/cpython:python'],
        ),
        scripts.ensure_file(
            parameters['//base:root'] /
            'shipyard/scripts/render-template',
        ),
    ]
    for dir_path in template_dirs:
        cmd.append('--template-dir')
        cmd.append(dir_path)
    for name, value in template_vars:
        cmd.append('--template-var')
        cmd.append(name)
        cmd.append(json.dumps(value))
    cmd.append('--output')
    cmd.append(output_path.absolute())
    cmd.append(template_path.absolute())
    scripts.execute(cmd)
<commit_msg>Separate render_template and render_template_to_path<commit_after>"""Helpers for writing templates (not for build rules directly)."""

__all__ = [
    'parse_common_args',
    'tapeout_files',
    'write_json_to',
    'render_template',
    'render_template_to_path',
]

import functools
import json
import tempfile

from garage import scripts

def parse_common_args(template):
    """Parse template arguments by the convention."""

    # Don't use setdefault() in parsers since arguments may be
    # None-valued.

    def parse_root(kwargs, arg):
        root = kwargs.get(arg)
        kwargs[arg] = root or '//base:root'

    def parse_name(kwargs, arg):
        name = kwargs.get(arg)
        if not name:
            name = ''
        elif not name.endswith('/'):
            name += '/'
        kwargs[arg] = name

    parsers = []
    for arg, anno in template.__annotations__.items():
        if anno == 'root':
            parsers.append(functools.partial(parse_root, arg=arg))
        elif anno == 'name':
            parsers.append(functools.partial(parse_name, arg=arg))
        else:
            raise AssertionError('cannot parse: %s' % anno)

    @functools.wraps(template)
    def wrapper(*args, **kwargs):
        for parser in parsers:
            parser(kwargs)
        return template(*args, **kwargs)

    return wrapper

def tapeout_files(parameters, paths, excludes=()):
    with scripts.using_sudo():
        rootfs = parameters['//base:drydock/rootfs']
        scripts.rsync(paths, rootfs, relative=True, excludes=excludes)

def write_json_to(obj, path):
    if scripts.is_dry_run():
        return
    with scripts.ensure_path(path).open('w') as json_file:
        json_file.write(json.dumps(obj, indent=4, sort_keys=True))
        json_file.write('\n')

def render_template(parameters, **kwargs):
    with tempfile.NamedTemporaryFile() as output:
        output_path = scripts.ensure_path(output.name)
        render_template_to_path(parameters, output_path=output_path, **kwargs)
        return output_path.read_text()

def render_template_to_path(
        parameters, *,
        template_path,
        template_dirs=(),
        template_vars=(),
        output_path):
    """Render a template file (with Mako)."""
    cmd = [
        scripts.ensure_file(
            parameters['//host/cpython:python'],
        ),
        scripts.ensure_file(
            parameters['//base:root'] /
            'shipyard/scripts/render-template',
        ),
    ]
    for dir_path in template_dirs:
        cmd.append('--template-dir')
        cmd.append(dir_path)
    for name, value in template_vars:
        cmd.append('--template-var')
        cmd.append(name)
        cmd.append(json.dumps(value))
    cmd.append('--output')
    cmd.append(output_path)
    cmd.append(template_path)
    scripts.execute(cmd)
<|endoftext|>


<commit_before>"""
Tinman command line interface
"""
__author__ = 'Gavin M. Roy'
__email__ = 'gmr@myyearbook.com'
__since__ = '2011-06-06'

# Tinman imports
from . import application
from . import utils
from . import __version__

# General imports
import logging
import multiprocessing
import optparse
import os
import signal
import socket
import sys
import time

# Tornado imports
from tornado import httpserver
from tornado import ioloop
from tornado import version as tornado_version

class TinmanProcess(object):
    """
    Manages the tinman master process when run from the cli
    """

    def __init__(self):

        # Set a default IOloop to None
        self._ioloop = None

        # Get our logger
        self._logger = logging.getLogger('tinman')

    def _build_connections(self, config):
        """Build and attach our supported connections to our IOLoop and
         application object.

        Full configuration file
        """
        # Build a RabbitMQ connection if it exists
        self._build_rabbitmq_connection(config.get('RabbitMQ'))

        # Build a Redis connection if it exists
        self._build_redis_connection(config.get('Redis'))

    def _build_rabbitmq_connection(self, config):
        """Create a connection to RabbitMQ if we have it configured in our
        configuration file.

        :param config: RabbitMQ node of the configuration file dictionary
        :type config: dict
        """
        if not config:
            return

        # Import RabbitMQ only if we need it
        from clients import rabbitmq

        # Create our RabbitMQ instance, it will auto-connect and setup based on
        # this
        rabbitmq = rabbitmq.RabbitMQ(config.get('host'),
                                     config.get('port'),
                                     config.get('virtual_host'),
                                     config.get('username'),
                                     config.get('password'))

        # Add it to our tinman attribute at the application scope
        self._application.tinman.add('rabbitmq', rabbitmq)

    def _build_redis_connection(self, config):
        """Create a connection to Redis if we have it configured in our
        configuration file.

        :param config: Redis node of the configuration file dictionary
        :type config: dict
        """
        if not config:
            return

        # Import Redis only if we need it
        from clients import redis

        # Create our Redis instance, it will auto-connect and setup
        redis = redis.Redis(config.get('host'),
                            config.get('port'),
                            config.get('db'),
                            config.get('password'))

        # Add it to our tinman attribute at the application scope
        self._application.tinman.add('redis', redis)

    def _check_required_configuration_parameters(self, config, options):
        """Validates that the required configuration parameters are set.

        :raises: AttributeError

        """
        # Required sections
        if 'Application' not in config:
            raise AttributeError("Missing Application section in configuration")

        if 'HTTPServer' not in config:
            raise AttributeError("Missing HTTPServer section in configuration")

        if 'Logging' not in config:
            raise AttributeError("Missing Logging section in configuration")

        if not options.route_decorator and 'Routes' not in config:
            raise AttributeError("Missing Routes section in configuration")

        if not options.route_decorator and \
           not isinstance(config['Routes'], list):
            raise AttributeError("Error in Routes section in configuration")

    def _process_options(self):
        """Process the cli options returning the options and arguments"""
        usage = "usage: %prog -c <configfile> [options]"
        version_string = "%%prog v%s" % __version__
        description = "Tinman's Tornado application runner"

        # Create our parser and setup our command line options
        parser = optparse.OptionParser(usage=usage,
                                       version=version_string,
                                       description=description)

        parser.add_option("-c", "--config",
                          action="store",
                          dest="config",
                          help="Specify the configuration file for use")
        parser.add_option("-f", "--foreground",
                          action="store_true",
                          dest="foreground",
                          default=False,
                          help="Run interactively in console")
        parser.add_option("-r", "--route-decorator",
                          action="store_true",
                          dest="route_decorator",
                          default=False,
                          help="Utilize the route decorator instead of the Routes\
                                section of the config")

        # Parse our options and arguments
        return parser.parse_args()

    def _shutdown_signal_handler(self, signum, frame):
        """
        Called on SIGTERM to shutdown the sub-process
        """
        if self._ioloop:
            self._ioloop.stop()

    def _start_processes(self, config):

        # Loop through and kick off our processes
        self._children = []
        for port in config['HTTPServer']['ports']:
            self._logger.info("Starting Tinman v%s process for port %i",
                              __version__, port)
            # Kick off the child process
            child = multiprocessing.Process(target=self._start_server,
                                            name="tinman-%i" % port,
                                            args=(config, port))
            self._children.append(child)
            child.start()

        # Log our completion
        self._logger.debug("All children spawned")

    def _start_server(self, config, port):

        # Setup our signal handler
        signal.signal(signal.SIGTERM, self._shutdown_signal_handler)

        # Start our application
        self._application = \
            application.TinmanApplication(config.get('Routes', None),
                                          **config.get('Application', dict()))

        # Try and build any connection types we automatically support
        self._build_connections(config)

        # Start the HTTP Server
        self._logger.info("Starting Tornado v%s HTTPServer on port %i",
                          tornado_version, port)
        http_server = httpserver.HTTPServer(self._application)
        try:
            http_server.listen(port)
        except socket.error as error:
            # If we couldn't bind to IPv6 (Tornado 2.0+)
            if str(error).find('bad family'):
                http_server.bind(port, family=socket.AF_INET)
                http_server.start(1)

        # Get a handle to the instance of IOLoop
        self._ioloop = ioloop.IOLoop.instance()

        # Start the IOLoop
        try:
            self._ioloop.start()
        except KeyboardInterrupt:
            self._logger.info('KeyboardInterrupt received, shutting down.')

    def run(self):
        """Run the Tinman Process"""

        # Process our command line options
        options, args = self._process_options()

        # No configuration file?
        if options.config is None:
            sys.stdout.write("\nConfiguration not specified, \
running Tinman Test Application\n")
            from . import test
            config = test.CONFIG

        # Load the configuration file
        else:
            config = utils.load_configuration_file(options.config)

        # Check our required options
        self._check_required_configuration_parameters(config, options)

        # If we have a base path set prepend it to our Python import path
        if 'base_path' in config:
            sys.path.insert(0, config['base_path'])

        # Setup our logging
        utils.setup_logging(config['Logging'], options.foreground)

        # Setup our signal handlers
        utils.setup_signals()

        # Daemonize if we need to
        if not options.foreground:
            utils.daemonize(pidfile=config.get("pidfile", None),
                            user=config.get("user", None),
                            group=config.get("group", None))
        else:
            config['Application']['debug'] = True

        self._start_processes(config)

        # Have a main event loop for dealing with stats
        utils.running = True
        while utils.running:
            try:
                time.sleep(1)
            except KeyboardInterrupt:
                self._logger.info("CTRL-C pressed, shutting down.")
                break

        # Tell all the children to shutdown
        for child in self._children:
            self._logger.debug("Sending terminate signal to %s", child.name)
            child.terminate()

        # Loop while the children are shutting down
        self._logger.debug("Waiting for children to shutdown")
        while True in [child.is_alive() for child in self._children]:
            time.sleep(0.5)

        # Remove our pidfile
        if 'pidfile' in config:
            if os.path.exists(config['pidfile']):
                try:
                    os.unlink(config['pidfile'])
                except OSError as e:
                    self._logger.error("Could not remove pidfile: %s", e)

        # Log that we're shutdown cleanly
        self._logger.info("tinman has shutdown")
<commit_msg>Remove the route decorator parts until its done<commit_after>"""
Tinman command line interface
"""
__author__ = 'Gavin M. Roy'
__email__ = 'gmr@myyearbook.com'
__since__ = '2011-06-06'

# Tinman imports
from . import application
from . import utils
from . import __version__

# General imports
import logging
import multiprocessing
import optparse
import os
import signal
import socket
import sys
import time

# Tornado imports
from tornado import httpserver
from tornado import ioloop
from tornado import version as tornado_version

class TinmanProcess(object):
    """
    Manages the tinman master process when run from the cli
    """

    def __init__(self):

        # Set a default IOloop to None
        self._ioloop = None

        # Get our logger
        self._logger = logging.getLogger('tinman')

    def _build_connections(self, config):
        """Build and attach our supported connections to our IOLoop and
         application object.

        Full configuration file
        """
        # Build a RabbitMQ connection if it exists
        self._build_rabbitmq_connection(config.get('RabbitMQ'))

        # Build a Redis connection if it exists
        self._build_redis_connection(config.get('Redis'))

    def _build_rabbitmq_connection(self, config):
        """Create a connection to RabbitMQ if we have it configured in our
        configuration file.

        :param config: RabbitMQ node of the configuration file dictionary
        :type config: dict
        """
        if not config:
            return

        # Import RabbitMQ only if we need it
        from clients import rabbitmq

        # Create our RabbitMQ instance, it will auto-connect and setup based on
        # this
        rabbitmq = rabbitmq.RabbitMQ(config.get('host'),
                                     config.get('port'),
                                     config.get('virtual_host'),
                                     config.get('username'),
                                     config.get('password'))

        # Add it to our tinman attribute at the application scope
        self._application.tinman.add('rabbitmq', rabbitmq)

    def _build_redis_connection(self, config):
        """Create a connection to Redis if we have it configured in our
        configuration file.

        :param config: Redis node of the configuration file dictionary
        :type config: dict
        """
        if not config:
            return

        # Import Redis only if we need it
        from clients import redis

        # Create our Redis instance, it will auto-connect and setup
        redis = redis.Redis(config.get('host'),
                            config.get('port'),
                            config.get('db'),
                            config.get('password'))

        # Add it to our tinman attribute at the application scope
        self._application.tinman.add('redis', redis)

    def _check_required_configuration_parameters(self, config, options):
        """Validates that the required configuration parameters are set.

        :raises: AttributeError

        """
        # Required sections
        if 'Application' not in config:
            raise AttributeError("Missing Application section in configuration")

        if 'HTTPServer' not in config:
            raise AttributeError("Missing HTTPServer section in configuration")

        if 'Logging' not in config:
            raise AttributeError("Missing Logging section in configuration")

        if not isinstance(config['Routes'], list):
            raise AttributeError("Error in Routes section in configuration")

    def _process_options(self):
        """Process the cli options returning the options and arguments"""
        usage = "usage: %prog -c <configfile> [options]"
        version_string = "%%prog v%s" % __version__
        description = "Tinman's Tornado application runner"

        # Create our parser and setup our command line options
        parser = optparse.OptionParser(usage=usage,
                                       version=version_string,
                                       description=description)

        parser.add_option("-c", "--config",
                          action="store",
                          dest="config",
                          help="Specify the configuration file for use")
        parser.add_option("-f", "--foreground",
                          action="store_true",
                          dest="foreground",
                          default=False,
                          help="Run interactively in console")

        # Parse our options and arguments
        return parser.parse_args()

    def _shutdown_signal_handler(self, signum, frame):
        """
        Called on SIGTERM to shutdown the sub-process
        """
        if self._ioloop:
            self._ioloop.stop()

    def _start_processes(self, config):

        # Loop through and kick off our processes
        self._children = []
        for port in config['HTTPServer']['ports']:
            self._logger.info("Starting Tinman v%s process for port %i",
                              __version__, port)
            # Kick off the child process
            child = multiprocessing.Process(target=self._start_server,
                                            name="tinman-%i" % port,
                                            args=(config, port))
            self._children.append(child)
            child.start()

        # Log our completion
        self._logger.debug("All children spawned")

    def _start_server(self, config, port):

        # Setup our signal handler
        signal.signal(signal.SIGTERM, self._shutdown_signal_handler)

        # Start our application
        self._application = \
            application.TinmanApplication(config.get('Routes', None),
                                          **config.get('Application', dict()))

        # Try and build any connection types we automatically support
        self._build_connections(config)

        # Start the HTTP Server
        self._logger.info("Starting Tornado v%s HTTPServer on port %i",
                          tornado_version, port)
        http_server = httpserver.HTTPServer(self._application)
        try:
            http_server.listen(port)
        except socket.error as error:
            # If we couldn't bind to IPv6 (Tornado 2.0+)
            if str(error).find('bad family'):
                http_server.bind(port, family=socket.AF_INET)
                http_server.start(1)

        # Get a handle to the instance of IOLoop
        self._ioloop = ioloop.IOLoop.instance()

        # Start the IOLoop
        try:
            self._ioloop.start()
        except KeyboardInterrupt:
            self._logger.info('KeyboardInterrupt received, shutting down.')

    def run(self):
        """Run the Tinman Process"""

        # Process our command line options
        options, args = self._process_options()

        # No configuration file?
        if options.config is None:
            sys.stdout.write("\nConfiguration not specified, \
running Tinman Test Application\n")
            from . import test
            config = test.CONFIG

        # Load the configuration file
        else:
            config = utils.load_configuration_file(options.config)

        # Check our required options
        self._check_required_configuration_parameters(config, options)

        # If we have a base path set prepend it to our Python import path
        if 'base_path' in config:
            sys.path.insert(0, config['base_path'])

        # Setup our logging
        utils.setup_logging(config['Logging'], options.foreground)

        # Setup our signal handlers
        utils.setup_signals()

        # Daemonize if we need to
        if not options.foreground:
            utils.daemonize(pidfile=config.get("pidfile", None),
                            user=config.get("user", None),
                            group=config.get("group", None))
        else:
            config['Application']['debug'] = True

        self._start_processes(config)

        # Have a main event loop for dealing with stats
        utils.running = True
        while utils.running:
            try:
                time.sleep(1)
            except KeyboardInterrupt:
                self._logger.info("CTRL-C pressed, shutting down.")
                break

        # Tell all the children to shutdown
        for child in self._children:
            self._logger.debug("Sending terminate signal to %s", child.name)
            child.terminate()

        # Loop while the children are shutting down
        self._logger.debug("Waiting for children to shutdown")
        while True in [child.is_alive() for child in self._children]:
            time.sleep(0.5)

        # Remove our pidfile
        if 'pidfile' in config:
            if os.path.exists(config['pidfile']):
                try:
                    os.unlink(config['pidfile'])
                except OSError as e:
                    self._logger.error("Could not remove pidfile: %s", e)

        # Log that we're shutdown cleanly
        self._logger.info("tinman has shutdown")
<|endoftext|>


<commit_before>"""
Problem migration operations for the shell manager.
"""

import json, os

from sys import stdin
from copy import deepcopy
from os.path import join
from re import findall

from shell_manager.util import PROBLEM_ROOT, sanitize_name, get_problem_root, get_problem

#More in-depth validation should occur with some sort of linting step.
PROBLEM_FIELDS = [
    ["author", {"required": True, "validation": str}],
    ["score", {"required": True, "validation": lambda score: type(score) == int and score >= 0}],
    ["name", {"required": True, "validation": str}],
    ["description", {"required": True, "validation": str}],
    ["category", {"required": True, "validation": str}],
    ["tags", {"required": False, "validation": list}],
    ["hints", {"required": False, "validation": list}],
    ["organization", {"required": False, "validation": str}],
    ["pkg_description", {"required": False, "validation": str}],
    ["pkg_version", {"required": False, "validation": str}],
    ["pkg_name", {"required": False, "validation": str}],
    ["pkg_dependencies", {"required": False, "validation": str}],
]

PROBLEM_DEFAULTS = {
    "version": lambda problem: "1.0-0",
    "pkg_dependencies": lambda problem: [],
    "tags": lambda problem: [],
    "hints": lambda problem: [],
    "organization": "ctf"
}

def translate_problem_fields(translation_table, problem):
    """
    Migrate old problem fields to newer names through the translation table.

    Args:
        translation_table: dict with keys which specify old fields and values which correspond to the new ones.
        problem: the problem object.

    Example:
        Given a translation table of {"name": "display_name"}, the problem would change the "name" key to "display_name".
    """

    for old_field, new_field in translation_table.items():
        if old_field in problem:
            value = problem.pop(old_field)
            problem[new_field] = value

def set_problem_defaults(problem, additional_defaults={}):
    """
    Set the default fields for a given problem.

    Args:
        problem: the problem object.
        additional_defaults: use case specific defaults that should be included.
                             In the same form as defaults, {field: lambda problem: "default"}
    """

    total_defaults = deepcopy(PROBLEM_DEFAULTS)
    total_defaults.update(**additional_defaults)

    for field, setter in total_defaults.items():
        if field not in problem:
            #Use the associated default function
            problem[field] = setter(problem)

def set_problem(problem_path, problem):
    """
    Overwrite the problem.json of a given problem.

    Args:
        problem_path: path to the root of the problem's directory.
        problem: the problem object.
    """

    serialized_problem = json.dumps(problem, indent=True)
    json_path = join(problem_path, "problem.json")

    with open(json_path, "w") as problem_file:
        problem_file.write(serialized_problem)

def migrate_cs2014_problem(problem_path, problem, overrides={}):
    """
    Convert a Cyberstakes 2014 problem to the updated format.

    Args:
        problem_path: path to the root of the problem's directory.
        problem: the problem object to modify.

    Returns:
        A new problem object that is consistent with the current spec.
    """

    field_table = {
        "basescore": "score",
        "displayname": "name",
        "desc": "description",
        "categories": "category"
    }

    category_translations = {
        "binary": "Binary Exploitation",
        "crypto": "Cryptography",
        "web": "Web Exploitation",
        "forensics": "Forensics",
        "reversing": "Reverse Engineering",
        "tutorial": "Tutorial"
    }

    new_defaults = {
        "author": lambda problem: overrides.get("author", "Nihil"),
        "organization": lambda problem: overrides.get("organization", "ctf")
    }

    def get_dependencies(problem_path):
        """
        Retrieve problem dependencies from challenge.py

        Args:
            problem_path: path to the root of the problem's directory.
        """

        challenge_path = join(problem_path, "challenge.py")
        with open(challenge_path, "r") as challenge_file:
            challenge = challenge_file.read()
            dependencies = []
            for requirements in ["local_requirements", "remote_requirements"]:
                requirements_index = challenge.find(requirements)

                #Could not find it, we shouldn't do anything.
                if requirements_index == -1:
                    continue

                requirements_bound = challenge.find("]", requirements_index)

                dependencies_text = challenge[requirements_index:requirements_bound]
                dependencies.extend(findall(r"'([a-z0-9-\+\.]+)'\s*(?:,)?", dependencies_text))

            #Remove any duplicates
            dependencies = list(set(dependencies))

            new_defaults["pkg_dependencies"] = lambda problem: dependencies

    get_dependencies(problem_path)
    translate_problem_fields(field_table, problem)

    #Get the first, primary category.
    problem["category"] = problem["category"][0] if type(problem["category"]) == list else problem["category"]
    if problem["category"] in category_translations:
        problem["category"] = category_translations[problem["category"]]

    set_problem_defaults(problem, additional_defaults=new_defaults)

    return problem

# Corresponds to possible migration formats.
MIGRATION_TABLE = {
    "cyberstakes2014": migrate_cs2014_problem
}

def migrate_problems(args, config):
    """ Main entrypoint for problem migration. """

    additional_defaults = {}
    for default_pair in args.set_defaults:
        if ":" in default_pair:
            field, value = default_pair.split(":")
            additional_defaults[field] = value

    for problem_path in args.problem_paths:
        problem = get_problem(problem_path)
        problem_copy = deepcopy(problem)

        print("Migrating '{}'...".format(problem_path))

        migrater = MIGRATION_TABLE[args.legacy_format]
        updated_problem = migrater(problem_path, problem_copy,
                                   overrides=additional_defaults)

        if args.dry:
            print(updated_problem)
        else:
            print("Updated '{}' to the new problem format.".format(updated_problem["name"]))
            set_problem(problem_path, updated_problem)
<commit_msg>Updated problem migration features with new logging.<commit_after>"""
Problem migration operations for the shell manager.
"""

import json, os, logging

from sys import stdin
from copy import deepcopy
from os.path import join
from re import findall

from shell_manager.util import PROBLEM_ROOT, sanitize_name, get_problem_root, get_problem
from shell_manager.util import FatalException

logger = logging.getLogger(__name__)

#More in-depth validation should occur with some sort of linting step.
PROBLEM_FIELDS = [
    ["author", {"required": True, "validation": str}],
    ["score", {"required": True, "validation": lambda score: type(score) == int and score >= 0}],
    ["name", {"required": True, "validation": str}],
    ["description", {"required": True, "validation": str}],
    ["category", {"required": True, "validation": str}],
    ["tags", {"required": False, "validation": list}],
    ["hints", {"required": False, "validation": list}],
    ["organization", {"required": False, "validation": str}],
    ["pkg_description", {"required": False, "validation": str}],
    ["pkg_version", {"required": False, "validation": str}],
    ["pkg_name", {"required": False, "validation": str}],
    ["pkg_dependencies", {"required": False, "validation": str}],
]

PROBLEM_DEFAULTS = {
    "version": lambda problem: "1.0-0",
    "pkg_dependencies": lambda problem: [],
    "tags": lambda problem: [],
    "hints": lambda problem: [],
    "organization": "ctf"
}

def translate_problem_fields(translation_table, problem):
    """
    Migrate old problem fields to newer names through the translation table.

    Args:
        translation_table: dict with keys which specify old fields and values which correspond to the new ones.
        problem: the problem object.

    Example:
        Given a translation table of {"name": "display_name"}, the problem would change the "name" key to "display_name".
    """

    for old_field, new_field in translation_table.items():
        if old_field in problem:
            value = problem.pop(old_field)
            problem[new_field] = value

def set_problem_defaults(problem, additional_defaults={}):
    """
    Set the default fields for a given problem.

    Args:
        problem: the problem object.
        additional_defaults: use case specific defaults that should be included.
                             In the same form as defaults, {field: lambda problem: "default"}
    """

    total_defaults = deepcopy(PROBLEM_DEFAULTS)
    total_defaults.update(**additional_defaults)

    for field, setter in total_defaults.items():
        if field not in problem:
            #Use the associated default function
            problem[field] = setter(problem)

def set_problem(problem_path, problem):
    """
    Overwrite the problem.json of a given problem.

    Args:
        problem_path: path to the root of the problem's directory.
        problem: the problem object.
    """

    serialized_problem = json.dumps(problem, indent=True)
    json_path = join(problem_path, "problem.json")

    with open(json_path, "w") as problem_file:
        problem_file.write(serialized_problem)

def migrate_cs2014_problem(problem_path, problem, overrides={}):
    """
    Convert a Cyberstakes 2014 problem to the updated format.

    Args:
        problem_path: path to the root of the problem's directory.
        problem: the problem object to modify.

    Returns:
        A new problem object that is consistent with the current spec.
    """

    field_table = {
        "basescore": "score",
        "displayname": "name",
        "desc": "description",
        "categories": "category"
    }

    category_translations = {
        "binary": "Binary Exploitation",
        "crypto": "Cryptography",
        "web": "Web Exploitation",
        "forensics": "Forensics",
        "reversing": "Reverse Engineering",
        "tutorial": "Tutorial"
    }

    new_defaults = {
        "author": lambda problem: overrides.get("author", "Nihil"),
        "organization": lambda problem: overrides.get("organization", "ctf")
    }

    def get_dependencies(problem_path):
        """
        Retrieve problem dependencies from challenge.py

        Args:
            problem_path: path to the root of the problem's directory.
        """

        challenge_path = join(problem_path, "challenge.py")
        with open(challenge_path, "r") as challenge_file:
            challenge = challenge_file.read()
            dependencies = []
            for requirements in ["local_requirements", "remote_requirements"]:
                requirements_index = challenge.find(requirements)

                #Could not find it, we shouldn't do anything.
                if requirements_index == -1:
                    continue

                requirements_bound = challenge.find("]", requirements_index)

                dependencies_text = challenge[requirements_index:requirements_bound]
                dependencies.extend(findall(r"'([a-z0-9-\+\.]+)'\s*(?:,)?", dependencies_text))

            #Remove any duplicates
            dependencies = list(set(dependencies))

            new_defaults["pkg_dependencies"] = lambda problem: dependencies

    get_dependencies(problem_path)
    translate_problem_fields(field_table, problem)

    #Get the first, primary category.
    problem["category"] = problem["category"][0] if type(problem["category"]) == list else problem["category"]
    if problem["category"] in category_translations:
        problem["category"] = category_translations[problem["category"]]

    set_problem_defaults(problem, additional_defaults=new_defaults)

    return problem

# Corresponds to possible migration formats.
MIGRATION_TABLE = {
    "cyberstakes2014": migrate_cs2014_problem
}

def migrate_problems(args, config):
    """ Main entrypoint for problem migration. """

    additional_defaults = {}
    for default_pair in args.set_defaults:
        if ":" in default_pair:
            field, value = default_pair.split(":")
            additional_defaults[field] = value

    for problem_path in args.problem_paths:
        problem = get_problem(problem_path)
        problem_copy = deepcopy(problem)

        logger.debug("Migrating '%s' from legacy %s format.", problem["name"], args.legacy_format)

        migrater = MIGRATION_TABLE[args.legacy_format]
        updated_problem = migrater(problem_path, problem_copy,
                                   overrides=additional_defaults)

        if args.dry:
            print(updated_problem)
        else:
            logger.info("Updated '%s' to the new problem format.", problem["name"])
            set_problem(problem_path, updated_problem)
<|endoftext|>