<commit_before>#!/usr/bin/env python3
"""
Communicate with the UPS.

Store data from a supported UPS in an sqlite3 database.
"""

import configparser
import datetime as dt
import os
import sqlite3
import subprocess
import sys
import syslog
import time
import traceback

import mausy5043funcs.fileops3 as mf
import mausy5043libs.libsignals3 as ml

import constants

from hanging_threads import start_monitoring
anti_freeze = constants.UPS['report_time'] * 2

# constants
DEBUG = False
HERE = os.path.realpath(__file__).split('/')
# runlist id :
MYID = HERE[-1]
# app_name :
MYAPP = HERE[-3]
MYROOT = "/".join(HERE[0:-3])
# host_name :
NODE = os.uname()[1]

# example values:
# HERE: ['', 'home', 'pi', 'upsdiagdd', 'bin', 'ups.py']
# MYID: 'ups.py
# MYAPP: upsdiagd
# MYROOT: /home/pi
# NODE: rbups

def main():
    """Execute main loop."""
    global DEBUG
    global MYAPP
    global MYID
    global MYROOT
    killer = ml.GracefulKiller()
    start_monitoring(seconds_frozen=anti_freeze, test_interval=1000)
    iniconf = configparser.ConfigParser()
    iniconf.read(f'{MYROOT}/{MYAPP}/config.ini')
    # report_time = iniconf.getint(MYID, 'reporttime')
    report_time = constants.UPS['report_time']
    # fdatabase = f"{MYROOT}/{iniconf.get('DEFAULT', 'databasefile')}"
    fdatabase = constants.UPS['database']
    # sqlcmd = iniconf.get(MYID, 'sqlcmd')
    sqlcmd = constants.UPS['sql_command']
    # samples_averaged = iniconf.getint(MYID, 'samplespercycle') \
    #                   * iniconf.getint(MYID, 'cycles')
    # sample_time = report_time / iniconf.getint(MYID, 'samplespercycle')
    samples_averaged = int(constants.UPS['samplespercycle']) \
                       * int(constants.UPS['cycles'])
    sample_time = report_time / int(constants.UPS['samplespercycle'])
    data = []

    test_db_connection(fdatabase)

    pause_time = time.time()
    while not killer.kill_now:
        if time.time() > pause_time:
            start_time = time.time()
            result = do_work()
            mf.syslog_trace(f"Result   : {result}", False, DEBUG)
            data.append([float(d) for d in result])
            if len(data) > samples_averaged:
                data.pop(0)

            # report sample average
            if start_time % report_time < sample_time:
                # somma       = list(map(sum, zip(*data)))
                somma = [sum(d) for d in zip(*data)]
                # not all entries should be float
                # ['234.000', '13.700', '100.000', '20.000', '1447.000']
                averages = [float(format(d / len(data), '.3f')) for d in somma]
                mf.syslog_trace("Averages : {0}".format(averages),
                                False,
                                DEBUG)
                do_add_to_database(averages, fdatabase, sqlcmd)

            pause_time = (sample_time - (time.time() - start_time) -
                          (start_time % sample_time) + time.time())
            if pause_time > 0:
                mf.syslog_trace(f"Waiting  : {pause_time - time.time():.1f}s", False, DEBUG)
                mf.syslog_trace("................................", False, DEBUG)
            else:
                mf.syslog_trace(f"Behind   : {pause_time - time.time():.1f}s", False, DEBUG)
                mf.syslog_trace("................................", False, DEBUG)
        else:
            time.sleep(1.0)

def do_work():
    """Do the thing.
    Example:
    *2*  battery.charge: 100
        battery.charge.low: 20
    *4*  battery.runtime: 1875
        battery.type: PbAc
        device.mfr: EATON
        device.model: Protection Station 650
        device.serial: AN2E49008
        device.type: ups
        driver.name: usbhid-ups
        driver.parameter.pollfreq: 30
        driver.parameter.pollinterval: 2
        driver.parameter.port: auto
        driver.parameter.synchronous: no
        driver.version: 2.7.4
        driver.version.data: MGE HID 1.39
        driver.version.internal: 0.41
        input.transfer.high: 264
        input.transfer.low: 184
        outlet.1.desc: PowerShare Outlet 1
        outlet.1.id: 2
        outlet.1.status: on
        outlet.1.switchable: no
        outlet.2.desc: PowerShare Outlet 2
        outlet.2.id: 3
        outlet.2.status: on
        outlet.2.switchable: no
        outlet.desc: Main Outlet
        outlet.id: 1
        outlet.power: 25
        outlet.switchable: no
        output.frequency.nominal: 50
    *0*  output.voltage: 230.0
        output.voltage.nominal: 230
        ups.beeper.status: enabled
        ups.delay.shutdown: 20
        ups.delay.start: 30
        ups.firmware: 1.13
    *3*  ups.load: 2
        ups.mfr: EATON
        ups.model: Protection Station 650
        ups.power.nominal: 650
        ups.productid: ffff
        ups.serial: AN2E49008
        ups.status: OL
        ups.timer.shutdown: -1
        ups.timer.start: -1
        ups.vendorid: 0463
    """
    # 5 datapoints gathered here
    try:
        upsc = str(subprocess.check_output(['upsc', 'ups@localhost'],
                                           stderr=subprocess.STDOUT),
                   'utf-8').splitlines()
    except subprocess.CalledProcessError:
        syslog.syslog(syslog.LOG_ALERT, "Waiting 10s ...")

        time.sleep(10)  # wait to let the driver crash properly
        mf.syslog_trace("*** RESTARTING nut-server.service ***",
                        syslog.LOG_ALERT, DEBUG)
        redo = str(subprocess.check_output(['sudo',
                                            'systemctl',
                                            'restart',
                                            'nut-server.service']
                                           ),
                   'utf-8').splitlines()
        mf.syslog_trace("Returned : {0}".format(redo), False, DEBUG)

        time.sleep(15)
        mf.syslog_trace("!!! Retrying communication with UPS !!!",
                        syslog.LOG_ALERT,
                        DEBUG)
        upsc = str(subprocess.check_output(['upsc',
                                            'ups@localhost'
                                            ],
                                           stderr=subprocess.STDOUT),
                   'utf-8').splitlines()

    ups_data = [-1.0, -1.0, -1.0, -1.0, -1.0]
    for element in upsc:
        var = element.split(': ')
        # if var[0] == 'input.voltage':
        if var[0] == 'output.voltage':
            ups_data[0] = float(var[1])
        # not available on Eaton Protection Station:
        if var[0] == 'battery.voltage':
            ups_data[1] = float(var[1])
        if var[0] == 'battery.charge':
            ups_data[2] = float(var[1])
        if var[0] == 'ups.load':
            ups_data[3] = float(var[1])
        if var[0] == 'battery.runtime':
            ups_data[4] = float(var[1])

    return ups_data

def do_add_to_database(result, fdatabase, sql_cmd):
    """Commit the results to the database."""
    # Get the time and date in human-readable form and UN*X-epoch...
    conn = None
    cursor = None
    dt_format = '%Y-%m-%d %H:%M:%S'
    out_date = dt.datetime.now()  # time.strftime('%Y-%m-%dT%H:%M:%S')
    out_epoch = int(out_date.timestamp())
    results = (out_date.strftime(dt_format),
               out_epoch,
               result[0],
               result[1],
               result[2],
               result[3],
               result[4]
               )
    mf.syslog_trace(f"   @: {out_date.strftime(dt_format)}", False, DEBUG)
    mf.syslog_trace(f"    : {results}", False, DEBUG)

    retries = 10
    while retries:
        retries -= 1
        try:
            conn = create_db_connection(fdatabase)
            cursor = conn.cursor()
            cursor.execute(sql_cmd, results)
            cursor.close()
            conn.commit()
            conn.close()
            retries = 0
        except sqlite3.OperationalError:
            if cursor:
                cursor.close()
            if conn:
                conn.close()
            if retries:
                raise

def create_db_connection(database_file):
    """
    Create a database connection to the SQLite3 database
    specified by database_file.
    """
    consql = None
    mf.syslog_trace(f"Connecting to: {database_file}", False, DEBUG)
    try:
        consql = sqlite3.connect(database_file, timeout=9000)
        return consql
    except sqlite3.Error:
        mf.syslog_trace("Unexpected SQLite3 error when connecting to server.",
                        syslog.LOG_CRIT,
                        DEBUG
                        )
        mf.syslog_trace(traceback.format_exc(), syslog.LOG_CRIT, DEBUG)
        if consql:  # attempt to close connection to SQLite3 server
            consql.close()
            mf.syslog_trace(" ** Closed SQLite3 connection. **",
                            syslog.LOG_CRIT,
                            DEBUG
                            )
        raise

def test_db_connection(fdatabase):
    """
    Test & log database engine connection.
    """
    try:
        conn = create_db_connection(fdatabase)
        cursor = conn.cursor()
        cursor.execute("SELECT sqlite_version();")
        versql = cursor.fetchone()
        cursor.close()
        conn.commit()
        conn.close()
        syslog.syslog(syslog.LOG_INFO, f"Attached to SQLite3 server: {versql} using {fdatabase}")
    except sqlite3.Error:
        mf.syslog_trace("Unexpected SQLite3 error during test.",
                        syslog.LOG_CRIT,
                        DEBUG
                        )
        mf.syslog_trace(traceback.format_exc(), syslog.LOG_CRIT, DEBUG)
        raise

if __name__ == "__main__":
    # initialise logging
    syslog.openlog(ident=f'{MYAPP}.{MYID.split(".")[0]}',
                   facility=syslog.LOG_LOCAL0)

    if len(sys.argv) == 2:
        if sys.argv[1] == 'start':
            main()
        elif sys.argv[1] == 'restart':
            main()
        elif sys.argv[1] == 'debug':
            # assist with debugging.
            DEBUG = True
            mf.syslog_trace("Debug-mode started.", syslog.LOG_DEBUG, DEBUG)
            print("Use <Ctrl>+C to stop.")
            main()
        else:
            print("Unknown command")
            sys.exit(2)
    else:
        print("usage: {0!s} start|restart|debug".format(sys.argv[0]))
        sys.exit(2)
    print("And it's goodnight from him")
    sys.exit(0)
<commit_msg>monitor main()<commit_after>#!/usr/bin/env python3
"""
Communicate with the UPS.

Store data from a supported UPS in an sqlite3 database.
"""

import configparser
import datetime as dt
import os
import sqlite3
import subprocess
import sys
import syslog
import time
import traceback

import mausy5043funcs.fileops3 as mf
import mausy5043libs.libsignals3 as ml

import constants

from hanging_threads import start_monitoring
anti_freeze = constants.UPS['report_time'] * 2

