<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
import sys
import csv
import json
import requests
import urllib
import re
from datetime import datetime
from datetime import timedelta
from HTMLParser import HTMLParser

# TODO:
#      ( ) monitor mode
#      ( ) cmd of monitor mode
#      ( ) colorful print
#      ( ) important status change
#      ( ) limit up/down print

class world_index:
    def __init__(self):
        self.google_url = 'http://www.google.com/finance/info?q='
        self.index_list = [['TPE:TAIEX'        , 'TAIEX' , '加權指數'], # Format: google finance id, nickname, timezone, name
                           ['INDEXDJX:.DJI'    , 'DOW'   , '道瓊指數'],
                           ['INDEXNASDAQ:.IXIC', 'NASDAQ', '那斯達克'],
                           ['INDEXNASDAQ:SOX'  , 'PHLX'  , '費半PHLX'],
                           ['INDEXDB:DAX'      , 'DAX'   , '德國DAX'],
                           ['INDEXFTSE:UKX'    , 'FTSE'  , '英國FTSE'],
                           ['INDEXEURO:PX1'    , 'CAC40' , '法國指數'],
                           ['INDEXNIKKEI:NI225', 'N225'  , '日本指數'],
                           ['KRX:KOSPI'        , 'KOSPI' , '韓國指數'],
                           ['SHA:000001'       , 'SHCOMP', '上證指數'],
                           ['INDEXHANGSENG:HSI', 'HK'    , '香港恆生']]
        self.json_data = ''
        self.query_url = ''
        self.data = []

    def create_query_url(self):
        stock_str = self.google_url
        for index in self.index_list:
            stock_str = stock_str + index[0] + ','
        self.query_url = stock_str

    def query_stock_info(self):
        result = requests.get(self.query_url)
        # remove '//' and other special char in json data
        self.json_data = re.sub(r'[^\x00-\x7F]','', result.content.replace('// ', '', 1))

    def parse_json_data(self):
        now = datetime.now()
        json_data = json.loads(self.json_data)
        for i in range(len(self.index_list)):
            result = []
            j = json_data[i]
            id = self.index_list[i][1]

            ratio = j["cp"]
            if ratio.find('-'):
                ratio = '+' + ratio

            result_time = datetime.strptime(j["lt_dts"], '%Y-%m-%dT%H:%M:%SZ') + timedelta(hours = self.timezone_diff(j['lt'][15:]))
            if (now-result_time).total_seconds() < 1800:
                time_str = str(result_time.strftime('%H:%M:%S        '))
            elif now.month == result_time.month and now.day == result_time.day:
                time_str = str(result_time.strftime('%H:%M:%S (today)'))
            else:
                time_str = str(result_time.strftime('%H:%M:%S (%m/%d)'))

            result = {'id':'', 'name':'', 'price':'', 'change':'', 'ratio':'', 'time': ''}
            result['id']     = id
            result['name']   = self.index_list[i][2]
            result['price']  = j["l"].replace(',','')
            result['change'] = j["c"]
            result['ratio']  = ratio
            result['time']   = time_str
            self.data.append(result)

    def timezone_diff(self, timezone):
        timezone = timezone.replace(' ', '')
        if timezone[:3] == 'GMT':
            return 8 - int(timezone[3:])
        elif timezone[:3] == 'EDT':
            return 8 - (-4)
        else:
            print 'unknown timezone: ' + timezone
            return 0

    def print_stock_info(self):
        print ' 代號      指數          點數         漲跌     百分比    資料時間'
        print '---------------------------------------------------------------------------'

        for stock in self.data:
            print ' {0:8s}' .format(stock['id']) + \
                  '{0:<10s}'.format(stock['name']) + \
                  '{0:>13s}'.format(stock['price']) + ' ' + \
                  '{0:>11s}'.format(stock['change']) + ' ' + \
                  '{0:>9s}%'.format(stock['ratio']) + ' ' + \
                  '{0:>19s}'.format(stock['time'])

    def run(self):
        self.create_query_url()
        self.query_stock_info()
        self.parse_json_data()
        self.print_stock_info()

class taiwan_stock:
    def __init__(self):
        self.user_stock_list = ['00632r','4127','2330','2317']
        self.stock_list = []
        self.stock_query_str = ''
        self.data = []
        self.twse_url = 'http://220.229.103.179/stock/api/getStockInfo.jsp?ex_ch='
        self.json_data = ''

    def create_stock_list(self):
        self.stock_list = []
        tmp_list = []

        # add stock in user list
        for i in self.user_stock_list:
            tmp_list.append(str(i).upper())

        # add stock in argv
        for i in sys.argv:
            if i.find('-') == -1:
                tmp_list.append(i.upper())

        # check duplicate and create stock `list
        seen = set()
        for i in tmp_list:
            if i not in seen:
                self.stock_list.append(i)
                seen.add(i)

    def create_query_url(self):
        stock_str = 'tse_t00.tw|otc_o00.tw|'

        for stock_no in self.stock_list:
            found = False

            f = open('tse.csv','r')
            for row in csv.reader(f):
                if not found and row[0] == str(stock_no):
                    stock_str = stock_str + 'tse_' + stock_no + '.tw|'
                    found = True
            f.close()

            f = open('otc.csv','r')
            for row in csv.reader(f):
                if not found and row[0] == str(stock_no):
                    stock_str = stock_str + 'otc_' + stock_no + '.tw|'
                    found = True
            f.close()

        self.stock_query_str = stock_str

    def query_stock_info(self):
        query_url = self.twse_url + self.stock_query_str
        r = requests.session()

        # we need access the website before query data
        r.get('http://220.229.103.179/')

        # query data
        self.json_data = r.get(query_url).content

    def parse_json_data(self):
        now = datetime.now()
        json_data = json.loads(self.json_data)

        for i in range(len(json_data['msgArray'])):
            j = json_data['msgArray'][i]
            status = ''
            price = j["z"]
            stock_no = j["c"]
            name     = j["n"]
            volume   = j["v"]
            highest = float(j['h'])
            lowest = float(j['l'])

            diff = float(j["z"]) - float(j["y"])
            if diff > 0:
                sign = '+'
            elif diff < 0:
                sign = ''
            else:
                sign = ' '

            change_str = sign + '{0:.2f}'.format(diff)
            change_str_p = sign + '{0:.2f}'.format(diff / float(j["y"]) *100)

            if stock_no == 't00':
                stock_no = 'TWSE'
                name = u'上市'
                price = '{0:.0f} '.format(float(price))
                change_str = sign + '{0:.0f} '.format(diff)
                volume = '{0:d}'.format(int(volume)/100)
            elif stock_no == 'o00':
                stock_no = 'OTC'
                name = u'上櫃'
                volume = '{0:d}'.format(int(volume)/100)

            # status
            if (float(price) == highest):
                status = u' 最高'
            elif (float(price) == lowest):
                status = u' 最低'

            if stock_no != 'TWSE' and stock_no != 'OTC':
                h_limit = float(j['u'])
                l_limit = float(j['w'])
                if (float(price) == h_limit):
                    status = u' 漲停！'
                elif (float(price) == l_limit):
                    status = u' 跌停！'

            # check data time
            date = datetime.strptime(j["d"], '%Y%m%d')
            result_time_str = j["d"] + ' ' + j["t"]
            result_time = datetime.strptime(result_time_str, '%Y%m%d %H:%M:%S')

            if now.month == date.month and now.day == date.day:
                if (now.hour > 13) or (now.hour == 13 and now.minute > 30):
                    time_str = j["t"] + ' (today)'
                else:
                    time_str = j["t"]
            else:
                time_str = j["t"] + date.strftime(' (%m/%d)')

            # FIXME: To avoid too long stock name
            if len(name) > 4:
                name = stock_no

            # save to self.data
            result = {'id':'', 'name':'', 'price':'', 'change':'', 'ratio':'', 'volume':'', 'time': '', 'status': ''}
            result['id']     = stock_no
            result['name']   = name
            result['price']  = price
            result['change'] = change_str
            result['ratio']  = change_str_p
            result['volume'] = volume
            result['time']   = time_str
            result['status']   = status
            self.data.append(result)

    def print_stock_info(self):
        print ' 股號     股名     成交價     漲跌    百分比   成交量    資料時間'
        print '---------------------------------------------------------------------------'

        for stock in self.data:
            print ' ' + '{0:s}'   .format(stock['id']) + \
                  '\t ' + stock['name'] + '\t' + \
                  '{0:>9s}' .format(stock['price']) + ' ' + \
                  '{0:>8s}' .format(stock['change']) + ' ' + \
                  '{0:>8s}%'.format(stock['ratio']) + ' ' + \
                  '{0:>8s}' .format(stock['volume']) + ' ' + \
                  '   {0:19s}'.format(stock['time']) + ' ' + \
                  stock['status']

    def add_tw_future(self):
        tw_future = taiwan_future()
        self.data.append(tw_future.get_data())

    def run(self):
        self.add_tw_future()
        self.create_stock_list()
        self.create_query_url()
        self.query_stock_info()
        self.parse_json_data()
        self.print_stock_info()

class taiwan_future(HTMLParser):
    def __init__(self):
        HTMLParser.__init__(self)
        self.data = []
        self.item_limit = 15
        self.item_index = 0

    def handle_starttag(self, tag, attrs):
        if tag == 'td':
            for attr in attrs:
                if attr[1] == 'bu13':
                    self.item_index = self.item_index + 1
                    break

    def handle_data(self, data):
        if self.item_index == 3 or self.item_index == 4:
            if self.item_limit and len(data) < 14:
                self.item_limit = self.item_limit - 1
                self.data.append(data)

    def read_data(self):
        if (self.data[1] == '收盤'):
            i = 6
            time_str = self.data[i+8] + ' (close)'
        else:
            i = 5
            time_str = self.data[i+8] + '        '

        price   = float(self.data[i].replace(',',''))
        change  = float(self.data[i+1])
        volume  = self.data[i+3].replace(',','')
        last_day_price = float(self.data[i+7].replace(',',''))
        ratio   = '{0:.02f}'.format(change / last_day_price * 100)
        highest = float(self.data[i+5].replace(',',''))
        lowest = float(self.data[i+6].replace(',',''))

        status = ''
        if price == highest:
            status = u' 最高'
        elif price == lowest:
            status = u' 最低'

        if change > 0:
            sign = '+'
        else:
            sign = ''

        change_str = sign + '{0:.0f} '.format(change)
        ratio_str = sign + ratio
        result = {'id':'', 'name':'', 'price':'', 'change':'', 'ratio':'', 'volume':'', 'time': '', 'status': ''}
        result['id']     = 'WTX'
        result['name']   = '台指期'
        result['price']  = '{0:.0f} '.format(price)
        result['change'] = change_str
        result['ratio']  = ratio_str
        result['volume'] = volume
        result['time']   = time_str
        result['status'] = status
        return result

    def get_data(self):
        self.feed(urllib.urlopen("http://info512.taifex.com.tw/Future/FusaQuote_Norl.aspx").read())
        self.close()
        return self.read_data()

