<commit_before>'''
conatcts.py
----------
Contact Book Application to CRUD personal contacts via a shelve DB
    - Functions
    - Objects
    - Modules
----------
Python Demo Workshop, March 22nd 2017
'''
### Import Modules
import shelve # Database Management Module [https://docs.python.org/2/library/shelve.html]
import os

# Database Name
DB_NAME = '.Contacts.db'

### Contact Class
# Python Class definiton of model used to represent personal contacts
class Contact():

    # Constructor Method
    def __init__(self, name, phone, address, email, nickname):
        self.name = name
        self.phone = phone
        self.address = address
        self.email = email
        self.nickname = nickname

# open_DB
# Returns open instance of shelve object
def open_DB():
    try:
        db = shelve.open(DB_NAME)

        # Syntax to check if db was opened properly
        # Verbose condition needed, as first time DB is opened it is empty
        #   -> Checking <if db> would not display confirmation, even though its true
        if db is not None:
            print('Opened DB to read: {0}'.format(DB_NAME))

        return db

    # ERROR CASE[1]: Unable to open Database
    except Exception:
        print('ERROR[1]: Unable to open Database')
        print('Exception information: {0}'.format(e))

        exit(1) # !! IF WE GOT HERE WE MESSED UP !! *#

# close_DB
# Closes instance of shelve object
def close_DB(db):
    try:
        db.close()
        return True

    # ERROR CASE[1]: Unable to close Database
    except Exception:
        print('ERROR[2]: Unable to open Database')
        print('Exception information: {0}'.format(e))

        exit(2) # !! IF WE GOT HERE WE MESSED UP !! *#

# store_contact
# CREATE
# Creates and saves new Contact object to DB
def store_contact(db):
    print('You are going to add a contact to your book...')

    # Get new contact information
    name = raw_input('Please enter the persons name: ')
    phone = raw_input('Please enter the telephone number: ')
    address = raw_input('Please enter the address: ')
    email = raw_input('Please enter the email: ')
    nick = raw_input('Please enter the nickname (or leave blank for none): ')

    # Account for no nickname
    if nick == '':
        nick = None

    # Create new Contact Instance
    contact = Contact(name, phone, address, email, nick)

    # Save to DB
    db[phone] = contact # Store contact as value with key being phone number

    return True

# print_DB
# READ
# Prints all currently stored in DB
def print_DB(db):
    print('Printing database...')

    # Loop through dictionary to print values
    for value in db.values():
        print('Name: {}'.format(value.name))
        print('Phone: {}'.format(value.phone))
        print('Address: {}'.format(value.address))
        print('Email: {}'.format(value.email))

        if value.nickname:
            print('Nickname: {}'.format(value.nickname))

        print '-' * 75

# edit_contact
# UPDATE
# Edits a contact accessed through key [phone number]
def edit_contact(db):
    print('Going to update a record in the database...')

    # Print out key/value pairs
    # iteritems() will return a list of tuples that contain the key in the first position
    # and value in the second
    for key, value in db.iteritems():
        print '-' * 75
        print('Contact name: {}'.format(value.name))
        print('Contact phone number: {}'.format(key))
        print '-' * 75

    idx = raw_input('Please enter the phone number of the contact to edit: ')

    try:
        contact = db[idx]
        print('Going to update {}...'.format(contact.name))

        # Phone number is not reassigned due to its use as primary key
        # -> Is there a way around this?
        contact.name = raw_input('Please enter new name of contact: ')
        contact.address = raw_input('Please enter the new address: ')
        contact.email = raw_input('Please enter the new email: ')
        contact.nickname = raw_input('Please enter the new nickname (or leave blank for none): ')

        if contact.nickname == '':
            contact.nickname = None

        db[idx] = contact

        return True

    # ERROR CASE[3]: Contact with matching key not found in DB
    except KeyError:
        print('No contact found with phone number to edit: {}'.format(idx))
        print('Exiting update function...')

        return False

#delet_contact
# DELETE
# Deletes a particular contact from the DB
def delete_contact(db):
    print('Going to delete a record in the database...')

    # Print out key/value pairs
    for key, value in db.iteritems():
        print '-' * 75
        print('Contact name: {}'.format(value.name))
        print('Contact phone number: {}'.format(key))
        print '-' * 75

    idx = raw_input('Please enter the phone number of the contact to delete: ')

    try:
        contact = db[idx]

        # Confirm user wants to delete contact
        confirm = raw_input('Are you SURE you want to delete {} from your contacts? (y/n):'.format(contact.name))

        if confirm.lower() == 'y':
            del db[idx] # Deletes key/value pair in DB on key
            print('Deleted {} OK'.format(contact.name))
            return True

        else:
            print('{} will stay in DB...'.format(contact.name))

    # ERROR CASE[4]: Contact with matching key not found in DB
    except KeyError:
        print('No contact found with phone number to delete: {}'.format(idx))
        print('Exiting delete function...')

        return False

# main
# Driver of program; Calls other functions that do stuff
def main():
    db = open_DB()

    # Main Loop to control flow of program
    while True:
        print '*' * 75
        print('Main Menu')
        print '*' * 75
        print('Please select a command: ')
        print('1. Add a contact')
        print('2. Print the database')
        print('3. Edit a contact information')
        print('4. Delete a contact')
        print('5. Exit program')

        # Get user choice to run desired function
        choice = raw_input('> ')

        try:
            choice = int(choice)
            flag = True

            ### No switch/case in Python, so gonna have to do this the long way...
            # EXIT CASE
            if choice == 5:
                print('Thanks for stopping by!')
                break

            # Store a new contact
            elif choice == 1:
                flag = store_contact(db)

            # Print DB
            elif choice == 2:
                print_DB(db)

            # Edit Contact
            elif choice == 3:
                flag = edit_contact(db)

            # Delete Contact
            elif choice == 4:
                flag = delete_contact(db)

            # ERROR[5]: User enters in option not covered in menu
            else:
                print('ERROR[5]: Please enter a valid command (1-5)')
                flag = False

            if flag:
                print('[Operation successful...]')

        # ERROR CASE[4]: User enters string in Main Menu
        except ValueError:
            print('ERROR[4]: Menu choice must be an int, not a string')

    isClosed = close_DB(db)

    if isClosed:
        print('Closed {} OK'.format(DB_NAME))

# Entry Point of Script
if __name__ == '__main__':
    main()
<commit_msg>Removed unused import<commit_after>'''
conatcts.py
----------
Contact Book Application to CRUD personal contacts via a shelve DB
    - Functions
    - Objects
    - Modules
----------
Python Demo Workshop, March 22nd 2017
'''
### Import Modules
import shelve # Database Management Module [https://docs.python.org/2/library/shelve.html]

# Database Name
DB_NAME = '.Contacts.db'

### Contact Class
# Python Class definiton of model used to represent personal contacts
class Contact():

    # Constructor Method
    def __init__(self, name, phone, address, email, nickname):
        self.name = name
        self.phone = phone
        self.address = address
        self.email = email
        self.nickname = nickname

# open_DB
# Returns open instance of shelve object
def open_DB():
    try:
        db = shelve.open(DB_NAME)

        # Syntax to check if db was opened properly
        # Verbose condition needed, as first time DB is opened it is empty
        #   -> Checking <if db> would not display confirmation, even though its true
        if db is not None:
            print('Opened DB to read: {0}'.format(DB_NAME))

        return db

    # ERROR CASE[1]: Unable to open Database
    except Exception:
        print('ERROR[1]: Unable to open Database')
        print('Exception information: {0}'.format(e))

        exit(1) # !! IF WE GOT HERE WE MESSED UP !! *#

# close_DB
# Closes instance of shelve object
def close_DB(db):
    try:
        db.close()
        return True

    # ERROR CASE[1]: Unable to close Database
    except Exception:
        print('ERROR[2]: Unable to open Database')
        print('Exception information: {0}'.format(e))

        exit(2) # !! IF WE GOT HERE WE MESSED UP !! *#

# store_contact
# CREATE
# Creates and saves new Contact object to DB
def store_contact(db):
    print('You are going to add a contact to your book...')

    # Get new contact information
    name = raw_input('Please enter the persons name: ')
    phone = raw_input('Please enter the telephone number: ')
    address = raw_input('Please enter the address: ')
    email = raw_input('Please enter the email: ')
    nick = raw_input('Please enter the nickname (or leave blank for none): ')

    # Account for no nickname
    if nick == '':
        nick = None

    # Create new Contact Instance
    contact = Contact(name, phone, address, email, nick)

    # Save to DB
    db[phone] = contact # Store contact as value with key being phone number

    return True