# constants
DEBUG = False
HERE = os.path.realpath(__file__).split('/')
# runlist id :
MYID = HERE[-1]
# app_name :
MYAPP = HERE[-3]
MYROOT = "/".join(HERE[0:-3])
# host_name :
NODE = os.uname()[1]

# example values:
# HERE: ['', 'home', 'pi', 'upsdiagdd', 'bin', 'ups.py']
# MYID: 'ups.py
# MYAPP: upsdiagd
# MYROOT: /home/pi
# NODE: rbups

def main():
    """Execute main loop."""
    global DEBUG
    global MYAPP
    global MYID
    global MYROOT
    killer = ml.GracefulKiller()
    start_monitoring(seconds_frozen=anti_freeze, test_interval=2000)
    iniconf = configparser.ConfigParser()
    iniconf.read(f'{MYROOT}/{MYAPP}/config.ini')
    # report_time = iniconf.getint(MYID, 'reporttime')
    report_time = constants.UPS['report_time']
    # fdatabase = f"{MYROOT}/{iniconf.get('DEFAULT', 'databasefile')}"
    fdatabase = constants.UPS['database']
    # sqlcmd = iniconf.get(MYID, 'sqlcmd')
    sqlcmd = constants.UPS['sql_command']
    # samples_averaged = iniconf.getint(MYID, 'samplespercycle') \
    #                   * iniconf.getint(MYID, 'cycles')
    # sample_time = report_time / iniconf.getint(MYID, 'samplespercycle')
    samples_averaged = int(constants.UPS['samplespercycle']) \
                       * int(constants.UPS['cycles'])
    sample_time = report_time / int(constants.UPS['samplespercycle'])
    data = []

    test_db_connection(fdatabase)

    pause_time = time.time()
    while not killer.kill_now:
        if time.time() > pause_time:
            start_time = time.time()
            result = do_work()
            mf.syslog_trace(f"Result   : {result}", False, DEBUG)
            data.append([float(d) for d in result])
            if len(data) > samples_averaged:
                data.pop(0)

            # report sample average
            if start_time % report_time < sample_time:
                # somma       = list(map(sum, zip(*data)))
                somma = [sum(d) for d in zip(*data)]
                # not all entries should be float
                # ['234.000', '13.700', '100.000', '20.000', '1447.000']
                averages = [float(format(d / len(data), '.3f')) for d in somma]
                mf.syslog_trace("Averages : {0}".format(averages),
                                False,
                                DEBUG)
                do_add_to_database(averages, fdatabase, sqlcmd)

            pause_time = (sample_time - (time.time() - start_time) -
                          (start_time % sample_time) + time.time())
            if pause_time > 0:
                mf.syslog_trace(f"Waiting  : {pause_time - time.time():.1f}s", False, DEBUG)
                mf.syslog_trace("................................", False, DEBUG)
            else:
                mf.syslog_trace(f"Behind   : {pause_time - time.time():.1f}s", False, DEBUG)
                mf.syslog_trace("................................", False, DEBUG)
        else:
            time.sleep(1.0)

def do_work():
    """Do the thing.
    Example:
    *2*  battery.charge: 100
        battery.charge.low: 20
    *4*  battery.runtime: 1875
        battery.type: PbAc
        device.mfr: EATON
        device.model: Protection Station 650
        device.serial: AN2E49008
        device.type: ups
        driver.name: usbhid-ups
        driver.parameter.pollfreq: 30
        driver.parameter.pollinterval: 2
        driver.parameter.port: auto
        driver.parameter.synchronous: no
        driver.version: 2.7.4
        driver.version.data: MGE HID 1.39
        driver.version.internal: 0.41
        input.transfer.high: 264
        input.transfer.low: 184
        outlet.1.desc: PowerShare Outlet 1
        outlet.1.id: 2
        outlet.1.status: on
        outlet.1.switchable: no
        outlet.2.desc: PowerShare Outlet 2
        outlet.2.id: 3
        outlet.2.status: on
        outlet.2.switchable: no
        outlet.desc: Main Outlet
        outlet.id: 1
        outlet.power: 25
        outlet.switchable: no
        output.frequency.nominal: 50
    *0*  output.voltage: 230.0
        output.voltage.nominal: 230
        ups.beeper.status: enabled
        ups.delay.shutdown: 20
        ups.delay.start: 30
        ups.firmware: 1.13
    *3*  ups.load: 2
        ups.mfr: EATON
        ups.model: Protection Station 650
        ups.power.nominal: 650
        ups.productid: ffff
        ups.serial: AN2E49008
        ups.status: OL
        ups.timer.shutdown: -1
        ups.timer.start: -1
        ups.vendorid: 0463
    """
    # 5 datapoints gathered here
    try:
        upsc = str(subprocess.check_output(['upsc', 'ups@localhost'],
                                           stderr=subprocess.STDOUT),
                   'utf-8').splitlines()
    except subprocess.CalledProcessError:
        syslog.syslog(syslog.LOG_ALERT, "Waiting 10s ...")

        time.sleep(10)  # wait to let the driver crash properly
        mf.syslog_trace("*** RESTARTING nut-server.service ***",
                        syslog.LOG_ALERT, DEBUG)
        redo = str(subprocess.check_output(['sudo',
                                            'systemctl',
                                            'restart',
                                            'nut-server.service']
                                           ),
                   'utf-8').splitlines()
        mf.syslog_trace("Returned : {0}".format(redo), False, DEBUG)

        time.sleep(15)
        mf.syslog_trace("!!! Retrying communication with UPS !!!",
                        syslog.LOG_ALERT,
                        DEBUG)
        upsc = str(subprocess.check_output(['upsc',
                                            'ups@localhost'
                                            ],
                                           stderr=subprocess.STDOUT),
                   'utf-8').splitlines()

    ups_data = [-1.0, -1.0, -1.0, -1.0, -1.0]
    for element in upsc:
        var = element.split(': ')
        # if var[0] == 'input.voltage':
        if var[0] == 'output.voltage':
            ups_data[0] = float(var[1])
        # not available on Eaton Protection Station:
        if var[0] == 'battery.voltage':
            ups_data[1] = float(var[1])
        if var[0] == 'battery.charge':
            ups_data[2] = float(var[1])
        if var[0] == 'ups.load':
            ups_data[3] = float(var[1])
        if var[0] == 'battery.runtime':
            ups_data[4] = float(var[1])

    return ups_data

def do_add_to_database(result, fdatabase, sql_cmd):
    """Commit the results to the database."""
    # Get the time and date in human-readable form and UN*X-epoch...
    conn = None
    cursor = None
    dt_format = '%Y-%m-%d %H:%M:%S'
    out_date = dt.datetime.now()  # time.strftime('%Y-%m-%dT%H:%M:%S')
    out_epoch = int(out_date.timestamp())
    results = (out_date.strftime(dt_format),
               out_epoch,
               result[0],
               result[1],
               result[2],
               result[3],
               result[4]
               )
    mf.syslog_trace(f"   @: {out_date.strftime(dt_format)}", False, DEBUG)
    mf.syslog_trace(f"    : {results}", False, DEBUG)

    retries = 10
    while retries:
        retries -= 1
        try:
            conn = create_db_connection(fdatabase)
            cursor = conn.cursor()
            cursor.execute(sql_cmd, results)
            cursor.close()
            conn.commit()
            conn.close()
            retries = 0
        except sqlite3.OperationalError:
            if cursor:
                cursor.close()
            if conn:
                conn.close()
            if retries:
                raise

def create_db_connection(database_file):
    """
    Create a database connection to the SQLite3 database
    specified by database_file.
    """
    consql = None
    mf.syslog_trace(f"Connecting to: {database_file}", False, DEBUG)
    try:
        consql = sqlite3.connect(database_file, timeout=9000)
        return consql
    except sqlite3.Error:
        mf.syslog_trace("Unexpected SQLite3 error when connecting to server.",
                        syslog.LOG_CRIT,
                        DEBUG
                        )
        mf.syslog_trace(traceback.format_exc(), syslog.LOG_CRIT, DEBUG)
        if consql:  # attempt to close connection to SQLite3 server
            consql.close()
            mf.syslog_trace(" ** Closed SQLite3 connection. **",
                            syslog.LOG_CRIT,
                            DEBUG
                            )
        raise

def test_db_connection(fdatabase):
    """
    Test & log database engine connection.
    """
    try:
        conn = create_db_connection(fdatabase)
        cursor = conn.cursor()
        cursor.execute("SELECT sqlite_version();")
        versql = cursor.fetchone()
        cursor.close()
        conn.commit()
        conn.close()
        syslog.syslog(syslog.LOG_INFO, f"Attached to SQLite3 server: {versql} using {fdatabase}")
    except sqlite3.Error:
        mf.syslog_trace("Unexpected SQLite3 error during test.",
                        syslog.LOG_CRIT,
                        DEBUG
                        )
        mf.syslog_trace(traceback.format_exc(), syslog.LOG_CRIT, DEBUG)
        raise

if __name__ == "__main__":
    # initialise logging
    syslog.openlog(ident=f'{MYAPP}.{MYID.split(".")[0]}',
                   facility=syslog.LOG_LOCAL0)

    if len(sys.argv) == 2:
        if sys.argv[1] == 'start':
            main()
        elif sys.argv[1] == 'restart':
            main()
        elif sys.argv[1] == 'debug':
            # assist with debugging.
            DEBUG = True
            mf.syslog_trace("Debug-mode started.", syslog.LOG_DEBUG, DEBUG)
            print("Use <Ctrl>+C to stop.")
            main()
        else:
            print("Unknown command")
            sys.exit(2)
    else:
        print("usage: {0!s} start|restart|debug".format(sys.argv[0]))
        sys.exit(2)
    print("And it's goodnight from him")
    sys.exit(0)
<|endoftext|>


<commit_before>#
# Copyright (c) 2004 Specifix, Inc.
# All rights reserved
#
import package
import files
import shutil
import string

# version is a short package version, not an SRS version string
def finalCommit(cfg, pkgName, version, root, fileList):
    pkgSet = package.PackageSet(cfg.reppath, pkgName)

    if pkgSet.hasVersion(version):
	raise KeyError, ("package %s version %s is already installed" %
		    (pkgName, version))
    p = pkgSet.createVersion(version)

    fileDB = cfg.reppath + "/files"

    for file in fileList:
	infoFile = files.FileDB(cfg.reppath, file.pathInRep(cfg.reppath))

	existing = infoFile.findVersion(file)
	if not existing:
	    file.version(version)
	    infoFile.addVersion(version, file)

	    if file.__class__ == files.SourceFile:
		p.addSource("/" + file.fileName(),
			    file.version())
	    else:
		p.addFile(file.path(), file.version())

	    infoFile.write()
	else:
	    if file.__class__ == files.SourceFile:
		p.addSource(file.path(), existing[0])
	    else:
		p.addFile(file.path(), existing[0])

	file.archive(cfg.reppath, root)

    pkgSet.write()