def main():
    # remote program name
    argv = sys.argv
    del argv[0]

    # read parameters
    for x in argv:
        if str(x) == '-c':
            color_print = True
        elif str(x) == '-a':
            show_world = True
            show_twse = True
            show_stock_list = True
        elif str(x) == '-w':
            show_world = True
        elif str(x) == '-i':
            show_twse = True
        elif str(x) == '-q':
            show_stock_list = True
        elif str(x) == '-s':
            show_simple = True
            show_stock_list = True
        elif str(x) == '-d':
            monitor_mode = True
        elif str(x) == '-h' or str(x) == '--help':
            usage()
            exit()

    world = world_index()
    world.run()
    print ''
    tw_stock = taiwan_stock()
    tw_stock.run()

if __name__ == '__main__':
    main()
<commit_msg>fix status mistake and add some code for monitor mode<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
import sys
import csv
import json
import requests
import urllib
import re
from datetime import datetime
from datetime import timedelta
from HTMLParser import HTMLParser

# TODO:
#      ( ) monitor mode
#      ( ) cmd of monitor mode
#      ( ) colorful print
#      ( ) important status change
#      ( ) limit up/down print

AUTO_UPDATE_SECOND = 20
PROFILE = {'color_print'   : False,
           'show_twse'     : False,
           'show_world'    : False,
           'show_user_list': False,
           'monitor_mode'  : False,
           'monitor_help'  : True}

class world_index:
    def __init__(self):
        self.google_url = 'http://www.google.com/finance/info?q='
        self.index_list = [['TPE:TAIEX'        , 'TAIEX' , '加權指數'], # Format: google finance id, nickname, timezone, name
                           ['INDEXDJX:.DJI'    , 'DOW'   , '道瓊指數'],
                           ['INDEXNASDAQ:.IXIC', 'NASDAQ', '那斯達克'],
                           ['INDEXNASDAQ:SOX'  , 'PHLX'  , '費半PHLX'],
                           ['INDEXDB:DAX'      , 'DAX'   , '德國DAX'],
                           ['INDEXFTSE:UKX'    , 'FTSE'  , '英國FTSE'],
                           ['INDEXEURO:PX1'    , 'CAC40' , '法國指數'],
                           ['INDEXNIKKEI:NI225', 'N225'  , '日本指數'],
                           ['KRX:KOSPI'        , 'KOSPI' , '韓國指數'],
                           ['SHA:000001'       , 'SHCOMP', '上證指數'],
                           ['INDEXHANGSENG:HSI', 'HK'    , '香港恆生']]
        self.json_data = ''
        self.query_url = ''
        self.data = []

    def create_query_url(self):
        stock_str = self.google_url
        for index in self.index_list:
            stock_str = stock_str + index[0] + ','
        self.query_url = stock_str

    def query_stock_info(self):
        result = requests.get(self.query_url)
        # remove '//' and other special char in json data
        self.json_data = re.sub(r'[^\x00-\x7F]','', result.content.replace('// ', '', 1))

    def parse_json_data(self):
        now = datetime.now()
        json_data = json.loads(self.json_data)
        for i in range(len(self.index_list)):
            result = []
            j = json_data[i]
            id = self.index_list[i][1]

            ratio = j["cp"]
            if ratio.find('-'):
                ratio = '+' + ratio

            result_time = datetime.strptime(j["lt_dts"], '%Y-%m-%dT%H:%M:%SZ') + timedelta(hours = self.timezone_diff(j['lt'][15:]))
            if (now-result_time).total_seconds() < 1800:
                time_str = str(result_time.strftime('%H:%M:%S        '))
            elif now.month == result_time.month and now.day == result_time.day:
                time_str = str(result_time.strftime('%H:%M:%S (today)'))
            else:
                time_str = str(result_time.strftime('%H:%M:%S (%m/%d)'))

            result = {'id':'', 'name':'', 'price':'', 'change':'', 'ratio':'', 'time': ''}
            result['id']     = id
            result['name']   = self.index_list[i][2]
            result['price']  = j["l"].replace(',','')
            result['change'] = j["c"]
            result['ratio']  = ratio
            result['time']   = time_str
            self.data.append(result)

    def timezone_diff(self, timezone):
        timezone = timezone.replace(' ', '')
        if timezone[:3] == 'GMT':
            return 8 - int(timezone[3:])
        elif timezone[:3] == 'EDT':
            return 8 - (-4)
        else:
            print 'unknown timezone: ' + timezone
            return 0

    def print_stock_info(self):
        print ' 代號      指數          點數         漲跌     百分比    資料時間'
        print '---------------------------------------------------------------------------'

        for stock in self.data:
            print ' {0:8s}' .format(stock['id']) + \
                  '{0:<10s}'.format(stock['name']) + \
                  '{0:>13s}'.format(stock['price']) + ' ' + \
                  '{0:>11s}'.format(stock['change']) + ' ' + \
                  '{0:>9s}%'.format(stock['ratio']) + ' ' + \
                  '{0:>19s}'.format(stock['time'])

    def run(self):
        self.create_query_url()
        self.query_stock_info()
        self.parse_json_data()
        self.print_stock_info()

class taiwan_stock:
    def __init__(self):
        self.user_stock_list = ['00632r','4127','2330','2317']
        self.stock_list = []
        self.stock_query_str = ''
        self.data = []
        self.twse_url = 'http://220.229.103.179/stock/api/getStockInfo.jsp?ex_ch='
        self.json_data = ''

    def create_stock_list(self):
        self.stock_list = []
        tmp_list = []

        # add stock in user list
        for i in self.user_stock_list:
            tmp_list.append(str(i).upper())

        # add stock in argv
        for i in sys.argv:
            if i.find('-') == -1:
                tmp_list.append(i.upper())

        # check duplicate and create stock `list
        seen = set()
        for i in tmp_list:
            if i not in seen:
                self.stock_list.append(i)
                seen.add(i)

    def create_query_url(self):
        stock_str = 'tse_t00.tw|otc_o00.tw|'

        for stock_no in self.stock_list:
            found = False

            f = open('tse.csv','r')
            for row in csv.reader(f):
                if not found and row[0] == str(stock_no):
                    stock_str = stock_str + 'tse_' + stock_no + '.tw|'
                    found = True
            f.close()

            f = open('otc.csv','r')
            for row in csv.reader(f):
                if not found and row[0] == str(stock_no):
                    stock_str = stock_str + 'otc_' + stock_no + '.tw|'
                    found = True
            f.close()

        self.stock_query_str = stock_str

    def query_stock_info(self):
        query_url = self.twse_url + self.stock_query_str
        r = requests.session()

        # we need access the website before query data
        r.get('http://220.229.103.179/')

        # query data
        self.json_data = r.get(query_url).content

    def parse_json_data(self):
        now = datetime.now()
        json_data = json.loads(self.json_data)

        for i in range(len(json_data['msgArray'])):
            j = json_data['msgArray'][i]
            price = j["z"]
            stock_no = j["c"]
            name     = j["n"]
            volume   = j["v"]
            highest = float(j['h'])
            lowest = float(j['l'])

            diff = float(j["z"]) - float(j["y"])
            if diff > 0:
                sign = '+'
            elif diff < 0:
                sign = ''
            else:
                sign = ' '

            change_str = sign + '{0:.2f}'.format(diff)
            change_str_p = sign + '{0:.2f}'.format(diff / float(j["y"]) *100)

            # status
            status = ''
            if (float(price) == highest):
                status = u'最高'
            elif (float(price) == lowest):
                status = u'最低'

            if stock_no != 't00' and stock_no != 'o00':
                h_limit = float(j['u'])
                l_limit = float(j['w'])
                if (float(price) == h_limit):
                    status = u'漲停！'
                elif (float(price) == l_limit):
                    status = u'跌停！'

            if stock_no == 't00':
                stock_no = 'TWSE'
                name = u'上市'
                price = '{0:.0f} '.format(float(price))
                change_str = sign + '{0:.0f} '.format(diff)
                volume = '{0:d}'.format(int(volume)/100)
            elif stock_no == 'o00':
                stock_no = 'OTC'
                name = u'上櫃'
                volume = '{0:d}'.format(int(volume)/100)

            # check data time
            date = datetime.strptime(j["d"], '%Y%m%d')
            result_time_str = j["d"] + ' ' + j["t"]
            result_time = datetime.strptime(result_time_str, '%Y%m%d %H:%M:%S')

            if now.month == date.month and now.day == date.day:
                if (now.hour > 13) or (now.hour == 13 and now.minute > 30):
                    time_str = j["t"] + ' (today)'
                else:
                    time_str = j["t"]
            else:
                time_str = j["t"] + date.strftime(' (%m/%d)')

            # FIXME: To avoid too long stock name
            if len(name) > 4:
                name = stock_no

            # save to self.data
            result = {'id':'', 'name':'', 'price':'', 'change':'', 'ratio':'', 'volume':'', 'time': '', 'status': ''}
            result['id']     = stock_no
            result['name']   = name
            result['price']  = price
            result['change'] = change_str
            result['ratio']  = change_str_p
            result['volume'] = volume
            result['time']   = time_str
            result['status']   = status
            self.data.append(result)

    def print_stock_info(self):
        print ' 股號     股名     成交價     漲跌    百分比   成交量    資料時間 & 狀態'
        print '--------------------------------------------------------------------------------'

        for stock in self.data:
            print ' ' + '{0:s}'   .format(stock['id']) + \
                  '\t ' + stock['name'] + '\t' + \
                  '{0:>9s}' .format(stock['price']) + ' ' + \
                  '{0:>8s}' .format(stock['change']) + ' ' + \
                  '{0:>8s}%'.format(stock['ratio']) + ' ' + \
                  '{0:>8s}' .format(stock['volume']) + ' ' + \
                  '   ' + stock['time'] + ' ' + stock['status']

    def add_tw_future(self):
        tw_future = taiwan_future()
        self.data.append(tw_future.get_data())

    def run(self):
        self.add_tw_future()
        self.create_stock_list()
        self.create_query_url()
        self.query_stock_info()
        self.parse_json_data()
        self.print_stock_info()