# print_DB
# READ
# Prints all currently stored in DB
def print_DB(db):
    print('Printing database...')

    # Loop through dictionary to print values
    for value in db.values():
        print('Name: {}'.format(value.name))
        print('Phone: {}'.format(value.phone))
        print('Address: {}'.format(value.address))
        print('Email: {}'.format(value.email))

        if value.nickname:
            print('Nickname: {}'.format(value.nickname))

        print '-' * 75

# edit_contact
# UPDATE
# Edits a contact accessed through key [phone number]
def edit_contact(db):
    print('Going to update a record in the database...')

    # Print out key/value pairs
    # iteritems() will return a list of tuples that contain the key in the first position
    # and value in the second
    for key, value in db.iteritems():
        print '-' * 75
        print('Contact name: {}'.format(value.name))
        print('Contact phone number: {}'.format(key))
        print '-' * 75

    idx = raw_input('Please enter the phone number of the contact to edit: ')

    try:
        contact = db[idx]
        print('Going to update {}...'.format(contact.name))

        # Phone number is not reassigned due to its use as primary key
        # -> Is there a way around this?
        contact.name = raw_input('Please enter new name of contact: ')
        contact.address = raw_input('Please enter the new address: ')
        contact.email = raw_input('Please enter the new email: ')
        contact.nickname = raw_input('Please enter the new nickname (or leave blank for none): ')

        if contact.nickname == '':
            contact.nickname = None

        db[idx] = contact

        return True

    # ERROR CASE[3]: Contact with matching key not found in DB
    except KeyError:
        print('No contact found with phone number to edit: {}'.format(idx))
        print('Exiting update function...')

        return False

#delet_contact
# DELETE
# Deletes a particular contact from the DB
def delete_contact(db):
    print('Going to delete a record in the database...')

    # Print out key/value pairs
    for key, value in db.iteritems():
        print '-' * 75
        print('Contact name: {}'.format(value.name))
        print('Contact phone number: {}'.format(key))
        print '-' * 75

    idx = raw_input('Please enter the phone number of the contact to delete: ')

    try:
        contact = db[idx]

        # Confirm user wants to delete contact
        confirm = raw_input('Are you SURE you want to delete {} from your contacts? (y/n):'.format(contact.name))

        if confirm.lower() == 'y':
            del db[idx] # Deletes key/value pair in DB on key
            print('Deleted {} OK'.format(contact.name))
            return True

        else:
            print('{} will stay in DB...'.format(contact.name))

    # ERROR CASE[4]: Contact with matching key not found in DB
    except KeyError:
        print('No contact found with phone number to delete: {}'.format(idx))
        print('Exiting delete function...')

        return False

# main
# Driver of program; Calls other functions that do stuff
def main():
    db = open_DB()

    # Main Loop to control flow of program
    while True:
        print '*' * 75
        print('Main Menu')
        print '*' * 75
        print('Please select a command: ')
        print('1. Add a contact')
        print('2. Print the database')
        print('3. Edit a contact information')
        print('4. Delete a contact')
        print('5. Exit program')

        # Get user choice to run desired function
        choice = raw_input('> ')

        try:
            choice = int(choice)
            flag = True

            ### No switch/case in Python, so gonna have to do this the long way...
            # EXIT CASE
            if choice == 5:
                print('Thanks for stopping by!')
                break

            # Store a new contact
            elif choice == 1:
                flag = store_contact(db)

            # Print DB
            elif choice == 2:
                print_DB(db)

            # Edit Contact
            elif choice == 3:
                flag = edit_contact(db)

            # Delete Contact
            elif choice == 4:
                flag = delete_contact(db)

            # ERROR[5]: User enters in option not covered in menu
            else:
                print('ERROR[5]: Please enter a valid command (1-5)')
                flag = False

            if flag:
                print('[Operation successful...]')

        # ERROR CASE[4]: User enters string in Main Menu
        except ValueError:
            print('ERROR[4]: Menu choice must be an int, not a string')

    isClosed = close_DB(db)

    if isClosed:
        print('Closed {} OK'.format(DB_NAME))

# Entry Point of Script
if __name__ == '__main__':
    main()
<|endoftext|>


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

import os
import sys
import re
import socket
import getopt
import threading
import subprocess
import shlex
import time
import select

blackhole = (
'10::2222',
'101::1234',
'2001::212',
'2001:da8:112::21ae',
'2003:ff:1:2:3:4:5fff:6',
'2003:ff:1:2:3:4:5fff:7',
'2003:ff:1:2:3:4:5fff:8',
'2003:ff:1:2:3:4:5fff:9',
'2003:ff:1:2:3:4:5fff:10',
'2003:ff:1:2:3:4:5fff:11',
'2003:ff:1:2:3:4:5fff:12',
'21:2::2',
'2123::3e12')

dns = {
'google_a':'2001:4860:4860::8888',
'google_b':'2001:4860:4860::8844',
'he_net':'2001:470:20::2',
'lax_he_net':'2001:470:0:9d::2'
}

config = {
'dns':dns['google_b'],
'infile':'hosts',
'outfile':'hosts.out',
'querytype':'aaaa',
'threadnum':10
}

hosts = []
done_num = 0
thread_lock = threading.Lock()
running = True

class worker_thread(threading.Thread):
    def __init__(self, start_pt, end_pt):
        threading.Thread.__init__(self)
        self.start_pt = start_pt
        self.end_pt = end_pt

    def run(self):
        global hosts, done_num
        for i in range(self.start_pt, self.end_pt):
            if not running: break

            line = hosts[i].strip()

            with thread_lock:
                done_num += 1

            if line == '' or line[0:2] == '##':
                hosts[i] = line + '\r\n'
                continue

            arr = line.lstrip('#').split()

            if len(arr) == 1:
                domain = arr[0]
            else:
                domain = arr[1]

            flag = False
            if validate_domain(domain):
                ret = query_domain(domain, False)

                if ret in blackhole or ret == '':
                    ret = query_domain(domain, True)

                if ret:
                    flag = True
                    arr[0] = ret

            if flag:
                if len(arr) == 1:
                   arr.append(domain)
            else:
                arr[0] = '#' + arr[0]

            hosts[i] = ' '.join(arr)
            hosts[i] += '\r\n'

class watcher_thread(threading.Thread):
    def run(self):
        total_num = len(hosts)

        wn = int(config['threadnum'])
        if wn > total_num:
            wn = total_num
        print "There are %d threads working..." % wn
        print "Press 'Enter' to exit.\n"

        while True:
            if sys.stdin in select.select([sys.stdin], [], [], 0)[0]:
                t = raw_input()
                global running
                with thread_lock:
                    running = False
                print 'Waiting threads to exit...'
                break

            with thread_lock:
                dn = done_num

            outbuf = "Total: %d lines, Done: %d lines, Ratio: %d %%.\r"\
                   % (total_num, dn, dn * 100 / total_num)
            print outbuf,
            sys.stdout.flush()

            if done_num == total_num:
                print outbuf
                break

            time.sleep(1)

def query_domain(domain, tcp):
    cmd = "dig +short +time=2 -6 %s @'%s' '%s'"\
        % (config['querytype'], config['dns'], domain)

    if tcp:
        cmd = cmd[:3] + ' +tcp' + cmd[3:]

    proc = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE)
    out, err = proc.communicate()

    outarr = out.splitlines()

    if len(outarr) == 0:
        ret = ''
    else:
        if validate_ip_addr(outarr[-1]):
            ret = outarr[-1]
        else:
            ret = ''

    return ret

def validate_domain(domain):
    pattern = '^((?!-)[*A-Za-z0-9-]{1,63}(?<!-)\\.)+[A-Za-z]{2,6}$'
    p = re.compile(pattern)
    m = p.match(domain)
    if m:
        return True
    else:
        return False

def validate_ip_addr(ip_addr):
    if ':' in ip_addr:
        try:
            socket.inet_pton(socket.AF_INET6, ip_addr)
            return True
        except socket.error:
            return False
    else:
        try:
            socket.inet_pton(socket.AF_INET, ip_addr)
            return True
        except socket.error:
            return False

def print_help():
    print('''usage: update_hosts [OPTIONS] FILE
A simple multi-threading tool used to update hosts file.

Options:
  -h, --help             show this help message and exit
  -s DNS                 set another dns server, default: 2001:4860:4860::8844
  -o OUT_FILE            ouput file, default: inputfilename.out
  -t QUERY_TYPE          dig command query type, defalut: aaaa
  -n THREAD_NUM          set the number of worker threads, default: 10
''')

