<commit_before>from __future__ import print_function
import os, time
from collections import defaultdict
from ..database import get_db
from ..errors import UsageError

def abbrev(t):
    if t is None:
        return "-"
    if t > 1.0:
        return "%.3fs" % t
    if t > 1e-3:
        return "%.1fms" % (t*1e3)
    return "%.1fus" % (t*1e6)

def abbreviate_space(s, SI=True):
    if s is None:
        return "-"
    if SI:
        U = 1000.0
        isuffix = "B"
    else:
        U = 1024.0
        isuffix = "iB"
    def r(count, suffix):
        return "%.2f %s%s" % (count, suffix, isuffix)

    if s < 1024: # 1000-1023 get emitted as bytes, even in SI mode
        return "%d B" % s
    if s < U*U:
        return r(s/U, "k")
    if s < U*U*U:
        return r(s/(U*U), "M")
    if s < U*U*U*U:
        return r(s/(U*U*U), "G")
    if s < U*U*U*U*U:
        return r(s/(U*U*U*U), "T")
    if s < U*U*U*U*U*U:
        return r(s/(U*U*U*U*U), "P")
    return r(s/(U*U*U*U*U*U), "E")

def print_event(event):
    event_type, started, result, total_bytes, waiting_time, total_time = event
    followthrough = None
    if waiting_time and total_time:
        followthrough = total_time - waiting_time
    print(" %16s: total=%7s wait=%7s ft=%7s size=%s (%s)" %
          ("%s-%s" % (event_type, result),
           abbrev(total_time),
           abbrev(waiting_time),
           abbrev(followthrough),
           abbreviate_space(total_bytes),
           time.ctime(started),
          ))

def show_usage(args):
    if not os.path.exists("relay.sqlite"):
        raise UsageError("cannot find relay.sqlite, please run from the server directory")
    oldest = None
    newest = None
    rendezvous_counters = defaultdict(int)
    transit_counters = defaultdict(int)
    total_transit_bytes = 0
    db = get_db("relay.sqlite")
    c = db.execute("SELECT * FROM `usage`"
                   " ORDER BY `started` ASC LIMIT ?",
                   (args.n,))
    for row in c.fetchall():
        if row["type"] == u"rendezvous":
            counters = rendezvous_counters
        elif row["type"] == u"transit":
            counters = transit_counters
            total_transit_bytes += row["total_bytes"]
        else:
            continue
        counters["total"] += 1
        counters[row["result"]] += 1
        if oldest is None or row["started"] < oldest:
            oldest = row["started"]
        if newest is None or row["started"] > newest:
            newest = row["started"]
        event = (row["type"], row["started"], row["result"],
                 row["total_bytes"], row["waiting_time"], row["total_time"])
        print_event(event)
    if rendezvous_counters["total"] or transit_counters["total"]:
        print("---")
        print("(most recent started %s ago)" % abbrev(time.time() - newest))
    if rendezvous_counters["total"]:
        print("rendezvous events:")
        counters = rendezvous_counters
        elapsed = time.time() - oldest
        total = counters["total"]
        print(" %d events in %s (%.2f per hour)" % (total, abbrev(elapsed),
                                                    (3600 * total / elapsed)))
        print("", ", ".join(["%s=%d (%d%%)" %
                             (k, counters[k], (100.0 * counters[k] / total))
                             for k in sorted(counters)
                             if k != "total"]))
    if transit_counters["total"]:
        print("transit events:")
        counters = transit_counters
        elapsed = time.time() - oldest
        total = counters["total"]
        print(" %d events in %s (%.2f per hour)" % (total, abbrev(elapsed),
                                                    (3600 * total / elapsed)))
        rate = total_transit_bytes / elapsed
        print(" %s total bytes, %sps" % (abbreviate_space(total_transit_bytes),
                                         abbreviate_space(rate)))
        print("", ", ".join(["%s=%d (%d%%)" %
                             (k, counters[k], (100.0 * counters[k] / total))
                             for k in sorted(counters)
                             if k != "total"]))
    return 0

def tail_usage(args):
    if not os.path.exists("relay.sqlite"):
        raise UsageError("cannot find relay.sqlite, please run from the server directory")
    db = get_db("relay.sqlite")
    # we don't seem to have unique row IDs, so this is an inaccurate and
    # inefficient hack
    seen = set()
    try:
        while True:
            old = time.time() - 2*60*60
            c = db.execute("SELECT * FROM `usage`"
                           " WHERE `started` > ?"
                           " ORDER BY `started` ASC", (old,))
            for row in c.fetchall():
                event = (row["type"], row["started"], row["result"],
                         row["total_bytes"], row["waiting_time"],
                         row["total_time"])
                if event not in seen:
                    print_event(event)
                    seen.add(event)
            time.sleep(2)
    except KeyboardInterrupt:
        return 0
    return 0
<commit_msg>show-usage: give more room for "rendezvous-lonely"<commit_after>from __future__ import print_function
import os, time
from collections import defaultdict
from ..database import get_db
from ..errors import UsageError

def abbrev(t):
    if t is None:
        return "-"
    if t > 1.0:
        return "%.3fs" % t
    if t > 1e-3:
        return "%.1fms" % (t*1e3)
    return "%.1fus" % (t*1e6)

def abbreviate_space(s, SI=True):
    if s is None:
        return "-"
    if SI:
        U = 1000.0
        isuffix = "B"
    else:
        U = 1024.0
        isuffix = "iB"
    def r(count, suffix):
        return "%.2f %s%s" % (count, suffix, isuffix)

    if s < 1024: # 1000-1023 get emitted as bytes, even in SI mode
        return "%d B" % s
    if s < U*U:
        return r(s/U, "k")
    if s < U*U*U:
        return r(s/(U*U), "M")
    if s < U*U*U*U:
        return r(s/(U*U*U), "G")
    if s < U*U*U*U*U:
        return r(s/(U*U*U*U), "T")
    if s < U*U*U*U*U*U:
        return r(s/(U*U*U*U*U), "P")
    return r(s/(U*U*U*U*U*U), "E")

def print_event(event):
    event_type, started, result, total_bytes, waiting_time, total_time = event
    followthrough = None
    if waiting_time and total_time:
        followthrough = total_time - waiting_time
    print("%17s: total=%7s wait=%7s ft=%7s size=%s (%s)" %
          ("%s-%s" % (event_type, result),
           abbrev(total_time),
           abbrev(waiting_time),
           abbrev(followthrough),
           abbreviate_space(total_bytes),
           time.ctime(started),
          ))

def show_usage(args):
    if not os.path.exists("relay.sqlite"):
        raise UsageError("cannot find relay.sqlite, please run from the server directory")
    oldest = None
    newest = None
    rendezvous_counters = defaultdict(int)
    transit_counters = defaultdict(int)
    total_transit_bytes = 0
    db = get_db("relay.sqlite")
    c = db.execute("SELECT * FROM `usage`"
                   " ORDER BY `started` ASC LIMIT ?",
                   (args.n,))
    for row in c.fetchall():
        if row["type"] == u"rendezvous":
            counters = rendezvous_counters
        elif row["type"] == u"transit":
            counters = transit_counters
            total_transit_bytes += row["total_bytes"]
        else:
            continue
        counters["total"] += 1
        counters[row["result"]] += 1
        if oldest is None or row["started"] < oldest:
            oldest = row["started"]
        if newest is None or row["started"] > newest:
            newest = row["started"]
        event = (row["type"], row["started"], row["result"],
                 row["total_bytes"], row["waiting_time"], row["total_time"])
        print_event(event)
    if rendezvous_counters["total"] or transit_counters["total"]:
        print("---")
        print("(most recent started %s ago)" % abbrev(time.time() - newest))
    if rendezvous_counters["total"]:
        print("rendezvous events:")
        counters = rendezvous_counters
        elapsed = time.time() - oldest
        total = counters["total"]
        print(" %d events in %s (%.2f per hour)" % (total, abbrev(elapsed),
                                                    (3600 * total / elapsed)))
        print("", ", ".join(["%s=%d (%d%%)" %
                             (k, counters[k], (100.0 * counters[k] / total))
                             for k in sorted(counters)
                             if k != "total"]))
    if transit_counters["total"]:
        print("transit events:")
        counters = transit_counters
        elapsed = time.time() - oldest
        total = counters["total"]
        print(" %d events in %s (%.2f per hour)" % (total, abbrev(elapsed),
                                                    (3600 * total / elapsed)))
        rate = total_transit_bytes / elapsed
        print(" %s total bytes, %sps" % (abbreviate_space(total_transit_bytes),
                                         abbreviate_space(rate)))
        print("", ", ".join(["%s=%d (%d%%)" %
                             (k, counters[k], (100.0 * counters[k] / total))
                             for k in sorted(counters)
                             if k != "total"]))
    return 0

def tail_usage(args):
    if not os.path.exists("relay.sqlite"):
        raise UsageError("cannot find relay.sqlite, please run from the server directory")
    db = get_db("relay.sqlite")
    # we don't seem to have unique row IDs, so this is an inaccurate and
    # inefficient hack
    seen = set()
    try:
        while True:
            old = time.time() - 2*60*60
            c = db.execute("SELECT * FROM `usage`"
                           " WHERE `started` > ?"
                           " ORDER BY `started` ASC", (old,))
            for row in c.fetchall():
                event = (row["type"], row["started"], row["result"],
                         row["total_bytes"], row["waiting_time"],
                         row["total_time"])
                if event not in seen:
                    print_event(event)
                    seen.add(event)
            time.sleep(2)
    except KeyboardInterrupt:
        return 0
    return 0
<|endoftext|>


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

import os
import subprocess
import tempfile
from nose.tools import assert_true
from nose_parameterized import parameterized

try:
    from shlex import quote
except ImportError:  # <3.3 fallback
    from shellescape import quote

sample_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'files')

if os.name == 'nt':
    # Sigh, shlex.quote quotes incorrectly on Windows
    quote = lambda x: x

    # Comes from appveyor
    cur_path = os.environ['PATH']
    python_path = os.environ.get('PYTHON')
    if python_path:
        new_path = r'{py};{py}\Scripts;{path}'.format(
            py=python_path, path=cur_path,
        )
        os.environ['PATH'] = new_path