class taiwan_future(HTMLParser):
    def __init__(self):
        HTMLParser.__init__(self)
        self.data = []
        self.item_limit = 15
        self.item_index = 0

    def handle_starttag(self, tag, attrs):
        if tag == 'td':
            for attr in attrs:
                if attr[1] == 'bu13':
                    self.item_index = self.item_index + 1
                    break

    def handle_data(self, data):
        if self.item_index == 3 or self.item_index == 4:
            if self.item_limit and len(data) < 14:
                self.item_limit = self.item_limit - 1
                self.data.append(data)

    def read_data(self):
        if (self.data[1] == '收盤'):
            i = 6
            time_str = self.data[i+8] + ' (close)'
        else:
            i = 5
            time_str = self.data[i+8]

        price   = float(self.data[i].replace(',',''))
        change  = float(self.data[i+1])
        volume  = self.data[i+3].replace(',','')
        last_day_price = float(self.data[i+7].replace(',',''))
        ratio   = '{0:.02f}'.format(change / last_day_price * 100)
        highest = float(self.data[i+5].replace(',',''))
        lowest = float(self.data[i+6].replace(',',''))

        status = ''
        if price == highest:
            status = u' 最高'
        elif price == lowest:
            status = u' 最低'

        if change > 0:
            sign = '+'
        else:
            sign = ''

        change_str = sign + '{0:.0f} '.format(change)
        ratio_str = sign + ratio
        result = {'id':'', 'name':'', 'price':'', 'change':'', 'ratio':'', 'volume':'', 'time': '', 'status': ''}
        result['id']     = 'WTX'
        result['name']   = '台指期'
        result['price']  = '{0:.0f} '.format(price)
        result['change'] = change_str
        result['ratio']  = ratio_str
        result['volume'] = volume
        result['time']   = time_str
        result['status'] = status
        return result

    def get_data(self):
        self.feed(urllib.urlopen("http://info512.taifex.com.tw/Future/FusaQuote_Norl.aspx").read())
        self.close()
        return self.read_data()

def usage():
    print 'stockcmd - get stock information for TWSE'
    print ''
    print 'Usage:'
    print '    stockcmd.py [Options] [stock numbers]'
    print ''
    print 'Options:'
    print '    -a: list all stock infomation (include TSE, OTC, other stock)'
    print '    -s: list information with simple format'
    print '    -i: list include TSE index and OTC index'
    print '    -w: list International Stock Indexes'
    print '    -q: list the stocks predefined'
    print '    -c: list with color'
    print '    -d: continue update information every 20 seconds'
    print '    -h: show this page'
    print ''
    print 'Example:'
    print '    stockcmd.py -i'
    print '    stockcmd.py -w'
    print '    stockcmd.py -q'
    print '    stockcmd.py 2330 2317 3008'
    print ''

def read_option(opt):
    global PROFILE
    for x in opt:
        if str(x) == '-c':
            PROFILE['color_print'] = True
        elif str(x) == '-a':
            PROFILE['show_world'] = True
            PROFILE['show_twse'] = True
            PROFILE['show_stock_list'] = True
        elif str(x) == '-w':
            PROFILE['show_world'] = True
        elif str(x) == '-i':
            PROFILE['show_twse'] = True
        elif str(x) == '-q':
            PROFILE['show_stock_list'] = True
        elif str(x) == '-s':
            PROFILE['show_simple'] = True
            PROFILE['show_stock_list'] = True
        elif str(x) == '-d':
            PROFILE['monitor_mode'] = True
        elif str(x) == '-h' or str(x) == '--help':
            usage()
            exit()

def monitor_read_setting():
    global PROFILE
    for x in range(1, AUTO_UPDATE_SECOND, 1):
        input_cmd = ''
        i, o, e = select.select([sys.stdin], [], [], 1)
        if not i:
            continue
        input = sys.stdin.readline().strip().upper()
        if input == 'Q':
            exit()
        elif input == 'C':
            PROFILE.color_print = not PROFILE.color_print
        elif input == 'S':
            PROFILE.show_simple = not PROFILE.show_simple
        elif input == 'I':
            PROFILE.show_twse = not PROFILE.show_twse
        elif input == 'W':
            PROFILE.show_world = not PROFILE.show_world
        elif input == 'U':
            PROFILE.show_user_list = not PROFILE.show_user_list
        elif input == 'H':
            PROFILE.monitor_help = not PROFILE.monitor_help
        elif input == 'X':
            hide_closed_index = not hide_closed_index
        elif input[:1] == '+' and len(input) > 1:
            append_stock = True
            for i in USER_STOCK_LIST:
                if i == input[1:]:
                       append_stock = False;
            if append_stock:
                USER_STOCK_LIST.append(input[1:])

        elif input[:1] == '-' and len(input) > 1:
            for i in USER_STOCK_LIST:
                if i == input[1:]:
                    USER_STOCK_LIST.remove(i)
                    break

def main():
    # remove program name
    argv = sys.argv
    del argv[0]

    read_option(argv)

    world = world_index()
    world.run()
    print ''
    tw_stock = taiwan_stock()
    tw_stock.run()

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


<commit_before>"""
Matplotlib offers support for custom classes, such as unyt_array, allowing customization
of axis information and unit conversion. In the case of unyt, the axis label is set
based on the unyt_array.name and unyt_array.units attributes. It is also possible to
convert the plotted units.

This feature is optional and has to be enabled using the matplotlib_support context
manager.
"""

# -----------------------------------------------------------------------------
# Copyright (c) 2020, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the LICENSE file, distributed with this software.
# -----------------------------------------------------------------------------

try:
    from matplotlib.units import ConversionInterface, AxisInfo, registry
except ImportError:
    pass
else:
    from weakref import WeakKeyDictionary
    from unyt import unyt_array, unyt_quantity, Unit

    __all__ = ["matplotlib_support"]

    class unyt_arrayConverter(ConversionInterface):
        """Matplotlib interface for unyt_array"""

        _instance = None
        _labelstyle = "()"
        _axisnames = WeakKeyDictionary()

        # ensure that unyt_arrayConverter is a singleton
        def __new__(cls):
            if unyt_arrayConverter._instance is None:
                unyt_arrayConverter._instance = super().__new__(cls)
            return unyt_arrayConverter._instance

        # When matplotlib first encounters a type in its units.registry, it will
        # call default_units to obtain the units. Then it calls axisinfo to
        # customize the axis - in our case, just set the label. Then matplotlib calls
        # convert.

        @staticmethod
        def axisinfo(unit, axis):
            """Set the axis label based on unit

            Parameters
            ----------

            unit : Unit object, string, or tuple
                This parameter comes from unyt_arrayConverter.default_units() or from
                user code such as Axes.plot(), Axis.set_units(), etc. In user code, it
                is possible to convert the plotted units by specifing the new unit as
                a string, such as "ms", or as a tuple, such as ("J", "thermal")
                following the call signature of unyt_array.convert_to_units().
            axis : Axis object

            Returns
            -------

            AxisInfo object with the label formatted as in-line math latex
            """
            if isinstance(unit, tuple):
                unit = unit[0]
            unit_obj = unit if isinstance(unit, Unit) else Unit(unit)
            name = unyt_arrayConverter._axisnames.get(axis, "")
            if unit_obj.is_dimensionless:
                label = name
            else:
                name += " "
                unit_str = unit_obj.latex_representation()
                if unyt_arrayConverter._labelstyle == "[]":
                    label = name + "$\\left[" + unit_str + "\\right]$"
                elif unyt_arrayConverter._labelstyle == "/":
                    axsym = "$q_{\\rm" + axis.axis_name + "}$"
                    name = axsym if name == " " else name
                    if "/" in unit_str:
                        label = name + "$\\;/\\;\\left(" + unit_str + "\\right)$"
                    else:
                        label = name + "$\\;/\\;" + unit_str + "$"
                else:
                    label = name + "$\\left(" + unit_str + "\\right)$"
            return AxisInfo(label=label.strip())

        @staticmethod
        def default_units(x, axis):
            """Return the Unit object of the unyt_array x

            Parameters
            ----------

            x : unyt_array
            axis : Axis object

            Returns
            -------

            Unit object
            """
            name = getattr(x, "name", "")
            # maintain a mapping between Axis and name since Axis does not point to
            # its underlying data and we want to propagate the name to the axis
            # label in the subsequent call to axisinfo
            unyt_arrayConverter._axisnames[axis] = name if name is not None else ""
            return x.units

        @staticmethod
        def convert(value, unit, axis):
            """Convert the units of value to unit

            Parameters
            ----------

            value : unyt_array, unyt_quantity, or sequence there of
            unit : Unit, string or tuple
                This parameter comes from unyt_arrayConverter.default_units() or from
                user code such as Axes.plot(), Axis.set_units(), etc. In user code, it
                is possible to convert the plotted units by specifing the new unit as
                a string, such as "ms", or as a tuple, such as ("J", "thermal")
                following the call signature of unyt_array.convert_to_units().
            axis : Axis object

            Returns
            -------

            unyt_array

            Raises
            ------

            UnitConversionError if unit does not have the same dimensions as value or
            if we don't know how to convert value.
            """
            converted_value = value
            if isinstance(unit, str) or isinstance(unit, Unit):
                unit = (unit,)
            if isinstance(value, (unyt_array, unyt_quantity)):
                converted_value = value.to(*unit)
            else:
                value_type = type(value)
                converted_value = []
                for obj in value:
                    converted_value.append(obj.to(*unit))
                converted_value = value_type(converted_value)
            return converted_value

    class matplotlib_support:
        """Context manager for enabling the feature

        When used in a with statement, the feature is enabled during the context and
        then disabled after it exits.

        Parameters
        ----------

        label_style : str
          One of the following set, ``{'()', '[]', '/'}``. These choices
          correspond to the following unit labels:

            * ``'()'`` -> ``'(unit)'``
            * ``'[]'`` -> ``'[unit]'``
            * ``'/'`` -> ``'q_x / unit'``
        """

        def __init__(self, label_style="()"):
            self._labelstyle = label_style
            unyt_arrayConverter._labelstyle = label_style
            self._enabled = False

        def __call__(self):
            self.__enter__()

        @property
        def label_style(self):
            """str: One of the following set, ``{'()', '[]', '/'}``.
               These choices correspond to the following unit labels:

                 * ``'()'`` -> ``'(unit)'``
                 * ``'[]'`` -> ``'[unit]'``
                 * ``'/'`` -> ``'q_x / unit'``
            """
            return self._labelstyle

        @label_style.setter
        def label_style(self, label_style="()"):
            self._labelstyle = label_style
            unyt_arrayConverter._labelstyle = label_style

        def __enter__(self):
            registry[unyt_array] = unyt_arrayConverter()
            registry[unyt_quantity] = unyt_arrayConverter()
            self._enabled = True

        def __exit__(self, exc_type, exc_val, exc_tb):
            registry.pop(unyt_array)
            registry.pop(unyt_quantity)
            self._enabled = False

        def enable(self):
            self.__enter__()

        def disable(self):
            if self._enabled:
                self.__exit__(None, None, None)