def get_config():
    shortopts = 'hs:o:t:n:'
    longopts = ['help']

    try:
        optlist, args = getopt.gnu_getopt(sys.argv[1:], shortopts, longopts)
    except getopt.GetoptError as e:
        print e, '\n'
        print_help()
        sys.exit(1)

    global config
    if len(args) != 1:
        print "You must specify the input hosts file (only one)."
        sys.exit(1)

    config['infile'] = args[0]
    config['outfile'] = args[0] + '.out'

    for key, value in optlist:
        if key == '-s':
            config['dns'] = value
        elif key == '-o':
            config['outfile'] = value
        elif key == '-t':
            config['querytype'] = value
        elif key == '-m':
            config['method'] = value
        elif key == '-n':
            config['threadnum'] = int(value)
        elif key in ('-h', '--help'):
            print_help()
            sys.exit(0)

def main():
    get_config()

    global hosts

    try:
        with open(config['infile'], 'r') as infile:
            hosts = infile.readlines()
    except IOError as e:
        print e
        sys.exit(e.errno)

    if os.path.exists(config['outfile']):
        config['outfile'] += '.new'

    try:
        outfile = open(config['outfile'], 'w')
    except IOError as e:
        print e
        sys.exit(e.errno)

    print "Input: %s    Output: %s\n" % (config['infile'], config['outfile'])

    threads = []

    t = watcher_thread()
    t.start()
    threads.append(t)

    worker_num = config['threadnum']
    lines_num = len(hosts)

    lines_per_thread = lines_num / worker_num
    lines_remain = lines_num % worker_num

    start_pt = 0

    for i in range(worker_num):
        if not running: break

        lines_for_thread = lines_per_thread

        if lines_for_thread == 0 and lines_remain == 0:
            break

        if lines_remain > 0:
            lines_for_thread += 1
            lines_remain -= 1

        t = worker_thread(start_pt, start_pt + lines_for_thread)
        start_pt += lines_for_thread

        t.start()
        threads.append(t)

    for t in threads:
        t.join()

    try:
        outfile.writelines(hosts)
    except IOError as e:
        print e
        sys.exit(e.errno)

    sys.exit(0)

if __name__ == '__main__':
    main()

<commit_msg>fix the bug when only -h option is given<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import re
import socket
import getopt
import threading
import subprocess
import shlex
import time
import select

blackhole = (
'10::2222',
'101::1234',
'2001::212',
'2001:da8:112::21ae',
'2003:ff:1:2:3:4:5fff:6',
'2003:ff:1:2:3:4:5fff:7',
'2003:ff:1:2:3:4:5fff:8',
'2003:ff:1:2:3:4:5fff:9',
'2003:ff:1:2:3:4:5fff:10',
'2003:ff:1:2:3:4:5fff:11',
'2003:ff:1:2:3:4:5fff:12',
'21:2::2',
'2123::3e12')

dns = {
'google_a':'2001:4860:4860::8888',
'google_b':'2001:4860:4860::8844',
'he_net':'2001:470:20::2',
'lax_he_net':'2001:470:0:9d::2'
}

config = {
'dns':dns['google_b'],
'infile':'',
'outfile':'',
'querytype':'aaaa',
'threadnum':10
}

hosts = []
done_num = 0
thread_lock = threading.Lock()
running = True

class worker_thread(threading.Thread):
    def __init__(self, start_pt, end_pt):
        threading.Thread.__init__(self)
        self.start_pt = start_pt
        self.end_pt = end_pt

    def run(self):
        global hosts, done_num
        for i in range(self.start_pt, self.end_pt):
            if not running: break

            line = hosts[i].strip()

            with thread_lock:
                done_num += 1

            if line == '' or line[0:2] == '##':
                hosts[i] = line + '\r\n'
                continue

            arr = line.lstrip('#').split()

            if len(arr) == 1:
                domain = arr[0]
            else:
                domain = arr[1]

            flag = False
            if validate_domain(domain):
                ret = query_domain(domain, False)

                if ret in blackhole or ret == '':
                    ret = query_domain(domain, True)

                if ret:
                    flag = True
                    arr[0] = ret

            if flag:
                if len(arr) == 1:
                   arr.append(domain)
            else:
                arr[0] = '#' + arr[0]

            hosts[i] = ' '.join(arr)
            hosts[i] += '\r\n'

class watcher_thread(threading.Thread):
    def run(self):
        total_num = len(hosts)

        wn = int(config['threadnum'])
        if wn > total_num:
            wn = total_num
        print "There are %d threads working..." % wn
        print "Press 'Enter' to exit.\n"

        while True:
            if sys.stdin in select.select([sys.stdin], [], [], 0)[0]:
                t = raw_input()
                global running
                with thread_lock:
                    running = False
                print 'Waiting threads to exit...'
                break

            with thread_lock:
                dn = done_num

            outbuf = "Total: %d lines, Done: %d lines, Ratio: %d %%.\r"\
                   % (total_num, dn, dn * 100 / total_num)
            print outbuf,
            sys.stdout.flush()

            if done_num == total_num:
                print outbuf
                break

            time.sleep(1)

def query_domain(domain, tcp):
    cmd = "dig +short +time=2 -6 %s @'%s' '%s'"\
        % (config['querytype'], config['dns'], domain)

    if tcp:
        cmd = cmd[:3] + ' +tcp' + cmd[3:]

    proc = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE)
    out, err = proc.communicate()

    outarr = out.splitlines()

    if len(outarr) == 0:
        ret = ''
    else:
        if validate_ip_addr(outarr[-1]):
            ret = outarr[-1]
        else:
            ret = ''

    return ret

def validate_domain(domain):
    pattern = '^((?!-)[*A-Za-z0-9-]{1,63}(?<!-)\\.)+[A-Za-z]{2,6}$'
    p = re.compile(pattern)
    m = p.match(domain)
    if m:
        return True
    else:
        return False

def validate_ip_addr(ip_addr):
    if ':' in ip_addr:
        try:
            socket.inet_pton(socket.AF_INET6, ip_addr)
            return True
        except socket.error:
            return False
    else:
        try:
            socket.inet_pton(socket.AF_INET, ip_addr)
            return True
        except socket.error:
            return False

def print_help():
    print('''usage: update_hosts [OPTIONS] FILE
A simple multi-threading tool used to update hosts file.

Options:
  -h, --help             show this help message and exit
  -s DNS                 set another dns server, default: 2001:4860:4860::8844
  -o OUT_FILE            ouput file, default: inputfilename.out
  -t QUERY_TYPE          dig command query type, defalut: aaaa
  -n THREAD_NUM          set the number of worker threads, default: 10
''')

def get_config():
    shortopts = 'hs:o:t:n:'
    longopts = ['help']

    try:
        optlist, args = getopt.gnu_getopt(sys.argv[1:], shortopts, longopts)
    except getopt.GetoptError as e:
        print e, '\n'
        print_help()
        sys.exit(1)

    global config
    for key, value in optlist:
        if key == '-s':
            config['dns'] = value
        elif key == '-o':
            config['outfile'] = value
        elif key == '-t':
            config['querytype'] = value
        elif key == '-m':
            config['method'] = value
        elif key == '-n':
            config['threadnum'] = int(value)
        elif key in ('-h', '--help'):
            print_help()
            sys.exit(0)

    if len(args) != 1:
        print "You must specify the input hosts file (only one)."
        sys.exit(1)

    config['infile'] = args[0]
    if config['outfile'] == '':
        config['outfile'] = config['infile'] + '.out'

def main():
    get_config()

    global hosts

    try:
        with open(config['infile'], 'r') as infile:
            hosts = infile.readlines()
    except IOError as e:
        print e
        sys.exit(e.errno)

    if os.path.exists(config['outfile']):
        config['outfile'] += '.new'

    try:
        outfile = open(config['outfile'], 'w')
    except IOError as e:
        print e
        sys.exit(e.errno)

    print "Input: %s    Output: %s\n" % (config['infile'], config['outfile'])

    threads = []

    t = watcher_thread()
    t.start()
    threads.append(t)

    worker_num = config['threadnum']
    lines_num = len(hosts)

    lines_per_thread = lines_num / worker_num
    lines_remain = lines_num % worker_num

    start_pt = 0

    for i in range(worker_num):
        if not running: break

        lines_for_thread = lines_per_thread

        if lines_for_thread == 0 and lines_remain == 0:
            break

        if lines_remain > 0:
            lines_for_thread += 1
            lines_remain -= 1

        t = worker_thread(start_pt, start_pt + lines_for_thread)
        start_pt += lines_for_thread

        t.start()
        threads.append(t)

    for t in threads:
        t.join()

    try:
        outfile.writelines(hosts)
    except IOError as e:
        print e
        sys.exit(e.errno)

    sys.exit(0)

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before>from __future__ import absolute_import

import requests

from dataverse.dataset import Dataset
from dataverse.exceptions import (
    ConnectionError, MethodNotAllowedError, OperationFailedError,
)
from dataverse.utils import get_element, get_elements, sanitize