def run_srt_util(cmd, shell=False, encoding='ascii'):
    raw_out = subprocess.check_output(
        cmd, shell=shell, env={'PYTHONPATH': '.'},
    )
    return raw_out.decode(encoding)

def assert_supports_all_io_methods(cmd, exclude_output=False,
                                   exclude_stdin=False):
    cmd[0] = 'srt_tools/' + cmd[0]

    if os.name == 'nt':
        cmd.insert(0, 'python')

    in_file = os.path.join(sample_dir, 'ascii.srt')
    in_file_gb = os.path.join(sample_dir, 'gb2312.srt')
    fd, out_file = tempfile.mkstemp()

    # This is accessed by filename, not fd
    os.close(fd)

    outputs = []
    cmd_string = ' '.join(quote(x) for x in cmd)

    try:
        outputs.append(run_srt_util(cmd + ['-i', in_file]))
        if not exclude_stdin:
            outputs.append(run_srt_util(
                '%s < %s' % (cmd_string, quote(in_file)),
                shell=True,
            ))
        if not exclude_output:
            run_srt_util(cmd + ['-i', in_file, '-o', out_file])
            run_srt_util(
                cmd + ['-i', in_file_gb, '-o', out_file, '-e', 'gb2312'],
                encoding='gb2312',
            )
            if not exclude_stdin:
                run_srt_util(
                    '%s < %s > %s' % (
                        cmd_string, quote(in_file), quote(out_file),
                    ),
                    shell=True,
                )
        assert_true(len(set(outputs)) == 1, repr(outputs))
    finally:
        os.remove(out_file)

@parameterized([
    (['srt-fix-subtitle-indexing'], False),
    (['srt-fixed-timeshift', '--seconds', '5'], False),
    ([
        'srt-linear-timeshift',
        '--f1', '00:00:01,000',
        '--f2', '00:00:02,000',
        '--t1', '00:00:03,000',
        '--t2', '00:00:04,000',
    ], False),
    (['srt-lines-matching', '-f', 'lambda x: True'], False),
    (['srt-mux'], False, True),
    (['srt-strip-html'], False),

    # Need to sort out time/thread issues
    # (('srt-play'), True),
])
def test_tools_support(args, exclude_output=False, exclude_stdin=False):
    assert_supports_all_io_methods(args, exclude_output, exclude_stdin)
<commit_msg>Pass path as part of subprocess.check_output<commit_after>#!/usr/bin/env python

import os
import subprocess
import tempfile
from nose.tools import assert_true
from nose_parameterized import parameterized

try:
    from shlex import quote
except ImportError:  # <3.3 fallback
    from shellescape import quote

sample_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'files')

if os.name == 'nt':
    # Sigh, shlex.quote quotes incorrectly on Windows
    quote = lambda x: x

def run_srt_util(cmd, shell=False, encoding='ascii'):
    if os.name == 'nt':
        # Comes from appveyor
        cur_path = os.environ['PATH']
        python_path = os.environ.get('PYTHON')
        if python_path:
            new_path = r'{py};{py}\Scripts;{path}'.format(
                py=python_path, path=cur_path,
            )
        else:
            new_path = cur_path

    raw_out = subprocess.check_output(
        cmd, shell=shell, env={'PYTHONPATH': '.', 'PATH': new_path},
    )
    return raw_out.decode(encoding)

def assert_supports_all_io_methods(cmd, exclude_output=False,
                                   exclude_stdin=False):
    cmd[0] = 'srt_tools/' + cmd[0]

    if os.name == 'nt':
        cmd.insert(0, 'python')

    in_file = os.path.join(sample_dir, 'ascii.srt')
    in_file_gb = os.path.join(sample_dir, 'gb2312.srt')
    fd, out_file = tempfile.mkstemp()

    # This is accessed by filename, not fd
    os.close(fd)

    outputs = []
    cmd_string = ' '.join(quote(x) for x in cmd)

    try:
        outputs.append(run_srt_util(cmd + ['-i', in_file]))
        if not exclude_stdin:
            outputs.append(run_srt_util(
                '%s < %s' % (cmd_string, quote(in_file)),
                shell=True,
            ))
        if not exclude_output:
            run_srt_util(cmd + ['-i', in_file, '-o', out_file])
            run_srt_util(
                cmd + ['-i', in_file_gb, '-o', out_file, '-e', 'gb2312'],
                encoding='gb2312',
            )
            if not exclude_stdin:
                run_srt_util(
                    '%s < %s > %s' % (
                        cmd_string, quote(in_file), quote(out_file),
                    ),
                    shell=True,
                )
        assert_true(len(set(outputs)) == 1, repr(outputs))
    finally:
        os.remove(out_file)

@parameterized([
    (['srt-fix-subtitle-indexing'], False),
    (['srt-fixed-timeshift', '--seconds', '5'], False),
    ([
        'srt-linear-timeshift',
        '--f1', '00:00:01,000',
        '--f2', '00:00:02,000',
        '--t1', '00:00:03,000',
        '--t2', '00:00:04,000',
    ], False),
    (['srt-lines-matching', '-f', 'lambda x: True'], False),
    (['srt-mux'], False, True),
    (['srt-strip-html'], False),

    # Need to sort out time/thread issues
    # (('srt-play'), True),
])
def test_tools_support(args, exclude_output=False, exclude_stdin=False):
    assert_supports_all_io_methods(args, exclude_output, exclude_stdin)
<|endoftext|>


<commit_before>#
# PLCAPI XML-RPC and SOAP interfaces
#
# Aaron Klingaman <alk@absarokasoft.com>
# Mark Huang <mlhuang@cs.princeton.edu>
#
# Copyright (C) 2004-2006 The Trustees of Princeton University
# $Id$
# $URL$
#

import sys
import traceback
import string

import xmlrpclib

# See "2.2 Characters" in the XML specification:
#
# #x9 | #xA | #xD | [#x20-#xD7FF] | [#xE000-#xFFFD]
# avoiding
# [#x7F-#x84], [#x86-#x9F], [#xFDD0-#xFDDF]

invalid_xml_ascii = map(chr, range(0x0, 0x8) + [0xB, 0xC] + range(0xE, 0x1F))
xml_escape_table = string.maketrans("".join(invalid_xml_ascii), "?" * len(invalid_xml_ascii))

def xmlrpclib_escape(s, replace = string.replace):
    """
    xmlrpclib does not handle invalid 7-bit control characters. This
    function augments xmlrpclib.escape, which by default only replaces
    '&', '<', and '>' with entities.
    """

    # This is the standard xmlrpclib.escape function
    s = replace(s, "&", "&")
    s = replace(s, "<", "<")
    s = replace(s, ">", ">",)

    # Replace invalid 7-bit control characters with '?'
    return s.translate(xml_escape_table)

def xmlrpclib_dump(self, value, write):
    """
    xmlrpclib cannot marshal instances of subclasses of built-in
    types. This function overrides xmlrpclib.Marshaller.__dump so that
    any value that is an instance of one of its acceptable types is
    marshalled as that type.

    xmlrpclib also cannot handle invalid 7-bit control characters. See
    above.
    """

    # Use our escape function
    args = [self, value, write]
    if isinstance(value, (str, unicode)):
        args.append(xmlrpclib_escape)

    try:
        # Try for an exact match first
        f = self.dispatch[type(value)]
    except KeyError:
        # Try for an isinstance() match
        for Type, f in self.dispatch.iteritems():
            if isinstance(value, Type):
                f(*args)
                return
        raise TypeError, "cannot marshal %s objects" % type(value)
    else:
        f(*args)

# You can't hide from me!
xmlrpclib.Marshaller._Marshaller__dump = xmlrpclib_dump

# SOAP support is optional
try:
    import SOAPpy
    from SOAPpy.Parser import parseSOAPRPC
    from SOAPpy.Types import faultType
    from SOAPpy.NS import NS
    from SOAPpy.SOAPBuilder import buildSOAP
except ImportError:
    SOAPpy = None

from PLC.Config import Config
from PLC.Faults import *
import PLC.Methods
import PLC.Accessors

def import_deep(name):
    mod = __import__(name)
    components = name.split('.')
    for comp in components[1:]:
        mod = getattr(mod, comp)
    return mod

class PLCAPI:

    # flat list of method names
    native_methods = PLC.Methods.native_methods

    # other_methods_map : dict {methodname: fullpath}
    # e.g. 'Accessors' -> 'PLC.Accessors.Accessors'
    other_methods_map={}
    for subdir in [ 'Accessors' ]:
        path="PLC."+subdir
        # scan e.g. PLC.Accessors.__all__
        pkg = __import__(path).__dict__[subdir]
        for modulename in getattr(pkg,"__all__"):
            fullpath=path+"."+modulename
            for method in getattr(import_deep(fullpath),"methods"):
                other_methods_map[method] = fullpath

    all_methods = native_methods + other_methods_map.keys()

    def __init__(self, config = "/etc/planetlab/plc_config", encoding = "utf-8"):
        self.encoding = encoding

        # Better just be documenting the API
        if config is None:
            return

        # Load configuration
        self.config = Config(config)

	# Initialize database connection
        if self.config.PLC_DB_TYPE == "postgresql":
            from PLC.PostgreSQL import PostgreSQL
            self.db = PostgreSQL(self)
        else:
            raise PLCAPIError, "Unsupported database type " + self.config.PLC_DB_TYPE

        # Aspects modify the API injecting code before/after method
        # calls. As of now we only have aspects for OMF integration,
        # that's why we enable aspects only if PLC_OMF is set to true.
        if self.config.PLC_OMF:
            from aspects import apply_aspects; apply_aspects()

    def callable(self, method):
        """
        Return a new instance of the specified method.
        """

        # Look up method
        if method not in self.all_methods:
            raise PLCInvalidAPIMethod, method

        # Get new instance of method
        try:
            classname = method.split(".")[-1]
            if method in self.native_methods:
                fullpath="PLC.Methods." + method
            else:
                fullpath=self.other_methods_map[method]
            module = __import__(fullpath, globals(), locals(), [classname])
            return getattr(module, classname)(self)
        except ImportError, AttributeError:
            raise PLCInvalidAPIMethod, "import error %s for %s" % (AttributeError,fullpath)

    def call(self, source, method, *args):
        """
        Call the named method from the specified source with the
        specified arguments.
        """

        function = self.callable(method)
        function.source = source
        return function(*args)

    def handle(self, source, data):
        """
        Handle an XML-RPC or SOAP request from the specified source.
        """

        # Parse request into method name and arguments
        try:
            interface = xmlrpclib
            (args, method) = xmlrpclib.loads(data)
            methodresponse = True
        except Exception, e:
            if SOAPpy is not None:
                interface = SOAPpy
                (r, header, body, attrs) = parseSOAPRPC(data, header = 1, body = 1, attrs = 1)
                method = r._name
                args = r._aslist()
                # XXX Support named arguments
            else:
                raise e

        try:
            result = self.call(source, method, *args)
        except PLCFault, fault:
            # Handle expected faults
            if interface == xmlrpclib:
                result = fault
                methodresponse = None
            elif interface == SOAPpy:
                result = faultParameter(NS.ENV_T + ":Server", "Method Failed", method)
                result._setDetail("Fault %d: %s" % (fault.faultCode, fault.faultString))

        # Return result
        if interface == xmlrpclib:
            if not isinstance(result, PLCFault):
                result = (result,)
            data = xmlrpclib.dumps(result, methodresponse = True, encoding = self.encoding, allow_none = 1)
        elif interface == SOAPpy:
            data = buildSOAP(kw = {'%sResponse' % method: {'Result': result}}, encoding = self.encoding)

        return data