<commit_msg>New black formatting<commit_after>"""
Matplotlib offers support for custom classes, such as unyt_array, allowing customization
of axis information and unit conversion. In the case of unyt, the axis label is set
based on the unyt_array.name and unyt_array.units attributes. It is also possible to
convert the plotted units.

This feature is optional and has to be enabled using the matplotlib_support context
manager.
"""

# -----------------------------------------------------------------------------
# Copyright (c) 2020, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the LICENSE file, distributed with this software.
# -----------------------------------------------------------------------------

try:
    from matplotlib.units import ConversionInterface, AxisInfo, registry
except ImportError:
    pass
else:
    from weakref import WeakKeyDictionary
    from unyt import unyt_array, unyt_quantity, Unit

    __all__ = ["matplotlib_support"]

    class unyt_arrayConverter(ConversionInterface):
        """Matplotlib interface for unyt_array"""

        _instance = None
        _labelstyle = "()"
        _axisnames = WeakKeyDictionary()

        # ensure that unyt_arrayConverter is a singleton
        def __new__(cls):
            if unyt_arrayConverter._instance is None:
                unyt_arrayConverter._instance = super().__new__(cls)
            return unyt_arrayConverter._instance

        # When matplotlib first encounters a type in its units.registry, it will
        # call default_units to obtain the units. Then it calls axisinfo to
        # customize the axis - in our case, just set the label. Then matplotlib calls
        # convert.

        @staticmethod
        def axisinfo(unit, axis):
            """Set the axis label based on unit

            Parameters
            ----------

            unit : Unit object, string, or tuple
                This parameter comes from unyt_arrayConverter.default_units() or from
                user code such as Axes.plot(), Axis.set_units(), etc. In user code, it
                is possible to convert the plotted units by specifing the new unit as
                a string, such as "ms", or as a tuple, such as ("J", "thermal")
                following the call signature of unyt_array.convert_to_units().
            axis : Axis object

            Returns
            -------

            AxisInfo object with the label formatted as in-line math latex
            """
            if isinstance(unit, tuple):
                unit = unit[0]
            unit_obj = unit if isinstance(unit, Unit) else Unit(unit)
            name = unyt_arrayConverter._axisnames.get(axis, "")
            if unit_obj.is_dimensionless:
                label = name
            else:
                name += " "
                unit_str = unit_obj.latex_representation()
                if unyt_arrayConverter._labelstyle == "[]":
                    label = name + "$\\left[" + unit_str + "\\right]$"
                elif unyt_arrayConverter._labelstyle == "/":
                    axsym = "$q_{\\rm" + axis.axis_name + "}$"
                    name = axsym if name == " " else name
                    if "/" in unit_str:
                        label = name + "$\\;/\\;\\left(" + unit_str + "\\right)$"
                    else:
                        label = name + "$\\;/\\;" + unit_str + "$"
                else:
                    label = name + "$\\left(" + unit_str + "\\right)$"
            return AxisInfo(label=label.strip())

        @staticmethod
        def default_units(x, axis):
            """Return the Unit object of the unyt_array x

            Parameters
            ----------

            x : unyt_array
            axis : Axis object

            Returns
            -------

            Unit object
            """
            name = getattr(x, "name", "")
            # maintain a mapping between Axis and name since Axis does not point to
            # its underlying data and we want to propagate the name to the axis
            # label in the subsequent call to axisinfo
            unyt_arrayConverter._axisnames[axis] = name if name is not None else ""
            return x.units

        @staticmethod
        def convert(value, unit, axis):
            """Convert the units of value to unit

            Parameters
            ----------

            value : unyt_array, unyt_quantity, or sequence there of
            unit : Unit, string or tuple
                This parameter comes from unyt_arrayConverter.default_units() or from
                user code such as Axes.plot(), Axis.set_units(), etc. In user code, it
                is possible to convert the plotted units by specifing the new unit as
                a string, such as "ms", or as a tuple, such as ("J", "thermal")
                following the call signature of unyt_array.convert_to_units().
            axis : Axis object

            Returns
            -------

            unyt_array

            Raises
            ------

            UnitConversionError if unit does not have the same dimensions as value or
            if we don't know how to convert value.
            """
            converted_value = value
            if isinstance(unit, str) or isinstance(unit, Unit):
                unit = (unit,)
            if isinstance(value, (unyt_array, unyt_quantity)):
                converted_value = value.to(*unit)
            else:
                value_type = type(value)
                converted_value = []
                for obj in value:
                    converted_value.append(obj.to(*unit))
                converted_value = value_type(converted_value)
            return converted_value

    class matplotlib_support:
        """Context manager for enabling the feature

        When used in a with statement, the feature is enabled during the context and
        then disabled after it exits.

        Parameters
        ----------

        label_style : str
          One of the following set, ``{'()', '[]', '/'}``. These choices
          correspond to the following unit labels:

            * ``'()'`` -> ``'(unit)'``
            * ``'[]'`` -> ``'[unit]'``
            * ``'/'`` -> ``'q_x / unit'``
        """

        def __init__(self, label_style="()"):
            self._labelstyle = label_style
            unyt_arrayConverter._labelstyle = label_style
            self._enabled = False

        def __call__(self):
            self.__enter__()

        @property
        def label_style(self):
            """str: One of the following set, ``{'()', '[]', '/'}``.
            These choices correspond to the following unit labels:

              * ``'()'`` -> ``'(unit)'``
              * ``'[]'`` -> ``'[unit]'``
              * ``'/'`` -> ``'q_x / unit'``
            """
            return self._labelstyle

        @label_style.setter
        def label_style(self, label_style="()"):
            self._labelstyle = label_style
            unyt_arrayConverter._labelstyle = label_style

        def __enter__(self):
            registry[unyt_array] = unyt_arrayConverter()
            registry[unyt_quantity] = unyt_arrayConverter()
            self._enabled = True

        def __exit__(self, exc_type, exc_val, exc_tb):
            registry.pop(unyt_array)
            registry.pop(unyt_quantity)
            self._enabled = False

        def enable(self):
            self.__enter__()

        def disable(self):
            if self._enabled:
                self.__exit__(None, None, None)
<|endoftext|>


<commit_before>"""Writes mechanism header and output testing files"""

# Python 2 compatibility
from __future__ import division
from __future__ import print_function

# Standard libraries
import sys
import math

# Local imports
import chem_utilities as chem
import utils