class Dataverse(object):
    def __init__(self, connection, collection):
        self.connection = connection
        self.collection = collection

        self._collection_info = None
        self._contents_json = None

    @property
    def is_published(self):

        # Always check latest version
        collection_info = requests.get(
            self.collection.get('href'),
            auth=self.connection.auth,
        ).content

        status_tag = get_element(
            collection_info,
            namespace='http://purl.org/net/sword/terms/state',
            tag='dataverseHasBeenReleased',
        )
        status = status_tag.text

        return status.lower() == 'true'

    @property
    def alias(self):
        return self.collection.get('href').split('/')[-1]

    @property
    def title(self):
        return sanitize(get_element(
            self.collection,
            namespace='atom',
            tag='title',
        ).text)

    def get_contents(self, refresh=False):
        if not refresh and self._contents_json:
            return self._contents_json

        content_uri = '{0}/dataverses/{1}/contents'.format(
            self.connection.native_base_url, self.alias
        )
        resp = requests.get(
            content_uri,
            params={'key': self.connection.token}
        )

        if resp.status_code != 200:
            raise ConnectionError('Atom entry could not be retrieved.')

        self._contents_json = resp.json()['data']
        return self._contents_json

    def get_collection_info(self, refresh=False):
        if not refresh and self._collection_info:
            return self._collection_info

        self._collection_info = requests.get(
            self.collection.get('href'),
            auth=self.connection.auth,
        ).content
        return self._collection_info

    def publish(self):
        edit_uri = '{0}/edit/dataverse/{1}'.format(
            self.connection.sword_base_url, self.alias
        )
        resp = requests.post(
            edit_uri,
            headers={'In-Progress': 'false'},
            auth=self.connection.auth,
        )

        if resp.status_code != 200:
            raise OperationFailedError('The Dataverse could not be published.')

    def create_dataset(self, title, description, creator, **kwargs):
        dataset = Dataset(
            title=title,
            description=description,
            creator=creator,
            **kwargs
        )

        self._add_dataset(dataset)
        return dataset

    def _add_dataset(self, dataset):

        resp = requests.post(
            self.collection.get('href'),
            data=dataset.get_entry(),
            headers={'Content-type': 'application/atom+xml'},
            auth=self.connection.auth,
        )

        if resp.status_code != 201:
            raise OperationFailedError('This dataset could not be added.')

        dataset.dataverse = self
        dataset._refresh(receipt=resp.content)
        self.get_collection_info(refresh=True)

    def delete_dataset(self, dataset):
        if dataset.get_state() == 'DELETED' or dataset.get_state() == 'DEACCESSIONED':
            return

        resp = requests.delete(
            dataset.edit_uri,
            auth=self.connection.auth,
        )
        if resp.status_code == 405:
            raise MethodNotAllowedError(
                'Published datasets can only be deleted from the GUI. For '
                'more information, please refer to '
                'https://github.com/IQSS/dataverse/issues/778'
            )

        dataset.is_deleted = True
        self.get_collection_info(refresh=True)

    def get_datasets(self, refresh=False):
        collection_info = self.get_collection_info(refresh)
        entries = get_elements(collection_info, tag='entry')
        return [Dataset.from_dataverse(entry, self) for entry in entries]

    def get_dataset_by_doi(self, doi, refresh=False):
        return next((s for s in self.get_datasets(refresh) if s.doi == doi), None)

    def get_dataset_by_title(self, title, refresh=False):
        return next((s for s in self.get_datasets(refresh) if s.title == title), None)

    def get_dataset_by_string_in_entry(self, string, refresh=False):
        return next((s for s in self.get_datasets(refresh) if string in s.get_entry()), None)
<commit_msg>optional timeout for get dataset methods<commit_after>from __future__ import absolute_import

import requests

from dataverse.dataset import Dataset
from dataverse.exceptions import (
    ConnectionError, MethodNotAllowedError, OperationFailedError,
)
from dataverse.utils import get_element, get_elements, sanitize

class Dataverse(object):
    def __init__(self, connection, collection):
        self.connection = connection
        self.collection = collection

        self._collection_info = None
        self._contents_json = None

    @property
    def is_published(self):

        # Always check latest version
        collection_info = requests.get(
            self.collection.get('href'),
            auth=self.connection.auth,
        ).content

        status_tag = get_element(
            collection_info,
            namespace='http://purl.org/net/sword/terms/state',
            tag='dataverseHasBeenReleased',
        )
        status = status_tag.text

        return status.lower() == 'true'

    @property
    def alias(self):
        return self.collection.get('href').split('/')[-1]

    @property
    def title(self):
        return sanitize(get_element(
            self.collection,
            namespace='atom',
            tag='title',
        ).text)

    def get_contents(self, refresh=False):
        if not refresh and self._contents_json:
            return self._contents_json

        content_uri = '{0}/dataverses/{1}/contents'.format(
            self.connection.native_base_url, self.alias
        )
        resp = requests.get(
            content_uri,
            params={'key': self.connection.token}
        )

        if resp.status_code != 200:
            raise ConnectionError('Atom entry could not be retrieved.')

        self._contents_json = resp.json()['data']
        return self._contents_json

    def get_collection_info(self, refresh=False, timeout=None):
        if not refresh and self._collection_info:
            return self._collection_info

        self._collection_info = requests.get(
            self.collection.get('href'),
            auth=self.connection.auth,
            timeout=timeout,
        ).content
        return self._collection_info

    def publish(self):
        edit_uri = '{0}/edit/dataverse/{1}'.format(
            self.connection.sword_base_url, self.alias
        )
        resp = requests.post(
            edit_uri,
            headers={'In-Progress': 'false'},
            auth=self.connection.auth,
        )

        if resp.status_code != 200:
            raise OperationFailedError('The Dataverse could not be published.')

    def create_dataset(self, title, description, creator, **kwargs):
        dataset = Dataset(
            title=title,
            description=description,
            creator=creator,
            **kwargs
        )

        self._add_dataset(dataset)
        return dataset

    def _add_dataset(self, dataset):

        resp = requests.post(
            self.collection.get('href'),
            data=dataset.get_entry(),
            headers={'Content-type': 'application/atom+xml'},
            auth=self.connection.auth,
        )

        if resp.status_code != 201:
            raise OperationFailedError('This dataset could not be added.')

        dataset.dataverse = self
        dataset._refresh(receipt=resp.content)
        self.get_collection_info(refresh=True)

    def delete_dataset(self, dataset):
        if dataset.get_state() == 'DELETED' or dataset.get_state() == 'DEACCESSIONED':
            return

        resp = requests.delete(
            dataset.edit_uri,
            auth=self.connection.auth,
        )
        if resp.status_code == 405:
            raise MethodNotAllowedError(
                'Published datasets can only be deleted from the GUI. For '
                'more information, please refer to '
                'https://github.com/IQSS/dataverse/issues/778'
            )

        dataset.is_deleted = True
        self.get_collection_info(refresh=True)

    def get_datasets(self, refresh=False, timeout=None):
        collection_info = self.get_collection_info(refresh, timeout=timeout)
        entries = get_elements(collection_info, tag='entry')
        return [Dataset.from_dataverse(entry, self) for entry in entries]

    def get_dataset_by_doi(self, doi, refresh=False, timeout=None):
        return next((s for s in self.get_datasets(refresh, timeout=timeout) if s.doi == doi), None)

    def get_dataset_by_title(self, title, refresh=False, timeout=None):
        return next((s for s in self.get_datasets(refresh, timeout=timeout) if s.title == title), None)

    def get_dataset_by_string_in_entry(self, string, refresh=False, timeout=None):
        return next((s for s in self.get_datasets(refresh, timeout=timeout) if string in s.get_entry()), None)
<|endoftext|>


<commit_before>import sys

RIGHT = 0
LEFT = 1
INC = 2
DEC = 3
LOOP_START = 4
LOOP_END = 5
GETC = 6
PUTC = 7

def compile_file(source_file):
    program_n = 0

    try:
        with open(source_file, "r") as f:
            eof = False
            while not eof:
                c = f.read(1)
                if len(c) == 0:
                    eof = True
                elif c == '.':
                    program_n += 1
    except IOError:
        print "Error: could not open file."
        sys.exit(1)

    binary = bin(program_n)[3:]

    if len(binary) % 3 != 0:
        print "Error: either you have too many or too few dots."
        sys.exit(1)

    program = []
    for i in xrange(0, len(binary), 3):
        b = binary[i:i+3]
        inst = int(b, 2)
        program.append(inst)

    return program

def match_loop(text, start):
    if text[start] == LOOP_START:
        stop = len(text)
        dir = 1
        search = LOOP_END
    elif text[start] == LOOP_END:
        stop = -1
        dir = -1
        search = LOOP_START
    else:
        raise ValueError('Incorrect number of dots.')
        return -1

    depth = 0

    for i in range(start, stop, dir):
        c = text[i]
        if c == LOOP_START:
            depth += 1
        elif c == LOOP_END:
            depth -= 1
        if c == search and depth == 0:
            return i
    return -1