<commit_msg>fix setting name<commit_after>#
# PLCAPI XML-RPC and SOAP interfaces
#
# Aaron Klingaman <alk@absarokasoft.com>
# Mark Huang <mlhuang@cs.princeton.edu>
#
# Copyright (C) 2004-2006 The Trustees of Princeton University
# $Id$
# $URL$
#

import sys
import traceback
import string

import xmlrpclib

# See "2.2 Characters" in the XML specification:
#
# #x9 | #xA | #xD | [#x20-#xD7FF] | [#xE000-#xFFFD]
# avoiding
# [#x7F-#x84], [#x86-#x9F], [#xFDD0-#xFDDF]

invalid_xml_ascii = map(chr, range(0x0, 0x8) + [0xB, 0xC] + range(0xE, 0x1F))
xml_escape_table = string.maketrans("".join(invalid_xml_ascii), "?" * len(invalid_xml_ascii))

def xmlrpclib_escape(s, replace = string.replace):
    """
    xmlrpclib does not handle invalid 7-bit control characters. This
    function augments xmlrpclib.escape, which by default only replaces
    '&', '<', and '>' with entities.
    """

    # This is the standard xmlrpclib.escape function
    s = replace(s, "&", "&")
    s = replace(s, "<", "<")
    s = replace(s, ">", ">",)

    # Replace invalid 7-bit control characters with '?'
    return s.translate(xml_escape_table)

def xmlrpclib_dump(self, value, write):
    """
    xmlrpclib cannot marshal instances of subclasses of built-in
    types. This function overrides xmlrpclib.Marshaller.__dump so that
    any value that is an instance of one of its acceptable types is
    marshalled as that type.

    xmlrpclib also cannot handle invalid 7-bit control characters. See
    above.
    """

    # Use our escape function
    args = [self, value, write]
    if isinstance(value, (str, unicode)):
        args.append(xmlrpclib_escape)

    try:
        # Try for an exact match first
        f = self.dispatch[type(value)]
    except KeyError:
        # Try for an isinstance() match
        for Type, f in self.dispatch.iteritems():
            if isinstance(value, Type):
                f(*args)
                return
        raise TypeError, "cannot marshal %s objects" % type(value)
    else:
        f(*args)

# You can't hide from me!
xmlrpclib.Marshaller._Marshaller__dump = xmlrpclib_dump

# SOAP support is optional
try:
    import SOAPpy
    from SOAPpy.Parser import parseSOAPRPC
    from SOAPpy.Types import faultType
    from SOAPpy.NS import NS
    from SOAPpy.SOAPBuilder import buildSOAP
except ImportError:
    SOAPpy = None

from PLC.Config import Config
from PLC.Faults import *
import PLC.Methods
import PLC.Accessors

def import_deep(name):
    mod = __import__(name)
    components = name.split('.')
    for comp in components[1:]:
        mod = getattr(mod, comp)
    return mod

class PLCAPI:

    # flat list of method names
    native_methods = PLC.Methods.native_methods

    # other_methods_map : dict {methodname: fullpath}
    # e.g. 'Accessors' -> 'PLC.Accessors.Accessors'
    other_methods_map={}
    for subdir in [ 'Accessors' ]:
        path="PLC."+subdir
        # scan e.g. PLC.Accessors.__all__
        pkg = __import__(path).__dict__[subdir]
        for modulename in getattr(pkg,"__all__"):
            fullpath=path+"."+modulename
            for method in getattr(import_deep(fullpath),"methods"):
                other_methods_map[method] = fullpath

    all_methods = native_methods + other_methods_map.keys()

    def __init__(self, config = "/etc/planetlab/plc_config", encoding = "utf-8"):
        self.encoding = encoding

        # Better just be documenting the API
        if config is None:
            return

        # Load configuration
        self.config = Config(config)

	# Initialize database connection
        if self.config.PLC_DB_TYPE == "postgresql":
            from PLC.PostgreSQL import PostgreSQL
            self.db = PostgreSQL(self)
        else:
            raise PLCAPIError, "Unsupported database type " + self.config.PLC_DB_TYPE

        # Aspects modify the API injecting code before/after method
        # calls. As of now we only have aspects for OMF integration,
        # that's why we enable aspects only if PLC_OMF is set to true.
        if self.config.PLC_OMF_ENABLED:
            from aspects import apply_aspects; apply_aspects()

    def callable(self, method):
        """
        Return a new instance of the specified method.
        """

        # Look up method
        if method not in self.all_methods:
            raise PLCInvalidAPIMethod, method

        # Get new instance of method
        try:
            classname = method.split(".")[-1]
            if method in self.native_methods:
                fullpath="PLC.Methods." + method
            else:
                fullpath=self.other_methods_map[method]
            module = __import__(fullpath, globals(), locals(), [classname])
            return getattr(module, classname)(self)
        except ImportError, AttributeError:
            raise PLCInvalidAPIMethod, "import error %s for %s" % (AttributeError,fullpath)

    def call(self, source, method, *args):
        """
        Call the named method from the specified source with the
        specified arguments.
        """

        function = self.callable(method)
        function.source = source
        return function(*args)

    def handle(self, source, data):
        """
        Handle an XML-RPC or SOAP request from the specified source.
        """

        # Parse request into method name and arguments
        try:
            interface = xmlrpclib
            (args, method) = xmlrpclib.loads(data)
            methodresponse = True
        except Exception, e:
            if SOAPpy is not None:
                interface = SOAPpy
                (r, header, body, attrs) = parseSOAPRPC(data, header = 1, body = 1, attrs = 1)
                method = r._name
                args = r._aslist()
                # XXX Support named arguments
            else:
                raise e

        try:
            result = self.call(source, method, *args)
        except PLCFault, fault:
            # Handle expected faults
            if interface == xmlrpclib:
                result = fault
                methodresponse = None
            elif interface == SOAPpy:
                result = faultParameter(NS.ENV_T + ":Server", "Method Failed", method)
                result._setDetail("Fault %d: %s" % (fault.faultCode, fault.faultString))

        # Return result
        if interface == xmlrpclib:
            if not isinstance(result, PLCFault):
                result = (result,)
            data = xmlrpclib.dumps(result, methodresponse = True, encoding = self.encoding, allow_none = 1)
        elif interface == SOAPpy:
            data = buildSOAP(kw = {'%sResponse' % method: {'Result': result}}, encoding = self.encoding)

        return data
<|endoftext|>


<commit_before>#!/usr/bin/python
import interrupt_client, database, MCP342X, wind_direction, HTU21D, bmp085, tgs2600, ds18b20_therm

pressure = bmp085.BMP085()
temp_probe = ds18b20_therm.DS18B20()
air_qual = tgs2600.TGS2600(adc_channel = 0)
humidity = HTU21D.HTU21D()
wind_dir = wind_direction.wind_direction(adc_channel = 0, margin = 20)
interrupts = interrupt_client.interrupt_client(port = 49501)

db = database.weather_database() #Local MySQL db

wind_average = wind_dir.get_value(10) #ten seconds

print("Inserting...")
db.insert(humidity.read_temperature(), temp_probe.read_temp(), air_qual.get_value(), pressure.get_pressure(), humidity.read_humidity(), wind_average, interrupts.get_wind(), interrupts.get_wind_gust(), interrupts.get_rain())
print("done")

interrupts.reset()
<commit_msg>tidy up<commit_after>#!/usr/bin/python
import interrupt_client, MCP342X, wind_direction, HTU21D, bmp085, tgs2600, ds18b20_therm
import database # requires MySQLdb python 2 library which is not ported to python 3 yet

pressure = bmp085.BMP085()
temp_probe = ds18b20_therm.DS18B20()
air_qual = tgs2600.TGS2600(adc_channel = 0)
humidity = HTU21D.HTU21D()
wind_dir = wind_direction.wind_direction(adc_channel = 0, margin = 20)
interrupts = interrupt_client.interrupt_client(port = 49501)

db = database.weather_database() #Local MySQL db

wind_average = wind_dir.get_value(10) #ten seconds

print("Inserting...")
db.insert(humidity.read_temperature(), temp_probe.read_temp(), air_qual.get_value(), pressure.get_pressure(), humidity.read_humidity(), wind_average, interrupts.get_wind(), interrupts.get_wind_gust(), interrupts.get_rain())
print("done")