def write_mechanism_initializers(path, lang, specs, reacs):
    if lang in ['matlab', 'fortran']:
        raise NotImplementedError

    #the mechanism header defines a number of useful preprocessor defines, as well as defining method stubs for setting initial conditions
    with open(path + 'mechanism.h', 'w') as file:

        file.write('#ifndef MECHANISM_H\n'
                   '#define MECHANISM_H\n\n')

        #convience: write species indexes
        file.write('/* Species Indexes\n')
        file.write('\n'.join('{}  {}'.format(i, spec.name) for i, spec in enumerate(specs)))
        file.write('*/\n\n')

        #defines
        file.write("//Number of species\n")
        file.write('#define NSP {}\n'.format(len(specs)))
        file.write("//Number of variables. NN = NSP + 1 (temperature)\n")
        file.write('#define NN {}\n'.format(len(specs) + 1))
        file.write('//Number of forward reactions\n')
        file.write('#define FWD_RATES {}\n'.format(len(reacs)))
        file.write('//Number of reversible reactions\n')
        file.write('#define REV_RATES {}\n'.format(len([reac for reac in reacs if reac.rev])))
        file.write('//Number of reactions with pressure modified rates\n')
        file.write('#define PRES_MOD_RATES {}\n\n'.format(len([reac for reac in reacs if reac.pdep or reac.thd])))

        #to work with both c and cuda
        file.write('#ifdef __cplusplus\n'
                   'extern "C" {\n'
                   '#endif\n')

        file.write('    //Must be implemented by user on a per mechanism basis in mechanism.c\n'
                   '    #ifdef CONP\n'
                   '    void set_same_initial_conditions(int NUM, double* y_host, double* pres_host);\n'
                   '    #elif CONV\n'
                   '    void set_same_initial_conditions(int NUM, double* y_host, double* pres_host, double* rho_host);\n'
                   '    #endif\n'
                   '    #ifdef RATES_TEST\n'
                   '    void write_jacobian_and_rates_output();\n'
                   '    #endif\n'
                   )

        #close previous extern
        file.write('#ifdef __cplusplus\n'
                   '}\n'
                   '#endif\n\n')

        file.write('#endif\n\n')

    #now the mechanism file
    with open(path + 'mechanism.c', 'w') as file:
        file.write('#include <stdio.h>\n'
                   '#include "mass_mole.h"\n'
                   '#include "mechanism.h"\n'
                   '#ifdef RATES_TEST\n'
                   '    #include "rates.h"\n'
                   '    #include "derivs.h"\n'
                   '#endif\n')
        file.write('#ifdef CONP\n'
                   '    void set_same_initial_conditions(int NUM, double* y_host, double* pres_host) {\n'
                   '#elif CONV\n'
                   '    void set_same_initial_conditions(int NUM, double* y_host, double* pres_host, double* rho_host) {\n'
                   '#endif\n'
                   '        double Xi [NSP] = {0.0};\n'
                   '        //set initial mole fractions here\n\n'
                   '        //Normalize mole fractions to sum to one\n'
                   '        double Xsum = 0.0;\n'
                   '        for (int j = 0; j < NSP; ++ j) {\n'
                   '            Xsum += Xi[j];\n'
                   '        }\n'
                   '        if (Xsum == 0.0) {\n'
                   '            printf("Use of the set initial conditions function requires user implementation!");\n'
                   '            exit(-1);\n'
                   '        }\n'
                   '        for (int j = 0; j < NSP; ++ j) {\n'
                   '            Xi[j] /= Xsum;\n'
                   '        }\n\n'
                   '        //convert to mass fractions\n'
                   '        double Yi[NSP];\n'
                   '        mole2mass(Xi, Yi);\n\n'
                   '        //set initial pressure, units [dyn/cm^2]\n'
                   '        double pres = 1.01325e6;\n'
                   '        // set intial temperature, units [K]\n'
                   '        double T0 = 1600;\n\n'
                   '        //load temperature and mass fractions for all threads (cells)\n'
                   '        for (int i = 0; i < NUM; ++i) {\n'
                   '            y_host[i] = T0;\n'
                   '            //loop through species\n'
                   '            for (int j = 1; j < NN; ++j) {\n'
                   '                y_host[i + NUM * j] = Yi[j - 1];\n'
                   '            }\n'
                   '        }\n\n'
                   '        #ifdef CONV\n'
                   '        //calculate density\n'
                   '        double rho = getDensity(T0, pres, Xi);\n'
                   '        #endif\n\n'
                   '        for (int i = 0; i < NUM; ++i) {\n'
                   '            #ifdef CONV\n'
                   '            rho_host[i] = rho;\n'
                   '            #endif\n'
                   '            pres_host[i] = pres;\n'
                   '        }\n'
                   '    }\n\n'
                   )
        file.write('#ifdef RATES_TEST\n')
        #write utility function that finds concentrations at a given state
        file.write('    void get_concentrations(double pres, const double* y, double* conc) {\n'
                   '        double rho;\n'
                   )

        line = '        rho = '
        isfirst = True
        for sp in specs:
            if len(line) > 70:
                line += '\n'
                file.write(line)
                line = '         '

            if not isfirst: line += ' + '
            line += '(y[{}] / {})'.format(specs.index(sp) + 1, sp.mw)
            isfirst = False

        line += ';\n'
        file.write(line)
        line = '        rho = pres / ({:.8e} * y[0] * rho);\n\n'.format(chem.RU)
        file.write(line)

        # calculation of species molar concentrations
        file.write('        // species molar concentrations\n')
        # loop through species
        for sp in specs:
            isp = specs.index(sp)
            line = '        conc[{}] = rho * y[{}] / '.format(isp, isp)
            line += '{}'.format(sp.mw) + utils.line_end[lang]
            file.write(line)

        file.write('    }\n\n')

        #convience method to write rates to file
        file.write('    void write_rates(FILE* fp, const double* fwd_rates, const double* rev_rates, const double* pres_mod, const double* dy) {\n'
                   '        for(int i = 0; i < FWD_RATES; ++i) {\n'
                   '            fprintf(fp, "%.15le\\n", fwd_rates[i]);\n'
                   '        }\n'
                   '        for(int i = 0; i < REV_RATES; ++i) {\n'
                   '            fprintf(fp, "%.15le\\n", rev_rates[i]);\n'
                   '        }\n'
                   '        for(int i = 0; i < PRES_MOD_RATES; i++) {\n'
                   '            fprintf(fp, "%.15le\\n", pres_mod[i]);\n'
                   '        }\n'
                   '        for(int i = 0; i < NSP; i++) {\n'
                   '            fprintf(fp, "%.15le\\n", dy[i]);\n'
                   '        }\n'
                   '    }\n\n'
                  )

        file.write('    void write_jacob(FILE* fp, const double* jacob) {\n'
                   '        for (int i = 0; i < NN * NN; ++i) {\n'
                   '            fprintf(fp, "%.15le\\n", jacob[i]);\n'
                   '        }\n'
                   '    }\n\n'
                   )

        file.write('    void write_jacobian_and_rates_output() {\n'
                   '        //set mass fractions to unity to turn on all reactions\n'
                   '        double Yi[NN];\n'
                   '        for (int i = 1; i < NN; ++i) {\n'
                   '            Yi[i] = 1.0 / ((double)NSP);\n'
                   '        }\n'
                   '        double conc[NSP];\n'
                   '        double fwd_rates[FWD_RATES];\n'
                   '        double rev_rates[REV_RATES];\n'
                   '        double pres_mod[PRES_MOD_RATES];\n'
                   '        double dy[NSP];\n'
                   '        double jacob[NN * NN];\n'
                   '        //evaluate and write rates for various conditions\n'
                   '        FILE* fp = fopen ("rates_data.txt", "w");\n'
                   '        fprintf(fp, "800K, 1 atm\\n");\n'
                   '        Yi[0] = 800;\n'
                   '        get_concentrations(1.01325e6, Yi, conc);\n'
                   '        eval_rxn_rates(800, conc, fwd_rates, rev_rates);\n'
                   '        get_rxn_pres_mod (800, 1.01325e6, conc, pres_mod);\n'
                   '        eval_spec_rates (fwd_rates, rev_rates, pres_mod, dy);\n'
                   '        write_rates(fp, fwd_rates, rev_rates, pres_mod, dy);\n'
                   '        eval_jacob(0, 1.01325e6, Yi, jacob);\n'
                   '        write_jacob(fp, jacob);\n'
                   '        fprintf(fp, "1600K, 1 atm\\n");\n'
                   '        Yi[0] = 1600;\n'
                   '        get_concentrations(1.01325e6, Yi, conc);\n'
                   '        eval_rxn_rates(1600, conc, fwd_rates, rev_rates);\n'
                   '        get_rxn_pres_mod (1600, 1.01325e6, conc, pres_mod);\n'
                   '        eval_spec_rates (fwd_rates, rev_rates, pres_mod, dy);\n'
                   '        write_rates(fp, fwd_rates, rev_rates, pres_mod, dy);\n'
                   '        eval_jacob(0, 1.01325e6, Yi, jacob);\n'
                   '        write_jacob(fp, jacob);\n'
                   '        fprintf(fp, "800K, 10 atm\\n");\n'
                   '        Yi[0] = 800;\n'
                   '        get_concentrations(1.01325e7, Yi, conc);\n'
                   '        eval_rxn_rates(800, conc, fwd_rates, rev_rates);\n'
                   '        get_rxn_pres_mod (800, 1.01325e7, conc, pres_mod);\n'
                   '        eval_spec_rates (fwd_rates, rev_rates, pres_mod, dy);\n'
                   '        write_rates(fp, fwd_rates, rev_rates, pres_mod, dy);\n'
                   '        eval_jacob(0, 1.01325e7, Yi, jacob);\n'
                   '        write_jacob(fp, jacob);\n'
                   '        fclose(fp);\n'
                   '    }\n'
                   )
        file.write('#endif')

<commit_msg>added support for GPU/CPU versions of set_same_initial_conditions, as well as the mechanism file.  TODO: 1. test, 2. cuda macros, 3. read conditions from file<commit_after>"""Writes mechanism header and output testing files"""

# Python 2 compatibility
from __future__ import division
from __future__ import print_function

# Standard libraries
import sys
import math

# Local imports
import chem_utilities as chem
import utils
import CUDAParams

def __write_rate_evaluator(file, have_rev_rxns, have_pdep_rxns, T, P, Pretty_P):
    file.write('        fprintf(fp, "{}, {} atm\\n");\n'.format(T, Pretty_P) +
               '        Yi[0] = {};\n'.format(T) +
               '        get_concentrations({}, Yi, conc);\n'.format(P) +
               '        eval_rxn_rates({}, conc, fwd_rates{}'.format(T, 'rev_rates);\n' if have_rev_rxns else ');\n')
               )
    if have_pdep_rxns:
        file.write('        get_rxn_pres_mod ({}, {}, conc, pres_mod);\n'.format(T, P))
    file.write('        eval_spec_rates (fwd_rates,{}{} dy);\n'.format(' rev_rates,' if have_rev_rxns else '', ' pres_mod,' if have_pdep_rxns else '') +
               '        write_rates(fp,{}{} dy);\n'.format(' rev_rates,' if have_rev_rxns else '', ' pres_mod,' if have_pdep_rxns else '') +
               '        eval_jacob(0, {}, Yi, jacob);\n'.format(P) +
               '        write_jacob(fp, jacob);\n'
               )