def run(program):
    pc = 0

    mem = [0 for i in xrange(30000)]
    pointer = 0

    while pc < len(program):
        c = program[pc]
        old_pc = pc
        if c == RIGHT:
            pointer += 1
            while pointer <= len(mem):
                mem.append(0)
        elif c == LEFT:
            pointer -= 1
            if pointer < 0:
                print "Error: either you have too many or too few dots. (out of bounds)"
                break
        elif c == INC:
            mem[pointer] += 1
            mem[pointer] %= 256
        elif c == DEC:
            mem[pointer] -= 1
            mem[pointer] %= 256
        elif c == PUTC:
            sys.stdout.write(chr(mem[pointer]))
        elif c == GETC:
            mem[pointer] = ord(sys.stdin.read(1))
        elif c == LOOP_START:
            if mem[pointer] == 0:
                pc = match_loop(program, pc)
                if pc < 0:
                    print "Error: either you have too many or too few dots. (unmatched loop)"
                    break
        elif c == LOOP_END:
            if mem[pointer] != 0:
                pc = match_loop(program, pc)
                if pc < 0:
                    print "Error: either you have too many or too few dots. (unmatched loop)"
                    break
        else:
            pass
        pc += 1

def main():
    if len(sys.argv) <= 1:
        print "usage: %s sourcefile" % sys.argv[0]
        sys.exit(1)

    source_file = sys.argv[1]
    program = compile_file(source_file)
    run(program)

if __name__ == "__main__":
    main()

<commit_msg>Fix RIGHT Operator<commit_after>import sys

RIGHT = 0
LEFT = 1
INC = 2
DEC = 3
LOOP_START = 4
LOOP_END = 5
GETC = 6
PUTC = 7

def compile_file(source_file):
    program_n = 0

    try:
        with open(source_file, "r") as f:
            eof = False
            while not eof:
                c = f.read(1)
                if len(c) == 0:
                    eof = True
                elif c == '.':
                    program_n += 1
    except IOError:
        print "Error: could not open file."
        sys.exit(1)

    binary = bin(program_n)[3:]

    if len(binary) % 3 != 0:
        print "Error: either you have too many or too few dots."
        sys.exit(1)

    program = []
    for i in xrange(0, len(binary), 3):
        b = binary[i:i+3]
        inst = int(b, 2)
        program.append(inst)

    return program

def match_loop(text, start):
    if text[start] == LOOP_START:
        stop = len(text)
        dir = 1
        search = LOOP_END
    elif text[start] == LOOP_END:
        stop = -1
        dir = -1
        search = LOOP_START
    else:
        raise ValueError('Incorrect number of dots.')
        return -1

    depth = 0

    for i in range(start, stop, dir):
        c = text[i]
        if c == LOOP_START:
            depth += 1
        elif c == LOOP_END:
            depth -= 1
        if c == search and depth == 0:
            return i
    return -1

def run(program):
    pc = 0

    mem = [0 for i in xrange(30000)]
    pointer = 0

    while pc < len(program):
        c = program[pc]
        old_pc = pc
        if c == RIGHT:
            pointer += 1
            while pointer >= len(mem):
                mem.append(0)
        elif c == LEFT:
            pointer -= 1
            if pointer < 0:
                print "Error: either you have too many or too few dots. (out of bounds)"
                break
        elif c == INC:
            mem[pointer] += 1
            mem[pointer] %= 256
        elif c == DEC:
            mem[pointer] -= 1
            mem[pointer] %= 256
        elif c == PUTC:
            sys.stdout.write(chr(mem[pointer]))
        elif c == GETC:
            mem[pointer] = ord(sys.stdin.read(1))
        elif c == LOOP_START:
            if mem[pointer] == 0:
                pc = match_loop(program, pc)
                if pc < 0:
                    print "Error: either you have too many or too few dots. (unmatched loop)"
                    break
        elif c == LOOP_END:
            if mem[pointer] != 0:
                pc = match_loop(program, pc)
                if pc < 0:
                    print "Error: either you have too many or too few dots. (unmatched loop)"
                    break
        else:
            pass
        pc += 1

def main():
    if len(sys.argv) <= 1:
        print "usage: %s sourcefile" % sys.argv[0]
        sys.exit(1)

    source_file = sys.argv[1]
    program = compile_file(source_file)
    run(program)

if __name__ == "__main__":
    main()

<|endoftext|>


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

"""
A class object for generating and storing Toytree plotting coordinates.
"""

import numpy as np
from .utils import ToytreeError

"""
TODO:

- it may be useful to decompose update so plotting coords change but not idxs
-
"""