interrupts.reset()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""
Pong in Pygame (written against Python 3.5)
"""

import argparse
import math
import random
import sys
from collections import namedtuple

import pygame
from pygame import locals as consts

DEBUG = False
DEBUG_PATH = None
DEBUG_PATH_COLOR = (128, 0, 0,)

# Ye olde window constants
FRAMES_PER_SECOND = 40
BACKGROUND_COLOR = (0, 0, 0,)
SCREEN_HEIGHT = 1000
SCREEN_WIDTH = 1600

# Ball constants
BALL_START_ANGLE = 45
BALL_MIN_SPEED = 10
BALL_MAX_SPEED = 30
BALL_HEIGHT = 20
BALL_WIDTH = 20
BALL_COLOR = (127, 216, 127,)

# Paddle constants
PADDLE_MIN_SPEED = 10
PADDLE_MAX_SPEED = 20
PADDLE_HEIGHT = 150
PADDLE_WIDTH = 30
PADDLE_COLOR = (127, 216, 127,)
PADDLE_DOWN_DIR = 1
PADDLE_UP_DIR = -1

# Ye olde data structures
Sides = namedtuple('Sides', ['top', 'right', 'bottom', 'left'])
Cartesian = namedtuple('Cartesian', ['x', 'y'])

class Path:

    def __init__(self):
        self.points = []

    def __repr__(self):
        return '<Path ({})>'.format(','.join(self.points))

    def add(self, point):
        self.points.append(point)

    def clear(self):
        self.points = []

class Vector:

    def __init__(self, angle, magnitude):
        self.angle = angle
        self.magnitude = magnitude

    def __repr__(self):
        return '<Vector (a: {}, m: {} x: {}, y: {})>'.format(
            self.angle,
            self.magnitude,
            *self.cartesian)

    @classmethod
    def from_cartesian(x, y):
        self.magnitude = math.sqrt(x**2 + y**2)
        if x != 0:
            self.angle = math.degrees(math.atan(y / x))
        else:
            self.angle = 90 if y >= 0 else 270

    @property
    def cartesian(self):
        magnitude = self.magnitude
        angle = self.angle
        return Cartesian(
            x = magnitude * math.cos(math.radians(angle)),
            y = -magnitude * math.sin(math.radians(angle)),
        )

    def reflect(self, horizontally=False, vertically=False):
        normed_angle = self.angle % 360
        if horizontally:
            normed_angle = (360 - (normed_angle - 180)) % 360
        if vertically:
            normed_angle = 360 - normed_angle
        return self.__class__(normed_angle, self.magnitude)

class Ball:

    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((BALL_WIDTH, BALL_HEIGHT))
        self.image.fill(BALL_COLOR)
        center_vert = SCREEN_HEIGHT // 2 - (BALL_HEIGHT // 2)
        center_horiz = SCREEN_WIDTH // 2 - (BALL_WIDTH // 2)
        self.rect = self.image.get_rect().move(center_horiz, center_vert)
        angle = BALL_START_ANGLE or random.randint(0, 360)
        self.vector = Vector(angle, BALL_MAX_SPEED)
        self.touching_paddle = False

    def calc_sides_touched(self):
        return Sides(
            top = self.rect.top <= 0,
            right = self.rect.right >= SCREEN_WIDTH,
            bottom = self.rect.bottom >= SCREEN_HEIGHT,
            left = self.rect.left <= 0,
        )

    def check_paddle_collision(self, paddle):
        if pygame.sprite.collide_rect(self, paddle):
            self.touching_paddle = True

    def should_change_direction(self):
        screen_area = pygame.display.get_surface().get_rect()
        return not screen_area.contains(self.rect) or self.touching_paddle

    def update(self):
        if self.should_change_direction():
            sides_touched = self.calc_sides_touched()
            if any(sides_touched):
                reflect_h = sides_touched.left or sides_touched.right
                reflect_v = sides_touched.top or sides_touched.bottom
            else:
                reflect_h = True
                reflect_v = False
            self.vector = self.vector.reflect(reflect_h, reflect_v)
        self.rect.move_ip(*self.vector.cartesian)
        self.touching_paddle = False

class Paddle:

    def __init__(self, top, left):
        super().__init__()
        self.height = PADDLE_HEIGHT
        self.width = PADDLE_WIDTH
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(PADDLE_COLOR)
        self.rect = self.image.get_rect().move(left, top)
        self.speed = 0
        self.direction = 0
        self.is_moving = False

    def up(self):
        self.is_moving = True
        self.direction = PADDLE_UP_DIR
        self.speed = PADDLE_MAX_SPEED

    def down(self):
        self.is_moving = True
        self.direction = PADDLE_DOWN_DIR
        self.speed = PADDLE_MAX_SPEED

    def recenter(self):
        self.is_moving = True
        d_center = self.rect.centery - (SCREEN_HEIGHT // 2)
        self.direction = PADDLE_DOWN_DIR if d_center < 0 else PADDLE_UP_DIR
        self.speed = min(PADDLE_MAX_SPEED, abs(d_center))

    def update(self):
        if self.is_moving:
            self.rect.move_ip(0, self.speed * self.direction)
            self.is_moving = False

class MercilessAutomaton:

    def __init__(self, paddle):
        self.paddle = paddle
        self.sweet_spot_radius = self.paddle.rect.height // 4

    # TODO: don't assume we are on the right side
    def predict_intercept(self, start, vector, intercept_x, max_reflections=5):
        """
        Predict where the ball will intercept our paddle

        max_reflections specifies the maximum number of bounces off
        the top or bottom of the screen we are willing to calculate
        ahead of time. This estimate will be a little inaccurate if
        start and vector are calculated from the center of the ball
        because the bounce is calculated from the edges
        """
        DEBUG_PATH.add(start)
        x, y = start
        diff_x = intercept_x - x
        angle = vector.angle
        magnitude = diff_x / math.cos(math.radians(angle))
        projected = Vector(angle, magnitude)
        proj_y = y + projected.cartesian.y
        intercept = Cartesian(intercept_x, proj_y)

        if 0 <= proj_y <= SCREEN_HEIGHT:
            DEBUG_PATH.add(intercept)
            return intercept
        elif max_reflections == 0:
            return None

        if proj_y < 0:
            next_y = 0
        else:
            next_y = SCREEN_HEIGHT

        new_diff_y = abs(next_y - y)
        # kind of wonky looking math. here's what is intended:
        # 90 - (angle % 90) yields angle compared to y-axis
        #                   when we are moving to the right
        # Then SohCahToa...TAN is opposite over adjacent. we've defined
        #             angle such that length of adjacent leg is
        #             distance from starting coord to the top/bottom
        #             edge of the screen we are moving toward
        # this allows the computer player to predict what will happen
        # when the ball bounces off the top/bottom edge of the screen.
        new_diff_x = math.tan(math.radians(abs(90 - (angle % 180)))) * new_diff_y

        new_start = Cartesian(x + new_diff_x, next_y)
        new_vector = vector.reflect(vertically=True)
        return self.predict_intercept(new_start, new_vector, intercept_x,
                                      max_reflections-1)

    # TODO: don't assume we are on the right side
    # TODO: handle ricochets
    def play(self, ball):
        """
        Implement a basic strategy

        Basic strategy:
        ===============
        1) Use geometry (naively - doesn't account for ricochets yet)
        2) Never get tired, never get hungry. Never feel pity. Never
           be moved by a great symphony.
        3) Crush hoomuns.
        """
        DEBUG_PATH.clear()
        angle = ball.vector.angle
        its_coming_right_for_us = not (90 <= angle <= 270)
        is_in_play = (its_coming_right_for_us and
                      ball.rect.left <= self.paddle.rect.left)
        if is_in_play:
            ball_location = Cartesian(*ball.rect.center)
            intercept = self.predict_intercept(
                ball_location, ball.vector, self.paddle.rect.left)
            sweet_spot = Sides(
                top=self.paddle.rect.centery - self.sweet_spot_radius,
                bottom=self.paddle.rect.centery + self.sweet_spot_radius,
                left=None,
                right=None,
            )
            if intercept and intercept.y < sweet_spot.top:
                self.paddle.up()
            if intercept and intercept.y > sweet_spot.bottom:
                self.paddle.down()
        else:
            self.paddle.recenter()

def is_quit_event(e):
    return (
        e.type == consts.QUIT or
        (e.type == consts.KEYDOWN and e.key == consts.K_ESCAPE)
    )

def main():
    pygame.init()
    screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
    pygame.display.set_caption('pong')
    pygame.mouse.set_visible(False)

    background = pygame.Surface(screen.get_size())
    background.fill(BACKGROUND_COLOR)
    screen.blit(background, (0, 0))

    paddle1 = Paddle(50, 100)
    paddle2 = Paddle(50, screen.get_rect().width - 110)
    computer = MercilessAutomaton(paddle2)
    ball = Ball()
    all_sprites = (paddle1, paddle2, ball,)
    clock = pygame.time.Clock()

    if DEBUG:
        import pdb

    while True:
        pygame.event.pump()
        keys = pygame.key.get_pressed()

        for event in pygame.event.get():
            if is_quit_event(event):
                pygame.quit()
                sys.exit()

        if DEBUG and keys[ord('d')]:
            pdb.set_trace()

        if keys[consts.K_UP]:
            paddle1.up()
        elif keys[consts.K_DOWN]:
            paddle1.down()

        computer.play(ball)
        ball.check_paddle_collision(paddle1)
        ball.check_paddle_collision(paddle2)

        screen.blit(background, (0, 0))
        for sprite in all_sprites:
            sprite.update()
            screen.blit(sprite.image, sprite.rect)

        if DEBUG and len(DEBUG_PATH.points) > 1:
            pygame.draw.lines(screen, DEBUG_PATH_COLOR, False, DEBUG_PATH.points)
        pygame.display.flip()
        clock.tick(FRAMES_PER_SECOND)

    pygame.quit()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--debug', action='store_true')
    args = parser.parse_args()
    global DEBUG
    global DEBUG_PATH
    DEBUG = args.debug
    DEBUG_PATH = Path()
    main()
<commit_msg>Fix obsolete comments<commit_after>#!/usr/bin/env python3
"""
Pong in Pygame (written against Python 3.5)
"""

import argparse
import math
import random
import sys
from collections import namedtuple

import pygame
from pygame import locals as consts

DEBUG = False
DEBUG_PATH = None
DEBUG_PATH_COLOR = (128, 0, 0,)

# Ye olde window constants
FRAMES_PER_SECOND = 40
BACKGROUND_COLOR = (0, 0, 0,)
SCREEN_HEIGHT = 1000
SCREEN_WIDTH = 1600

# Ball constants
BALL_START_ANGLE = 45
BALL_MIN_SPEED = 10
BALL_MAX_SPEED = 30
BALL_HEIGHT = 20
BALL_WIDTH = 20
BALL_COLOR = (127, 216, 127,)

# Paddle constants
PADDLE_MIN_SPEED = 10
PADDLE_MAX_SPEED = 20
PADDLE_HEIGHT = 150
PADDLE_WIDTH = 30
PADDLE_COLOR = (127, 216, 127,)
PADDLE_DOWN_DIR = 1
PADDLE_UP_DIR = -1

# Ye olde data structures
Sides = namedtuple('Sides', ['top', 'right', 'bottom', 'left'])
Cartesian = namedtuple('Cartesian', ['x', 'y'])

class Path:

    def __init__(self):
        self.points = []

    def __repr__(self):
        return '<Path ({})>'.format(','.join(self.points))

    def add(self, point):
        self.points.append(point)

    def clear(self):
        self.points = []

class Vector:

    def __init__(self, angle, magnitude):
        self.angle = angle
        self.magnitude = magnitude

    def __repr__(self):
        return '<Vector (a: {}, m: {} x: {}, y: {})>'.format(
            self.angle,
            self.magnitude,
            *self.cartesian)

    @classmethod
    def from_cartesian(x, y):
        self.magnitude = math.sqrt(x**2 + y**2)
        if x != 0:
            self.angle = math.degrees(math.atan(y / x))
        else:
            self.angle = 90 if y >= 0 else 270

    @property
    def cartesian(self):
        magnitude = self.magnitude
        angle = self.angle
        return Cartesian(
            x = magnitude * math.cos(math.radians(angle)),
            y = -magnitude * math.sin(math.radians(angle)),
        )

    def reflect(self, horizontally=False, vertically=False):
        normed_angle = self.angle % 360
        if horizontally:
            normed_angle = (360 - (normed_angle - 180)) % 360
        if vertically:
            normed_angle = 360 - normed_angle
        return self.__class__(normed_angle, self.magnitude)

class Ball:

    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((BALL_WIDTH, BALL_HEIGHT))
        self.image.fill(BALL_COLOR)
        center_vert = SCREEN_HEIGHT // 2 - (BALL_HEIGHT // 2)
        center_horiz = SCREEN_WIDTH // 2 - (BALL_WIDTH // 2)
        self.rect = self.image.get_rect().move(center_horiz, center_vert)
        angle = BALL_START_ANGLE or random.randint(0, 360)
        self.vector = Vector(angle, BALL_MAX_SPEED)
        self.touching_paddle = False

    def calc_sides_touched(self):
        return Sides(
            top = self.rect.top <= 0,
            right = self.rect.right >= SCREEN_WIDTH,
            bottom = self.rect.bottom >= SCREEN_HEIGHT,
            left = self.rect.left <= 0,
        )

    def check_paddle_collision(self, paddle):
        if pygame.sprite.collide_rect(self, paddle):
            self.touching_paddle = True

    def should_change_direction(self):
        screen_area = pygame.display.get_surface().get_rect()
        return not screen_area.contains(self.rect) or self.touching_paddle

    def update(self):
        if self.should_change_direction():
            sides_touched = self.calc_sides_touched()
            if any(sides_touched):
                reflect_h = sides_touched.left or sides_touched.right
                reflect_v = sides_touched.top or sides_touched.bottom
            else:
                reflect_h = True
                reflect_v = False
            self.vector = self.vector.reflect(reflect_h, reflect_v)
        self.rect.move_ip(*self.vector.cartesian)
        self.touching_paddle = False

class Paddle:

    def __init__(self, top, left):
        super().__init__()
        self.height = PADDLE_HEIGHT
        self.width = PADDLE_WIDTH
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(PADDLE_COLOR)
        self.rect = self.image.get_rect().move(left, top)
        self.speed = 0
        self.direction = 0
        self.is_moving = False

    def up(self):
        self.is_moving = True
        self.direction = PADDLE_UP_DIR
        self.speed = PADDLE_MAX_SPEED

    def down(self):
        self.is_moving = True
        self.direction = PADDLE_DOWN_DIR
        self.speed = PADDLE_MAX_SPEED

    def recenter(self):
        self.is_moving = True
        d_center = self.rect.centery - (SCREEN_HEIGHT // 2)
        self.direction = PADDLE_DOWN_DIR if d_center < 0 else PADDLE_UP_DIR
        self.speed = min(PADDLE_MAX_SPEED, abs(d_center))

    def update(self):
        if self.is_moving:
            self.rect.move_ip(0, self.speed * self.direction)
            self.is_moving = False

class MercilessAutomaton:

    def __init__(self, paddle):
        self.paddle = paddle
        self.sweet_spot_radius = self.paddle.rect.height // 4

    # TODO: don't assume we are on the right side
    def predict_intercept(self, start, vector, intercept_x, max_reflections=5):
        """
        Predict where the ball will intercept our paddle

        max_reflections specifies the maximum number of bounces off
        the top or bottom of the screen we are willing to calculate
        ahead of time. This estimate will be a little inaccurate if
        start and vector are calculated from the center of the ball
        because the bounce is calculated from the edges
        """
        DEBUG_PATH.add(start)
        x, y = start
        diff_x = intercept_x - x
        angle = vector.angle
        magnitude = diff_x / math.cos(math.radians(angle))
        projected = Vector(angle, magnitude)
        proj_y = y + projected.cartesian.y
        intercept = Cartesian(intercept_x, proj_y)

        if 0 <= proj_y <= SCREEN_HEIGHT:
            DEBUG_PATH.add(intercept)
            return intercept
        elif max_reflections == 0:
            return None

        if proj_y < 0:
            next_y = 0
        else:
            next_y = SCREEN_HEIGHT

        new_diff_y = abs(next_y - y)
        # kind of wonky looking math. here's what is intended:
        # 90 - (angle % 180) yields angle compared to y-axis
        #                    when we are moving to the right
        # Then SohCahToa...TAN is opposite over adjacent. we've defined
        #             angle such that length of adjacent leg is
        #             distance from starting coord to the top/bottom
        #             edge of the screen we are moving toward
        # this allows the computer player to predict what will happen
        # when the ball bounces off the top/bottom edge of the screen.
        new_diff_x = math.tan(math.radians(abs(90 - (angle % 180)))) * new_diff_y

        new_start = Cartesian(x + new_diff_x, next_y)
        new_vector = vector.reflect(vertically=True)
        return self.predict_intercept(new_start, new_vector, intercept_x,
                                      max_reflections-1)

    # TODO: don't assume we are on the right side
    def play(self, ball):
        """
        Implement a basic strategy

        Basic strategy:
        ===============
        1) Use geometry
        2) Never get tired, never get hungry. Never feel pity. Never
           be moved by a great symphony.
        3) Crush hoomuns.
        """
        DEBUG_PATH.clear()
        angle = ball.vector.angle
        its_coming_right_for_us = not (90 <= angle <= 270)
        is_in_play = (its_coming_right_for_us and
                      ball.rect.left <= self.paddle.rect.left)
        if is_in_play:
            ball_location = Cartesian(*ball.rect.center)
            intercept = self.predict_intercept(
                ball_location, ball.vector, self.paddle.rect.left)
            sweet_spot = Sides(
                top=self.paddle.rect.centery - self.sweet_spot_radius,
                bottom=self.paddle.rect.centery + self.sweet_spot_radius,
                left=None,
                right=None,
            )
            if intercept and intercept.y < sweet_spot.top:
                self.paddle.up()
            if intercept and intercept.y > sweet_spot.bottom:
                self.paddle.down()
        else:
            self.paddle.recenter()

def is_quit_event(e):
    return (
        e.type == consts.QUIT or
        (e.type == consts.KEYDOWN and e.key == consts.K_ESCAPE)
    )

def main():
    pygame.init()
    screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
    pygame.display.set_caption('pong')
    pygame.mouse.set_visible(False)

    background = pygame.Surface(screen.get_size())
    background.fill(BACKGROUND_COLOR)
    screen.blit(background, (0, 0))

    paddle1 = Paddle(50, 100)
    paddle2 = Paddle(50, screen.get_rect().width - 110)
    computer = MercilessAutomaton(paddle2)
    ball = Ball()
    all_sprites = (paddle1, paddle2, ball,)
    clock = pygame.time.Clock()

    if DEBUG:
        import pdb

    while True:
        pygame.event.pump()
        keys = pygame.key.get_pressed()

        for event in pygame.event.get():
            if is_quit_event(event):
                pygame.quit()
                sys.exit()

        if DEBUG and keys[ord('d')]:
            pdb.set_trace()

        if keys[consts.K_UP]:
            paddle1.up()
        elif keys[consts.K_DOWN]:
            paddle1.down()

        computer.play(ball)
        ball.check_paddle_collision(paddle1)
        ball.check_paddle_collision(paddle2)

        screen.blit(background, (0, 0))
        for sprite in all_sprites:
            sprite.update()
            screen.blit(sprite.image, sprite.rect)

        if DEBUG and len(DEBUG_PATH.points) > 1:
            pygame.draw.lines(screen, DEBUG_PATH_COLOR, False, DEBUG_PATH.points)
        pygame.display.flip()
        clock.tick(FRAMES_PER_SECOND)

    pygame.quit()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--debug', action='store_true')
    args = parser.parse_args()
    global DEBUG
    global DEBUG_PATH
    DEBUG = args.debug
    DEBUG_PATH = Path()
    main()
<|endoftext|>


<commit_before>
import db
import seaborn as sns

def follower_botness_histogram(username):

    scorelist=[]
    followers = db.getFollowers(toName=username)
    for f in followers:
        follower = f['_from'].split('/')[1]
        score=db.getUser(follower)['botness']['score']
        scorelist.append(score)
    ax = sns.distplot(scorelist, name='probability of follower bot')
    fig = ax.get_figure()
    fig.savefig('plots/testfig.png')

<commit_msg>followers_botness now also returns percentage of follower bots<commit_after>import db
import seaborn as sns
import numpy as np

# takes data (that should come from the frontend) and creates the output we would like to have
# on the front end.

def follower_botness(username):
"""  given a username, it creates the histogram of the botness of the followers
    and saves it in plots (for now)  it also returns the probable percentage of follower bots
    (cutoff needs to be defined, for now it is 0.7)"""
    cutoff = 0.7
    scorelist = []
    followers = db.getFollowers(toName=username)
    for f in followers:
        follower = f['_from'].split('/')[1]
        score=db.getUser(follower)['botness']['score']
        scorelist.append(score)
    ax = sns.distaplot(scorelist, name='probability of follower bot')
    fig = ax.get_figure()
    fig.savefig('plots/testfig.png')
    botpercent = sum(np.array(scorelist)>cutoff)/len(scorelist)
    return botpercent

<|endoftext|>


<commit_before>import argparse
import os
import sys
import pysam

__author__ = 'Rob Edwards'

"""
Extract the PCR regions from indexed bam alignments. These are regions that we have asked people to pull out, and
we will extract the sequences from those regions