def write_mechanism_initializers(path, lang, specs, reacs):
    if lang in ['matlab', 'fortran']:
        raise NotImplementedError

    #some information variables
    have_rev_rxns = any(reac.rev for reac in reacs)
    have_pdep_rxns = any(reac.thd or reac.pdep for reac in reacs)

    #the mechanism header defines a number of useful preprocessor defines, as well as defining method stubs for setting initial conditions
    with open(path + 'mechanism.{}h'.format('cu' if lang == 'cuda' else ''), 'w') as file:

        file.write('#ifndef MECHANISM_{}H\n'.format('cu' if lang == 'cuda' else '') +
                   '#define MECHANISM_{}H\n\n'.format('cu' if lang == 'cuda' else ''))

        #convience: write species indexes
        file.write('/* Species Indexes\n')
        file.write('\n'.join('{}  {}'.format(i, spec.name) for i, spec in enumerate(specs)))
        file.write('*/\n\n')

        #defines
        file.write("//Number of species\n")
        file.write('#define NSP {}\n'.format(len(specs)))
        file.write("//Number of variables. NN = NSP + 1 (temperature)\n")
        file.write('#define NN {}\n'.format(len(specs) + 1))
        file.write('//Number of forward reactions\n')
        file.write('#define FWD_RATES {}\n'.format(len(reacs)))
        file.write('//Number of reversible reactions\n')
        file.write('#define REV_RATES {}\n'.format(len([reac for reac in reacs if reac.rev])))
        file.write('//Number of reactions with pressure modified rates\n')
        file.write('#define PRES_MOD_RATES {}\n\n'.format(len([reac for reac in reacs if reac.pdep or reac.thd])))

        if lang == 'cuda':
          file.write('#ifdef __cplusplus\n'
                     'extern "C" {\n'
                     '#endif\n')

        file.write('    //Must be implemented by user on a per mechanism basis in mechanism.c\n'
                   '    #ifdef CONP\n'
                   '    {} set_same_initial_conditions(int NUM,{} double* y_host, double* pres_host);\n'\
                        .format(' int block_size, int grid_size, ' if lang == 'cuda' else '', 'int' if lang == 'cuda' else 'void') +
                   '    #elif CONV\n'
                   '    {} set_same_initial_conditions(int NUM,{} double* y_host, double* pres_host, double* rho_host);\n'\
                        .format(' int block_size, int grid_size, ' if lang == 'cuda' else '', 'int' if lang == 'cuda' else 'void') +
                   '    #endif\n'
                   '    #ifdef RATES_TEST\n'
                   '    void write_jacobian_and_rates_output();\n'
                   '    #endif\n'
                   )

        if lang == 'cuda':
          #close previous extern
          file.write('#ifdef __cplusplus\n'
                     '}\n'
                     '#endif\n\n')

        file.write('#endif\n\n')

    #now the mechanism file
    with open(path + 'mechanism.c{}'.format('u' if lang == 'cuda' else ''), 'w') as file:
        file.write('#include <stdio.h>\n'
                   '#include "mass_mole.h"\n'
                   '#include "mechanism.{}h"\n'.format('cu' if lang == 'cuda' else '') +
                   '#ifdef RATES_TEST\n'
                   '    #include "rates.h"\n'
                   '    #include "derivs.h"\n'
                   '#endif\n')
        if lang == 'cuda':
            file.write('#include "gpu_memory.cuh"\n'
                       '#include "gpu_macros.cuh"\n'
                       'extern __device__ double* y;\n'
                       'extern __device__ double* pres;\n'
                       '#ifdef CONV\n'
                       '    extern __device__ double* rho;\n'
                       '#endif\n'
                       )
        file.write('#ifdef CONP\n'
                   '    {} set_same_initial_conditions(int NUM, double* y_host, double* pres_host) {{\n'.format('int' if lang == 'cuda' else 'void') +
                   '#elif CONV\n'
                   '    {} set_same_initial_conditions(int NUM, double* y_host, double* pres_host, double* rho_host) {{\n'.format('int' if lang == 'cuda' else 'void') +
                   '#endif\n'
                   )
        if lang == 'cuda':
            #do cuda mem init and copying
            file.write('        int size = initialize_gpu_memory(NUM, block_size, grid_size);\n')
        else:
            file.write('        int size = NUM;\n')
        file.write('        y_host = (double*)malloc(NN * size * sizeof(double));\n'
                   '        pres_host = (double*)malloc(size * sizeof(double));\n'
                   '#ifdef CONV\n'
                   '        rho_host = (double*)malloc(size * sizeof(double));\n'
                   '#endif'
                   '        double Xi [NSP] = {0.0};\n'
                   '        //set initial mole fractions here\n\n'
                   '        //Normalize mole fractions to sum to one\n'
                   '        double Xsum = 0.0;\n'
                   '        for (int j = 0; j < NSP; ++ j) {\n'
                   '            Xsum += Xi[j];\n'
                   '        }\n'
                   '        if (Xsum == 0.0) {\n'
                   '            printf("Use of the set initial conditions function requires user implementation!");\n'
                   '            exit(-1);\n'
                   '        }\n'
                   '        for (int j = 0; j < NSP; ++ j) {\n'
                   '            Xi[j] /= Xsum;\n'
                   '        }\n\n'
                   '        //convert to mass fractions\n'
                   '        double Yi[NSP];\n'
                   '        mole2mass(Xi, Yi);\n\n'
                   '        //set initial pressure, units [dyn/cm^2]\n'
                   '        double pres = 1.01325e6;\n'
                   '        // set intial temperature, units [K]\n'
                   '        double T0 = 1600;\n\n'
                   '        //load temperature and mass fractions for all threads (cells)\n'
                   '        for (int i = 0; i < size; ++i) {\n'
                   '            y_host[i] = T0;\n'
                   '            //loop through species\n'
                   '            for (int j = 1; j < NN; ++j) {\n'
                   '                y_host[i + NUM * j] = Yi[j - 1];\n'
                   '            }\n'
                   '        }\n\n'
                   '        #ifdef CONV\n'
                   '        //calculate density\n'
                   '        double rho = getDensity(T0, pres, Xi);\n'
                   '        #endif\n\n'
                   '        for (int i = 0; i < size; ++i) {\n'
                   '            #ifdef CONV\n'
                   '            rho_host[i] = rho;\n'
                   '            #endif\n'
                   '            pres_host[i] = pres;\n'
                   '        }\n'
                   )
        if lang == 'cuda': #copy memory over
            file.write('        cudaMemcpy(y, y_host, size * NN * sizeof(double), cudaMemcpyHostToDevice);\n'
                       '        cudaMemcpy(pres, pres_host, size * sizeof(double), cudaMemcpyHostToDevice);\n'
                       '#ifdef CONV\n'
                       '        cudaMemcpy(rho, rho_host, size * sizeof(double), cudaMemcpyHostToDevice);\n'
                       '#endif\n')

        file.write('    }\n\n')
        file.write('#ifdef RATES_TEST\n')
        #write utility function that finds concentrations at a given state
        file.write('    void get_concentrations(double pres, const double* y, double* conc) {\n'
                   '        double rho;\n'
                   )

        line = '        rho = '
        isfirst = True
        for sp in specs:
            if len(line) > 70:
                line += '\n'
                file.write(line)
                line = '         '

            if not isfirst: line += ' + '
            line += '(y[{}] / {})'.format(specs.index(sp) + 1, sp.mw)
            isfirst = False

        line += ';\n'
        file.write(line)
        line = '        rho = pres / ({:.8e} * y[0] * rho);\n\n'.format(chem.RU)
        file.write(line)

        # calculation of species molar concentrations
        file.write('        // species molar concentrations\n')
        # loop through species
        for sp in specs:
            isp = specs.index(sp)
            line = '        conc[{}] = rho * y[{}] / '.format(isp, isp)
            line += '{}'.format(sp.mw) + utils.line_end[lang]
            file.write(line)

        file.write('    }\n\n')

        line = 'void write_rates(FILE* fp, const double* fwd_rates'
        if have_rev_rxns:
            line += ', const double* rev_rates'
        if have_pdep_rxns:
            line += ', const double* pres_mod'
        line += 'const double dy*) {\n'
        file.write(line)
        #convience method to write rates to file
        file.write('        for(int i = 0; i < FWD_RATES; ++i) {\n'
                   '            fprintf(fp, "%.15le\\n", fwd_rates[i]);\n'
                   '        }\n'
                   )
        if have_rev_rxns:
            file.write('        for(int i = 0; i < REV_RATES; ++i) {\n'
                       '            fprintf(fp, "%.15le\\n", rev_rates[i]);\n'
                       '        }\n'
                       )
        if have_pdep_rxns:
            file.write('        for(int i = 0; i < PRES_MOD_RATES; i++) {\n'
                       '            fprintf(fp, "%.15le\\n", pres_mod[i]);\n'
                       '        }\n'
                       )
        file.write('        for(int i = 0; i < NSP; i++) {\n'
                   '            fprintf(fp, "%.15le\\n", dy[i]);\n'
                   '        }\n'
                   '    }\n\n'
                  )

        file.write('    void write_jacob(FILE* fp, const double* jacob) {\n'
                   '        for (int i = 0; i < NN * NN; ++i) {\n'
                   '            fprintf(fp, "%.15le\\n", jacob[i]);\n'
                   '        }\n'
                   '    }\n\n'
                   )

        file.write('    void write_jacobian_and_rates_output() {\n'
                   '        //set mass fractions to unity to turn on all reactions\n'
                   '        double Yi[NN];\n'
                   '        for (int i = 1; i < NN; ++i) {\n'
                   '            Yi[i] = 1.0 / ((double)NSP);\n'
                   '        }\n'
                   '        double conc[NSP];\n'
                   '        double fwd_rates[FWD_RATES];\n'
                   )
        if have_rev_rxns:
            file.write('        double rev_rates[REV_RATES];\n')
        if have_pdep_rxns:
            file.write('        double pres_mod[PRES_MOD_RATES];\n')

        file.write('        double dy[NSP];\n'
                   '        double jacob[NN * NN];\n'
                   '        //evaluate and write rates for various conditions\n'
                   '        FILE* fp = fopen ("rates_data.txt", "w");\n'
                   )
        __write_rate_evaluator(file, have_rev_rxns, have_pdep_rxns, '800', '1.01325e6', '1')
        __write_rate_evaluator(file, have_rev_rxns, have_pdep_rxns, '1600', '1.01325e6', '1')
        __write_rate_evaluator(file, have_rev_rxns, have_pdep_rxns, '800', '1.01325e7', '10')
        file.write('        fclose(fp);\n'
                   '    }\n'
                   )
        file.write('#endif')

    if lang == 'cuda' and CUDAParams.is_global():
        with open(path + 'gpu_memory.cuh', 'w') as file:
            file.write('#ifndef GPU_MEMORY_CUH\n'
                       '#define GPU_MEMORY_CUH\n'
                       '\n'
                       '#include "mechanism.cuh"\n'
                       '\n'
                       'int initialize_gpu_memory(int NUM, int block_size, int grid_size);\n'
                       'void free_gpu_memory();\n'
                       '\n'
                       '#endif\n')

        with open(path + 'gpu_memory.cu', 'w') as file:
            file.write('//define global memory pointers')
            template = '__device__ double* {};'
            host_template = 'double* d_{};'
            conp = ['h', 'cp']
            conv = ['u', 'cv']
            arrays = ['y', 'dy', 'conc', 'fwd_rates' if have_rev_rxns else '', 'rev_rates' if have_rev_rxns else '', 'rates' if have_rev_rxns else '',\
             'pres_mod' if have_pdep_rxns else '', 'sp_rates']
            arrays = [arr for arr in arrays if arr.strip()]

            all_arrays_conv = arrays + conv
            all_arrays_conp = arrays + conp
            method_template = 'double* {}_in'
            setter_template = '{} = {}_in;'

            file.write('\n#ifdef CONP\n' +
                        '\n'.join([template.format(arr) for arr in conp]) +
                        '\n#else\n' +
                        '\n'.join([template.format(arr) for arr in conv]) +
                        '\n#endif\n'
                        )
            file.write('\n'.join([template.format(arr) for arr in arrays]))
            file.write('\n\n'
                       '#ifdef CONP\n'
                       '__global__ pointer_set_kernel(' +
                       ', '.join([method_template.format(arr) for arr in all_arrays_conp]) +
                       ') {\n  ' +
                       '\n  '.join([setter_template.format(arr, arr) for arr in all_arrays_conp]) +
                       '\n}\n'
                       '#else\n'
                       '__global__ pointer_set_kernel(' +
                       ', '.join([method_template.format(arr) for arr in all_arrays_conv]) +
                       ') {\n  ' +
                       '\n  '.join([setter_template.format(arr, arr) for arr in all_arrays_conv]) +
                       '\n}\n'
                       '#endif\n'
                       )

            file.write('\n#ifdef CONP\n' +
            '\n'.join([host_template.format(arr) for arr in conp]) +
            '\n#else\n' +
            '\n'.join([host_template.format(arr) for arr in conv]) +
            '\n#endif\n' +
            '\n'.join([host_template.format(arr) for arr in arrays])
            )
            file.write('\n')
            malloc_template = 'cudaMalloc((void**)&{}, {});'
            free_template = 'cudaFree({});'
            file.write('\nvoid initialize_gpu_memory(int NUM, int block_size, int grid_size) {\n'
                       '  int padded = grid_size * block_size;\n'
                       '  padded = padded > NUM ? padded : NUM;\n'
                       '#ifdef CONP\n  ' +
                       '\n  '.join([malloc_template.format('d_' + arr, 'NN * padded * sizeof(double)') for arr in all_arrays_conp[:2]]) +
                       '\n  ' +
                       '\n  '.join([malloc_template.format('d_' + arr, 'NSP * padded * sizeof(double)') for arr in all_arrays_conp[2:]]) +
                       '\n  pointer_set_kernel<<<1, 1>>>({});'.format(', '.join(all_arrays_conp)) +
                       '\n#else\n  ' +
                       '\n  '.join([malloc_template.format('d_' + arr, 'NN * padded * sizeof(double)') for arr in all_arrays_conv[:2]]) +
                       '\n  ' +
                       '\n  '.join([malloc_template.format('d_' + arr, 'NSP * padded * sizeof(double)') for arr in all_arrays_conv[2:]]) +
                       '\n  pointer_set_kernel<<<1, 1>>>({});'.format(', '.join(all_arrays_conv)) +
                       '\n#endif\n'
                       '  return padded;\n'
                       '}'
                       )
            file.write('\n')
            file.write('\n\nvoid free_gpu_memory() {\n'
                       '#ifdef CONP\n  ' +
                       '\n  '.join([free_template.format('d_' + arr) for arr in all_arrays_conp]) +
                       '\n#else\n  ' +
                       '\n  '.join([free_template.format('d_' + arr) for arr in all_arrays_conv]) +
                       '\n#endif\n'
                       '}\n'
                       )<|endoftext|>