class Coords:
    """
    Generates and stores plotting coordinates for nodes and edges of a tree.
    Uses the toytree _style information (e.g., layout, use_edge_lengths).

    Attributes:
    -----------
    edges: ndarray
        2-d array of pairs showing (parent, child) relationships.
        Rows are ordered in postorder traversal order for plotting.
    verts: ndarray
        2-d array with relative cartesian coordinate positions of nodes (n)
        Rows are ordered by the numbered 'idx' feature of nodes.
        ...
    """
    def __init__(self, ttree):

        # the toytree
        self.ttree = ttree

        # the vertices and edge tuples for normal layouts ('n')
        self.edges = None
        self.verts = None

    def update(self, layout=None):
        """
        Updates cartesian coordinates for drawing tree graph
        """
        # store tree dimensions on update
        self.ttree.nnodes = 0
        self.ttree.ntips = 0
        for node in self.ttree.treenode.traverse():
            self.ttree.nnodes += 1
            if node.is_leaf():
                self.ttree.ntips += 1

        # updates idxs and fixed_idx for any tree manipulations
        self.update_idxs()
        self.update_fixed_order()      # in case ntips changed

        # get new shape and clear for attrs
        self.edges = None  # np.zeros((self.ttree.nnodes - 1, 2), dtype=int)
        self.verts = None  # np.zeros((self.ttree.nnodes, 2), dtype=float)

        # get edges and verts (node locations)
        layout = self.ttree.style.layout
        if layout == 'c':
            self.circ = Circle(self.ttree)
            self.verts, self.edges = self.get_radial_coords()
        else:
            # returns with verts oriented by layout
            self.verts, self.edges = self.get_linear_coords(layout=layout)

        # orientation can reorder dimensions
        # self.verts = self.reorient_coords(self.verts)

    def update_idxs(self):
        "set root idx highest, tip idxs lowest ordered as ladderized"

        # internal nodes: root is highest idx
        idx = self.ttree.nnodes - 1

        for node in self.ttree.treenode.traverse("levelorder"):
            if not node.is_leaf():
                node.add_feature("idx", idx)
                self.ttree.idx_dict[idx] = node
                if not node.name:
                    node.name = str(idx)
                idx -= 1

        # external nodes: lowest numbers are for tips (0-N)
        for node in self.ttree.treenode.get_leaves():
            node.add_feature("idx", idx)
            self.ttree.idx_dict[idx] = node
            if not node.name:
                node.name = str(idx)
            idx -= 1

    def update_fixed_order(self):
        "after pruning fixed order needs update to match new nnodes/ntips."
        # set tips order if fixing for multi-tree plotting (default None)
        fixed_order = self.ttree._fixed_order
        self.ttree_fixed_order = None
        self.ttree_fixed_idx = list(range(self.ttree.ntips))

        # check if fixed_order changed:
        if fixed_order:
            fixed_order = [
                i for i in fixed_order if i in self.ttree.get_tip_labels()
            ]
            self.ttree._set_fixed_order(fixed_order)
        # else:
            # self.ttree._fixed_idx = list(range(self.ttree.ntips))

    def get_radial_coords(self, use_edge_lengths=True):
        """
        Assign .edges and .verts for node positions in a fan tree.
        The farthest tip aligns at the circumference.
        """
        # shortname
        if not use_edge_lengths:
            nbits = self.ttree.treenode.get_farthest_leaf(True)[1]

        # use cache to fill edges array
        for idx in range(self.ttree.nnodes):
            node = self.ttree.idx_dict[idx]
            if node.up:
                self.edges[idx, :] = (node.up.idx, idx)

            # leaves: x positions are evenly spaced around circumference
            if node.is_leaf() and (not node.is_root()):

                # store radians (how far around from zero to 2pi)
                node.radians = self.circ.tip_radians[idx]

                # get positions of tips using radians and radius
                if use_edge_lengths:
                    node.radius = self.circ.radius - node.height
                    node.x, node.y = self.circ.get_node_coords(node)
                else:
                    node.radius = nbits
                    node.x, node.y = self.circ.get_node_coords(node)

            # internal nodes comes after tips. Inherit position from children.
            else:

                # height is either distance or nodes from root
                if use_edge_lengths:
                    node.radius = self.circ.radius - node.height
                else:
                    node.radius = sum(1 for i in node.iter_ancestors())
                    # max([i.radius for i in node.children]) - 1

                # x position is halfway between children x-pos
                node.radians = np.mean([i.radians for i in node.children])
                node.x, node.y = self.circ.get_node_coords(node)

            # store the x,y vertex positions
            self.verts[node.idx] = [node.x, node.y]

        # scale so that node idx 0 (or fixed_order x) is at (0, 0)
        # adjust = self.ttree._coords.verts[:, 1].min()
        # self.ttree._coords.verts[:, 1] -= adjust

    def get_linear_coords(self, layout=None, use_edge_lengths=True):
        """
        Sets .edges, .verts for node positions.
        X and Y positions here refer to base assumption that tree is right
        facing, reorient_coordinates() will handle re-translating this.
        """
        edges = np.zeros((self.ttree.nnodes - 1, 2), dtype=int)
        verts = np.zeros((self.ttree.nnodes, 2), dtype=float)

        # use cache to fill edges array
        for idx in range(self.ttree.nnodes - 1):
            parent = self.ttree.idx_dict[idx].up
            if parent:
                edges[idx, :] = (parent.idx, idx)

        # store verts array with x,y positions of nodes (lengths of branches)
        # we want tips to align at the right face (larger axis number)
        fartip = self.ttree.treenode.get_farthest_leaf()[0]
        hgt = sum([
            fartip.dist,
            sum(i.dist for i in fartip.iter_ancestors() if not i.is_root())
        ])

        # finding tip distance from root is a little slow, it would be faster
        # just progress forward from root, but this wouldn't easily allow for
        # spacing fixed order tips. So instead we start from tips.
        for idx in range(self.ttree.nnodes):
            node = self.ttree.idx_dict[idx]

            if node.is_leaf() and (not node.is_root()):
                if use_edge_lengths:
                    toroot = sum(
                        i.dist for i in node.iter_ancestors()
                        if not i.is_root()
                    )
                    node.y = hgt - (node.dist + toroot)
                else:
                    node.y = 0.
                if self.ttree._fixed_order:
                    node.x = self.ttree._fixed_order.index(node.name)
                else:
                    node.x = node.idx

            else:
                nch = node.children
                # empty trees only
                if not nch:
                    node.x = node.y = 0

                # all other trees
                else:
                    if use_edge_lengths:
                        node.y = nch[0].y + nch[0].dist
                    else:
                        node.y = max((i.y for i in nch)) + 1

                    # internal node at midpoint of children
                    node.x = sum(i.x for i in nch) / float(len(nch))

            # store the vertex
            verts[idx] = [node.x, node.y]

        # scale so that node idx 0 (or fixed_order x) is at (0, 0)
        if use_edge_lengths:
            adjust = verts[:, 1].min()
            verts[:, 1] -= adjust

        # orient for layout
        verts = self.reorient_coords(verts, layout)
        return verts, edges

    def reorient_coords(self, verts, layout=None):
        """
        Returns a modified .verts array with new coordinates for nodes.
        This does not need to modify .edges. The order of nodes, and therefore
        of verts rows is still the same because it is still based on the tree
        branching order (ladderized usually).
        """

        # select layout
        if layout is None:
            layout = self.ttree.style.layout

        # if tree is empty then bail out
        if len(self.ttree) < 2:
            return verts

        # this is done later in drawing.
        if layout == 'c':
            return verts

        # default: Down-facing tips align at y=0, first ladderized tip at x=0
        elif layout == 'd':
            return verts

        # right-facing tips align at x=0, last ladderized tip at y=0
        elif layout == 'r':

            # verts swap x and ys and make xs 0 to negative
            tmp = np.zeros(verts.shape)
            tmp[:, 1] = verts[:, 0]
            tmp[:, 0] = verts[:, 1] * -1
            return tmp

        elif layout == 'l':

            # verts swap x and ys and make xs 0 to negative
            tmp = np.zeros(verts.shape)
            tmp[:, 1] = verts[:, 0]
            tmp[:, 0] = verts[:, 1]
            return tmp

        elif layout == 'u':
            # verts swap x and ys and make xs 0 to negative
            tmp = np.zeros(verts.shape)
            tmp[:, 1] = verts[:, 1] * -1
            tmp[:, 0] = verts[:, 0] * -1
            return tmp

        else:
            raise ToytreeError("layout not recognized")

    def get_radii(self, layout):
        """
        radius of every node relative to 0,0 origin.
        """

        # custom request layout
        if layout is None:
            layout = self.ttree.layout

        # get circular layout
        if layout == 'c':
            self.radii = np.array(
                [self.ttree.idx_dict[i].radius for i in range(self.nnodes)])
        else:
            self.radii = np.repeat(0, self.ntips)

class Circle:
    """
    When init from a toytree it can return coordinates based on
    distributing tips equally in radians from 0-2pi (TODO: allow
    setting start/end).

    The radius for farthest tip-ends is set to the tree height,
    and origin is at 0.0.
    """
    def __init__(self, tre):

        # set radius
        self.tre = tre
        self.radius = self.tre.treenode.height
        # get_distance(self.tre.treenode.get_farthest_leaf()[0])

        # origin
        self.o = (0, 0)
        # self.tre.style.xbaseline, self.tre.style.ybaseline)  # -self.radius, 0)

        # tips (bottom to top) are evenly spread from 0 to -2pi (counter clock)
        self.tip_radians = np.linspace(0, -np.pi * 2, self.tre.ntips + 1)[:-1]

    def get_node_coords(self, node):
        """
        get node coord. Node has radius assigned on the fly.
        """
        x = self.o[0] + node.radius * np.cos(node.radians)
        y = self.o[1] - node.radius * np.sin(node.radians)
        return x, y

    # def get_node_lines(self, node):
    #     x = self.o[0] + node.up.radius * np.cos(node.radians)
    #     y = self.o[1] - node.up.radius * np.sin(node.radians)
    #     return x, y

    def get_tip_end_angles(self):
        """
        node radians for printing tip labels should be from the root (0,0)
        """
        return np.array([np.rad2deg(abs(i)) for i in self.tip_radians])

    def get_tip_end_coords(self):
        """
        node tip coords must calculate new radian angle relative to parent
        node and then add offset amount to radius when calculating (x, y).
        """
        xs = [self.o[0] + self.radius * np.cos(i) for i in self.tip_radians]
        ys = [self.o[1] - self.radius * np.sin(i) for i in self.tip_radians]
        return np.stack([xs, ys], axis=1)

    # def get_parent_coords(self, node):
    #     """
    #     get x,y coordinates of parent based on radians of children
    #     """
    #     # calculate x, y coords
    #     x = self.o[0] + node.radius * np.cos(midradian)
    #     y = self.o[1] + node.radius * np.sin(midradian)
    #     return x, y

<commit_msg>all return arrays<commit_after>#!/usr/bin/env python

"""
A class object for generating and storing Toytree plotting coordinates.
"""

import numpy as np
from .utils import ToytreeError

"""
TODO:

- it may be useful to decompose update so plotting coords change but not idxs
-
"""