<commit_msg>use new packageset interface handle version strings properly<commit_after>#
# Copyright (c) 2004 Specifix, Inc.
# All rights reserved
#
import package
import files
import shutil
import string
import copy

# version is a short package version, not an SRS version string
def finalCommit(cfg, pkgName, simpleVersion, root, fileList):
    pkgSet = package.PackageSet(cfg.reppath, pkgName)

    version = pkgSet.getLatestVersion(cfg.defaultbranch)
    if version and simpleVersion == version.trailingVersion():
	# yes, increment the revision
	version = copy.deepcopy(version)
	version.incrementVersionRelease()
    else:
	# no, make a new version
	version = copy.deepcopy(cfg.defaultbranch)
	version.appendVersionRelease(simpleVersion, 1)

    p = package.Package(version)

    fileDB = cfg.reppath + "/files"

    for file in fileList:
	infoFile = files.FileDB(cfg.reppath, file.pathInRep(cfg.reppath))

	duplicateVersion = infoFile.checkBranchForDuplicate(cfg.defaultbranch,
					file)
	if not duplicateVersion:
	    file.version(version)
	    infoFile.addVersion(version, file)

	    if file.__class__ == files.SourceFile:
		p.addSource("/" + file.fileName(),
			    file.version())
	    else:
		p.addFile(file.path(), file.version())

	    infoFile.close()
	else:
	    if file.__class__ == files.SourceFile:
		p.addSource(file.path(), duplicateVersion)
	    else:
		p.addFile(file.path(),  duplicateVersion)

	file.archive(cfg.reppath, root)

    pkgSet.addVersion(version, p)
    pkgSet.close()
<|endoftext|>


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

# author:       me@itzo.org
# version:      2.1
# description:  Get current data from bittrex for market exchange rates and amounts in order book

import urllib
import json
import datetime
import sys
import getopt
import sqlite3 as db

# print usage
def usage():
    print "usage: ./bittrex.py [-iph] -m market"
    print "   eg: ./bittrex.py -m BTC-ETH"
    print
    print "     -i [--init]     initializes a new sqlite database 'market.db'"
    print "     -p [--print]    prints out history for given market"
    print "     -m [--market]   specifies the market to use"
    print "     -h [--help]     prints this menu"

# initialize the market.db if requested
def create_db():
    try:
        con = db.connect('market.db')
        cur = con.cursor()
        cur.execute('DROP TABLE IF EXISTS history')
        cur.executescript("""
            CREATE TABLE history(  market TEXT,
                date TIMESTAMP,
                buyq INT,
                sellq INT,
                bid REAL,
                ask REAL,
                buy_orders INT,
                sell_orders INT,
                buy_min REAL,
                sell_max REAL);""")
        con.commit()
    except db.Error, e:
        if con:
            con.rollback()
        print "Error %s:" % e.args[0]
        sys.exit(1)
    finally:
        if con:
            con.close()

# update the db with the latest info
def db_insert(market,buyq,sellq,bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max):
    timestamp = int(datetime.datetime.now().strftime("%s"))
    con = db.connect('market.db')
    cur = con.cursor()
    cur.execute('INSERT INTO history VALUES(?,?,?,?,?,?,?,?,?,?)', \
        (market,timestamp,int(buyq),int(sellq),bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max));
    con.commit()

# get the market data
def get_data(market):
    # load JSON data from Bittrex API
    url = 'https://bittrex.com/api/v1.1/public/getorderbook?market='+market+'&type=both&depth=50'
    json_obj = urllib.urlopen(url)
    data = json.load(json_obj)
    # get the sum of all orders that we care about
    buyq = 0
    sellq = 0
    # only top 'percent' % of each orderbook considered relevant data
    percent = 25
    bid = data['result']['buy'][0]['Rate']
    total_buy_orders = len(data['result']['buy'])
    buy_index = percent * total_buy_orders/100
    buy_min = data['result']['buy'][buy_index]['Rate']
    ask = data['result']['sell'][0]['Rate']
    total_sell_orders = len(data['result']['sell'])
    sell_index = percent * total_sell_orders/100
    sell_max = data['result']['sell'][sell_index]['Rate']
    for item in data['result']['buy']:
        if item['Rate'] > buy_min:
            buyq += item['Quantity']
    for item in data['result']['sell']:
        if item['Rate'] < sell_max:
            sellq += item['Quantity']
    db_insert(market,buyq,sellq,bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max)

    print "----------------------------------------"
    print "buy orders:\t\t"+str(total_buy_orders)
    print "buy_index:\t\t"+str(buy_index)
    print "MIN buy allowed:\t"+str(buy_min)
    print "----------------------------------------"
    print "sell orders:\t\t"+str(total_sell_orders)
    print "sell_index:\t\t"+str(sell_index)
    print "MAX sell allowed:\t"+str(sell_max)
    print "----------------------------------------"

# print market history
def print_history(market):
    con = db.connect('market.db')
    with con:
        cur = con.cursor()
        cur.execute('SELECT * from history')
        con.commit()
        rows = cur.fetchall()
        for row in rows:
            print "%-8s\t%s\t%+8s\t%+8s\t%0.8f\t%0.8f" % (row[0], row[1], row[2], row[3], row[4], row[5])
            #market,timestamp,int(buyq),int(sellq),bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max

# get cmd line options a.k.a. main... ;\
required_m = False
try:
    opts, args = getopt.getopt(sys.argv[1:], 'ipm:h', ['init', 'print', 'market=', 'help'])
except getopt.GetoptError:
    usage()
    sys.exit(2)

for opt, arg in opts:
    if opt in ('-h', '--help'):
        usage()
        sys.exit(2)
    elif opt in ('-m', '--market'):
        market = arg
        get_data(market)
        required_m = True
    elif opt in ('-i', '--init'):
        create_db()
    elif opt in ('-p', '--print'):
        print_history(market)
    else:
        usage()
        sys.exit(2)
if required_m == False:
    usage()
    sys.exit(2)
<commit_msg>fixing some formatting on print<commit_after>#!/usr/bin/env python

# author:       me@itzo.org
# version:      2.1
# description:  Get current data from bittrex for market exchange rates and amounts in order book

import urllib
import json
import datetime
import sys
import getopt
import sqlite3 as db

# print usage
def usage():
    print "usage: ./bittrex.py [-iph] -m market"
    print "   eg: ./bittrex.py -m BTC-ETH"
    print
    print "     -i [--init]     initializes a new sqlite database 'market.db'"
    print "     -p [--print]    prints out history for given market"
    print "     -m [--market]   specifies the market to use"
    print "     -h [--help]     prints this menu"

# initialize the market.db if requested
def create_db():
    try:
        con = db.connect('market.db')
        cur = con.cursor()
        cur.execute('DROP TABLE IF EXISTS history')
        cur.executescript("""
            CREATE TABLE history(  market TEXT,
                date TIMESTAMP,
                buyq INT,
                sellq INT,
                bid REAL,
                ask REAL,
                buy_orders INT,
                sell_orders INT,
                buy_min REAL,
                sell_max REAL);""")
        con.commit()
    except db.Error, e:
        if con:
            con.rollback()
        print "Error %s:" % e.args[0]
        sys.exit(1)
    finally:
        if con:
            con.close()

# update the db with the latest info
def db_insert(market,buyq,sellq,bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max):
    timestamp = int(datetime.datetime.now().strftime("%s"))
    con = db.connect('market.db')
    cur = con.cursor()
    cur.execute('INSERT INTO history VALUES(?,?,?,?,?,?,?,?,?,?)', \
        (market,timestamp,int(buyq),int(sellq),bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max));
    con.commit()

# get the market data
def get_data(market):
    # load JSON data from Bittrex API
    url = 'https://bittrex.com/api/v1.1/public/getorderbook?market='+market+'&type=both&depth=50'
    json_obj = urllib.urlopen(url)
    data = json.load(json_obj)
    # get the sum of all orders that we care about
    buyq = 0
    sellq = 0
    # only top 'percent' % of each orderbook considered relevant data
    percent = 25
    bid = data['result']['buy'][0]['Rate']
    total_buy_orders = len(data['result']['buy'])
    buy_index = percent * total_buy_orders/100
    buy_min = data['result']['buy'][buy_index]['Rate']
    ask = data['result']['sell'][0]['Rate']
    total_sell_orders = len(data['result']['sell'])
    sell_index = percent * total_sell_orders/100
    sell_max = data['result']['sell'][sell_index]['Rate']
    for item in data['result']['buy']:
        if item['Rate'] > buy_min:
            buyq += item['Quantity']
    for item in data['result']['sell']:
        if item['Rate'] < sell_max:
            sellq += item['Quantity']
    db_insert(market,buyq,sellq,bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max)

    print "----------------------------------------"
    print "buy orders:\t\t"+str(total_buy_orders)
    print "buy_index:\t\t"+str(buy_index)
    print "MIN buy allowed:\t"+str(buy_min)
    print "----------------------------------------"
    print "sell orders:\t\t"+str(total_sell_orders)
    print "sell_index:\t\t"+str(sell_index)
    print "MAX sell allowed:\t"+str(sell_max)
    print "----------------------------------------"

# print market history
def print_history(market):
    con = db.connect('market.db')
    with con:
        cur = con.cursor()
        cur.execute('SELECT * from history')
        con.commit()
        rows = cur.fetchall()
        print "%-8s %-10s %+8s %+8s %-10s  %-10s %+5s %+5s  %+10s %+10s" % \
        ('Market', 'Time', 'Buyq', 'Sellq', 'Bid', 'Ask', 'Buy#', 'Sell#', 'Buy min', 'Sell max')
        for row in rows:
            if row[0] == market:
                print "%-8s %s %+8s %+8s %0.8f  %0.8f %+5s %+5s  %0.8f %0.8f" % \
                (row[0], row[1], row[2], row[3], row[4], row[5], row[6], row[7], row[8], row[9])
                #market,timestamp,int(buyq),int(sellq),bid,ask,total_buy_orders,total_sell_orders,buy_min,sell_max

# get cmd line options a.k.a. main... ;\
required_m = False
try:
    opts, args = getopt.getopt(sys.argv[1:], 'ipm:h', ['init', 'print', 'market=', 'help'])
except getopt.GetoptError:
    usage()
    sys.exit(2)

for opt, arg in opts:
    if opt in ('-h', '--help'):
        usage()
        sys.exit(2)
    elif opt in ('-m', '--market'):
        market = arg
        get_data(market)
        required_m = True
    elif opt in ('-i', '--init'):
        create_db()
    elif opt in ('-p', '--print'):
        print_history(market)
    else:
        usage()
        sys.exit(2)