Current regions:

Primer sequences:			Position on JQ995537 (the original genbank record)
Primer A:
	Fwd: CTGATAGTATGATTGGTAAT	Position 25634 .. 25653
	Rev: ATAAGTTCTCCAACTATCTT	Position complement(26945 .. 26964)

Primer B:
	Fwd: CCAGTATCTCCATAAGCATC	Position 33709 .. 33728
	Rev: GTGAGGGCGGAATAGCTA		Position complement(35045 .. 35062)

Primer C:
	Fwd: GCAACAGGAGTAGTAAAATCTC	Position 43820 .. 43841
	Rev: GCTCCTGTTAATCCTGATGTTA	Position complement(45036 .. 45057)

"""

locations = {
    'JQ995537': {
        'PrimerA': (25633, 26964),
        'PrimerB': (33708, 35062),
        'PrimerC': (43819, 45057)
    }
}

def print_query_regions(bam):
    """
    Print the regions from the query sequences that overlap our regions of interest

    :param bam: The bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            for read in bam.fetch(reference=template, start=start, end=end):
                # this is an AlignedSegment: http://pysam.readthedocs.org/en/latest/api.html#pysam.AlignedSegment
                # sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #     primer, start, end, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #     read.reference_start, read.reference_end
                # ))

                # this checks for sequences that overlap the start and end (none do in the Ondrej data set
                # if read.reference_start <= start and read.reference_end >= stop:
                #     sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #         primer, start, stop, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #         read.reference_start, read.reference_end
                #     ))

                # get just the sequence that maps to the region
                seq = read.query_sequence
                beg_offset = None
                end_offset = None
                if read.reference_start < start:
                    beg_offset = start - read.reference_start - 1
                if read.reference_end > end:
                    end_offset = len(seq) - (read.reference_end - end)

                if beg_offset and end_offset:
                    seq = seq[beg_offset:end_offset]
                elif beg_offset:
                    seq = seq[beg_offset:]
                elif end_offset:
                    seq = seq[:end_offset]

                print(">{} {} {} {}\n{}".format(read.query_name, primer, read.reference_start, read.reference_end, seq))