<commit_before>import json
from json.decoder import JSONDecodeError

try:
    import xattr2
except ImportError:
    print("===")
    print("pyxattr library not found (https://pypi.python.org/pypi/pyxattr)")
    print("===")
    raise

from mediafs.fs import File, Directory, CachedRootDirectory

class XAttrMetadata(object):
    """
    A dict-like class that uses the ``xattr`` module to reflect key/value pairs from
    extended filesystem attributes.

    Keys must be strings.

    Values are fed through the ``json`` module for better type support.

    No caching is done, which means that all values will reflect the status of the file
    at the time of the call. This is useful if you have multiple applications operating
    on the same attribute data.

    If you are utilizing this metadata with tools other than this library, be aware that
    all attribute keys are in the user namespace.

    For example:

    | >>> fs['asdf2.txt'].metadata['author'] = "John Smith"

    | $ getfattr -n user.author asdf2.txt
    | # file: asdf2.txt
    | user.author="\"John Smith\""
    """

    def __init__(self, path):
        self._path = path

    def _decodeVal(self, val):
        """
        Helper function for decoding values we get from ``xattr.get()``.
        Converts the value to a string, then passes it through ``json.loads()`` for type conversion.
        """
        val = val.decode()
        try:
            return json.loads(val)
        except JSONDecodeError:
            return val

    def __getitem__(self, key):
        """
        Retrieves an extended filesystem attribute. Raises ``KeyError`` if the file does not
        have an attribute with that name.
        """
        try:
            return self._decodeVal(xattr.get(self._path, key, namespace=xattr.NS_USER))
        except OSError:
            pass
        raise KeyError(key)

    def __setitem__(self, key, val):
        """
        Sets a filesystem attribute
        """
        xattr.set(self._path, key, json.dumps(val), namespace=xattr.NS_USER)

    def __delitem__(self, key):
        """
        Removes a filesystem attribute
        """
        xattr.remove(self._path, key, namespace=xattr.NS_USER)

    def __contains__(self, key):
        """
        Checks if a filesystem attribute exists on this file
        """
        return key in self.keys()

    def __iter__(self):
        """
        An interator on all values stored in filesystem attributes
        """
        for val in self.values():
            yield val

    def __len__(self):
        """
        Returns the number of attributes stored
        """
        return len(self.keys())

    def keys(self):
        """
        Returns a list of all filesystem attributes
        """
        return [ key.decode() for key in xattr.list(self._path, namespace=xattr.NS_USER) ]

    def values(self):
        """
        Returns a list of all filesystem attribute values
        """
        return [ self[key] for key in self.keys() ]

    def items(self):
        """
        Returns a list of key/value pairs for all filesystem attributes
        """
        allAttrs = xattr.get_all(self._path, namespace=xattr.NS_USER)
        return [ (key.decode(), self._decodeVal(val)) for key, val in allAttrs ]

    def pop(self, key):
        """
        Removes a filesystem attribute by its name and returns its value
        """
        val = self[key]
        del self[key]
        return val

    def copy(self):
        """
        Returns a dict with all filesystem attributes as key/value pairs
        """
        return { k:v for k, v in self.items() }

    def __str__(self):
        return "{%s}" % ", ".join("%r: %r" % (key, val) for key, val in self.items())
    __repr__ = __str__

class XAttrRootDirectory(CachedRootDirectory):
    """
    A RootDirectory object that stores file and directory metadata as extended filesystem
    attributes on Linux systems. Depends on the ``pyxattrs`` Python module.

    This class is derived from ``CachedRootDirectory``, so you will still need to call
    ``save()`` to write the directory tree cache. If you do not need this functionality
    at all, pass ``treeFile=None`` as an argument to the constructor.

    Using extended filesystem attributes for metadata can situationally be faster than
    using a single JSON file to store all metadata because it doesn't have the overhead
    of hashing (even partially) files before doing the metadata lookup. It also has the
    advantage of working on the filesystem attributes directly - calling
    ``file.metadata['asdf'] = 5`` immediately serializes the value, which is then available
    to other applications without needing to call ``save()`` first.

    The primary disadvantage is portability. Many system tools such as ``cp`` and ``rsync``
    will ignore extended attributes and you can lose data this way. It can also be tricky
    to get your data out if you move it to another system.

    If you have issues using this class, try using the command line tools ``attr``,
    ``getfattr``, and ``setfattr``. Please note that this class uses the "user" namespace,
    which means that if you set a metadata key called ``title``, then the actual key name
    will be ``user.title`` when using the command line tools.

    If the command line tools do not work for you, then there is probably something wrong
    with your OS configuration. Some filesystems (like ``tmpfs``) do not support extended
    attributes at all, and others may need to have particular mount options to work correctly.
    """

    def __init__(self, path, treeFile=".tree.json"):
        CachedRootDirectory.__init__(self, path, metadataFile=None, treeFile=treeFile)

    def _getMetadataForObject(self, obj):
        """
        Return a wrapper around the ``xattr`` module for metadata access
        """
        return XAttrMetadata(obj.abspath)

<commit_msg>Fixed typo<commit_after>import json
from json.decoder import JSONDecodeError

try:
    import xattr
except ImportError:
    print("===")
    print("pyxattr library not found (https://pypi.python.org/pypi/pyxattr)")
    print("===")
    raise

from mediafs.fs import File, Directory, CachedRootDirectory