if required_m == False:
    usage()
    sys.exit(2)
<|endoftext|>


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

"""
Created on 27 Feb 2017

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)

DESCRIPTION
The afe_calib utility is used to retrieve or install the calibration sheet for the Alphasense analogue front-end (AFE)
board or single sensor installed on the host system.

Alphasense electrochemical sensors are calibrated in the factory when fitted to their AFE board, or individually.
The calibration values are provided in a structured document, either on paper or - for sensors provided by South Coast
Science - in electronic form. The afe_calib utility is used to retrieve this JSON document via a web API.

The afe_calib utility may also be used to set a "test" calibration sheet, for use in a manufacturing environment.

Note that the scs_dev/gasses_sampler process must be restarted for changes to take effect.

SYNOPSIS
afe_calib.py [{ -a SERIAL_NUMBER | -s SERIAL_NUMBER YYYY-MM-DD | -t }] [-v]

EXAMPLES
./afe_calib.py -s 212810465 2019-08-22

DOCUMENT EXAMPLE
{"serial_number": null, "type": "ISI", "calibrated_on": "2019-09-09", "dispatched_on": null, "pt1000_v20": null,
"sn1": {"serial_number": "212810464", "sensor_type": "NOGA4", "we_electronic_zero_mv": 300, "we_sensor_zero_mv": 0,
"we_total_zero_mv": 300, "ae_electronic_zero_mv": 300, "ae_sensor_zero_mv": 0, "ae_total_zero_mv": 300,
"we_sensitivity_na_ppb": -0.445, "we_cross_sensitivity_no2_na_ppb": "n/a", "pcb_gain": -0.7,
"we_sensitivity_mv_ppb": 0.325, "we_cross_sensitivity_no2_mv_ppb": 0.324}}

FILES
~/SCS/conf/afe_calib.json

SEE ALSO
scs_dev/gases_sampler
scs_mfr/afe_baseline

RESOURCES
https://www.alphasense-technology.co.uk/
"""

import json
import sys

from scs_core.data.json import JSONify

from scs_core.gas.afe_calib import AFECalib
from scs_core.gas.sensor_calib import SensorCalib

from scs_core.sys.http_exception import HTTPException

from scs_host.client.http_client import HTTPClient
from scs_host.sys.host import Host

from scs_mfr.cmd.cmd_afe_calib import CmdAFECalib

# TODO: change name to gas_calib

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

if __name__ == '__main__':

    calib = None
    calibrated_on = None
    we_sens_mv = None

    # ----------------------------------------------------------------------------------------------------------------
    # cmd...

    cmd = CmdAFECalib()

    if cmd.verbose:
        print("afe_calib: %s" % cmd, file=sys.stderr)
        sys.stderr.flush()

    # ----------------------------------------------------------------------------------------------------------------
    # resources...

    calib = AFECalib.load(Host)

    # ----------------------------------------------------------------------------------------------------------------
    # run...

    if cmd.set():
        if cmd.test:
            jdict = json.loads(AFECalib.TEST_LOAD)
            calib = AFECalib.construct_from_jdict(jdict)

        elif cmd.afe_serial_number:
            client = HTTPClient()
            client.connect(AFECalib.ALPHASENSE_HOST)

            try:
                path = AFECalib.ALPHASENSE_PATH + cmd.afe_serial_number
                jdict = json.loads(client.get(path, None, AFECalib.ALPHASENSE_HEADER))
                calib = AFECalib.construct_from_jdict(jdict)

            except HTTPException as ex:
                print("afe_calib: %s" % ex, file=sys.stderr)
                exit(1)

            finally:
                client.close()

        else:
            client = HTTPClient()
            client.connect(SensorCalib.ALPHASENSE_HOST)

            try:
                path = SensorCalib.ALPHASENSE_PATH + cmd.sensor_serial_number
                jstr = client.get(path, None, SensorCalib.ALPHASENSE_HEADER)

                sensor_calib = SensorCalib.construct_from_jdict(json.loads(jstr))
                sensor_calib.set_defaults()
                sensor_calib.set_sens_mv_from_sens_na()

                calib = AFECalib.construct_for_sensor(cmd.sensor_calibration_date, sensor_calib)

            except HTTPException as ex:
                print("afe_calib: %s" % ex, file=sys.stderr)
                exit(1)

            finally:
                client.close()

        if calib is not None:
            calib.save(Host)

    elif cmd.delete:
        calib.delete(Host)
        calib = None

    if calib:
        print(JSONify.dumps(calib))
<commit_msg>Added AFECalib download method<commit_after>#!/usr/bin/env python3

"""
Created on 27 Feb 2017

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)

DESCRIPTION
The afe_calib utility is used to retrieve or install the calibration sheet for the Alphasense analogue front-end (AFE)
board or single sensor installed on the host system.

Alphasense electrochemical sensors are calibrated in the factory when fitted to their AFE board, or individually.
The calibration values are provided in a structured document, either on paper or - for sensors provided by South Coast
Science - in electronic form. The afe_calib utility is used to retrieve this JSON document via a web API.

The afe_calib utility may also be used to set a "test" calibration sheet, for use in a manufacturing environment.

Note that the scs_dev/gasses_sampler process must be restarted for changes to take effect.

SYNOPSIS
afe_calib.py [{ -a SERIAL_NUMBER | -s SERIAL_NUMBER YYYY-MM-DD | -t }] [-v]

EXAMPLES
./afe_calib.py -s 212810465 2019-08-22

DOCUMENT EXAMPLE
{"serial_number": null, "type": "ISI", "calibrated_on": "2019-09-09", "dispatched_on": null, "pt1000_v20": null,
"sn1": {"serial_number": "212810464", "sensor_type": "NOGA4", "we_electronic_zero_mv": 300, "we_sensor_zero_mv": 0,
"we_total_zero_mv": 300, "ae_electronic_zero_mv": 300, "ae_sensor_zero_mv": 0, "ae_total_zero_mv": 300,
"we_sensitivity_na_ppb": -0.445, "we_cross_sensitivity_no2_na_ppb": "n/a", "pcb_gain": -0.7,
"we_sensitivity_mv_ppb": 0.325, "we_cross_sensitivity_no2_mv_ppb": 0.324}}

FILES
~/SCS/conf/afe_calib.json

SEE ALSO
scs_dev/gases_sampler
scs_mfr/afe_baseline

RESOURCES
https://www.alphasense-technology.co.uk/
"""

import json
import sys

from scs_core.data.json import JSONify

from scs_core.gas.afe_calib import AFECalib
from scs_core.gas.sensor_calib import SensorCalib

from scs_core.sys.http_exception import HTTPException

from scs_host.client.http_client import HTTPClient
from scs_host.sys.host import Host

from scs_mfr.cmd.cmd_afe_calib import CmdAFECalib

# TODO: change name to gas_calib

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

if __name__ == '__main__':

    calib = None
    calibrated_on = None
    we_sens_mv = None

    # ----------------------------------------------------------------------------------------------------------------
    # cmd...

    cmd = CmdAFECalib()

    if cmd.verbose:
        print("afe_calib: %s" % cmd, file=sys.stderr)
        sys.stderr.flush()

    # ----------------------------------------------------------------------------------------------------------------
    # resources...

    calib = AFECalib.load(Host)

    # ----------------------------------------------------------------------------------------------------------------
    # run...

    if cmd.set():
        if cmd.test:
            jdict = json.loads(AFECalib.TEST_LOAD)
            calib = AFECalib.construct_from_jdict(jdict)

        elif cmd.afe_serial_number:
            try:
                calib = AFECalib.download(HTTPClient(), cmd.afe_serial_number)

            except HTTPException as ex:
                print("afe_calib: %s" % ex, file=sys.stderr)
                exit(1)

        else:
            client = HTTPClient()
            client.connect(SensorCalib.ALPHASENSE_HOST)

            try:
                path = SensorCalib.ALPHASENSE_PATH + cmd.sensor_serial_number
                jstr = client.get(path, None, SensorCalib.ALPHASENSE_HEADER)

                sensor_calib = SensorCalib.construct_from_jdict(json.loads(jstr))
                sensor_calib.set_defaults()
                sensor_calib.set_sens_mv_from_sens_na()

                calib = AFECalib.construct_for_sensor(cmd.sensor_calibration_date, sensor_calib)

            except HTTPException as ex:
                print("afe_calib: %s" % ex, file=sys.stderr)
                exit(1)

            finally:
                client.close()

        if calib is not None:
            calib.save(Host)

    elif cmd.delete:
        calib.delete(Host)
        calib = None

    if calib:
        print(JSONify.dumps(calib))
<|endoftext|>


<commit_before>#!/usr/bin/python