def print_pileup(bam):
    """
    Print the information about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end  = locations[template][primer]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                    else:
                        bp = '-'
                    bases[bp] = bases.get(bp, 0) + 1
                sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, str(bases)))

def print_consensus(bam):
    """
    Print the consensus sequence about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            cseq = []
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                        bases[bp] = bases.get(bp, 0) + 1
                # if we don't have any bases at this position, add an N
                if not bases:
                    bases['N'] = 1
                bps = sorted(bases, key=bases.get, reverse=True)
                # text = ""
                # for b in bps:
                #     text += " " + b + ": " + str(bases[b])
                # sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, text))

                # make the consensus seq
                cseq.append(bps[0])
            print(">{} {} {} {}\n{}".format(primer, template, start, end, ''.join(cseq)))

def print_alignment(bam, output_dir, filename):
    """
    Print an alignment of all matching pileups

    :param output_dir: The output directory to store the results in. We will make a sub directory for each PCR region
    :type output_dir: str
    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    for template in locations:
        for primer in locations[template]:
            if not os.path.exists(os.path.join(output_dir, primer)):
                os.mkdir(os.path.join(output_dir, primer))
            start, end = locations[template][primer]
            # print("ALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            # This is a failed attempt to get the reference sequence for this region, but I am not sure that this
            # is even possible from a BAM file, since each read will have a unique alignment to the reference
            # refseq = ['-' for i in range(start, end)]
            # for aln in bam.fetch(reference=template, start=start, end=end, until_eof=True):
            #     posns = aln.get_reference_positions()
            #     seq = aln.get_reference_sequence()
            #     if len(posns) > len(seq):
            #         sys.stderr.write("There are more positions {} than sequences {}\n".format(len(posns), len(seq)))
            #         continue
            #     for i in range(len(posns)):
            #         if posns[i] - start > len(refseq) -1:
            #             sys.stderr.write("Too many positions\n")
            #         if i > len(seq)-1:
            #             sys.stderr.write("Too many seq\n")
            #         refseq[posns[i]-start] = seq[i]
            #
            # print("{}_{}     {}".format(template, primer, ''.join(refseq)))

            alignment = {}
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                for pilups in p.pileups:
                    if pilups.alignment.query_name not in alignment:
                        alignment[pilups.alignment.query_name] = ['-' for idx in range(start, end+1)]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                rp = p.reference_pos
                idx = rp - start
                for pilups in p.pileups:
                    if pilups.query_position:
                        posn = pilups.query_position - start
                        # sys.stderr.write("Posn: {} Q.position: {} start: {} end: {} len: {}\n".format(posn, pilups.query_position, start, end, end-start))
                        alignment[pilups.alignment.query_name][idx] = pilups.alignment.query_sequence[pilups.query_position]

            # find the longest name
            longest_name = 0
            for n in alignment:
                if len(n) > longest_name:
                    longest_name = len(n)
            longest_name += 5

            # I want to sort by the number of -'s at the beginning of the sequence
            beginning_gaps = {}
            for n in alignment:
                gap = 0
                while (gap < len(alignment[n]) and alignment[n][gap] == '-'):
                    gap += 1
                beginning_gaps[n] = gap

            out = open(os.path.join(output_dir, primer, filename), 'w')

            for n in sorted(alignment.keys(), key=beginning_gaps.get):
                # this is for "pylip style output"
                # out.write(n)
                # out.write(" " * (longest_name - len(n)))
                # out.write(''.join(alignment[n]) + "\n")
                out.write(">{}\n{}".format(n, ''.join(alignment[n])))
            out.close()
        #print("\n\n")

def list_sequences(bam):
    """
    List the sequences involved and whether they are forward or reverse

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            print("\nALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            for read in bam.fetch(reference=template, start=start, end=end):
                print("{}\t{}".format(read.query_name, read.is_reverse))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract PCRd regions from BAM files')
    parser.add_argument('-b', help='bam file', required=True)
    parser.add_argument('-q', help='print query regions. This is a fasta output of all sequences that are in the region', action='store_true')
    parser.add_argument('-p', help='print pileup. Prints debug information about each position in the pileup', action='store_true')
    parser.add_argument('-c', help='print consensus sequence. Prints a single sequence for each region', action='store_true')
    parser.add_argument('-a', help='print alignment. Prints an alignment for each region.', action='store_true')
    parser.add_argument('-l', help='list read ids and whether they are reversed', action='store_true')
    parser.add_argument('-d', help='output directory to write the alignments to (required with -a)')
    parser.add_argument('-v', help='verbose output')
    args = parser.parse_args()

    bam = pysam.AlignmentFile(args.b, 'rb')
    filename = args.b.split(os.pathsep)[-1]
    filename = filename.replace('.bam', '.aln')

    if args.q:
        print_query_regions(bam)
    if args.p:
        print_pileup(bam)
    if args.c:
        print_consensus(bam)
    if args.a:
        if not args.d:
            sys.exit("You must provide an output directory with -a flag")
        print_alignment(bam, args.d, filename)
    if args.l:
        list_sequences(bam)<commit_msg>updating output<commit_after>import argparse
import os
import sys
import pysam

__author__ = 'Rob Edwards'

"""
Extract the PCR regions from indexed bam alignments. These are regions that we have asked people to pull out, and
we will extract the sequences from those regions

Current regions:

Primer sequences:			Position on JQ995537 (the original genbank record)
Primer A:
	Fwd: CTGATAGTATGATTGGTAAT	Position 25634 .. 25653
	Rev: ATAAGTTCTCCAACTATCTT	Position complement(26945 .. 26964)

Primer B:
	Fwd: CCAGTATCTCCATAAGCATC	Position 33709 .. 33728
	Rev: GTGAGGGCGGAATAGCTA		Position complement(35045 .. 35062)

Primer C:
	Fwd: GCAACAGGAGTAGTAAAATCTC	Position 43820 .. 43841
	Rev: GCTCCTGTTAATCCTGATGTTA	Position complement(45036 .. 45057)

"""

locations = {
    'JQ995537': {
        'PrimerA': (25633, 26964),
        'PrimerB': (33708, 35062),
        'PrimerC': (43819, 45057)
    }
}

def print_query_regions(bam):
    """
    Print the regions from the query sequences that overlap our regions of interest

    :param bam: The bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            for read in bam.fetch(reference=template, start=start, end=end):
                # this is an AlignedSegment: http://pysam.readthedocs.org/en/latest/api.html#pysam.AlignedSegment
                # sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #     primer, start, end, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #     read.reference_start, read.reference_end
                # ))

                # this checks for sequences that overlap the start and end (none do in the Ondrej data set
                # if read.reference_start <= start and read.reference_end >= stop:
                #     sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #         primer, start, stop, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #         read.reference_start, read.reference_end
                #     ))

                # get just the sequence that maps to the region
                seq = read.query_sequence
                beg_offset = None
                end_offset = None
                if read.reference_start < start:
                    beg_offset = start - read.reference_start - 1
                if read.reference_end > end:
                    end_offset = len(seq) - (read.reference_end - end)

                if beg_offset and end_offset:
                    seq = seq[beg_offset:end_offset]
                elif beg_offset:
                    seq = seq[beg_offset:]
                elif end_offset:
                    seq = seq[:end_offset]

                print(">{} {} {} {}\n{}".format(read.query_name, primer, read.reference_start, read.reference_end, seq))

def print_pileup(bam):
    """
    Print the information about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end  = locations[template][primer]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                    else:
                        bp = '-'
                    bases[bp] = bases.get(bp, 0) + 1
                sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, str(bases)))