class Coords:
    """
    Generates and stores plotting coordinates for nodes and edges of a tree.
    Uses the toytree _style information (e.g., layout, use_edge_lengths).

    Attributes:
    -----------
    edges: ndarray
        2-d array of pairs showing (parent, child) relationships.
        Rows are ordered in postorder traversal order for plotting.
    verts: ndarray
        2-d array with relative cartesian coordinate positions of nodes (n)
        Rows are ordered by the numbered 'idx' feature of nodes.
        ...
    """
    def __init__(self, ttree):

        # the toytree
        self.ttree = ttree

        # the vertices and edge tuples for normal layouts ('n')
        self.edges = None
        self.verts = None

    def update(self, layout=None):
        """
        Updates cartesian coordinates for drawing tree graph
        """
        # store tree dimensions on update
        self.ttree.nnodes = 0
        self.ttree.ntips = 0
        for node in self.ttree.treenode.traverse():
            self.ttree.nnodes += 1
            if node.is_leaf():
                self.ttree.ntips += 1

        # updates idxs and fixed_idx for any tree manipulations
        self.update_idxs()
        self.update_fixed_order()      # in case ntips changed
        self.circ = Circle(self.ttree)

        # get new shape and clear for attrs
        self.edges = self.get_edges()

        # get edges and verts (node locations)
        layout = self.ttree.style.layout
        if layout == 'c':
            self.verts = self.get_radial_coords()
        else:
            # returns with verts oriented by layout
            self.verts = self.get_linear_coords(layout=layout)

        # orientation can reorder dimensions
        # self.verts = self.reorient_coords(self.verts)

    def update_idxs(self):
        "set root idx highest, tip idxs lowest ordered as ladderized"

        # internal nodes: root is highest idx
        idx = self.ttree.nnodes - 1

        for node in self.ttree.treenode.traverse("levelorder"):
            if not node.is_leaf():
                node.add_feature("idx", idx)
                self.ttree.idx_dict[idx] = node
                if not node.name:
                    node.name = str(idx)
                idx -= 1

        # external nodes: lowest numbers are for tips (0-N)
        for node in self.ttree.treenode.get_leaves():
            node.add_feature("idx", idx)
            self.ttree.idx_dict[idx] = node
            if not node.name:
                node.name = str(idx)
            idx -= 1

    def update_fixed_order(self):
        "after pruning fixed order needs update to match new nnodes/ntips."
        # set tips order if fixing for multi-tree plotting (default None)
        fixed_order = self.ttree._fixed_order
        self.ttree_fixed_order = None
        self.ttree_fixed_idx = list(range(self.ttree.ntips))

        # check if fixed_order changed:
        if fixed_order:
            fixed_order = [
                i for i in fixed_order if i in self.ttree.get_tip_labels()
            ]
            self.ttree._set_fixed_order(fixed_order)
        # else:
            # self.ttree._fixed_idx = list(range(self.ttree.ntips))

    def get_edges(self):
        """
        This could change if a tree dropped nodes.
        """
        # use cache to fill edges array
        edges = np.zeros((self.ttree.nnodes - 1, 2), dtype=int)
        for idx in range(self.ttree.nnodes - 1):
            parent = self.ttree.idx_dict[idx].up
            if parent:
                edges[idx, :] = (parent.idx, idx)
        return edges

    def get_radial_coords(self, use_edge_lengths=True):
        """
        Assign .edges and .verts for node positions in a fan tree.
        The farthest tip aligns at the circumference.
        """

        verts = np.zeros((self.ttree.nnodes, 2), dtype=float)

        # shortname
        if not use_edge_lengths:
            nbits = self.ttree.treenode.get_farthest_leaf(True)[1]

        # use cache to fill edges array
        for idx in range(self.ttree.nnodes):
            node = self.ttree.idx_dict[idx]

            # leaves: x positions are evenly spaced around circumference
            if node.is_leaf() and (not node.is_root()):

                # store radians (how far around from zero to 2pi)
                node.radians = self.circ.tip_radians[idx]

                # get positions of tips using radians and radius
                if use_edge_lengths:
                    node.radius = self.circ.radius - node.height
                    node.x, node.y = self.circ.get_node_coords(node)
                else:
                    node.radius = nbits
                    node.x, node.y = self.circ.get_node_coords(node)

            # internal nodes comes after tips. Inherit position from children.
            else:

                # height is either distance or nodes from root
                if use_edge_lengths:
                    node.radius = self.circ.radius - node.height
                else:
                    node.radius = sum(1 for i in node.iter_ancestors())
                    # max([i.radius for i in node.children]) - 1

                # x position is halfway between children x-pos
                node.radians = np.mean([i.radians for i in node.children])
                node.x, node.y = self.circ.get_node_coords(node)

            # store the x,y vertex positions
            verts[node.idx] = [node.x, node.y]
        return verts

        # scale so that node idx 0 (or fixed_order x) is at (0, 0)
        # adjust = self.ttree._coords.verts[:, 1].min()
        # self.ttree._coords.verts[:, 1] -= adjust

    def get_linear_coords(self, layout=None, use_edge_lengths=True):
        """
        Sets .edges, .verts for node positions.
        X and Y positions here refer to base assumption that tree is right
        facing, reorient_coordinates() will handle re-translating this.
        """
        verts = np.zeros((self.ttree.nnodes, 2), dtype=float)

        # store verts array with x,y positions of nodes (lengths of branches)
        # we want tips to align at the right face (larger axis number)
        fartip = self.ttree.treenode.get_farthest_leaf()[0]
        hgt = sum([
            fartip.dist,
            sum(i.dist for i in fartip.iter_ancestors() if not i.is_root())
        ])

        # finding tip distance from root is a little slow, it would be faster
        # just progress forward from root, but this wouldn't easily allow for
        # spacing fixed order tips. So instead we start from tips.
        for idx in range(self.ttree.nnodes):
            node = self.ttree.idx_dict[idx]

            if node.is_leaf() and (not node.is_root()):
                if use_edge_lengths:
                    toroot = sum(
                        i.dist for i in node.iter_ancestors()
                        if not i.is_root()
                    )
                    node.y = hgt - (node.dist + toroot)
                else:
                    node.y = 0.
                if self.ttree._fixed_order:
                    node.x = self.ttree._fixed_order.index(node.name)
                else:
                    node.x = node.idx

            else:
                nch = node.children
                # empty trees only
                if not nch:
                    node.x = node.y = 0

                # all other trees
                else:
                    if use_edge_lengths:
                        node.y = nch[0].y + nch[0].dist
                    else:
                        node.y = max((i.y for i in nch)) + 1

                    # internal node at midpoint of children
                    node.x = sum(i.x for i in nch) / float(len(nch))

            # store the vertex
            verts[idx] = [node.x, node.y]

        # scale so that node idx 0 (or fixed_order x) is at (0, 0)
        if use_edge_lengths:
            adjust = verts[:, 1].min()
            verts[:, 1] -= adjust

        # default: Down-facing tips align at y=0, first ladderized tip at x=0
        if layout == 'd':
            return verts

        # right-facing tips align at x=0, last ladderized tip at y=0
        elif layout == 'r':

            # verts swap x and ys and make xs 0 to negative
            tmp = np.zeros(verts.shape)
            tmp[:, 1] = verts[:, 0]
            tmp[:, 0] = verts[:, 1] * -1
            return tmp

        elif layout == 'l':

            # verts swap x and ys and make xs 0 to negative
            tmp = np.zeros(verts.shape)
            tmp[:, 1] = verts[:, 0]
            tmp[:, 0] = verts[:, 1]
            return tmp

        elif layout == 'u':
            # verts swap x and ys and make xs 0 to negative
            tmp = np.zeros(verts.shape)
            tmp[:, 1] = verts[:, 1] * -1
            tmp[:, 0] = verts[:, 0] * -1
            return tmp

        else:
            raise ToytreeError("layout not recognized")

    def get_radii(self, layout):
        """
        radius of every node relative to 0,0 origin.
        """

        # custom request layout
        if layout is None:
            layout = self.ttree.layout

        # get circular layout
        if layout == 'c':
            self.radii = np.array(
                [self.ttree.idx_dict[i].radius for i in range(self.nnodes)])
        else:
            self.radii = np.repeat(0, self.ntips)

class Circle:
    """
    When init from a toytree it can return coordinates based on
    distributing tips equally in radians from 0-2pi (TODO: allow
    setting start/end).

    The radius for farthest tip-ends is set to the tree height,
    and origin is at 0.0.
    """
    def __init__(self, tre):

        # set radius
        self.tre = tre
        self.radius = self.tre.treenode.height
        # get_distance(self.tre.treenode.get_farthest_leaf()[0])

        # origin
        self.o = (0, 0)
        # self.tre.style.xbaseline, self.tre.style.ybaseline)  # -self.radius, 0)

        # tips (bottom to top) are evenly spread from 0 to -2pi (counter clock)
        self.tip_radians = np.linspace(0, -np.pi * 2, self.tre.ntips + 1)[:-1]

    def get_node_coords(self, node):
        """
        get node coord. Node has radius assigned on the fly.
        """
        x = self.o[0] + node.radius * np.cos(node.radians)
        y = self.o[1] - node.radius * np.sin(node.radians)
        return x, y

    # def get_node_lines(self, node):
    #     x = self.o[0] + node.up.radius * np.cos(node.radians)
    #     y = self.o[1] - node.up.radius * np.sin(node.radians)
    #     return x, y

    def get_tip_end_angles(self):
        """
        node radians for printing tip labels should be from the root (0,0)
        """
        return np.array([np.rad2deg(abs(i)) for i in self.tip_radians])

    def get_tip_end_coords(self):
        """
        node tip coords must calculate new radian angle relative to parent
        node and then add offset amount to radius when calculating (x, y).
        """
        xs = [self.o[0] + self.radius * np.cos(i) for i in self.tip_radians]
        ys = [self.o[1] - self.radius * np.sin(i) for i in self.tip_radians]
        return np.stack([xs, ys], axis=1)

    # def get_parent_coords(self, node):
    #     """
    #     get x,y coordinates of parent based on radians of children
    #     """
    #     # calculate x, y coords
    #     x = self.o[0] + node.radius * np.cos(midradian)
    #     y = self.o[1] + node.radius * np.sin(midradian)
    #     return x, y