"""
Copyright 2016 EMBL-European Bioinformatics Institute

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import argparse, urllib2, gzip, shutil, shlex, subprocess, os.path

from socket import error as SocketError
import errno

try :
    import pysam
except ImportError :
    print "[Error] Cannot import \"pysam\" package. Have you installed it?"
    exit(-1)

class common:
    """
    Functions for downloading and processing Mnase-seq FastQ files. Files are
    downloaded from the European Nucleotide Archive (ENA), then aligned,
    filtered and analysed for peak calling
    """

    def __init__ (self):
        """
        Initialise the module
        """
        self.ready = ""

    def getGenomeFile(self, data_dir, species, assembly, index = True):
        """
        Function for downloading and extracting the DNA files from the ensembl FTP
        """

        file_name = data_dir + species + '_' + assembly + '/' + species + '.' + assembly + '.dna.toplevel.fa.gz'
        file_name_unzipped = file_name.replace('.fa.gz', '.fa')
        print file_name

        if os.path.isfile(file_name) == False:
            ftp_url = 'ftp://ftp.ensembl.org/pub/current_fasta/' + species.lower() + '/dna/' + species[0].upper() + species[1:] + '.' + assembly + '.dna.toplevel.fa.gz'

            self.download_file(file_name, ftp_url)

        if os.path.isfile(file_name_unzipped) == False:
            print "Unzipping"
            with gzip.open(file_name, 'rb') as f_in, open(file_name_unzipped, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)

        indexes = {}
        if index == True:
            indexes = self.run_indexers(file_name_unzipped)

        return {'zipped': file_name, 'unzipped': file_name_unzipped, 'index' : indexes}

    def getGenomeFromENA(self, data_dir, species, assembly, index = True):
        """
        Obtain the Assembly from the ENA. This means that it is possible to get
        the exact assembly matching that defined by the user.
        """

        file_name = data_dir + species + '_' + assembly + '/' + species + '.' + assembly + '.fa'

        if os.path.isfile(file_name) == False:
          ftp_list_url = 'ftp://ftp.ebi.ac.uk/pub/databases/ena/assembly/' + assembly[0:7] + '/' + assembly[0:10] + '/' + assembly + '_sequence_report.txt'
          res_list = urllib2.urlopen(ftp_list_url)
          table = res_list.read()
          table = table.split("\n")

          chr_list = []
          for row in table:
              col = row.split("\t")
              if len(col) > 5 and col[3] == 'assembled-molecule':
                  chr_list.append(col[0])

          ftp_url = 'http://www.ebi.ac.uk/ena/data/view/' + ','.join(chr_list) + '&display=fasta'
          self.download_file(file_name, ftp_url)
          self.replaceENAHeader(file_name)

        indexes = {}
        if index == True:
            indexes = self.run_indexers(file_name)

        return {'unzipped': file_name, 'index' : indexes}

    def replaceENAHeader(file_path):
        """
        The ENA header has pipes in the header as part of teh stable_id. This
        function removes the ENA stable_id and replaces it with the final
        section after splitting the stable ID on the pipe.
        """

        #Create temp file
        fh, abs_path = mkstemp()
        with open(abs_path,'w') as new_file:
            with open(file_path) as old_file:
                for line in old_file:
                    if line[0] == '>':
                        space_line = line.split(" ")
                        new_file.write(">" + space_line[0].split("|")[-1].replace(">", "") + " " + " ".join(space_line[1:]))
                    else:
                        new_file.write(line)
        close(fh)
        #Remove original file
        remove(file_path)
        #Move new file
        move(abs_path, file_path)

    def getcDNAFiles(self, data_dir, species, assembly, e_release):
        """
        Function for downloading and extracting the CDNA files from the ensembl FTP
        """

        file_name = data_dir + species + '_' + assembly + '/' + species + '.' + assembly + '.cdna.all.fa.gz'

        if os.path.isfile(file_name) == False:
            ftp_url = 'ftp://ftp.ensembl.org/pub/release-' + e-release + '/' + species.lower() + '/cdna/' + species[0].upper() + species[1:] + '.' + assembly + '.cdna.all.fa.gz'

            self.download_file(file_name, ftp_url)

        return file_name

    def getFastqFiles(self, ena_err_id, data_dir, ena_srr_id = None):
        """
        Function for downloading and extracting the FastQ files from the ENA
        """

        f_index = urllib2.urlopen(
        'http://www.ebi.ac.uk/ena/data/warehouse/filereport?accession=' + str(ena_err_id) + '&result=read_run&fields=study_accession,run_accession,tax_id,scientific_name,instrument_model,library_layout,fastq_ftp&download=txt')
        data = f_index.read()
        rows = data.split("\n")
        row_count = 0
        files = []
        gzfiles  = []
        for row in rows:
            if row_count == 0:
                row_count += 1
                continue

            row = row.rstrip()
            row = row.split("\t")

            if len(row) < 6:
                continue

            project = row[0]
            srr_id = row[1]
            if (ena_srr_id != None and srr_id != ena_srr_id):
                continue
            fastq_files = row[6].split(';')
            row_count += 1

            for fastq_file in fastq_files:
                file_name = fastq_file.split("/")

                print data_dir + project + "/" + file_name[-1]

                if os.path.isfile(data_dir + '/' + project + "/" + file_name[-1]) == False and os.path.isfile(data_dir + '/' + project + "/" + file_name[-1].replace('.fastq.gz', '.fastq')) == False:
                    file_location = data_dir + '/' + project + "/" + file_name[-1]
                    ftp_url = "ftp://" + fastq_file
                    self.download_file(file_location, ftp_url)

                if os.path.isfile(data_dir + '/' + project + "/" + file_name[-1]) == True:
                    gzfiles.append(data_dir + '/' + project + "/" + file_name[-1])
                files.append(data_dir + '/' + project + "/" + file_name[-1].replace('.fastq.gz', '.fastq'))

        for gzf in gzfiles:
            with gzip.open(gzf, 'rb') as f_in, open(gzf.replace('.fastq.gz', '.fastq'), 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
            os.remove(gzf)

        return files

    def download_file(self, file_location, url):
        """
        Function to download a file to a given location and file name. Will
        attempt to restart the download up to 5 times if the connection is
        closed by the remote server.
        """

        restart_counter = 0

        while True:
            meta = urllib2.urlopen(url)
            info = meta.info()
            if restart_counter > 0:
                byte_start = int(os.stat(file_location).st_size) + 1
                request = urllib2.Request(url, headers={"Range" : "bytes=" + str(byte_start) + "-" + str(info["content-length"])})
            else:
                request = urllib2.Request(url)

            req = urllib2.urlopen(request)
            CHUNK = 16 * 1024

            try:
                if restart_counter == 0 :
                    with open(file_location, 'wb') as fp:
                        while True:
                            chunk = req.read(CHUNK)
                            if not chunk: break
                            fp.write(chunk)
                else:
                    with open(file_location, 'ab') as fp:
                        while True:
                            chunk = req.read(CHUNK)
                            if not chunk: break
                            fp.write(chunk)

                req.close()
            except SocketError as e:
                if e.errno != errno.ECONNRESET:
                    raise
                print e
                print "Attempting to restart download ..."
                restart_counter += 1

            if restart_counter >= 5 :
                return False
            break

        return True

    def run_indexers(self, file_name):
        """

        """

        file_name_unzipped = file_name.replace('.fa.gz', '.fa')
        file_name_nofa = file_name_unzipped.replace('.fa', '')

        if os.path.isfile(file_name_unzipped) == False:
            print "Unzipping"
            with gzip.open(file_name, 'rb') as f_in, open(file_name_unzipped, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)

        if os.path.isfile(file_name_unzipped + '.bwt') == False:
            print "Indexing - BWA"
            self.bwa_index_genome(file_name)

        if os.path.isfile(file_name_nofa + '.1.bt2') == False:
            print "Indexing - Bowtie"
            self.bowtie_index_genome(file_name_unzipped)

        if os.path.isfile(file_name_nofa + '.gem') == False:
            print "Indexing - GEM"
            self.bowtie_index_genome(file_name_unzipped)

        return {'bowtie' : file_name_unzipped + '.1.bt2', 'bwa' : file_name_unzipped + '.bwt', 'gem' : file_name_unzipped + '.gem'}

    def gem_index_genome(self, genome_file):
        """
        Create an index of the genome FASTA file with GEM
        """
        file_name = genome_file.split("/")
        file_name[-1].replace('.fa', '')

        command_line = 'gem-indexer -i ' + genome_file + ' -o ' + file_name[-1].replace('.fa', '')

        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        p.wait()

    def bowtie_index_genome(self, genome_file):
        """
        Create an index of the genome FASTA file with Bowtie2
        """
        file_name = genome_file.split("/")
        file_name[-1].replace('.fa', '')

        command_line = 'bowtie2-build ' + genome_file + ' ' + file_name[-1].replace('.fa', '')

        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        p.wait()

    def bwa_index_genome(self, genome_file):
        """
        Create an index of the genome FASTA file with BWA
        """
        command_line = 'bwa index ' + genome_file

        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        p.wait()

    def bwa_align_reads(self, genome_file, data_dir, project_id, run_id):
        """
        Map the reads to the genome
        """

        reads_file = data_dir + project_id + '/' + run_id + '.fastq'
        intermediate_file = data_dir + project_id + '/' + run_id + '.sai'
        intermediate_sam_file = data_dir + project_id + '/' + run_id + '.sam'
        output_bam_file = data_dir + project_id + '/' + run_id + '.bam'

        command_lines = [
            'bwa aln -q 5 -f ' + intermediate_file + ' ' + genome_file + ' ' + reads_file,
            'bwa samse -f ' + intermediate_sam_file  + ' ' + genome_file + ' ' + intermediate_file + ' ' + reads_file,
            'samtools view -b -o ' + output_bam_file + ' ' + intermediate_sam_file
        ]

        print command_lines

        for command_line in command_lines:
            args = shlex.split(command_line)
            p = subprocess.Popen(args)
            p.wait()

    def merge_bam(self, data_dir, project_id, final_id, run_ids=[]):
        """
        Merge together all the bams in a directory and sort to create the final
        bam ready to be filtered

        If run_ids is blank then the function looks for all bam files in the
        data_dir
        """
        out_bam_file = data_dir + project_id + '/' + final_id + '.bam'

        if len(run_ids) == 0:
            bam_files = [f for f in listdir(data_dir + project_id) if f.endswith(("sai"))]
        else:
            bam_files = [f + ".bam" for f in run_ids]

        bam_sort_files = []
        bam_merge_files = []
        for bam in bam_files:
            bam_loc = data_dir + project_id + '/' + bam
            bam_sort_files.append(bam_loc)
            bam_merge_files.append(bam_loc)

        for bam_sort_file in bam_sort_files:
            print bam_sort_file
            pysam.sort("-o", str(bam_sort_file), str(bam_sort_file))

        if len(bam_sort_files) == 1:
            pysam.sort("-o", str(out_bam_file), str(bam_sort_files[0]))
        else:
            pysam.merge(out_bam_file, *bam_merge_files)
            pysam.sort("-o", str(out_bam_file), "-T", str(out_bam_file) + ".bam_sort", str(out_bam_file))

        pysam.index(str(out_bam_file))
<commit_msg>Moved the replace to run before the indexer runs<commit_after>#!/usr/bin/python

"""
Copyright 2016 EMBL-European Bioinformatics Institute

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import argparse, urllib2, gzip, shutil, shlex, subprocess, os.path

from socket import error as SocketError
import errno

try :
    import pysam
except ImportError :
    print "[Error] Cannot import \"pysam\" package. Have you installed it?"
    exit(-1)