def print_consensus(bam):
    """
    Print the consensus sequence about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            cseq = []
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                        bases[bp] = bases.get(bp, 0) + 1
                # if we don't have any bases at this position, add an N
                if not bases:
                    bases['N'] = 1
                bps = sorted(bases, key=bases.get, reverse=True)
                # text = ""
                # for b in bps:
                #     text += " " + b + ": " + str(bases[b])
                # sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, text))

                # make the consensus seq
                cseq.append(bps[0])
            print(">{} {} {} {}\n{}".format(primer, template, start, end, ''.join(cseq)))

def print_alignment(bam, output_dir, filename):
    """
    Print an alignment of all matching pileups

    :param output_dir: The output directory to store the results in. We will make a sub directory for each PCR region
    :type output_dir: str
    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    for template in locations:
        for primer in locations[template]:
            if not os.path.exists(os.path.join(output_dir, primer)):
                os.mkdir(os.path.join(output_dir, primer))
            start, end = locations[template][primer]
            # print("ALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            # This is a failed attempt to get the reference sequence for this region, but I am not sure that this
            # is even possible from a BAM file, since each read will have a unique alignment to the reference
            # refseq = ['-' for i in range(start, end)]
            # for aln in bam.fetch(reference=template, start=start, end=end, until_eof=True):
            #     posns = aln.get_reference_positions()
            #     seq = aln.get_reference_sequence()
            #     if len(posns) > len(seq):
            #         sys.stderr.write("There are more positions {} than sequences {}\n".format(len(posns), len(seq)))
            #         continue
            #     for i in range(len(posns)):
            #         if posns[i] - start > len(refseq) -1:
            #             sys.stderr.write("Too many positions\n")
            #         if i > len(seq)-1:
            #             sys.stderr.write("Too many seq\n")
            #         refseq[posns[i]-start] = seq[i]
            #
            # print("{}_{}     {}".format(template, primer, ''.join(refseq)))

            alignment = {}
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                for pilups in p.pileups:
                    if pilups.alignment.query_name not in alignment:
                        alignment[pilups.alignment.query_name] = ['-' for idx in range(start, end+1)]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                rp = p.reference_pos
                idx = rp - start
                for pilups in p.pileups:
                    if pilups.query_position:
                        posn = pilups.query_position - start
                        # sys.stderr.write("Posn: {} Q.position: {} start: {} end: {} len: {}\n".format(posn, pilups.query_position, start, end, end-start))
                        alignment[pilups.alignment.query_name][idx] = pilups.alignment.query_sequence[pilups.query_position]

            # find the longest name
            longest_name = 0
            for n in alignment:
                if len(n) > longest_name:
                    longest_name = len(n)
            longest_name += 5

            # I want to sort by the number of -'s at the beginning of the sequence
            beginning_gaps = {}
            for n in alignment:
                gap = 0
                while (gap < len(alignment[n]) and alignment[n][gap] == '-'):
                    gap += 1
                beginning_gaps[n] = gap

            out = open(os.path.join(output_dir, primer, filename), 'w')

            for n in sorted(alignment.keys(), key=beginning_gaps.get):
                # this is for "pylip style output"
                # out.write(n)
                # out.write(" " * (longest_name - len(n)))
                # out.write(''.join(alignment[n]) + "\n")
                out.write(">{}\n{}\n".format(n, ''.join(alignment[n])))
            out.close()
        #print("\n\n")

def list_sequences(bam):
    """
    List the sequences involved and whether they are forward or reverse

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            print("\nALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            for read in bam.fetch(reference=template, start=start, end=end):
                print("{}\t{}".format(read.query_name, read.is_reverse))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract PCRd regions from BAM files')
    parser.add_argument('-b', help='bam file', required=True)
    parser.add_argument('-q', help='print query regions. This is a fasta output of all sequences that are in the region', action='store_true')
    parser.add_argument('-p', help='print pileup. Prints debug information about each position in the pileup', action='store_true')
    parser.add_argument('-c', help='print consensus sequence. Prints a single sequence for each region', action='store_true')
    parser.add_argument('-a', help='print alignment. Prints an alignment for each region.', action='store_true')
    parser.add_argument('-l', help='list read ids and whether they are reversed', action='store_true')
    parser.add_argument('-d', help='output directory to write the alignments to (required with -a)')
    parser.add_argument('-v', help='verbose output')
    args = parser.parse_args()

    filename = os.path.split(args.b)[-1]
    filename = filename.replace('.bam', '.aln')
    bam = pysam.AlignmentFile(args.b, 'rb')

    if args.q:
        print_query_regions(bam)
    if args.p:
        print_pileup(bam)
    if args.c:
        print_consensus(bam)
    if args.a:
        if not args.d:
            sys.exit("You must provide an output directory with -a flag")
        print_alignment(bam, args.d, filename)
    if args.l:
        list_sequences(bam)
<|endoftext|>


<commit_before>import argparse
import os
import sys
import pysam

__author__ = 'Rob Edwards'

"""
Extract the PCR regions from indexed bam alignments. These are regions that we have asked people to pull out, and
we will extract the sequences from those regions

Current regions:

Primer sequences:			Position on JQ995537 (the original genbank record)
Primer A:
	Fwd: CTGATAGTATGATTGGTAAT	Position 25634 .. 25653
	Rev: ATAAGTTCTCCAACTATCTT	Position complement(26945 .. 26964)

Primer B:
	Fwd: CCAGTATCTCCATAAGCATC	Position 33709 .. 33728
	Rev: GTGAGGGCGGAATAGCTA		Position complement(35045 .. 35062)

Primer C:
	Fwd: GCAACAGGAGTAGTAAAATCTC	Position 43820 .. 43841
	Rev: GCTCCTGTTAATCCTGATGTTA	Position complement(45036 .. 45057)

"""

locations = {
    'JQ995537': {
        'PrimerA': (25633, 26964),
        'PrimerB': (33708, 35062),
        'PrimerC': (43819, 45057)
    }
}

def print_query_regions(bam):
    """
    Print the regions from the query sequences that overlap our regions of interest

    :param bam: The bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            for read in bam.fetch(reference=template, start=start, end=end):
                # this is an AlignedSegment: http://pysam.readthedocs.org/en/latest/api.html#pysam.AlignedSegment
                # sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #     primer, start, end, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #     read.reference_start, read.reference_end
                # ))

                # this checks for sequences that overlap the start and end (none do in the Ondrej data set
                # if read.reference_start <= start and read.reference_end >= stop:
                #     sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #         primer, start, stop, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #         read.reference_start, read.reference_end
                #     ))

                # get just the sequence that maps to the region
                seq = read.query_sequence
                beg_offset = None
                end_offset = None
                if read.reference_start < start:
                    beg_offset = start - read.reference_start - 1
                if read.reference_end > end:
                    end_offset = len(seq) - (read.reference_end - end)

                if beg_offset and end_offset:
                    seq = seq[beg_offset:end_offset]
                elif beg_offset:
                    seq = seq[beg_offset:]
                elif end_offset:
                    seq = seq[:end_offset]

                print(">{} {} {} {}\n{}".format(read.query_name, primer, read.reference_start, read.reference_end, seq))

def print_pileup(bam):
    """
    Print the information about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end  = locations[template][primer]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                    else:
                        bp = '-'
                    bases[bp] = bases.get(bp, 0) + 1
                sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, str(bases)))

def print_consensus(bam):
    """
    Print the consensus sequence about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            cseq = []
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                        bases[bp] = bases.get(bp, 0) + 1
                # if we don't have any bases at this position, add an N
                if not bases:
                    bases['N'] = 1
                bps = sorted(bases, key=bases.get, reverse=True)
                # text = ""
                # for b in bps:
                #     text += " " + b + ": " + str(bases[b])
                # sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, text))

                # make the consensus seq
                cseq.append(bps[0])
            print(">{} {} {} {}\n{}".format(primer, template, start, end, ''.join(cseq)))