class XAttrMetadata(object):
    """
    A dict-like class that uses the ``xattr`` module to reflect key/value pairs from
    extended filesystem attributes.

    Keys must be strings.

    Values are fed through the ``json`` module for better type support.

    No caching is done, which means that all values will reflect the status of the file
    at the time of the call. This is useful if you have multiple applications operating
    on the same attribute data.

    If you are utilizing this metadata with tools other than this library, be aware that
    all attribute keys are in the user namespace.

    For example:

    | >>> fs['asdf2.txt'].metadata['author'] = "John Smith"

    | $ getfattr -n user.author asdf2.txt
    | # file: asdf2.txt
    | user.author="\"John Smith\""
    """

    def __init__(self, path):
        self._path = path

    def _decodeVal(self, val):
        """
        Helper function for decoding values we get from ``xattr.get()``.
        Converts the value to a string, then passes it through ``json.loads()`` for type conversion.
        """
        val = val.decode()
        try:
            return json.loads(val)
        except JSONDecodeError:
            return val

    def __getitem__(self, key):
        """
        Retrieves an extended filesystem attribute. Raises ``KeyError`` if the file does not
        have an attribute with that name.
        """
        try:
            return self._decodeVal(xattr.get(self._path, key, namespace=xattr.NS_USER))
        except OSError:
            pass
        raise KeyError(key)

    def __setitem__(self, key, val):
        """
        Sets a filesystem attribute
        """
        xattr.set(self._path, key, json.dumps(val), namespace=xattr.NS_USER)

    def __delitem__(self, key):
        """
        Removes a filesystem attribute
        """
        xattr.remove(self._path, key, namespace=xattr.NS_USER)

    def __contains__(self, key):
        """
        Checks if a filesystem attribute exists on this file
        """
        return key in self.keys()

    def __iter__(self):
        """
        An interator on all values stored in filesystem attributes
        """
        for val in self.values():
            yield val

    def __len__(self):
        """
        Returns the number of attributes stored
        """
        return len(self.keys())

    def keys(self):
        """
        Returns a list of all filesystem attributes
        """
        return [ key.decode() for key in xattr.list(self._path, namespace=xattr.NS_USER) ]

    def values(self):
        """
        Returns a list of all filesystem attribute values
        """
        return [ self[key] for key in self.keys() ]

    def items(self):
        """
        Returns a list of key/value pairs for all filesystem attributes
        """
        allAttrs = xattr.get_all(self._path, namespace=xattr.NS_USER)
        return [ (key.decode(), self._decodeVal(val)) for key, val in allAttrs ]

    def pop(self, key):
        """
        Removes a filesystem attribute by its name and returns its value
        """
        val = self[key]
        del self[key]
        return val

    def copy(self):
        """
        Returns a dict with all filesystem attributes as key/value pairs
        """
        return { k:v for k, v in self.items() }

    def __str__(self):
        return "{%s}" % ", ".join("%r: %r" % (key, val) for key, val in self.items())
    __repr__ = __str__

class XAttrRootDirectory(CachedRootDirectory):
    """
    A RootDirectory object that stores file and directory metadata as extended filesystem
    attributes on Linux systems. Depends on the ``pyxattrs`` Python module.

    This class is derived from ``CachedRootDirectory``, so you will still need to call
    ``save()`` to write the directory tree cache. If you do not need this functionality
    at all, pass ``treeFile=None`` as an argument to the constructor.

    Using extended filesystem attributes for metadata can situationally be faster than
    using a single JSON file to store all metadata because it doesn't have the overhead
    of hashing (even partially) files before doing the metadata lookup. It also has the
    advantage of working on the filesystem attributes directly - calling
    ``file.metadata['asdf'] = 5`` immediately serializes the value, which is then available
    to other applications without needing to call ``save()`` first.

    The primary disadvantage is portability. Many system tools such as ``cp`` and ``rsync``
    will ignore extended attributes and you can lose data this way. It can also be tricky
    to get your data out if you move it to another system.

    If you have issues using this class, try using the command line tools ``attr``,
    ``getfattr``, and ``setfattr``. Please note that this class uses the "user" namespace,
    which means that if you set a metadata key called ``title``, then the actual key name
    will be ``user.title`` when using the command line tools.

    If the command line tools do not work for you, then there is probably something wrong
    with your OS configuration. Some filesystems (like ``tmpfs``) do not support extended
    attributes at all, and others may need to have particular mount options to work correctly.
    """

    def __init__(self, path, treeFile=".tree.json"):
        CachedRootDirectory.__init__(self, path, metadataFile=None, treeFile=treeFile)

    def _getMetadataForObject(self, obj):
        """
        Return a wrapper around the ``xattr`` module for metadata access
        """
        return XAttrMetadata(obj.abspath)

<|endoftext|>


<commit_before>from os import environ
from logging import getLogger
from os.path import join, dirname
from argparse import ArgumentParser
from time import sleep

from . import jobs
from boto.ec2 import EC2Connection

parser = ArgumentParser(description='Run some source files.')

parser.add_argument('bucketname',
                    help='Required S3 bucket name.')

parser.add_argument('-a', '--access-key', default=environ.get('AWS_ACCESS_KEY_ID', None),
                    help='Optional AWS access key name. Defaults to value of AWS_ACCESS_KEY_ID environment variable.')

parser.add_argument('-s', '--secret-key', default=environ.get('AWS_SECRET_ACCESS_KEY', None),
                    help='Optional AWS secret key name. Defaults to value of AWS_SECRET_ACCESS_KEY environment variable.')

parser.add_argument('--ec2-access-key',
                    help='Optional AWS access key name for setting up EC2; distinct from access key for populating S3 bucket. Defaults to value of S3 access key.')

parser.add_argument('--ec2-secret-key',
                    help='Optional AWS secret key name for setting up EC2; distinct from secret key for populating S3 bucket. Defaults to value of S3 secret key.')

parser.add_argument('--instance-type', default='m3.xlarge',
                    help='EC2 instance type, defaults to m3.xlarge.')

parser.add_argument('--machine-image', default='ami-4ae27e22',
                    help='AMI identifier, defaults to Alestic Ubuntu 14.04 (ami-4ae27e22).')

def main():
    args = parser.parse_args()
    jobs.setup_logger(None)

    with open(join(dirname(__file__), 'templates', 'user-data.sh')) as file:
        user_data = file.read().format(**args.__dict__)

    getLogger('openaddr').info('Prepared {} bytes of instance user data\n\n{}'.format(len(user_data), user_data))

    ec2 = EC2Connection(args.ec2_access_key, args.ec2_secret_key)
    history = ec2.get_spot_price_history(instance_type=args.instance_type)
    median = sorted([h.price for h in history])[len(history)/2]
    bid = median + .01

    getLogger('openaddr').info('Bidding ${:.4f}/hour'.format(bid))

    spot_args = dict(instance_type=args.instance_type, user_data=user_data,
                     key_name='cfa-keypair-2013', security_groups=['default'])

    spot_req = ec2.request_spot_instances(bid, args.machine_image, **spot_args)[0]

    url = 'https://console.aws.amazon.com/ec2/v2/home?region=us-east-1#SpotInstances:search={}'.format(spot_req.id)

    getLogger('openaddr').info(url)

    while spot_req.state == 'open':
        getLogger('openaddr').debug('Spot request is {}'.format(spot_req.status))
        sleep(15)
        spot_req = ec2.get_all_spot_instance_requests(spot_req.id)[0]

    if spot_req.state != 'active':
        raise Exception('FUCK')

    while not spot_req.instance_id:
        getLogger('openaddr').debug('Waiting for instance ID')
        sleep(15)
        spot_req = ec2.get_all_spot_instance_requests(spot_req.id)[0]

    getLogger('openaddr').info('Looking up instance {}'.format(spot_req.instance_id))
    instance = ec2.get_only_instances(spot_req.instance_id)[0]

    getLogger('openaddr').info('Found instance {}'.format(instance.public_dns_name))
    sleep(10)
    instance.terminate()

if __name__ == '__main__':
    exit(main())
<commit_msg>Commented parts of openaddr.run main()<commit_after>from os import environ
from logging import getLogger
from os.path import join, dirname
from argparse import ArgumentParser
from time import sleep

from . import jobs
from boto.ec2 import EC2Connection

parser = ArgumentParser(description='Run some source files.')

parser.add_argument('bucketname',
                    help='Required S3 bucket name.')

parser.add_argument('-a', '--access-key', default=environ.get('AWS_ACCESS_KEY_ID', None),
                    help='Optional AWS access key name. Defaults to value of AWS_ACCESS_KEY_ID environment variable.')

parser.add_argument('-s', '--secret-key', default=environ.get('AWS_SECRET_ACCESS_KEY', None),
                    help='Optional AWS secret key name. Defaults to value of AWS_SECRET_ACCESS_KEY environment variable.')

parser.add_argument('--ec2-access-key',
                    help='Optional AWS access key name for setting up EC2; distinct from access key for populating S3 bucket. Defaults to value of S3 access key.')

parser.add_argument('--ec2-secret-key',
                    help='Optional AWS secret key name for setting up EC2; distinct from secret key for populating S3 bucket. Defaults to value of S3 secret key.')

parser.add_argument('--instance-type', default='m3.xlarge',
                    help='EC2 instance type, defaults to m3.xlarge.')

parser.add_argument('--machine-image', default='ami-4ae27e22',
                    help='AMI identifier, defaults to Alestic Ubuntu 14.04 (ami-4ae27e22).')

def main():
    args = parser.parse_args()
    jobs.setup_logger(None)

    #
    # Prepare init script for new EC2 instance to run.
    #
    with open(join(dirname(__file__), 'templates', 'user-data.sh')) as file:
        user_data = file.read().format(**args.__dict__)

    getLogger('openaddr').info('Prepared {} bytes of instance user data\n\n{}'.format(len(user_data), user_data))

    #
    # Figure out how much we're willing to bid on a spot instance.
    #
    ec2 = EC2Connection(args.ec2_access_key, args.ec2_secret_key)
    history = ec2.get_spot_price_history(instance_type=args.instance_type)
    median = sorted([h.price for h in history])[len(history)/2]
    bid = median + .01

    getLogger('openaddr').info('Bidding ${:.4f}/hour'.format(bid))

    #
    # Request a spot instance, then wait while it does its thing.
    #
    spot_args = dict(instance_type=args.instance_type, user_data=user_data,
                     key_name='cfa-keypair-2013', security_groups=['default'])

    spot_req = ec2.request_spot_instances(bid, args.machine_image, **spot_args)[0]

    url = 'https://console.aws.amazon.com/ec2/v2/home?region=us-east-1#SpotInstances:search={}'.format(spot_req.id)

    getLogger('openaddr').info(url)

    while spot_req.state == 'open':
        getLogger('openaddr').debug('Spot request is {}'.format(spot_req.status))
        sleep(15)
        spot_req = ec2.get_all_spot_instance_requests(spot_req.id)[0]

    if spot_req.state != 'active':
        raise Exception('FUCK')

    while not spot_req.instance_id:
        getLogger('openaddr').debug('Waiting for instance ID')
        sleep(15)
        spot_req = ec2.get_all_spot_instance_requests(spot_req.id)[0]

    getLogger('openaddr').info('Looking up instance {}'.format(spot_req.instance_id))
    instance = ec2.get_only_instances(spot_req.instance_id)[0]

    getLogger('openaddr').info('Found instance {}'.format(instance.public_dns_name))
    sleep(10)
    instance.terminate()

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