class common:
    """
    Functions for downloading and processing Mnase-seq FastQ files. Files are
    downloaded from the European Nucleotide Archive (ENA), then aligned,
    filtered and analysed for peak calling
    """

    def __init__ (self):
        """
        Initialise the module
        """
        self.ready = ""

    def getGenomeFile(self, data_dir, species, assembly, index = True):
        """
        Function for downloading and extracting the DNA files from the ensembl FTP
        """

        file_name = data_dir + species + '_' + assembly + '/' + species + '.' + assembly + '.dna.toplevel.fa.gz'
        file_name_unzipped = file_name.replace('.fa.gz', '.fa')
        print file_name

        if os.path.isfile(file_name) == False:
            ftp_url = 'ftp://ftp.ensembl.org/pub/current_fasta/' + species.lower() + '/dna/' + species[0].upper() + species[1:] + '.' + assembly + '.dna.toplevel.fa.gz'

            self.download_file(file_name, ftp_url)

        if os.path.isfile(file_name_unzipped) == False:
            print "Unzipping"
            with gzip.open(file_name, 'rb') as f_in, open(file_name_unzipped, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)

        indexes = {}
        if index == True:
            indexes = self.run_indexers(file_name_unzipped)

        return {'zipped': file_name, 'unzipped': file_name_unzipped, 'index' : indexes}

    def getGenomeFromENA(self, data_dir, species, assembly, index = True):
        """
        Obtain the Assembly from the ENA. This means that it is possible to get
        the exact assembly matching that defined by the user.
        """

        file_name = data_dir + species + '_' + assembly + '/' + species + '.' + assembly + '.fa'

        if os.path.isfile(file_name) == False:
          ftp_list_url = 'ftp://ftp.ebi.ac.uk/pub/databases/ena/assembly/' + assembly[0:7] + '/' + assembly[0:10] + '/' + assembly + '_sequence_report.txt'
          res_list = urllib2.urlopen(ftp_list_url)
          table = res_list.read()
          table = table.split("\n")

          chr_list = []
          for row in table:
              col = row.split("\t")
              if len(col) > 5 and col[3] == 'assembled-molecule':
                  chr_list.append(col[0])

          ftp_url = 'http://www.ebi.ac.uk/ena/data/view/' + ','.join(chr_list) + '&display=fasta'
          self.download_file(file_name, ftp_url)

        self.replaceENAHeader(file_name)

        indexes = {}
        if index == True:
            self.replaceENAHeader(file_name)
            indexes = self.run_indexers(file_name)

        return {'unzipped': file_name, 'index' : indexes}

    def replaceENAHeader(file_path):
        """
        The ENA header has pipes in the header as part of teh stable_id. This
        function removes the ENA stable_id and replaces it with the final
        section after splitting the stable ID on the pipe.
        """

        #Create temp file
        fh, abs_path = mkstemp()
        with open(abs_path,'w') as new_file:
            with open(file_path) as old_file:
                for line in old_file:
                    if line[0] == '>':
                        space_line = line.split(" ")
                        new_file.write(">" + space_line[0].split("|")[-1].replace(">", "") + " " + " ".join(space_line[1:]))
                    else:
                        new_file.write(line)
        close(fh)
        #Remove original file
        remove(file_path)
        #Move new file
        move(abs_path, file_path)

    def getcDNAFiles(self, data_dir, species, assembly, e_release):
        """
        Function for downloading and extracting the CDNA files from the ensembl FTP
        """

        file_name = data_dir + species + '_' + assembly + '/' + species + '.' + assembly + '.cdna.all.fa.gz'

        if os.path.isfile(file_name) == False:
            ftp_url = 'ftp://ftp.ensembl.org/pub/release-' + e-release + '/' + species.lower() + '/cdna/' + species[0].upper() + species[1:] + '.' + assembly + '.cdna.all.fa.gz'

            self.download_file(file_name, ftp_url)

        return file_name

    def getFastqFiles(self, ena_err_id, data_dir, ena_srr_id = None):
        """
        Function for downloading and extracting the FastQ files from the ENA
        """

        f_index = urllib2.urlopen(
        'http://www.ebi.ac.uk/ena/data/warehouse/filereport?accession=' + str(ena_err_id) + '&result=read_run&fields=study_accession,run_accession,tax_id,scientific_name,instrument_model,library_layout,fastq_ftp&download=txt')
        data = f_index.read()
        rows = data.split("\n")
        row_count = 0
        files = []
        gzfiles  = []
        for row in rows:
            if row_count == 0:
                row_count += 1
                continue

            row = row.rstrip()
            row = row.split("\t")

            if len(row) < 6:
                continue

            project = row[0]
            srr_id = row[1]
            if (ena_srr_id != None and srr_id != ena_srr_id):
                continue
            fastq_files = row[6].split(';')
            row_count += 1

            for fastq_file in fastq_files:
                file_name = fastq_file.split("/")

                print data_dir + project + "/" + file_name[-1]

                if os.path.isfile(data_dir + '/' + project + "/" + file_name[-1]) == False and os.path.isfile(data_dir + '/' + project + "/" + file_name[-1].replace('.fastq.gz', '.fastq')) == False:
                    file_location = data_dir + '/' + project + "/" + file_name[-1]
                    ftp_url = "ftp://" + fastq_file
                    self.download_file(file_location, ftp_url)

                if os.path.isfile(data_dir + '/' + project + "/" + file_name[-1]) == True:
                    gzfiles.append(data_dir + '/' + project + "/" + file_name[-1])
                files.append(data_dir + '/' + project + "/" + file_name[-1].replace('.fastq.gz', '.fastq'))

        for gzf in gzfiles:
            with gzip.open(gzf, 'rb') as f_in, open(gzf.replace('.fastq.gz', '.fastq'), 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
            os.remove(gzf)

        return files

    def download_file(self, file_location, url):
        """
        Function to download a file to a given location and file name. Will
        attempt to restart the download up to 5 times if the connection is
        closed by the remote server.
        """

        restart_counter = 0

        while True:
            meta = urllib2.urlopen(url)
            info = meta.info()
            if restart_counter > 0:
                byte_start = int(os.stat(file_location).st_size) + 1
                request = urllib2.Request(url, headers={"Range" : "bytes=" + str(byte_start) + "-" + str(info["content-length"])})
            else:
                request = urllib2.Request(url)

            req = urllib2.urlopen(request)
            CHUNK = 16 * 1024

            try:
                if restart_counter == 0 :
                    with open(file_location, 'wb') as fp:
                        while True:
                            chunk = req.read(CHUNK)
                            if not chunk: break
                            fp.write(chunk)
                else:
                    with open(file_location, 'ab') as fp:
                        while True:
                            chunk = req.read(CHUNK)
                            if not chunk: break
                            fp.write(chunk)

                req.close()
            except SocketError as e:
                if e.errno != errno.ECONNRESET:
                    raise
                print e
                print "Attempting to restart download ..."
                restart_counter += 1

            if restart_counter >= 5 :
                return False
            break

        return True

    def run_indexers(self, file_name):
        """

        """

        file_name_unzipped = file_name.replace('.fa.gz', '.fa')
        file_name_nofa = file_name_unzipped.replace('.fa', '')

        if os.path.isfile(file_name_unzipped) == False:
            print "Unzipping"
            with gzip.open(file_name, 'rb') as f_in, open(file_name_unzipped, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)

        if os.path.isfile(file_name_unzipped + '.bwt') == False:
            print "Indexing - BWA"
            self.bwa_index_genome(file_name)

        if os.path.isfile(file_name_nofa + '.1.bt2') == False:
            print "Indexing - Bowtie"
            self.bowtie_index_genome(file_name_unzipped)

        if os.path.isfile(file_name_nofa + '.gem') == False:
            print "Indexing - GEM"
            self.bowtie_index_genome(file_name_unzipped)

        return {'bowtie' : file_name_unzipped + '.1.bt2', 'bwa' : file_name_unzipped + '.bwt', 'gem' : file_name_unzipped + '.gem'}

    def gem_index_genome(self, genome_file):
        """
        Create an index of the genome FASTA file with GEM
        """
        file_name = genome_file.split("/")
        file_name[-1].replace('.fa', '')

        command_line = 'gem-indexer -i ' + genome_file + ' -o ' + file_name[-1].replace('.fa', '')

        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        p.wait()

    def bowtie_index_genome(self, genome_file):
        """
        Create an index of the genome FASTA file with Bowtie2
        """
        file_name = genome_file.split("/")
        file_name[-1].replace('.fa', '')

        command_line = 'bowtie2-build ' + genome_file + ' ' + file_name[-1].replace('.fa', '')

        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        p.wait()

    def bwa_index_genome(self, genome_file):
        """
        Create an index of the genome FASTA file with BWA
        """
        command_line = 'bwa index ' + genome_file

        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        p.wait()

    def bwa_align_reads(self, genome_file, data_dir, project_id, run_id):
        """
        Map the reads to the genome
        """

        reads_file = data_dir + project_id + '/' + run_id + '.fastq'
        intermediate_file = data_dir + project_id + '/' + run_id + '.sai'
        intermediate_sam_file = data_dir + project_id + '/' + run_id + '.sam'
        output_bam_file = data_dir + project_id + '/' + run_id + '.bam'

        command_lines = [
            'bwa aln -q 5 -f ' + intermediate_file + ' ' + genome_file + ' ' + reads_file,
            'bwa samse -f ' + intermediate_sam_file  + ' ' + genome_file + ' ' + intermediate_file + ' ' + reads_file,
            'samtools view -b -o ' + output_bam_file + ' ' + intermediate_sam_file
        ]

        print command_lines

        for command_line in command_lines:
            args = shlex.split(command_line)
            p = subprocess.Popen(args)
            p.wait()

    def merge_bam(self, data_dir, project_id, final_id, run_ids=[]):
        """
        Merge together all the bams in a directory and sort to create the final
        bam ready to be filtered

        If run_ids is blank then the function looks for all bam files in the
        data_dir
        """
        out_bam_file = data_dir + project_id + '/' + final_id + '.bam'

        if len(run_ids) == 0:
            bam_files = [f for f in listdir(data_dir + project_id) if f.endswith(("sai"))]
        else:
            bam_files = [f + ".bam" for f in run_ids]

        bam_sort_files = []
        bam_merge_files = []
        for bam in bam_files:
            bam_loc = data_dir + project_id + '/' + bam
            bam_sort_files.append(bam_loc)
            bam_merge_files.append(bam_loc)

        for bam_sort_file in bam_sort_files:
            print bam_sort_file
            pysam.sort("-o", str(bam_sort_file), str(bam_sort_file))

        if len(bam_sort_files) == 1:
            pysam.sort("-o", str(out_bam_file), str(bam_sort_files[0]))
        else:
            pysam.merge(out_bam_file, *bam_merge_files)
            pysam.sort("-o", str(out_bam_file), "-T", str(out_bam_file) + ".bam_sort", str(out_bam_file))

        pysam.index(str(out_bam_file))
<|endoftext|>


<commit_before>"""Common functionality."""

import logging
import shlex
from pathlib import Path

import jsonfile

import media

from misctypes import Flag

import pyutils.files
import pyutils.misc
from pyutils.misc import fmt_size

from synd import Feed

import util

log = logging.getLogger(__name__)
messager = util.Messager(__name__)

WRAPPER = util.MultiWrapper(
    # width=79,
    width=-2,
    max_lines=20,
    break_long_words=False,
    break_on_hyphens=False,

    # initial_indent=' ' * 10,
    # subsequent_indent=' ' * 10,
    # paragraph_indent='\n' + ' ' * 10,

    # initial_indent=' ' * 10,
    # subsequent_indent=' ' * 10,
    # paragraph_indent=' ' * 12,

    initial_indent=' ' * 8,
    subsequent_indent=' ' * 4,
    paragraph_indent=' ' * 8,
)

class TextDict(dict):
    """Dictionary persistent in a text file."""
    def __init__(self, path):
        super().__init__()
        self.path = Path(path)
        self.read()

    # def __del__(self):
    #     self.write()

    # def __setitem__(self, key, value):
    #     super().__setitem__(key, value)
    #     self.write()

    def read(self):
        """Read from file, if it exists."""
        try:
            self.update(read_data(self.path))
        except FileNotFoundError:
            pass

    def write(self):
        """Write to file."""
        write_data(self.path, self)

def write_data(path, data):
    """Write object into JSON file."""
    pyutils.files.ensure_dir(path)
    with pyutils.files.tempfile_and_backup(path, 'w') as fp:
        return jsonfile.write_json(fp, data)

def read_data(path):
    """Read JSON file into dictionary."""
    with path.open() as fp:
        return jsonfile.read_json(fp)

class View(util.AttrDict):
    """A view to a feedlist."""
    DEFAULTS = dict(flags='foin', number=1, sortkey='fpD', sortkey2='SD')

    def __init__(self, directory=None, flags=None, number=None, sortkey=None,
                 sortkey2=None):
        super().__init__()
        self.directory = directory
        if self.directory is not None:
            self.directory = [Path(x) for x in self.directory]
        if flags is None:
            flags = self.DEFAULTS['flags']
        if number is None:
            number = self.DEFAULTS['number']
        if sortkey is None:
            sortkey = self.DEFAULTS['sortkey']
        if sortkey2 is None:
            sortkey2 = self.DEFAULTS['sortkey2']
        self.flags = [Flag(x) for x in flags]
        self.number = number
        self.sortkey = sortkey
        self.sortkey2 = sortkey2

    def __str__(self):
        s = '(f={flags}, n={number}, s={sortkey}, S={sortkey2}, {n})'
        d = dict(self, flags=''.join(x.value for x in self.flags),
                 n=len(self.directory))
        return s.format(**d)

    def parse(self, s, sep=','):
        """Parse view string, update original view, and create a new one."""
        if not s:
            s = ',,,'
        lst = list(s.split(sep))
        if len(lst) < 4:
            lst += [''] * (4 - len(lst))  # Fill in missing parts.
        try:
            flags, number, sortkey, sortkey2 = lst
        except ValueError as e:
            log.error('Cannot parse view: "%s"', s)
            raise
        return self.__class__(
            directory=self.directory,
            flags=flags or self.flags,
            number=int(number or self.number),
            sortkey=sortkey or self.sortkey,
            sortkey2=sortkey2 or self.sortkey2,
            )

# def message(value, *args, **kwargs):
#     """Print message."""
#     # print(value, *args, *kwargs)
#     messager.msg(value, *args, *kwargs)
#
#
# def feedback(value, *args, **kwargs):
#     """Print feedback."""
#     # print(value, *args, *kwargs)
#     messager.msg(value, *args, *kwargs)

def check_feed(feed, verbose=False):
    """Check feed. Return list of orphaned files."""
    for msg in feed.check(verbose=verbose):
        messager.msg('{}: {}'.format(feed, msg))
    return feed.get_orphans()

def add_url(url):
    """Add feed. Return directory name."""
    feed = Feed(url)
    feed.refresh(force=True)
    # check_feed(feed)
    if feed.directory.exists():
        raise IOError('Feed directory exists: {}'.format(feed.directory))
    feed.write()
    check_feed(feed)
    return feed.directory

def show(header, seq, wrapper=WRAPPER):
    """Show bodies of text prettily, using a TextWrapper object."""
    if header:
        messager.msg(*header, sep=' ▶ ')
    for text, title in seq:
        messager.msg(wrapper.fills(text, title+':'))

def show_feed(feed, verbose=0):
    """Show feed."""
    header = [feed.directory, feed.url]
    lst = []
    if verbose:
        lst.extend([
            (util.time_fmt(feed.head.date, fmt='rfc2822'), 'Date'),
            (feed.priority, 'Priority'),
            (feed.head.link, 'Link'),
            (feed.head.title, 'Title'),
            (feed.head.subtitle, 'Subt'),
            (feed.head.summary, 'Summ'),
            # (feed.head.language or 'unknown', 'Language'),
            # (', '.join(feed.get_tags().as_strings()) or [], 'Tags'),
            (str(feed.get_tags()), 'Tags'),
            (feed.head.image or [], 'Image'),
            (len(feed.entries), 'Entries'),
            ])
    if verbose > 1:
        # List flags and time range.
        def n_flagged(flag):
            return len(feed.list_entries(flags=flag.value))
        first = util.time_fmt(feed.entries[0].date, fmt='isodate')
        last = util.time_fmt(feed.entries[-1].date, fmt='isodate')
        lst.extend([
            (', '.join('{} {}'.format(x.name, n_flagged(x)) for x in Flag),
             'Flags'),
            ('earliest {}, latest {}'.format(first, last), 'Range'),
            ])
    show(header, lst)

def show_entry(entry, verbose=0):
    """Show feed."""
    date_str = util.time_fmt(entry.date, fmt='isodate')
    header = [date_str, entry.feed.directory, entry.title]
    lst = []
    if verbose:
        lst.extend([
            (util.time_fmt(entry.date_published, fmt='isofull'), "Publ'd"),
            (util.time_fmt(entry.date_seen, fmt='isofull'), 'Seen'),
            # (entry.score, 'Score'),
            ((entry.flag.name, entry.progress, entry.score), 'FlPrSc'),
            (entry.guid, 'GUID'),
            (entry.link, 'Link'),
            (entry.title, 'Title'),
            (entry.subtitle, 'Subt'),
            (entry.summary, 'Summ'),
            (str(entry.get_tags()), 'Tags'),
            ])
        for enc in entry.encs():
            d = dict(href=enc.href, length=fmt_size(enc.length), typ=enc.typ,
                     name=enc.filename)
            try:
                d['size'] = fmt_size(enc.size())
                d['duration'] = util.fmt_duration(enc.duration())
                d['gain'] = media.get_gain(enc.path)
            except FileNotFoundError:
                d['size'] = d['duration'] = d['gain'] = '-'
            s = '{length}, {typ}, {href}'
            lst.append((s.format(**d), 'URL'))
            s = '{size}, {duration}, {gain}, {name}'
            lst.append((s.format(**d), 'File'))
    show(header, lst)

def show_files(entry, verbose=0):
    """Show files."""
    for enc in entry.encs():
        if verbose or enc.path.exists():
            print(shlex.quote(str(enc.path)))

def show_enclosure(enc):
    if enc.path.exists():
        messager.msg('Media info:')
        messager.pp(media.get_media_info(enc.path))
        messager.msg('Extended attributes:')
        messager.pp(sorted(pyutils.files.XAttrStr(enc.path).items()))
    else:
        log.error('No file: %s', enc.path)

def show_enclosures(entry):
    """Show enclosure file info."""
    for enc in entry.encs():
        show_enclosure(enc)

def download_enclosure(enc, maxsize=None):
    if enc.path.exists():
        log.debug('Already exists: %s', enc.path)
    elif enc.is_too_big(maxsize):
        log.warning('Download too big: %s: %s', fmt_size(enc.length), enc.path)
    else:
        # messager.msg(truncate('Downloading {}: {}'.format(
        #     fmt_size(enc.length), enc.path)))
        messager.msg('Downloading {}: {}'.format(fmt_size(enc.length),
                                                 enc.path), truncate=True)
        try:
            if not enc.download():
                log.error('Download failed: %s', enc.path)
        except NotImplementedError:
            log.debug('Download not implemented: %s', enc)
        except KeyboardInterrupt:
            log.error('Download interrupted: %s', enc)

def download_enclosures(entry, maxsize=None):
    """Download entry enclosures."""
    for enc in entry.encs():
        download_enclosure(enc, maxsize=maxsize)

def normalize_enclosure(enc, force=False):
    if enc.path.exists():
        if force or not enc.is_normalized():
            media.normalize_volume(enc.path)

def normalize_enclosures(entry, force=False):
    """Normalize loudness in enclosures."""
    for enc in entry.encs():
        normalize_enclosure(enc, force=force)

def play_enclosure(enc):
    try:
        if enc.path.exists():
            exit_code = enc.play()
        else:
            messager.msg('File does not exist: {}'.format(enc.path))
            messager.msg('Streaming: {}'.format(enc.href))
            exit_code = enc.stream()
        messager.msg('Exit code:', exit_code)
        return exit_code
    except NotImplementedError:
        log.error('Play/stream not implemented: %s', enc)
        return None

def play_enclosures(entry, set_flag=True):
    """Play entry enclosures."""
    exit_codes = (play_enclosure(x) for x in entry.encs())
    if set_flag and all(x == 0 for x in exit_codes):
        if entry.flag in [Flag.fresh, Flag.important, Flag.new]:
            messager.msg('Flagging entry as opened')
            entry.set_flag(Flag.opened)

def remove_enclosure(enc):
    """Remove enclosure from disk."""
    if enc.path.exists():
        messager.msg('Removing:', enc.path)
        # enc.path.unlink()
        pyutils.files.trash_or_rm(enc.path)

def remove_enclosures(entry, set_flag=True):
    """Remove entry enclosures from disk."""
    for enc in entry.encs():
        remove_enclosure(enc)
    if set_flag:
        if entry.flag == Flag.opened:
            entry.progress = 1
        entry.set_flag('d')
<commit_msg>use boltons; misc<commit_after>"""Common functionality."""

import logging
import shlex
from pathlib import Path

from blessings import Terminal
from boltons.strutils import bytes2human, html2text, is_uuid
from boltons.timeutils import relative_time

import pyutils.files
import pyutils.misc
# from pyutils.misc import fmt_size

import jsonfile
import media
import util
from misctypes import Flag
from synd import Feed
from util import fmt_strings, fmt_duration, fmt_table, time_fmt

log = logging.getLogger(__name__)
messager = util.Messager(__name__)
term = Terminal()

WRAPPER = util.MultiWrapper(
    width=79,
    # width=-2,
    max_lines=20,
    break_long_words=False,
    break_on_hyphens=False,

    # initial_indent=' ' * 10,
    # subsequent_indent=' ' * 10,
    # paragraph_indent='\n' + ' ' * 10,

    # initial_indent=' ' * 10,
    # subsequent_indent=' ' * 10,
    # paragraph_indent=' ' * 12,

    initial_indent=' ' * 8,
    subsequent_indent=' ' * 4,
    paragraph_indent=' ' * 8,
)

class TextDict(dict):
    """Dictionary persistent in a text file."""
    def __init__(self, path):
        super().__init__()
        self.path = Path(path)
        self.read()

    # def __del__(self):
    #     self.write()

    # def __setitem__(self, key, value):
    #     super().__setitem__(key, value)
    #     self.write()

    def read(self):
        """Read from file, if it exists."""
        try:
            self.update(read_data(self.path))
        except FileNotFoundError:
            pass

    def write(self):
        """Write to file."""
        write_data(self.path, self)

def write_data(path, data):
    """Write object into JSON file."""
    pyutils.files.ensure_dir(path)
    with pyutils.files.tempfile_and_backup(path, 'w') as fp:
        return jsonfile.write_json(fp, data)

def read_data(path):
    """Read JSON file into dictionary."""
    with path.open() as fp:
        return jsonfile.read_json(fp)

class View(util.AttrDict):
    """A view to a feedlist. Giving '.' as flags stands for all flags."""
    DEFAULTS = dict(flags='foin', number=1, sortkey='fpD', sortkey2='SD')

    def __init__(self, directory=None, flags=None, number=None, sortkey=None,
                 sortkey2=None):
        super().__init__()
        self.directory = directory
        if self.directory is not None:
            self.directory = [Path(x) for x in self.directory]
        if flags is None:
            flags = self.DEFAULTS['flags']
        if number is None:
            number = self.DEFAULTS['number']
        if sortkey is None:
            sortkey = self.DEFAULTS['sortkey']
        if sortkey2 is None:
            sortkey2 = self.DEFAULTS['sortkey2']
        if flags == '.':
            flags = list(Flag)
        self.flags = [Flag(x) for x in flags]
        self.number = number
        self.sortkey = sortkey
        self.sortkey2 = sortkey2

    def __str__(self):
        s = '(f={flags}, n={number}, s={sortkey}, S={sortkey2}, {n})'
        d = dict(self, flags=''.join(x.value for x in self.flags),
                 n=len(self.directory))
        return s.format(**d)

    def parse(self, s, sep=','):
        """Parse view string, update original view, and create a new one."""
        if not s:
            s = ',,,'
        lst = list(s.split(sep))
        if len(lst) < 4:
            lst += [''] * (4 - len(lst))  # Fill in missing parts.
        try:
            flags, number, sortkey, sortkey2 = lst
        except ValueError as e:
            log.error('Cannot parse view: "%s"', s)
            raise
        return self.__class__(
            directory=self.directory,
            flags=flags or self.flags,
            number=int(number or self.number),
            sortkey=sortkey or self.sortkey,
            sortkey2=sortkey2 or self.sortkey2,
            )

# def message(value, *args, **kwargs):
#     """Print message."""
#     # print(value, *args, *kwargs)
#     messager.msg(value, *args, *kwargs)
#
#
# def feedback(value, *args, **kwargs):
#     """Print feedback."""
#     # print(value, *args, *kwargs)
#     messager.msg(value, *args, *kwargs)

def check_feed(feed, verbose=False):
    """Check feed. Return list of orphaned files."""
    for msg in feed.check(verbose=verbose):
        messager.msg('{}: {}'.format(feed, msg))
    return feed.get_orphans()

def add_url(url):
    """Add feed. Return directory name."""
    feed = Feed(url)
    feed.refresh(force=True)
    # check_feed(feed)
    if feed.directory.exists():
        raise IOError('Feed directory exists: {}'.format(feed.directory))
    feed.write()
    check_feed(feed)
    return feed.directory

def show(header, seq, wrapper=WRAPPER):
    """Show bodies of text prettily, using a TextWrapper object."""
    if header:
        # messager.msg(term.on_bright_green, *header, term.normal, sep=' ▶ ')
        messager.msg(*header, sep=' ▶ ')
    # for text, title in seq:
    #     messager.msg(wrapper.fills(text, title+':'))
    rows = [[y+':', wrapper.fills(x)] for x, y in seq]
    messager.msg(fmt_table(rows))

def show_feed(feed, verbose=0):
    """Show feed."""
    t_bold = term.bold
    header = [feed.directory, feed.url]
    lst = []
    if verbose:
        lst.extend([
            (time_fmt(feed.head.date, fmt='rfc2822'), 'Date'),
            (feed.priority, 'Priority'),
            (feed.head.link, 'Link'),
            (t_bold(feed.head.title), 'Title'),
            (feed.head.subtitle, 'Subt'),
            (feed.head.summary, 'Summ'),
            # (feed.head.language or 'unknown', 'Language'),
            # (', '.join(feed.get_tags().as_strings()) or [], 'Tags'),
            (str(feed.get_tags()), 'Tags'),
            (feed.head.image or [], 'Image'),
            (len(feed.entries), 'Entries'),
            ])
    if verbose > 1:
        # List flags and time range.
        def n_flagged(flag):
            return len(feed.list_entries(flags=flag.value))
        first = time_fmt(feed.entries[0].date, fmt='isodate')
        last = time_fmt(feed.entries[-1].date, fmt='isodate')
        lst.extend([
            (fmt_strings('{} {}'.format(x.name, n_flagged(x)) for x in Flag),
             'Flags'),
            ('earliest {}, latest {}'.format(first, last), 'Range'),
            ])
    show(header, lst)

def fmt_time(t):
    return fmt_strings([time_fmt(t, fmt='isofull'),
                        relative_time(t, ndigits=1) if t else '-'])

def show_entry(entry, verbose=0):
    """Show feed."""
    t_bold = term.bold

    def term_bool(x):
        return term.green if x else term.red

    date_str = time_fmt(entry.date, fmt='isodate')
    header = [date_str, entry.feed.directory, entry.title]
    lst = []
    if verbose:
        lst.extend([
            (fmt_time(entry.date_published), "Publ'd"),
            (fmt_time(entry.date_seen), 'Seen'),
            # (entry.score, 'Score'),
            (fmt_strings([entry.flag.name, entry.progress, entry.score]),
             'FlPrSc'),
            (term_bool(is_uuid(entry.guid))(entry.guid), 'GUID'),
            (entry.link, 'Link'),
            (t_bold(entry.title), 'Title'),
            (entry.subtitle, 'Subt'),
            (html2text(entry.summary), 'Summ'),
            (str(entry.get_tags()), 'Tags'),
            ])
        for enc in entry.encs():
            lst.append((fmt_strings([bytes2human(enc.length or 0), enc.typ,
                                     enc.href]), 'URL'))
            try:
                lst.append((fmt_strings([bytes2human(enc.size() or 0),
                                         fmt_duration(enc.duration()),
                                         media.get_gain(enc.path),
                                         enc.filename]), 'File'))
            except FileNotFoundError:
                lst.append((enc.filename, 'File'))
            ##
            lst.append((enc.filename, '#Name'))
            # lst.append((enc.filename_slugified, '#Slug'))
            ##
    show(header, lst)
    # messager.msg(WRAPPER.fills(html2text(entry.summary)))

def show_files(entry, verbose=0):
    """Show files."""
    for enc in entry.encs():
        if verbose or enc.path.exists():
            print(shlex.quote(str(enc.path)))

def show_enclosure(enc):
    if enc.path.exists():
        messager.msg('Media info:')
        messager.pp(media.get_media_info(enc.path))
        messager.msg('Extended attributes:')
        messager.pp(sorted(pyutils.files.XAttrStr(enc.path).items()))
    else:
        log.error('No file: %s', enc.path)

def show_enclosures(entry):
    """Show enclosure file info."""
    for enc in entry.encs():
        show_enclosure(enc)

def download_enclosure(enc, maxsize=None):
    if enc.path.exists():
        log.debug('Already exists: %s', enc.path)
    elif enc.is_too_big(maxsize):
        log.warning('Download too big: %s: %s', bytes2human(enc.length or 0),
                    enc.path)
    else:
        # messager.msg(truncate('Downloading {}: {}'.format(
        #     fmt_size(enc.length), enc.path)))
        # messager.msg('Downloading {}: {}'.format(fmt_size(enc.length),
        #                                          enc.path), truncate=True)
        try:
            if not enc.download():
                log.error('Download failed: %s', enc.path)
        except NotImplementedError:
            log.debug('Download not implemented: %s', enc)
        except KeyboardInterrupt:
            log.error('Download interrupted: %s', enc)

def download_enclosures(entry, maxsize=None):
    """Download entry enclosures."""
    for enc in entry.encs():
        download_enclosure(enc, maxsize=maxsize)

def normalize_enclosure(enc, force=False):
    if enc.path.exists():
        if force or not enc.is_normalized():
            media.normalize_volume(enc.path)

def normalize_enclosures(entry, force=False):
    """Normalize loudness in enclosures."""
    tags = entry.feed.get_tags()
    if force or 'nonorm' not in tags:
        for enc in entry.encs():
            normalize_enclosure(enc, force=force)

def play_enclosure(enc):
    try:
        if enc.path.exists():
            exit_code = enc.play()
        else:
            messager.msg('File does not exist: {}'.format(enc.path))
            messager.msg('Streaming: {}'.format(enc.href))
            exit_code = enc.stream()
        messager.msg('Exit code:', exit_code)
        return exit_code
    except NotImplementedError:
        log.error('Play/stream not implemented: %s', enc)
        return None

def play_enclosures(entry, set_flag=True):
    """Play entry enclosures."""
    # exit_codes = (play_enclosure(x) for x in entry.encs())
    # if set_flag and all(x == 0 for x in exit_codes):
    #     if entry.flag in [Flag.fresh, Flag.important, Flag.new]:
    #         messager.msg('Flagging entry as opened')
    #         entry.set_flag(Flag.opened)
    #     entry.progress = 1  # TODO: Should be set correctly (embed mpv).
    if set_flag:
        if entry.flag in [Flag.fresh, Flag.important, Flag.new]:
            messager.msg('Flagging entry as opened')
            entry.set_flag(Flag.opened)
        entry.progress = 1  # TODO: Should be set correctly (embed mpv).
        entry.feed.write()
    exit_codes = [play_enclosure(x) for x in entry.encs()]
    return exit_codes

def remove_enclosure(enc):
    """Remove enclosure from disk."""
    if enc.path.exists():
        messager.msg('Removing:', enc.path)
        # enc.path.unlink()
        pyutils.files.trash_or_rm(enc.path)

def remove_enclosures(entry, set_flag=True):
    """Remove entry enclosures from disk."""
    for enc in entry.encs():
        remove_enclosure(enc)
    if set_flag:
        # if entry.flag == Flag.opened:
        #     entry.progress = 1
        entry.set_flag('d')

def drop_enc(entry):
    """Drop enclosure information from entry."""
    for enc in entry.encs():
        if enc.path.exists():
            raise FileExistsError('Enclosure exists: {}'.format(enc.path))
    del entry.enclosures
    entry.enclosures = []
    entry.feed.modified = True
<|endoftext|>