def print_alignment(bam, output_dir, filename):
    """
    Print an alignment of all matching pileups

    :param output_dir: The output directory to store the results in. We will make a sub directory for each PCR region
    :type output_dir: str
    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    for template in locations:
        for primer in locations[template]:
            if not os.path.exists(os.path.join(output_dir, primer)):
                os.mkdir(os.path.join(output_dir, primer))
            start, end = locations[template][primer]
            # print("ALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            # This is a failed attempt to get the reference sequence for this region, but I am not sure that this
            # is even possible from a BAM file, since each read will have a unique alignment to the reference
            # refseq = ['-' for i in range(start, end)]
            # for aln in bam.fetch(reference=template, start=start, end=end, until_eof=True):
            #     posns = aln.get_reference_positions()
            #     seq = aln.get_reference_sequence()
            #     if len(posns) > len(seq):
            #         sys.stderr.write("There are more positions {} than sequences {}\n".format(len(posns), len(seq)))
            #         continue
            #     for i in range(len(posns)):
            #         if posns[i] - start > len(refseq) -1:
            #             sys.stderr.write("Too many positions\n")
            #         if i > len(seq)-1:
            #             sys.stderr.write("Too many seq\n")
            #         refseq[posns[i]-start] = seq[i]
            #
            # print("{}_{}     {}".format(template, primer, ''.join(refseq)))

            alignment = {}
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                for pilups in p.pileups:
                    if pilups.alignment.query_name not in alignment:
                        alignment[pilups.alignment.query_name] = ['-' for idx in range(start, end+1)]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                rp = p.reference_pos
                idx = rp - start
                for pilups in p.pileups:
                    if pilups.query_position:
                        posn = pilups.query_position - start
                        # sys.stderr.write("Posn: {} Q.position: {} start: {} end: {} len: {}\n".format(posn, pilups.query_position, start, end, end-start))
                        alignment[pilups.alignment.query_name][idx] = pilups.alignment.query_sequence[pilups.query_position]

            # don't want to carry on if we have no sequences
            if len(alignment) == 0:
                continue

            # find the longest name
            longest_name = 0
            for n in alignment:
                if len(n) > longest_name:
                    longest_name = len(n)
            longest_name += 5

            # I want to sort by the number of -'s at the beginning of the sequence
            beginning_gaps = {}
            for n in alignment:
                gap = 0
                while (gap < len(alignment[n]) and alignment[n][gap] == '-'):
                    gap += 1
                beginning_gaps[n] = gap

            out = open(os.path.join(output_dir, primer, filename), 'w')

            for n in sorted(alignment.keys(), key=beginning_gaps.get):
                # this is for "pylip style output"
                # out.write(n)
                # out.write(" " * (longest_name - len(n)))
                # out.write(''.join(alignment[n]) + "\n")
                out.write(">{}\n{}\n".format(n, ''.join(alignment[n])))
            out.close()
        #print("\n\n")

def list_sequences(bam):
    """
    List the sequences involved and whether they are forward or reverse

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            print("\nALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            for read in bam.fetch(reference=template, start=start, end=end):
                print("{}\t{}".format(read.query_name, read.is_reverse))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract PCRd regions from BAM files')
    parser.add_argument('-b', help='bam file', required=True)
    parser.add_argument('-q', help='print query regions. This is a fasta output of all sequences that are in the region', action='store_true')
    parser.add_argument('-p', help='print pileup. Prints debug information about each position in the pileup', action='store_true')
    parser.add_argument('-c', help='print consensus sequence. Prints a single sequence for each region', action='store_true')
    parser.add_argument('-a', help='print alignment. Prints an alignment for each region.', action='store_true')
    parser.add_argument('-l', help='list read ids and whether they are reversed', action='store_true')
    parser.add_argument('-d', help='output directory to write the alignments to (required with -a)')
    parser.add_argument('-v', help='verbose output')
    args = parser.parse_args()

    filename = os.path.split(args.b)[-1]
    filename = filename.replace('.bam', '.aln')
    bam = pysam.AlignmentFile(args.b, 'rb')

    if args.q:
        print_query_regions(bam)
    if args.p:
        print_pileup(bam)
    if args.c:
        print_consensus(bam)
    if args.a:
        if not args.d:
            sys.exit("You must provide an output directory with -a flag")
        print_alignment(bam, args.d, filename)
    if args.l:
        list_sequences(bam)
<commit_msg>output to fasta format alignment<commit_after>import argparse
import os
import sys
import pysam

__author__ = 'Rob Edwards'

"""
Extract the PCR regions from indexed bam alignments. These are regions that we have asked people to pull out, and
we will extract the sequences from those regions

Current regions:

Primer sequences:			Position on JQ995537 (the original genbank record)
Primer A:
	Fwd: CTGATAGTATGATTGGTAAT	Position 25634 .. 25653
	Rev: ATAAGTTCTCCAACTATCTT	Position complement(26945 .. 26964)

Primer B:
	Fwd: CCAGTATCTCCATAAGCATC	Position 33709 .. 33728
	Rev: GTGAGGGCGGAATAGCTA		Position complement(35045 .. 35062)

Primer C:
	Fwd: GCAACAGGAGTAGTAAAATCTC	Position 43820 .. 43841
	Rev: GCTCCTGTTAATCCTGATGTTA	Position complement(45036 .. 45057)

"""

locations = {
    'JQ995537': {
        'PrimerA': (25633, 26964),
        'PrimerB': (33708, 35062),
        'PrimerC': (43819, 45057)
    }
}

def print_query_regions(bam):
    """
    Print the regions from the query sequences that overlap our regions of interest

    :param bam: The bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            for read in bam.fetch(reference=template, start=start, end=end):
                # this is an AlignedSegment: http://pysam.readthedocs.org/en/latest/api.html#pysam.AlignedSegment
                # sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #     primer, start, end, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #     read.reference_start, read.reference_end
                # ))

                # this checks for sequences that overlap the start and end (none do in the Ondrej data set
                # if read.reference_start <= start and read.reference_end >= stop:
                #     sys.stderr.write("Primer: {} ({} .. {}). Found a region for {} ({} .. {}) -> ({} .. {})\n".format(
                #         primer, start, stop, read.query_name, read.query_alignment_start, read.query_alignment_end,
                #         read.reference_start, read.reference_end
                #     ))

                # get just the sequence that maps to the region
                seq = read.query_sequence
                beg_offset = None
                end_offset = None
                if read.reference_start < start:
                    beg_offset = start - read.reference_start - 1
                if read.reference_end > end:
                    end_offset = len(seq) - (read.reference_end - end)

                if beg_offset and end_offset:
                    seq = seq[beg_offset:end_offset]
                elif beg_offset:
                    seq = seq[beg_offset:]
                elif end_offset:
                    seq = seq[:end_offset]

                print(">{} {} {} {}\n{}".format(read.query_name, primer, read.reference_start, read.reference_end, seq))

def print_pileup(bam):
    """
    Print the information about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    for template in locations:
        for primer in locations[template]:
            start, end  = locations[template][primer]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                    else:
                        bp = '-'
                    bases[bp] = bases.get(bp, 0) + 1
                sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, str(bases)))

def print_consensus(bam):
    """
    Print the consensus sequence about the pileup

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            cseq = []
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                bases = {}
                for pilups in p.pileups:
                    if pilups.query_position:
                        bp = pilups.alignment.query_sequence[pilups.query_position]
                        bases[bp] = bases.get(bp, 0) + 1
                # if we don't have any bases at this position, add an N
                if not bases:
                    bases['N'] = 1
                bps = sorted(bases, key=bases.get, reverse=True)
                # text = ""
                # for b in bps:
                #     text += " " + b + ": " + str(bases[b])
                # sys.stdout.write("{} : {} -> {}\n".format(p.reference_name, p.reference_pos, text))

                # make the consensus seq
                cseq.append(bps[0])
            print(">{} {} {} {}\n{}".format(primer, template, start, end, ''.join(cseq)))

def print_alignment(bam, output_dir, filename, min_read_length=0, min_alignment_frac=0):
    """
    Print an alignment of all matching pileups

    :param output_dir: The output directory to store the results in. We will make a sub directory for each PCR region
    :type output_dir: str
    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    for template in locations:
        for primer in locations[template]:
            if not os.path.exists(os.path.join(output_dir, primer)):
                os.mkdir(os.path.join(output_dir, primer))
            start, end = locations[template][primer]
            # print("ALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            # This is a failed attempt to get the reference sequence for this region, but I am not sure that this
            # is even possible from a BAM file, since each read will have a unique alignment to the reference
            # refseq = ['-' for i in range(start, end)]
            # for aln in bam.fetch(reference=template, start=start, end=end, until_eof=True):
            #     posns = aln.get_reference_positions()
            #     seq = aln.get_reference_sequence()
            #     if len(posns) > len(seq):
            #         sys.stderr.write("There are more positions {} than sequences {}\n".format(len(posns), len(seq)))
            #         continue
            #     for i in range(len(posns)):
            #         if posns[i] - start > len(refseq) -1:
            #             sys.stderr.write("Too many positions\n")
            #         if i > len(seq)-1:
            #             sys.stderr.write("Too many seq\n")
            #         refseq[posns[i]-start] = seq[i]
            #
            # print("{}_{}     {}".format(template, primer, ''.join(refseq)))

            alignment = {}
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                for pilups in p.pileups:
                    if pilups.alignment.query_name not in alignment:
                        alignment_frac = 1.0 * pilups.alignment.query_alignment_length/ pilups.alignment.query_length
                        if pilups.alignment.query_length > min_read_length and alignment_frac > min_alignment_frac:
                            alignment[pilups.alignment.query_name] = ['-' for idx in range(start, end+1)]
            for p in bam.pileup(reference=template, start=start, end=end, truncate=True):
                rp = p.reference_pos
                idx = rp - start
                for pilups in p.pileups:
                    if pilups.query_position:
                        posn = pilups.query_position - start
                        # sys.stderr.write("Posn: {} Q.position: {} start: {} end: {} len: {}\n".format(posn, pilups.query_position, start, end, end-start))
                        if pilups.alignment.query_name in alignment:
                            # if the read is too short, we don't include it here!
                            alignment[pilups.alignment.query_name][idx] = pilups.alignment.query_sequence[pilups.query_position]

            # don't want to carry on if we have no sequences
            if len(alignment) == 0:
                continue

            # find the longest name
            longest_name = 0
            for n in alignment:
                if len(n) > longest_name:
                    longest_name = len(n)
            longest_name += 5

            # I want to sort by the number of -'s at the beginning of the sequence
            beginning_gaps = {}
            for n in alignment:
                gap = 0
                while (gap < len(alignment[n]) and alignment[n][gap] == '-'):
                    gap += 1
                beginning_gaps[n] = gap

            out = open(os.path.join(output_dir, primer, filename), 'w')

            for n in sorted(alignment.keys(), key=beginning_gaps.get):
                # this is for "pylip style output"
                # out.write(n)
                # out.write(" " * (longest_name - len(n)))
                # out.write(''.join(alignment[n]) + "\n")
                out.write(">{}\n{}\n".format(n, ''.join(alignment[n])))
            out.close()
        #print("\n\n")

def list_sequences(bam):
    """
    List the sequences involved and whether they are forward or reverse

    :param bam: the bam object from pysam
    :type bam: pysam.AlignmentFile
    :return:
    :rtype:
    """
    for template in locations:
        for primer in locations[template]:
            start, end = locations[template][primer]
            print("\nALIGNMENT: {} FROM {} TO {}\n".format(primer, start, end))
            for read in bam.fetch(reference=template, start=start, end=end):
                print("{}\t{}".format(read.query_name, read.is_reverse))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract PCRd regions from BAM files')
    parser.add_argument('-b', help='bam file', required=True)
    parser.add_argument('-q', help='print query regions. This is a fasta output of all sequences that are in the region', action='store_true')
    parser.add_argument('-p', help='print pileup. Prints debug information about each position in the pileup', action='store_true')
    parser.add_argument('-c', help='print consensus sequence. Prints a single sequence for each region', action='store_true')
    parser.add_argument('-a', help='print alignment. Prints an alignment for each region.', action='store_true')
    parser.add_argument('-l', help='list read ids and whether they are reversed', action='store_true')
    parser.add_argument('-d', help='output directory to write the alignments to (required with -a)')
    parser.add_argument('-m', help='minimum alignment length for inclusion in output. Default=0', default=0, type=int)
    parser.add_argument('-f', help='minimum fraction (i.e. <=1) of read length in alignment. Default=0', default=0, type=float)
    parser.add_argument('-v', help='verbose output')
    args = parser.parse_args()

    filename = os.path.split(args.b)[-1]
    filename = filename.replace('.bam', '.aln')
    bam = pysam.AlignmentFile(args.b, 'rb')

    if args.q:
        print_query_regions(bam)
    if args.p:
        print_pileup(bam)
    if args.c:
        print_consensus(bam)
    if args.a:
        if not args.d:
            sys.exit("You must provide an output directory with -a flag")
        print_alignment(bam, args.d, filename, args.m, args.f)
    if args.l:
        list_sequences(bam)
<|endoftext|>