<|endoftext|>


<commit_before>"""
Transforms are SheetOperations that manipulate dataviews, typically
for the purposes of visualization. Such transformations often apply to
SheetViews or SheetStacks and compose the data together in ways that
can be viewed conveniently, often by creating or manipulating color
channels.
"""

import colorsys
import param

import numpy as np
import matplotlib
from imagen.analysis import ViewOperation
from sheetviews import SheetView

rgb_to_hsv = np.vectorize(colorsys.rgb_to_hsv)
hsv_to_rgb = np.vectorize(colorsys.hsv_to_rgb)

class HCS(ViewOperation):
    """
    Hue-Confidence-Strength plot.

    Accepts an overlay containing either 2 or 3 layers. The first two
    layers are hue and confidence and the third layer (if available)
    is the strength channel.
    """

    S_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the strength value.""")

    C_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the confidence value.""")

    flipSC = param.Boolean(default=False, doc="""
        Whether to flip the strength and confidence channels""")

    def _process(self, overlay):
        hue = overlay[0]
        confidence = overlay[1]

        strength_data = overlay[2].data if (len(overlay) == 3) else np.ones(hue.shape)

        if hue.shape != confidence.shape:
            raise Exception("Cannot combine plots with different shapes")

        (h,s,v)= (hue.N.data.clip(0.0, 1.0),
                  (confidence.data * self.p.C_multiplier).clip(0.0, 1.0),
                  (strength_data * self.p.S_multiplier).clip(0.0, 1.0))

        if self.p.flipSC:
            (h,s,v) = (h,v,s.clip(0,1.0))

        r, g, b = hsv_to_rgb(h, s, v)
        rgb = np.dstack([r,g,b])
        return [SheetView(rgb, hue.bounds, roi_bounds=overlay.roi_bounds, mode='rgb')]

class colorize(ViewOperation):
    """
    Given a SheetOverlay consisting of a grayscale colormap and a
    second Sheetview with some specified colour map, use the second
    layer to colorize the data of the first layer.

    Currently, colorize only support the 'hsv' color map and is just a
    shortcut to the HCS transform using a constant confidence
    value. Arbitrary colorization will be supported in future.
    """

    def _process(self, overlay):

         if len(overlay) != 2 and overlay[0].mode != 'gray':
             raise Exception("Can only colorize grayscale overlayed with colour map.")
         if [overlay[0].depth, overlay[1].depth ] != [1,1]:
             raise Exception("Depth one layers required.")
         if overlay[0].shape != overlay[1].shape:
             raise Exception("Shapes don't match.")

         if overlay[1].mode != 'hsv':
             raise NotImplementedError

         C = SheetView(np.ones(overlay[1].data.shape),
                       bounds=overlay.bounds)
         return [HCS(overlay[1] * C * overlay[0].N)]

class cmap2rgb(ViewOperation):
    """
    Convert SheetViews using colormaps to RGBA mode.
    """

    cmap = param.String(default=None, allow_None=True, doc="""
          Force the use of a specific color map. Otherwise, the mode
          property of the view is used instead.""")

    def _process(self, sheetview):
        if sheetview.depth != 1:
            raise Exception("Can only apply colour maps to SheetViews with depth of 1.")
        cmap = matplotlib.cm.get_cmap(sheetview.mode if self.p.cmap is None else self.p.cmap)
        return [SheetView(cmap(sheetview.data),
                         bounds=sheetview.bounds,
                         cyclic_range=sheetview.cyclic_range,
                         style=sheetview.style,
                         metadata=sheetview.metadata,
                         mode='rgba')]

class split(ViewOperation):
    """
    Given SheetViews in RGBA mode, return the R,G,B and A channels as
    a list of SheetViews.
    """
    def _process(self, sheetview):
        if sheetview.depth not in [3,4]:
            raise Exception("Can only split SheetViews with a depth of 3 or 4")
        return [SheetView(sheetview.data[:,:,i], bounds=sheetview.bounds)
                for i in range(sheetview.depth)]
<commit_msg>Transforms updated to use the view label and style system<commit_after>"""
Transforms are SheetOperations that manipulate dataviews, typically
for the purposes of visualization. Such transformations often apply to
SheetViews or SheetStacks and compose the data together in ways that
can be viewed conveniently, often by creating or manipulating color
channels.
"""

import colorsys
import param

import numpy as np
import matplotlib
from imagen.analysis import ViewOperation
from sheetviews import SheetView

from styles import Styles, Style

rgb_to_hsv = np.vectorize(colorsys.rgb_to_hsv)
hsv_to_rgb = np.vectorize(colorsys.hsv_to_rgb)

class HCS(ViewOperation):
    """
    Hue-Confidence-Strength plot.

    Accepts an overlay containing either 2 or 3 layers. The first two
    layers are hue and confidence and the third layer (if available)
    is the strength channel.
    """

    S_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the strength value.""")

    C_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the confidence value.""")

    flipSC = param.Boolean(default=False, doc="""
        Whether to flip the strength and confidence channels""")

    def _process(self, overlay):
        hue = overlay[0]
        confidence = overlay[1]

        strength_data = overlay[2].data if (len(overlay) == 3) else np.ones(hue.shape)

        if hue.shape != confidence.shape:
            raise Exception("Cannot combine plots with different shapes")

        (h,s,v)= (hue.N.data.clip(0.0, 1.0),
                  (confidence.data * self.p.C_multiplier).clip(0.0, 1.0),
                  (strength_data * self.p.S_multiplier).clip(0.0, 1.0))

        if self.p.flipSC:
            (h,s,v) = (h,v,s.clip(0,1.0))

        r, g, b = hsv_to_rgb(h, s, v)
        rgb = np.dstack([r,g,b])
        return [SheetView(rgb, hue.bounds, roi_bounds=overlay.roi_bounds,
                          mode='rgb', label=hue.label+' HCS Plot')]

class colorize(ViewOperation):
    """
    Given a SheetOverlay consisting of a grayscale colormap and a
    second Sheetview with some specified colour map, use the second
    layer to colorize the data of the first layer.

    Currently, colorize only support the 'hsv' color map and is just a
    shortcut to the HCS transform using a constant confidence
    value. Arbitrary colorization will be supported in future.
    """

    def _process(self, overlay):

         if len(overlay) != 2 and overlay[0].mode != 'gray':
             raise Exception("Can only colorize grayscale overlayed with colour map.")
         if [overlay[0].depth, overlay[1].depth ] != [1,1]:
             raise Exception("Depth one layers required.")
         if overlay[0].shape != overlay[1].shape:
             raise Exception("Shapes don't match.")

         if overlay[1].mode != 'hsv':
             raise NotImplementedError

         C = SheetView(np.ones(overlay[1].data.shape),
                       bounds=overlay.bounds)
         hcs = HCS(overlay[1] * C * overlay[0].N)
         return [SheetView(hsc.data, hsc.bounds,
                           mode='rgb',
                           roi_bounds=hcs.roi_bounds,
                           label=sheetview.label+' Colorized')]

class cmap2rgb(ViewOperation):
    """
    Convert SheetViews using colormaps to RGBA mode.  The colormap of
    the style is used, if available. Otherwise, the colormap may be
    forced.
    """

    cmap = param.String(default=None, allow_None=True, doc="""
          Force the use of a specific color map. Otherwise, the cmap
          property of the applicable style is used.""")

    def _process(self, sheetview):
        if sheetview.depth != 1:
            raise Exception("Can only apply colour maps to SheetViews with depth of 1.")

        style_cmap = Styles[sheetview][0].get('cmap', None)
        if not any([self.p.cmap, style_cmap]):
            raise Exception("No color map supplied and no cmap in the active style.")

        cmap = matplotlib.cm.get_cmap(style_cmap if self.p.cmap is None else self.p.cmap)
        return [SheetView(cmap(sheetview.data),
                         bounds=sheetview.bounds,
                         cyclic_range=sheetview.cyclic_range,
                         style=sheetview.style,
                         metadata=sheetview.metadata,
                         label = sheetview.label+' RGB',
                         mode='rgba')]

class split(ViewOperation):
    """
    Given SheetViews in RGBA mode, return the R,G,B and A channels as
    a GridLayout.
    """
    def _process(self, sheetview):
        if sheetview.depth not in [3,4]:
            raise Exception("Can only split SheetViews with a depth of 3 or 4")
        return [SheetView(sheetview.data[:,:,i],
                          bounds=sheetview.bounds,
                          label='RGBA'[i] + ' Channel')
                for i in range(sheetview.depth)]

Styles.R_Channel = Style(cmap='gray')
Styles.G_Channel = Style(cmap='gray')
Styles.B_Channel = Style(cmap='gray')
Styles.A_Channel = Style(cmap='gray')
<|endoftext|>