<commit_before>import optparse
from socket import *

def connscan(tgt_host, tgt_port):
    """Attempt to create a connection to the target host and port."""
    try:
        connsckt = socket(AF_INET, SOCK_STREAM)
        connsckt.connect((tgt_host, tgt_port))
        connsckt.send('ViolentPython\r\n')
        results = connsckt.recv(100)
        print('[+]%d/tcp open' % tgt_port)
        print('[+] ' + str(results))
        connsckt.close()
    except:
        print('[-]%d/tcp closed' % tgt_port)

def portscan(tgt_host, tgt_ports):
    """Attempt to resolve an IP Address to a friendly hostname, then print hostname or IP address."""
    try:
        tgtIP = gethostbyname(tgt_host)
    except:
        print("[-] Cannot resolve '%s': Unknown host" % tgt_host)
        return
    try:
        tgtName = gethostbyaddr(tgtIP)
        print('\n[+] Scan Results for: ' + tgtName[0])
    except:
        print('\n[+] Scan Results for: ' + tgtIP)
    setdefaulttimeout(1)
    for tgt_port in tgt_ports:
        print('Scanning port ' + tgt_port)
        connscan(tgt_host, int(tgt_port))

def main():
    """Set up an run above functions."""
    parser = optparse.OptionParser(
        'usage %prog -H' + '<target host> -p <target port>')
    parser.add_option('-H', dest='tgt_host', type='string',
                      help='specify target host')
    parser.add_option('-p', dest='tgt_port', type='int',
                      help='specify target port')
    (options, args) = parser.parse_args()
    tgt_host = options.tgt_host
    tgt_port = options.tgt_port.split(', ')
    if (tgt_host is None) or (tgt_port is None):
        print(parser.usage)
        exit(0)
        portscan(tgt_host, tgt_ports)

if __name__ == '__main__':
    main()<commit_msg>threading, parser needs work.<commit_after>"""A port scanning script from the Violent Python cookbook."""
import optparse
from socket import *
from threading import *

screen_lock = Semaphore(value=1)

def connscan(tgt_host, tgt_port):
    """Attempt to create a connection to the target host and port."""
    try:
        connsckt = socket(AF_INET, SOCK_STREAM)
        connsckt.connect((tgt_host, tgt_port))
        connsckt.send('ViolentPython\r\n')
        results = connsckt.recv(100)
        screen_lock.acquire()
        print('[+]%d/tcp open' % tgt_port)
        print('[+] ' + str(results))
    except:
        screen_lock.acquire()
        print('[-]%d/tcp closed' % tgt_port)
    finally:
        screen_lock.release()
        connsckt.close()

def portscan(tgt_host, tgt_ports):
    """Attempt to resolve an IP Address to a friendly hostname, then print hostname or IP address."""
    try:
        tgt_ip = gethostbyname(tgt_host)
    except:
        print("[-] Cannot resolve '%s': Unknown host" % tgt_host)
        return
    try:
        tgt_name = gethostbyaddr(tgt_ip)
        print('\n[+] Scan Results for: ' + tgt_name[0])
    except:
        print('\n[+] Scan Results for: ' + tgt_ip)
    setdefaulttimeout(1)
    for tgt_port in tgt_ports:
        t = Thread(target=connscan, args=(tgt_host, int(tgt_port)))
        t.start()

def main():
    """Parse -p and -H flags and scan the ip and port associated with it."""
    parser = optparse.OptionParser(
        'usage %prog -H' + '<target host> -p <target port>')
    parser.add_option('-H', dest='tgt_host', type='string',
                      help='specify target host')
    parser.add_option('-p', dest='tgt_port', type='string',
                      help='specify target port')
    (options, args) = parser.parse_args()
    tgt_host = options.tgt_host
    tgt_ports = str(options.tgt_port).split(', ')
    if (tgt_host is None) or (tgt_ports[0] is None):
        print(parser.usage)
        exit(0)
    portscan(tgt_host, tgt_ports)

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


<commit_before>"""Module for SymPy containers

    (SymPy objects that store other SymPy objects)

    The containers implemented in this module are subclassed to Basic.
    They are supposed to work seamlessly within the SymPy framework.
"""

from sympy.core.basic import Basic
from sympy.core.sympify import sympify, converter
from sympy.utilities.iterables import iterable

class Tuple(Basic):
    """
    Wrapper around the builtin tuple object

    The Tuple is a subclass of Basic, so that it works well in the
    SymPy framework.  The wrapped tuple is available as self.args, but
    you can also access elements or slices with [:] syntax.

    >>> from sympy import symbols
    >>> from sympy.core.containers import Tuple
    >>> a, b, c, d = symbols('a b c d')
    >>> Tuple(a, b, c)[1:]
    (b, c)
    >>> Tuple(a, b, c).subs(a, d)
    (d, b, c)

    """

    def __new__(cls, *args, **assumptions):
        args = [ sympify(arg) for arg in args ]
        obj = Basic.__new__(cls, *args, **assumptions)
        return obj

    def __getitem__(self,i):
        if isinstance(i,slice):
            indices = i.indices(len(self))
            return Tuple(*[self.args[i] for i in range(*indices)])
        return self.args[i]

    def __len__(self):
        return len(self.args)

    def __contains__(self, item):
        return item in self.args

    def __iter__(self):
        return iter(self.args)

    def __add__(self, other):
        if isinstance(other, Tuple):
            return Tuple(*(self.args + other.args))
        elif isinstance(other, tuple):
            return Tuple(*(self.args + other))
        else:
            return NotImplemented

    def __radd__(self, other):
        if isinstance(other, Tuple):
            return Tuple(*(other.args + self.args))
        elif isinstance(other, tuple):
            return Tuple(*(other + self.args))
        else:
            return NotImplemented

    def __eq__(self, other):
        if isinstance(other, Basic):
            return super(Tuple, self).__eq__(other)
        return self.args == other

    def __ne__(self, other):
        if isinstance(other, Basic):
            return super(Tuple, self).__ne__(other)
        return self.args != other

    def __hash__(self):
        return hash(self.args)

    def _to_mpmath(self, prec):
        return tuple([a._to_mpmath(prec) for a in self.args])

    def __lt__(self, other):
        return self.args < other.args

    def __le__(self, other):
        return self.args <= other.args

converter[tuple] = lambda tup: Tuple(*tup)

def tuple_wrapper(method):
    """
    Decorator that converts any tuple in the function arguments into a Tuple.

    The motivation for this is to provide simple user interfaces.  The user can
    call a function with regular tuples in the argument, and the wrapper will
    convert them to Tuples before handing them to the function.

    >>> from sympy.core.containers import tuple_wrapper, Tuple
    >>> def f(*args):
    ...    return args
    >>> g = tuple_wrapper(f)

    The decorated function g sees only the Tuple argument:

    >>> g(0, (1, 2), 3)
    (0, (1, 2), 3)

    """
    def wrap_tuples(*args, **kw_args):
        newargs=[]
        for arg in args:
            if type(arg) is tuple:
                newargs.append(Tuple(*arg))
            else:
                newargs.append(arg)
        return method(*newargs, **kw_args)
    return wrap_tuples

class Dict(Basic):
    """
    Wrapper around the builtin dict object

    The Dict is a subclass of Basic, so that it works well in the
    SymPy framework.  Because it is immutable, it may be included
    in sets, but its values must all be given at instantiation and
    cannot be changed afterwards.  Otherwise it behaves identically
    to the Python dict.

    >>> from sympy import S
    >>> from sympy.core.containers import Dict

    >>> D = Dict({1:'one', 2:'two'})
    >>> for key in D: print key, D[key]
    1 one
    2 two

    The args are sympified so the 1 and 2 are Integers and the values
    are Symbols. Queries automatically sympify args so the following work:

    >>> 1 in D
    True
    >>> D.has('one') # searches keys and values
    True
    >>> 'one' in D # not in the keys
    False
    >>> D[1]
    one

    """

    def __new__(cls, *args):
        if len(args)==1 and args[0].__class__ is dict:
            items = [Tuple(k, v) for k, v in args[0].items()]
        elif iterable(args) and all(len(arg) == 2 for arg in args):
            items = [Tuple(k, v) for k, v in args]
        else:
            raise TypeError('Pass Dict args as Dict((k1, v1), ...) or Dict({k1: v1, ...})')
        elements = frozenset(items)
        obj = Basic.__new__(cls, elements)
        obj.elements = elements
        obj._dict = dict(items) # In case Tuple decides it wants to sympify
        return obj

    def __getitem__(self, key):
        """x.__getitem__(y) <==> x[y]"""
        return self._dict[sympify(key)]

    def __setitem__(self, key, value):
        raise NotImplementedError("SymPy Dicts are Immutable")

    @property
    def args(self):
        return tuple(self.elements)

    def items(self):
        '''D.items() -> list of D's (key, value) pairs, as 2-tuples'''
        return self._dict.items()

    def keys(self):
        '''D.keys() -> list of D's keys'''
        return self._dict.keys()

    def values(self):
        '''D.values() -> list of D's values'''
        return self._dict.values()

    def __iter__(self):
        '''x.__iter__() <==> iter(x)'''
        return iter(self._dict)

    def __len__(self):
        '''x.__len__() <==> len(x)'''
        return self._dict.__len__()

    def get(self, key, default=None):
        '''D.get(k[,d]) -> D[k] if k in D, else d.  d defaults to None.'''
        return self._dict.get(sympify(key), default)

    def __contains__(self, key):
        '''D.__contains__(k) -> True if D has a key k, else False'''
        return sympify(key) in self._dict

    def __lt__(self, other):
        return self.args < other.args

class TableForm(Basic):
    """
    Allows to create nice table representation of data.

    Example::

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])

    Then use "print t" to print the table. You can use the SymPy's printing
    system to produce tables in any format (ascii, latex, html, ...).

    """

    def __new__(cls, data, headings=None, alignment="left",
            column_formats=None):
        """
        Creates a TableForm.

        Parameters:

            data ... 2D data to be put into the table
            headings ... gives the labels for entries in each dimension:
                         None ... no labels in any dimension
                         "automatic" ... gives successive integer labels
                         [[l1, l2, ...], ...] gives labels for each entry in
                             each dimension (can be None for some dimension)
            alignment ... "left", "center", "right" (alignment of the columns)

        Example:

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])

        """
        # We only support 2D data. Check the consistency:
        from sympy.core.numbers import Integer
        _w = Integer(len(data[0]))
        _h = Integer(len(data))
        for line in data:
            assert len(line) == _w
        _lines = Tuple(*data)

        if headings is None:
            _headings = [None, None]
        elif headings == "automatic":
            _headings = [range(1, _h + 1), range(1, _w + 1)]
        else:
            h1, h2 = headings
            if h1 == "automatic":
                h1 = range(1, _h + 1)
            if h2 == "automatic":
                h2 = range(1, _w + 1)
            _headings = [h1, h2]

        _alignment = alignment
        if column_formats:
            _column_formats = column_formats
        else:
            _column_formats = [None]*_w

        obj = Basic.__new__(cls, _w, _h, _lines)
        obj._w = _w
        obj._h = _h
        obj._lines = _lines
        obj._headings = _headings
        obj._alignment = _alignment
        obj._column_formats = _column_formats
        return obj

    def as_str(self):
        """
        Returns the string representation of 'self'.

        Example:

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])
        >>> s = t.as_str()

        """
        column_widths = [0] * self._w
        lines = []
        for line in self._lines:
            new_line = []
            for i in range(self._w):
                # Format the item somehow if needed:
                s = str(line[i])
                w = len(s)
                if w > column_widths[i]:
                    column_widths[i] = w
                new_line.append(s)
            lines.append(new_line)

        # Check heading:
        if self._headings[1]:
            new_line = []
            for i in range(self._w):
                # Format the item somehow if needed:
                s = str(self._headings[1][i])
                w = len(s)
                if w > column_widths[i]:
                    column_widths[i] = w
                new_line.append(s)
            self._headings[1] = new_line

        format_str = ""
        for w in column_widths:
            if self._alignment == "left":
                align = "-"
            elif self._alignment == "right":
                align = ""
            else:
                raise NotImplementedError()
            format_str += "%" + align + str(w) + "s "
        format_str += "\n"

        if self._headings[0]:
            self._headings[0] = [str(x) for x in self._headings[0]]
            heading_width = max([len(x) for x in self._headings[0]])
            format_str = "%" + str(heading_width) + "s | " + format_str

        s = ""
        if self._headings[1]:
            d = self._headings[1]
            if self._headings[0]:
                d = [""] + d
            first_line = format_str % tuple(d)
            s += first_line
            s += "-" * (len(first_line) - 2) + "\n"
        for i, line in enumerate(lines):
            d = line
            if self._headings[0]:
                d = [self._headings[0][i]] + d
            s += format_str % tuple(d)
        return s

    def as_latex(self):
        """
        Returns the string representation of 'self'.

        Example:

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])
        >>> s = t.as_latex()

        """
        # Check heading:
        if self._headings[1]:
            new_line = []
            for i in range(self._w):
                # Format the item somehow if needed:
                s = str(self._headings[1][i])
                w = len(s)
                new_line.append(s)
            self._headings[1] = new_line

        if self._headings[0]:
            self._headings[0] = [str(x) for x in self._headings[0]]

        s = r"\begin{tabular}{" + " ".join(["c" for x in self._lines[0]]) + "}\n"
        if self._headings[1]:
            d = self._headings[1]
            if self._headings[0]:
                d = [""] + d
            first_line = " & ".join(d) + r" \\" + "\n"
            s += first_line
            s += r"\hline" + "\n"
        for i, line in enumerate(self._lines):
            d = []
            for j, x in enumerate(line):
                if self._column_formats[j]:
                    d.append(self._column_formats[j] % x)
                else:
                    d.append(str(x))
            if self._headings[0]:
                d = [self._headings[0][i]] + d
            s += " & ".join(d) + r" \\" + "\n"
        s += r"\end{tabular}"
        return s
<commit_msg>TableForm: add alignment="center" support<commit_after>"""Module for SymPy containers

    (SymPy objects that store other SymPy objects)

    The containers implemented in this module are subclassed to Basic.
    They are supposed to work seamlessly within the SymPy framework.
"""

from sympy.core.basic import Basic
from sympy.core.sympify import sympify, converter
from sympy.utilities.iterables import iterable

class Tuple(Basic):
    """
    Wrapper around the builtin tuple object

    The Tuple is a subclass of Basic, so that it works well in the
    SymPy framework.  The wrapped tuple is available as self.args, but
    you can also access elements or slices with [:] syntax.

    >>> from sympy import symbols
    >>> from sympy.core.containers import Tuple
    >>> a, b, c, d = symbols('a b c d')
    >>> Tuple(a, b, c)[1:]
    (b, c)
    >>> Tuple(a, b, c).subs(a, d)
    (d, b, c)

    """

    def __new__(cls, *args, **assumptions):
        args = [ sympify(arg) for arg in args ]
        obj = Basic.__new__(cls, *args, **assumptions)
        return obj

    def __getitem__(self,i):
        if isinstance(i,slice):
            indices = i.indices(len(self))
            return Tuple(*[self.args[i] for i in range(*indices)])
        return self.args[i]

    def __len__(self):
        return len(self.args)

    def __contains__(self, item):
        return item in self.args

    def __iter__(self):
        return iter(self.args)

    def __add__(self, other):
        if isinstance(other, Tuple):
            return Tuple(*(self.args + other.args))
        elif isinstance(other, tuple):
            return Tuple(*(self.args + other))
        else:
            return NotImplemented

    def __radd__(self, other):
        if isinstance(other, Tuple):
            return Tuple(*(other.args + self.args))
        elif isinstance(other, tuple):
            return Tuple(*(other + self.args))
        else:
            return NotImplemented

    def __eq__(self, other):
        if isinstance(other, Basic):
            return super(Tuple, self).__eq__(other)
        return self.args == other

    def __ne__(self, other):
        if isinstance(other, Basic):
            return super(Tuple, self).__ne__(other)
        return self.args != other

    def __hash__(self):
        return hash(self.args)

    def _to_mpmath(self, prec):
        return tuple([a._to_mpmath(prec) for a in self.args])

    def __lt__(self, other):
        return self.args < other.args

    def __le__(self, other):
        return self.args <= other.args

converter[tuple] = lambda tup: Tuple(*tup)

def tuple_wrapper(method):
    """
    Decorator that converts any tuple in the function arguments into a Tuple.

    The motivation for this is to provide simple user interfaces.  The user can
    call a function with regular tuples in the argument, and the wrapper will
    convert them to Tuples before handing them to the function.

    >>> from sympy.core.containers import tuple_wrapper, Tuple
    >>> def f(*args):
    ...    return args
    >>> g = tuple_wrapper(f)

    The decorated function g sees only the Tuple argument:

    >>> g(0, (1, 2), 3)
    (0, (1, 2), 3)

    """
    def wrap_tuples(*args, **kw_args):
        newargs=[]
        for arg in args:
            if type(arg) is tuple:
                newargs.append(Tuple(*arg))
            else:
                newargs.append(arg)
        return method(*newargs, **kw_args)
    return wrap_tuples

class Dict(Basic):
    """
    Wrapper around the builtin dict object

    The Dict is a subclass of Basic, so that it works well in the
    SymPy framework.  Because it is immutable, it may be included
    in sets, but its values must all be given at instantiation and
    cannot be changed afterwards.  Otherwise it behaves identically
    to the Python dict.

    >>> from sympy import S
    >>> from sympy.core.containers import Dict

    >>> D = Dict({1:'one', 2:'two'})
    >>> for key in D: print key, D[key]
    1 one
    2 two

    The args are sympified so the 1 and 2 are Integers and the values
    are Symbols. Queries automatically sympify args so the following work:

    >>> 1 in D
    True
    >>> D.has('one') # searches keys and values
    True
    >>> 'one' in D # not in the keys
    False
    >>> D[1]
    one

    """

    def __new__(cls, *args):
        if len(args)==1 and args[0].__class__ is dict:
            items = [Tuple(k, v) for k, v in args[0].items()]
        elif iterable(args) and all(len(arg) == 2 for arg in args):
            items = [Tuple(k, v) for k, v in args]
        else:
            raise TypeError('Pass Dict args as Dict((k1, v1), ...) or Dict({k1: v1, ...})')
        elements = frozenset(items)
        obj = Basic.__new__(cls, elements)
        obj.elements = elements
        obj._dict = dict(items) # In case Tuple decides it wants to sympify
        return obj

    def __getitem__(self, key):
        """x.__getitem__(y) <==> x[y]"""
        return self._dict[sympify(key)]

    def __setitem__(self, key, value):
        raise NotImplementedError("SymPy Dicts are Immutable")

    @property
    def args(self):
        return tuple(self.elements)

    def items(self):
        '''D.items() -> list of D's (key, value) pairs, as 2-tuples'''
        return self._dict.items()

    def keys(self):
        '''D.keys() -> list of D's keys'''
        return self._dict.keys()

    def values(self):
        '''D.values() -> list of D's values'''
        return self._dict.values()

    def __iter__(self):
        '''x.__iter__() <==> iter(x)'''
        return iter(self._dict)

    def __len__(self):
        '''x.__len__() <==> len(x)'''
        return self._dict.__len__()

    def get(self, key, default=None):
        '''D.get(k[,d]) -> D[k] if k in D, else d.  d defaults to None.'''
        return self._dict.get(sympify(key), default)

    def __contains__(self, key):
        '''D.__contains__(k) -> True if D has a key k, else False'''
        return sympify(key) in self._dict

    def __lt__(self, other):
        return self.args < other.args

class TableForm(Basic):
    """
    Allows to create nice table representation of data.

    Example::

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])

    Then use "print t" to print the table. You can use the SymPy's printing
    system to produce tables in any format (ascii, latex, html, ...).

    """

    def __new__(cls, data, headings=None, alignment="left",
            column_formats=None):
        """
        Creates a TableForm.

        Parameters:

            data ... 2D data to be put into the table
            headings ... gives the labels for entries in each dimension:
                         None ... no labels in any dimension
                         "automatic" ... gives successive integer labels
                         [[l1, l2, ...], ...] gives labels for each entry in
                             each dimension (can be None for some dimension)
            alignment ... alignment of the columns controlled by values;
                          "left" or "<"
                          "center" or "^"
                          "right" or ">"

        Example:

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])

        """
        # We only support 2D data. Check the consistency:
        from sympy.core.numbers import Integer
        _w = Integer(len(data[0]))
        _h = Integer(len(data))
        for line in data:
            assert len(line) == _w
        _lines = Tuple(*data)

        if headings is None:
            _headings = [None, None]
        elif headings == "automatic":
            _headings = [range(1, _h + 1), range(1, _w + 1)]
        else:
            h1, h2 = headings
            if h1 == "automatic":
                h1 = range(1, _h + 1)
            if h2 == "automatic":
                h2 = range(1, _w + 1)
            _headings = [h1, h2]

        _alignment = {'<': 'left', '>': 'right', '^': 'center'}.get(alignment, alignment)
        if column_formats:
            _column_formats = column_formats
        else:
            _column_formats = [None]*_w

        obj = Basic.__new__(cls, _w, _h, _lines)
        obj._w = _w
        obj._h = _h
        obj._lines = _lines
        obj._headings = _headings
        obj._alignment = _alignment
        obj._column_formats = _column_formats
        return obj

    def as_str(self):
        """
        Returns the string representation of 'self'.

        Example:

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])
        >>> s = t.as_str()

        """
        column_widths = [0] * self._w
        lines = []
        for line in self._lines:
            new_line = []
            for i in range(self._w):
                # Format the item somehow if needed:
                s = str(line[i])
                w = len(s)
                if w > column_widths[i]:
                    column_widths[i] = w
                new_line.append(s)
            lines.append(new_line)

        # Check heading:
        if self._headings[1]:
            new_line = []
            for i in range(self._w):
                # Format the item somehow if needed:
                s = str(self._headings[1][i])
                w = len(s)
                if w > column_widths[i]:
                    column_widths[i] = w
                new_line.append(s)
            self._headings[1] = new_line

        format_str = []
        for w in column_widths:
            if self._alignment == "left":
                align = "-"
            elif self._alignment == "right":
                align = ""
            elif self._alignment == "center":
                align = ""
            else:
                raise NotImplementedError()
            format_str += ["%" + align + str(w) + "s"]
        format_str = ' '.join(format_str)
        format_str += "\n"

        if self._headings[0]:
            self._headings[0] = [str(x) for x in self._headings[0]]
            heading_width = max([len(x) for x in self._headings[0]])
            format_str = "%" + str(heading_width) + "s | " + format_str

        s = []
        if self._headings[1]:
            d = self._headings[1]
            if self._headings[0]:
                d = [""] + d
            first_line = format_str % tuple(d)
            s.append(first_line)
            s.append("-" * (len(first_line) - 2) + "\n")
        for i, line in enumerate(lines):
            d = [l if self._alignment != 'center' else
                 l.center(column_widths[j]) for j,l in enumerate(line)]
            if self._headings[0]:
                d = [self._headings[0][i]] + d
            s.append(format_str % tuple(d))
        return ''.join(s)

    def as_latex(self):
        """
        Returns the string representation of 'self'.

        Example:

        >>> from sympy import TableForm
        >>> t = TableForm([[5, 7], [4, 2], [10, 3]])
        >>> s = t.as_latex()

        """
        # Check heading:
        if self._headings[1]:
            new_line = []
            for i in range(self._w):
                # Format the item somehow if needed:
                s = str(self._headings[1][i])
                w = len(s)
                new_line.append(s)
            self._headings[1] = new_line

        if self._headings[0]:
            self._headings[0] = [str(x) for x in self._headings[0]]

        s = r"\begin{tabular}{" + " ".join(["c" for x in self._lines[0]]) + "}\n"
        if self._headings[1]:
            d = self._headings[1]
            if self._headings[0]:
                d = [""] + d
            first_line = " & ".join(d) + r" \\" + "\n"
            s += first_line
            s += r"\hline" + "\n"
        for i, line in enumerate(self._lines):
            d = []
            for j, x in enumerate(line):
                if self._column_formats[j]:
                    d.append(self._column_formats[j] % x)
                else:
                    d.append(str(x))
            if self._headings[0]:
                d = [self._headings[0][i]] + d
            s += " & ".join(d) + r" \\" + "\n"
        s += r"\end{tabular}"
        return s
<|endoftext|>


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

'''
todo

 * testing
     * check why top couple bins don't match in plot
     * check some unnormalized count rows
     * check some normalized count rows
     * check some summary table rows
'''

from bamliquidator_internal.normalize_plot_and_summarize import normalize_plot_and_summarize
from bamliquidator_internal.chromosome_list import chromosomes

import argparse
import os
import subprocess
import tables
import datetime
import csv
from time import time

# creates empty file, overwriting any prior existing files
def create_count_table(h5file):
    class BinCount(tables.IsDescription):
        bin_number = tables.UInt32Col(    pos=0)
        cell_type  = tables.StringCol(16, pos=1)
        chromosome = tables.StringCol(16, pos=2)
        count      = tables.UInt64Col(    pos=3)
        # todo: rename file_name to genome or line
        file_name  = tables.StringCol(64, pos=4)

    table = h5file.create_table("/", "bin_counts", BinCount, "bin counts")

    table.flush()

    return table

def create_regions_table(h5file):
    class Region(tables.IsDescription):
        cell_type   = tables.StringCol(16, pos=0)
        file_name   = tables.StringCol(64, pos=1)
        chromosome  = tables.StringCol(16, pos=2)
        region_name = tables.StringCol(64, pos=3)
        start       = tables.UInt64Col(    pos=4)
        stop        = tables.UInt64Col(    pos=5)
        strand      = tables.StringCol(1,  pos=6)
        count       = tables.UInt64Col(    pos=7)

    table = h5file.create_table("/", "region_counts", Region, "region counts")

    table.flush()

    return table

def all_bam_file_paths_in_directory(bam_directory):
    bam_file_paths = []
    for dirpath, _, files in os.walk(bam_directory, followlinks=True):
        for file_ in files:
            if file_.endswith(".bam"):
                bam_file_paths.append(os.path.join(dirpath, file_))
    return bam_file_paths

def bam_file_paths_with_no_counts(counts, bam_file_paths):
    with_no_counts = []

    for bam_file_path in bam_file_paths:
        count_found = False
        condition = "file_name == '%s'" % os.path.basename(bam_file_path)

        for _ in counts.where(condition):
            count_found = True

        if not count_found:
            with_no_counts.append(bam_file_path)

    return with_no_counts

# returns a tuple of two dictionaries:
# 1) (bam_file_name, chromosome) -> sequence length
# 2) base_bam_file_name -> total mapped count
def lengths_and_total_counts(bam_file_paths):
    file_chromosome_tuple_to_length = {}
    file_to_count = {}

    chr_col         = 0
    length_col      = 1
    mapped_read_col = 2

    for bam_file_path in bam_file_paths:
        args = ["samtools", "idxstats", bam_file_path]
        output = subprocess.check_output(args)
        # skip last two lines: the unmapped chromosome line and the empty line
        reader = csv.reader(output.split('\n')[:-2], delimiter='\t')
        file_name = os.path.basename(bam_file_path)
        file_count = 0
        for row in reader:
            file_count += int(row[mapped_read_col])
            file_chromosome_tuple_to_length[file_name, row[chr_col]] = int(row[length_col])
        file_to_count[file_name] = file_count

    return file_chromosome_tuple_to_length, file_to_count

def liquidate(bam_file_paths, output_directory, file_chromosome_tuple_to_length, file_to_count,
              counts_file_path, executable_path, bin_size = None, region_file = None):

    if (bin_size is not None and region_file is not None) or (bin_size is None and region_file is None):
        sys.exit("either bin_size or region_file must be provided, but not both")

    for i, bam_file_path in enumerate(bam_file_paths):
        print "Liquidating %s (file %d of %d, %s)" % (
            bam_file_path, i+1, len(bam_file_paths), datetime.datetime.now().strftime('%H:%M:%S'))
        cell_type = os.path.basename(os.path.dirname(bam_file_path))

        if bin_size:
            args = [executable_path, cell_type, str(bin_size), bam_file_path, counts_file_path]

            bam_file_name = os.path.basename(bam_file_path)
            for chromosome in chromosomes:
                args.append(chromosome)
                args.append(str(file_chromosome_tuple_to_length[bam_file_name, chromosome]))
        else:
            args = [executable_path, cell_type, region_file, bam_file_path, counts_file_path]

        start = time()
        return_code = subprocess.call(args)
        end = time()
        print "Liquidation completed in %f seconds" % (end - start)

        if return_code != 0:
            print "%s failed with exit code %d" % (executable_path, return_code)
            exit(return_code)

    if bin_size:
        counts_file = tables.open_file(counts_file_path, mode = "r")
        normalize_plot_and_summarize(counts_file.root.bin_counts, output_directory, bin_size, file_to_count)
        counts_file.close()

    # todo: normalize regions

def main():
    parser = argparse.ArgumentParser(description='Count the number of base pair reads in each bin or region '
                                                 'in the bam file(s) at the given directory, and then normalize, plot bins, '
                                                 'and summarize the counts in the output directory.  For additional '
                                                 'help, please see https://github.com/BradnerLab/pipeline/wiki')
    parser.add_argument('--output_directory', default='output',
                        help='Directory to create and output the h5 and/or html files to (aborts if already exists). '
                             'Default is "./output".')
    parser.add_argument('--counts_file',
                        help='HDF5 counts file from a prior run to be appended to.  If unspecified, defaults to '
                             'creating a new file "counts.h5" in the output directory.')
    parser.add_argument('--bin_size', type=int, default=100000,
                        help="Number of base pairs in each bin -- the smaller the bin size the longer the runtime and "
                             "the larger the data files (default is 100000). This argument is ignored if regions are provided.")
    parser.add_argument('ucsc_chrom_sizes_or_regions_file',
                        help='Tab delimited text file, interpretted as a ucsc chromosome size file if ending in .txt and a region file if ending in .gff (todo: support .bed regions as well). A ucsc chromosome size file must have the first column naming the chromosome (e.g. chr1), the third column naming the genome type (e.g. mm8), and the fifth column naming the number of base pairs in the reference chromosome. (todo add description of .gff and .bed files)')
    parser.add_argument('bam_file_path',
                        help='The directory to recursively search for .bam files for counting.  Every .bam file must '
                             'have a corresponding .bai file at the same location.  To count just a single file, '
                             'provide the .bam file path instead of a directory.  The parent directory of each .bam '
                             'file is interpreted as the cell type (e.g. mm1s might be an appropriate directory '
                             'name).  Bam files in the same directory are grouped together for plotting. Plots use '
                             'normalized counts, such that all .bam files in the same directory have bin '
                             'counts that add up to 1 for each chromosome.  The .bam file name is also required to '
                             'contain the genome type so that the corresponding entries in the ucsc_chrom_sizes file '
                             'can be used.  If your .bam files are not in this directory format, please consider '
                             'creating a directory of sym links to your actual .bam and .bai files. If the .bam file '
                             'already has 1 or more reads in the HDF5 counts file, then that .bam file is skipped.')
    args = parser.parse_args()

    assert(tables.__version__ >= '3.0.0')

    _, extension = os.path.splitext(args.ucsc_chrom_sizes_or_regions_file)
    if extension == ".gff":
        region_mode = True
        args.bin_size = None
    else:
        region_mode = False

    executable_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                   "bamliquidator_internal",
                                   "bamliquidator_regions" if region_mode else "bamliquidator_bins")
    if not os.path.isfile(executable_path):
        sys.exit("%s is missing -- try cd'ing into the directory and running 'make'" % executable_path)

    os.mkdir(args.output_directory)

    if args.counts_file is None:
        args.counts_file = os.path.join(args.output_directory, "counts.h5")

        counts_file = tables.open_file(args.counts_file, mode = "w",
                                       title = "bam liquidator genome bin read counts")
    else:
        counts_file = tables.open_file(args.counts_file, "r+")

    try:
        counts = counts_file.root.region_counts if region_mode else counts_file.root.bin_counts
    except:
        counts = create_regions_table(counts_file) if region_mode else create_count_table(counts_file)

    if os.path.isdir(args.bam_file_path):
        bam_file_paths = all_bam_files_in_directory(args.bam_file_path)
    else:
        bam_file_paths = [args.bam_file_path]

    bam_file_paths = bam_file_paths_with_no_counts(counts, bam_file_paths)

    counts_file.close() # bamliquidator_bins/bamliquidator_regions will open this file and
                        # modify it, so it is best that we not hold an out of sync reference

    file_chromosome_tuple_to_length, file_to_count = lengths_and_total_counts(bam_file_paths)

    liquidate(bam_file_paths, args.output_directory, file_chromosome_tuple_to_length, file_to_count,
              args.counts_file, executable_path, args.bin_size)

if __name__ == "__main__":
    main()

'''
   The MIT License (MIT)

   Copyright (c) 2013 John DiMatteo (jdimatteo@gmail.com)

   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
   copies of the Software, and to permit persons to whom the Software is
   furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in
   all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
   THE SOFTWARE.
'''
<commit_msg>Updated region file handling so usc file is no longer required<commit_after>#!/usr/bin/env python

'''
todo

 * testing
     * check why top couple bins don't match in plot
     * check some unnormalized count rows
     * check some normalized count rows
     * check some summary table rows
'''

from bamliquidator_internal.normalize_plot_and_summarize import normalize_plot_and_summarize
from bamliquidator_internal.chromosome_list import chromosomes

import argparse
import os
import subprocess
import tables
import datetime
import csv
from time import time

# creates empty file, overwriting any prior existing files
def create_count_table(h5file):
    class BinCount(tables.IsDescription):
        bin_number = tables.UInt32Col(    pos=0)
        cell_type  = tables.StringCol(16, pos=1)
        chromosome = tables.StringCol(16, pos=2)
        count      = tables.UInt64Col(    pos=3)
        # todo: rename file_name to genome or line
        file_name  = tables.StringCol(64, pos=4)

    table = h5file.create_table("/", "bin_counts", BinCount, "bin counts")

    table.flush()

    return table

def create_regions_table(h5file):
    class Region(tables.IsDescription):
        cell_type   = tables.StringCol(16, pos=0)
        file_name   = tables.StringCol(64, pos=1)
        chromosome  = tables.StringCol(16, pos=2)
        region_name = tables.StringCol(64, pos=3)
        start       = tables.UInt64Col(    pos=4)
        stop        = tables.UInt64Col(    pos=5)
        strand      = tables.StringCol(1,  pos=6)
        count       = tables.UInt64Col(    pos=7)

    table = h5file.create_table("/", "region_counts", Region, "region counts")

    table.flush()

    return table

def all_bam_file_paths_in_directory(bam_directory):
    bam_file_paths = []
    for dirpath, _, files in os.walk(bam_directory, followlinks=True):
        for file_ in files:
            if file_.endswith(".bam"):
                bam_file_paths.append(os.path.join(dirpath, file_))
    return bam_file_paths

def bam_file_paths_with_no_counts(counts, bam_file_paths):
    with_no_counts = []

    for bam_file_path in bam_file_paths:
        count_found = False
        condition = "file_name == '%s'" % os.path.basename(bam_file_path)

        for _ in counts.where(condition):
            count_found = True

        if not count_found:
            with_no_counts.append(bam_file_path)

    return with_no_counts

# returns a tuple of two dictionaries:
# 1) (bam_file_name, chromosome) -> sequence length
# 2) base_bam_file_name -> total mapped count
def lengths_and_total_counts(bam_file_paths):
    file_chromosome_tuple_to_length = {}
    file_to_count = {}

    chr_col         = 0
    length_col      = 1
    mapped_read_col = 2

    for bam_file_path in bam_file_paths:
        args = ["samtools", "idxstats", bam_file_path]
        output = subprocess.check_output(args)
        # skip last two lines: the unmapped chromosome line and the empty line
        reader = csv.reader(output.split('\n')[:-2], delimiter='\t')
        file_name = os.path.basename(bam_file_path)
        file_count = 0
        for row in reader:
            file_count += int(row[mapped_read_col])
            file_chromosome_tuple_to_length[file_name, row[chr_col]] = int(row[length_col])
        file_to_count[file_name] = file_count

    return file_chromosome_tuple_to_length, file_to_count

def liquidate(bam_file_paths, output_directory, file_chromosome_tuple_to_length, file_to_count,
              counts_file_path, executable_path, bin_size = None, region_file = None):

    if (bin_size is not None and region_file is not None) or (bin_size is None and region_file is None):
        sys.exit("either bin_size or region_file must be provided, but not both")

    for i, bam_file_path in enumerate(bam_file_paths):
        print "Liquidating %s (file %d of %d, %s)" % (
            bam_file_path, i+1, len(bam_file_paths), datetime.datetime.now().strftime('%H:%M:%S'))
        cell_type = os.path.basename(os.path.dirname(bam_file_path))

        if bin_size:
            args = [executable_path, cell_type, str(bin_size), bam_file_path, counts_file_path]

            bam_file_name = os.path.basename(bam_file_path)
            for chromosome in chromosomes:
                args.append(chromosome)
                args.append(str(file_chromosome_tuple_to_length[bam_file_name, chromosome]))
        else:
            args = [executable_path, cell_type, region_file, bam_file_path, counts_file_path]

        start = time()
        return_code = subprocess.call(args)
        end = time()
        print "Liquidation completed in %f seconds" % (end - start)

        if return_code != 0:
            print "%s failed with exit code %d" % (executable_path, return_code)
            exit(return_code)

    if bin_size:
        counts_file = tables.open_file(counts_file_path, mode = "r")
        normalize_plot_and_summarize(counts_file.root.bin_counts, output_directory, bin_size, file_to_count)
        counts_file.close()

    # todo: normalize regions

def main():
    parser = argparse.ArgumentParser(description='Count the number of base pair reads in each bin or region '
                                                 'in the bam file(s) at the given directory, and then normalize, plot bins, '
                                                 'and summarize the counts in the output directory.  For additional '
                                                 'help, please see https://github.com/BradnerLab/pipeline/wiki')
    parser.add_argument('--output_directory', default='output',
                        help='Directory to create and output the h5 and/or html files to (aborts if already exists). '
                             'Default is "./output".')
    parser.add_argument('--counts_file',
                        help='HDF5 counts file from a prior run to be appended to.  If unspecified, defaults to '
                             'creating a new file "counts.h5" in the output directory.')
    parser.add_argument('--bin_size', type=int, default=100000,
                        help="Number of base pairs in each bin -- the smaller the bin size the longer the runtime and "
                             "the larger the data files (default is 100000). This argument is ignored if regions are provided.")
    parser.add_argument('--regions_file',
                        help='a region file in either .gff or .bed format -- todo: this is not yet fully supported')
    parser.add_argument('bam_file_path',
                        help='The directory to recursively search for .bam files for counting.  Every .bam file must '
                             'have a corresponding .bai file at the same location.  To count just a single file, '
                             'provide the .bam file path instead of a directory.  The parent directory of each .bam '
                             'file is interpreted as the cell type (e.g. mm1s might be an appropriate directory '
                             'name).  Bam files in the same directory are grouped together for plotting. Plots use '
                             'normalized counts, such that all .bam files in the same directory have bin '
                             'counts that add up to 1 for each chromosome.  The .bam file name is also required to '
                             'contain the genome type so that the corresponding entries in the ucsc_chrom_sizes file '
                             'can be used.  If your .bam files are not in this directory format, please consider '
                             'creating a directory of sym links to your actual .bam and .bai files. If the .bam file '
                             'already has 1 or more reads in the HDF5 counts file, then that .bam file is skipped.')
    args = parser.parse_args()

    assert(tables.__version__ >= '3.0.0')

    if args.regions_file is None:
        region_mode = False
    else:
        region_mode = True
        args.bin_size = None

    executable_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                   "bamliquidator_internal",
                                   "bamliquidator_regions" if region_mode else "bamliquidator_bins")
    if not os.path.isfile(executable_path):
        sys.exit("%s is missing -- try cd'ing into the directory and running 'make'" % executable_path)

    os.mkdir(args.output_directory)

    if args.counts_file is None:
        args.counts_file = os.path.join(args.output_directory, "counts.h5")

        counts_file = tables.open_file(args.counts_file, mode = "w",
                                       title = "bam liquidator genome bin read counts")
    else:
        counts_file = tables.open_file(args.counts_file, "r+")

    try:
        counts = counts_file.root.region_counts if region_mode else counts_file.root.bin_counts
    except:
        counts = create_regions_table(counts_file) if region_mode else create_count_table(counts_file)

    if os.path.isdir(args.bam_file_path):
        bam_file_paths = all_bam_files_in_directory(args.bam_file_path)
    else:
        bam_file_paths = [args.bam_file_path]

    bam_file_paths = bam_file_paths_with_no_counts(counts, bam_file_paths)

    counts_file.close() # bamliquidator_bins/bamliquidator_regions will open this file and
                        # modify it, so it is best that we not hold an out of sync reference

    file_chromosome_tuple_to_length, file_to_count = lengths_and_total_counts(bam_file_paths)

    liquidate(bam_file_paths, args.output_directory, file_chromosome_tuple_to_length, file_to_count,
              args.counts_file, executable_path, args.bin_size, args.regions_file)

if __name__ == "__main__":
    main()

'''
   The MIT License (MIT)

   Copyright (c) 2013 John DiMatteo (jdimatteo@gmail.com)

   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
   copies of the Software, and to permit persons to whom the Software is
   furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in
   all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
   THE SOFTWARE.
'''
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import keyword as kw
import sympy
from repr import ReprPrinter
from str import StrPrinter

# A list of classes that should be printed using StrPrinter
STRPRINT = ("Add", "Infinity", "Integer", "Mul", "NegativeInfinity",
            "Pow", "Zero")

class PythonPrinter(ReprPrinter, StrPrinter):
    """A printer which converts an expression into its Python interpretation."""

    def __init__(self, settings=None):
        ReprPrinter.__init__(self)
        StrPrinter.__init__(self, settings)
        self.symbols = []
        self.functions = []

        # Create print methods for classes that should use StrPrinter instead
        # of ReprPrinter.
        for name in STRPRINT:
            f_name = "_print_%s"%name
            f = getattr(StrPrinter, f_name)
            setattr(PythonPrinter, f_name, f)

    def _print_Function(self, expr):
        func = expr.func.__name__
        if not hasattr(sympy, func) and not func in self.functions:
            self.functions.append(func)
        return StrPrinter._print_Function(self, expr)

    # procedure (!) for defining symbols which have be defined in print_python()
    def _print_Symbol(self, expr):
        symbol = self._str(expr)
        if symbol not in self.symbols:
            self.symbols.append(symbol)
        return StrPrinter._print_Symbol(self, expr)

    def _print_module(self, expr):
        raise ValueError('Modules in the expression are unacceptable')

def python(expr, **settings):
    """Return Python interpretation of passed expression
    (can be passed to the exec() function without any modifications)"""

    printer = PythonPrinter(settings)
    exprp = printer.doprint(expr)

    result = ''
    # Returning found symbols and functions
    renamings = {}
    for symbolname in printer.symbols:
        variablename = symbolname
        if kw.iskeyword(variablename):
            while True:
                variablename += "_"
                if variablename not in printer.symbols:
                    renamings[sympy.Symbol(symbolname)] = sympy.Symbol(variablename)
                    break
        result += variablename + ' = Symbol(\'' + variablename + '\')\n'

    for function in printer.functions:
        result += function + ' = Function(\'' + function + '\')\n'

    exprr = expr.xreplace(renamings)
    result += 'e = ' + printer._str(exprr)
    return result

def print_python(expr, **settings):
    """Print output of python() function"""
    print python(expr, **settings)
<commit_msg>Fix test failures<commit_after># -*- coding: utf-8 -*-

import keyword as kw
import sympy
from repr import ReprPrinter
from str import StrPrinter

# A list of classes that should be printed using StrPrinter
STRPRINT = ("Add", "Infinity", "Integer", "Mul", "NegativeInfinity",
            "Pow", "Zero")

class PythonPrinter(ReprPrinter, StrPrinter):
    """A printer which converts an expression into its Python interpretation."""

    def __init__(self, settings=None):
        ReprPrinter.__init__(self)
        StrPrinter.__init__(self, settings)
        self.symbols = []
        self.functions = []

        # Create print methods for classes that should use StrPrinter instead
        # of ReprPrinter.
        for name in STRPRINT:
            f_name = "_print_%s"%name
            f = getattr(StrPrinter, f_name)
            setattr(PythonPrinter, f_name, f)

    def _print_Function(self, expr):
        func = expr.func.__name__
        if not hasattr(sympy, func) and not func in self.functions:
            self.functions.append(func)
        return StrPrinter._print_Function(self, expr)

    # procedure (!) for defining symbols which have be defined in print_python()
    def _print_Symbol(self, expr):
        symbol = self._str(expr)
        if symbol not in self.symbols:
            self.symbols.append(symbol)
        return StrPrinter._print_Symbol(self, expr)

    def _print_module(self, expr):
        raise ValueError('Modules in the expression are unacceptable')

def python(expr, **settings):
    """Return Python interpretation of passed expression
    (can be passed to the exec() function without any modifications)"""

    printer = PythonPrinter(settings)
    exprp = printer.doprint(expr)

    result = ''
    # Returning found symbols and functions
    renamings = {}
    for symbolname in printer.symbols:
        variablename = symbolname
        if kw.iskeyword(variablename):
            while True:
                variablename += "_"
                if variablename not in printer.symbols:
                    renamings[sympy.Symbol(symbolname)] = sympy.Symbol(variablename)
                    break
        result += variablename + ' = Symbol(\'' + variablename + '\')\n'

    for function in printer.functions:
        result += function + ' = Function(\'' + function + '\')\n'

    if not len(renamings) == 0:
        exprp = expr.xreplace(renamings)
    result += 'e = ' + printer._str(exprp)
    return result

def print_python(expr, **settings):
    """Print output of python() function"""
    print python(expr, **settings)
<|endoftext|>


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

import logging

import base
from kcaa import screens
from kcaa import kcsapi

logger = logging.getLogger('kcaa.manipulators.rebuilding')

def compute_gain_with_bonus(gain):
    return gain + (gain + 1) / 5

def compute_rebuilding_gain(material_ships):
    firepower = sum(s.rebuilding_material.firepower for s in material_ships)
    thunderstroke = sum(s.rebuilding_material.thunderstroke for s in
                        material_ships)
    anti_air = sum(s.rebuilding_material.anti_air for s in material_ships)
    armor = sum(s.rebuilding_material.armor for s in material_ships)
    return kcsapi.AbilityEnhancement(
        firepower=compute_gain_with_bonus(firepower),
        thunderstroke=compute_gain_with_bonus(thunderstroke),
        anti_air=compute_gain_with_bonus(anti_air),
        armor=compute_gain_with_bonus(armor))

def compute_gain_cap(target_ship):
    return kcsapi.AbilityEnhancement(
        firepower=(target_ship.firepower.maximum -
                   (target_ship.firepower.baseline +
                    target_ship.enhanced_ability.firepower)),
        thunderstroke=(target_ship.thunderstroke.maximum -
                       (target_ship.thunderstroke.baseline +
                        target_ship.enhanced_ability.thunderstroke)),
        anti_air=(target_ship.anti_air.maximum -
                  (target_ship.anti_air.baseline +
                   target_ship.enhanced_ability.anti_air)),
        armor=(target_ship.armor.maximum -
               (target_ship.armor.baseline +
                target_ship.enhanced_ability.armor)))

def can_enhance(gain_cap, material_pool):
    return ((gain_cap.firepower > 0 and material_pool.firepower > 0) or
            (gain_cap.thunderstroke > 0 and material_pool.thunderstroke > 0) or
            (gain_cap.anti_air > 0 and material_pool.anti_air > 0) or
            (gain_cap.armor > 0 and material_pool.armor > 0))

def compute_capped_gain(gain, gain_cap):
    return kcsapi.AbilityEnhancement(
        firepower=min(gain.firepower, gain_cap.firepower),
        thunderstroke=min(gain.thunderstroke, gain_cap.thunderstroke),
        anti_air=min(gain.anti_air, gain_cap.anti_air),
        armor=min(gain.armor, gain_cap.armor))

def has_improvement(gain_a, gain_b):
    return (gain_a.firepower > gain_b.firepower or
            gain_a.thunderstroke > gain_b.thunderstroke or
            gain_a.anti_air > gain_b.anti_air or
            gain_a.armor > gain_b.armor)

class RebuildShip(base.Manipulator):

    def run(self, target_ship_id, material_ship_ids):
        target_ship_id = int(target_ship_id)
        if not isinstance(material_ship_ids, list):
            material_ship_ids = [int(ship_id) for ship_id in
                                 material_ship_ids.split(',')]
        if len(material_ship_ids) == 0 or len(material_ship_ids) > 5:
            logger.error('Needs 1 - 5 material ships.')
            return
        ship_list = self.objects.get('ShipList')
        if not ship_list:
            logger.error('No ship list was found. Giving up.')
            return
        fleet_list = self.objects.get('FleetList')
        if not fleet_list:
            logger.error('No fleet list was found. Giving up.')
            return
        target_ship = ship_list.ships[str(target_ship_id)]
        material_ships = [ship_list.ships[str(ship_id)] for ship_id in
                          material_ship_ids]
        logger.info('Trying to rebuild {} with material ships {}.'.format(
            target_ship.name.encode('utf8'),
            ', '.join([s.name.encode('utf8') for s in material_ships])))
        logger.info('Expected gain: {}'.format(
            compute_rebuilding_gain(material_ships).json()))
        if not target_ship.locked:
            logger.error('Target ship is not locked.')
            return
        if target_ship.is_under_repair:
            logger.error('Target ship is under repair.')
            return
        fleet = fleet_list.find_fleet_for_ship(target_ship.id)
        if fleet and fleet.mission_id:
            logger.error('Target ship is undertaking a mission.')
        for material_ship in material_ships:
            name = material_ship.name.encode('utf8')
            if material_ship.locked:
                logger.error('Material ship {} is locked.'.format(name))
                return
            fleet = fleet_list.find_fleet_for_ship(material_ship.id)
            if fleet and fleet.mission_id:
                logger.error('Material ship {} is undertaking a mission.'
                             .format(name))
                return
            if ship_list.is_unique(material_ship):
                logger.error('Material ship {} is unique.'.format(name))
                return
        yield self.screen.change_screen(screens.PORT_REBUILDING)
        fleet = fleet_list.find_fleet_for_ship(target_ship.id)
        if fleet:
            yield self.screen.select_fleet(fleet.id)
            yield self.screen.select_fleet_ship(
                fleet.ship_ids.index(target_ship_id))
        else:
            # First select the first fleet to cancel the effect of page
            # skipping of the last attempt.
            yield self.screen.select_fleet(1)
            yield self.screen.select_ship_list()
            page, in_page_index = (
                ship_list.get_ship_position_rebuilding_target(
                    target_ship_id, fleet_list))
            yield self.screen.select_page(page)
            yield self.screen.select_ship(in_page_index)
        yield self.screen.try_rebuilding()
        for i, material_ship in enumerate(material_ships):
            yield self.screen.select_slot(i)
            max_page = ship_list.max_page_rebuilding(material_ship_ids[:i])
            page, in_page_index = ship_list.get_ship_position_rebuilding(
                material_ship.id, material_ship_ids[:i])
            yield self.screen.select_material_page(page, max_page)
            yield self.screen.select_material_ship(in_page_index)
        yield self.screen.finalyze_rebuilding()
        yield self.screen.confirm_rebuilding()
        yield self.screen.check_rebuilding_result()
        yield self.screen.cancel()

class EnhanceBestShip(base.Manipulator):

    def run(self):
        logger.info(
            'Trying to enhance the best ship which has the room to grow.')
        ship_list = self.objects.get('ShipList')
        if not ship_list:
            logger.error('No ship list was found. Giving up.')
            return
        fleet_list = self.objects.get('FleetList')
        if not fleet_list:
            logger.error('No fleet list was found. Giving up.')
            return
        target_candidates = sorted(
            ship_list.rebuilding_enhanceable_ships(fleet_list),
            kcsapi.ShipSorter.kancolle_level, reverse=True)
        material_candidates = sorted(
            ship_list.rebuilding_available_material_ships(fleet_list),
            kcsapi.ShipSorter.rebuilding_rank)
        material_pool = compute_rebuilding_gain(material_candidates)
        logger.debug('Material pool: {}'.format(material_pool.json()))
        for target_ship in target_candidates:
            gain_cap = compute_gain_cap(target_ship)
            if not can_enhance(gain_cap, material_pool):
                continue
            material_ships = []
            last_gain = compute_rebuilding_gain([])
            for material_ship in material_candidates:
                gain = compute_rebuilding_gain(
                    material_ships + [material_ship])
                capped_gain = compute_capped_gain(gain, gain_cap)
                # There should be additional improvement.
                # Even that's true, anti air and firepower enhancements are
                # relatively rare. Try not to exceed them.
                gain_no_waste = (gain.firepower <= gain_cap.firepower and
                                 gain.anti_air <= gain_cap.anti_air)
                if has_improvement(capped_gain, last_gain) and gain_no_waste:
                    material_ships.append(material_ship)
                    last_gain = capped_gain
                # If the least useful material becomes a mere waste due to the
                # newly added material, remove it from the list.
                while len(material_ships) > 1:
                    gain_r = compute_rebuilding_gain(material_ships[1:])
                    capped_gain_r = compute_capped_gain(gain_r, gain_cap)
                    if has_improvement(capped_gain, capped_gain_r):
                        break
                    # Now the first material ship turned to be a mere waste;
                    # there is no drop in gain even it is removed.
                    del material_ships[0]
                if len(material_ships) == 5:
                    break
            else:
                # Using less than 5 ships is considered "mottainai".
                # It may be acceptable when the ship is reaching the enhance
                # limit.
                if (len(material_ships) < 5 and
                        has_improvement(gain_cap, last_gain)):
                    continue
            logger.info('{} has the room to grow: {}'.format(
                target_ship.name.encode('utf8'), gain_cap.json()))
            logger.info('Expected capped gain: {}'.format(last_gain.json()))
            yield self.do_manipulator(RebuildShip, target_ship.id,
                                      [s.id for s in material_ships])
            return
        else:
            logger.error('No ship has the room to grow.')

class AutoEnhanceBestShip(base.AutoManipulator):

    @classmethod
    def monitored_objects(cls):
        return ['ShipList', 'PlayerInfo']

    @classmethod
    def can_trigger(cls, owner):
        if (owner.screen_id != screens.PORT_MAIN and
                owner.screen_id != screens.PORT_REBUILDING):
            return
        ship_list = owner.objects.get('ShipList')
        player_info = owner.objects.get('PlayerInfo')
        if player_info.max_ships - len(ship_list.ships) < 5:
            return {}

    def run(self):
        yield self.do_manipulator(EnhanceBestShip)
<commit_msg>Fix the logic to detect the cap reach.<commit_after>#!/usr/bin/env python

import logging

import base
from kcaa import screens
from kcaa import kcsapi

logger = logging.getLogger('kcaa.manipulators.rebuilding')

def compute_gain_with_bonus(gain):
    return gain + (gain + 1) / 5

def compute_rebuilding_gain(material_ships):
    firepower = sum(s.rebuilding_material.firepower for s in material_ships)
    thunderstroke = sum(s.rebuilding_material.thunderstroke for s in
                        material_ships)
    anti_air = sum(s.rebuilding_material.anti_air for s in material_ships)
    armor = sum(s.rebuilding_material.armor for s in material_ships)
    return kcsapi.AbilityEnhancement(
        firepower=compute_gain_with_bonus(firepower),
        thunderstroke=compute_gain_with_bonus(thunderstroke),
        anti_air=compute_gain_with_bonus(anti_air),
        armor=compute_gain_with_bonus(armor))

def compute_gain_cap(target_ship):
    return kcsapi.AbilityEnhancement(
        firepower=(target_ship.firepower.maximum -
                   (target_ship.firepower.baseline +
                    target_ship.enhanced_ability.firepower)),
        thunderstroke=(target_ship.thunderstroke.maximum -
                       (target_ship.thunderstroke.baseline +
                        target_ship.enhanced_ability.thunderstroke)),
        anti_air=(target_ship.anti_air.maximum -
                  (target_ship.anti_air.baseline +
                   target_ship.enhanced_ability.anti_air)),
        armor=(target_ship.armor.maximum -
               (target_ship.armor.baseline +
                target_ship.enhanced_ability.armor)))

def can_enhance(gain_cap, material_pool):
    return ((gain_cap.firepower > 0 and material_pool.firepower > 0) or
            (gain_cap.thunderstroke > 0 and material_pool.thunderstroke > 0) or
            (gain_cap.anti_air > 0 and material_pool.anti_air > 0) or
            (gain_cap.armor > 0 and material_pool.armor > 0))

def compute_capped_gain(gain, gain_cap):
    return kcsapi.AbilityEnhancement(
        firepower=min(gain.firepower, gain_cap.firepower),
        thunderstroke=min(gain.thunderstroke, gain_cap.thunderstroke),
        anti_air=min(gain.anti_air, gain_cap.anti_air),
        armor=min(gain.armor, gain_cap.armor))

def has_improvement(gain_a, gain_b):
    """Returns True if gain_a has at least 1 improvement in any parameter
    compared to gain_b."""
    return (gain_a.firepower > gain_b.firepower or
            gain_a.thunderstroke > gain_b.thunderstroke or
            gain_a.anti_air > gain_b.anti_air or
            gain_a.armor > gain_b.armor)

def reached_cap(gain_cap, gain):
    """Returns True if the given gain has at least 1 non-zero gain that reaches
    the capped gain."""
    return ((gain.firepower > 0 and gain.firepower == gain_cap.firepower) or
            (gain.thunderstroke > 0 and
                gain.thunderstroke == gain_cap.thunderstroke) or
            (gain.anti_air > 0 and gain.anti_air == gain_cap.anti_air) or
            (gain.armor > 0 and gain.armor == gain_cap.armor))

class RebuildShip(base.Manipulator):

    def run(self, target_ship_id, material_ship_ids):
        target_ship_id = int(target_ship_id)
        if not isinstance(material_ship_ids, list):
            material_ship_ids = [int(ship_id) for ship_id in
                                 material_ship_ids.split(',')]
        if len(material_ship_ids) == 0 or len(material_ship_ids) > 5:
            logger.error('Needs 1 - 5 material ships.')
            return
        ship_list = self.objects.get('ShipList')
        if not ship_list:
            logger.error('No ship list was found. Giving up.')
            return
        fleet_list = self.objects.get('FleetList')
        if not fleet_list:
            logger.error('No fleet list was found. Giving up.')
            return
        target_ship = ship_list.ships[str(target_ship_id)]
        material_ships = [ship_list.ships[str(ship_id)] for ship_id in
                          material_ship_ids]
        logger.info('Trying to rebuild {} with material ships {}.'.format(
            target_ship.name.encode('utf8'),
            ', '.join([s.name.encode('utf8') for s in material_ships])))
        logger.info('Expected gain: {}'.format(
            compute_rebuilding_gain(material_ships).json()))
        if not target_ship.locked:
            logger.error('Target ship is not locked.')
            return
        if target_ship.is_under_repair:
            logger.error('Target ship is under repair.')
            return
        fleet = fleet_list.find_fleet_for_ship(target_ship.id)
        if fleet and fleet.mission_id:
            logger.error('Target ship is undertaking a mission.')
        for material_ship in material_ships:
            name = material_ship.name.encode('utf8')
            if material_ship.locked:
                logger.error('Material ship {} is locked.'.format(name))
                return
            fleet = fleet_list.find_fleet_for_ship(material_ship.id)
            if fleet and fleet.mission_id:
                logger.error('Material ship {} is undertaking a mission.'
                             .format(name))
                return
            if ship_list.is_unique(material_ship):
                logger.error('Material ship {} is unique.'.format(name))
                return
        yield self.screen.change_screen(screens.PORT_REBUILDING)
        fleet = fleet_list.find_fleet_for_ship(target_ship.id)
        if fleet:
            yield self.screen.select_fleet(fleet.id)
            yield self.screen.select_fleet_ship(
                fleet.ship_ids.index(target_ship_id))
        else:
            # First select the first fleet to cancel the effect of page
            # skipping of the last attempt.
            yield self.screen.select_fleet(1)
            yield self.screen.select_ship_list()
            page, in_page_index = (
                ship_list.get_ship_position_rebuilding_target(
                    target_ship_id, fleet_list))
            yield self.screen.select_page(page)
            yield self.screen.select_ship(in_page_index)
        yield self.screen.try_rebuilding()
        for i, material_ship in enumerate(material_ships):
            yield self.screen.select_slot(i)
            max_page = ship_list.max_page_rebuilding(material_ship_ids[:i])
            page, in_page_index = ship_list.get_ship_position_rebuilding(
                material_ship.id, material_ship_ids[:i])
            yield self.screen.select_material_page(page, max_page)
            yield self.screen.select_material_ship(in_page_index)
        yield self.screen.finalyze_rebuilding()
        yield self.screen.confirm_rebuilding()
        yield self.screen.check_rebuilding_result()
        yield self.screen.cancel()

class EnhanceBestShip(base.Manipulator):

    def run(self):
        logger.info(
            'Trying to enhance the best ship which has the room to grow.')
        ship_list = self.objects.get('ShipList')
        if not ship_list:
            logger.error('No ship list was found. Giving up.')
            return
        fleet_list = self.objects.get('FleetList')
        if not fleet_list:
            logger.error('No fleet list was found. Giving up.')
            return
        target_candidates = sorted(
            ship_list.rebuilding_enhanceable_ships(fleet_list),
            kcsapi.ShipSorter.kancolle_level, reverse=True)
        material_candidates = sorted(
            ship_list.rebuilding_available_material_ships(fleet_list),
            kcsapi.ShipSorter.rebuilding_rank)
        material_pool = compute_rebuilding_gain(material_candidates)
        logger.debug('Material pool: {}'.format(material_pool.json()))
        for target_ship in target_candidates:
            gain_cap = compute_gain_cap(target_ship)
            if not can_enhance(gain_cap, material_pool):
                continue
            material_ships = []
            last_gain = compute_rebuilding_gain([])
            for material_ship in material_candidates:
                gain = compute_rebuilding_gain(
                    material_ships + [material_ship])
                capped_gain = compute_capped_gain(gain, gain_cap)
                # There should be additional improvement.
                # Even that's true, anti air and firepower enhancements are
                # relatively rare. Try not to exceed them.
                gain_no_waste = (gain.firepower <= gain_cap.firepower and
                                 gain.anti_air <= gain_cap.anti_air)
                if has_improvement(capped_gain, last_gain) and gain_no_waste:
                    material_ships.append(material_ship)
                    last_gain = capped_gain
                # If the least useful material becomes a mere waste due to the
                # newly added material, remove it from the list.
                while len(material_ships) > 1:
                    gain_r = compute_rebuilding_gain(material_ships[1:])
                    capped_gain_r = compute_capped_gain(gain_r, gain_cap)
                    if has_improvement(capped_gain, capped_gain_r):
                        break
                    # Now the first material ship turned to be a mere waste;
                    # there is no drop in gain even it is removed.
                    del material_ships[0]
                if len(material_ships) == 5:
                    break
            else:
                # Using less than 5 ships is considered "mottainai".
                # It may be acceptable when the ship is reaching the enhance
                # limit.
                if (len(material_ships) < 5 and
                        not reached_cap(gain_cap, last_gain)):
                    continue
            logger.info('{} has the room to grow: {}'.format(
                target_ship.name.encode('utf8'), gain_cap.json()))
            logger.info('Expected capped gain: {}'.format(last_gain.json()))
            yield self.do_manipulator(RebuildShip, target_ship.id,
                                      [s.id for s in material_ships])
            return
        else:
            logger.error('No ship has the room to grow.')

class AutoEnhanceBestShip(base.AutoManipulator):

    @classmethod
    def monitored_objects(cls):
        return ['ShipList', 'PlayerInfo']

    @classmethod
    def can_trigger(cls, owner):
        if (owner.screen_id != screens.PORT_MAIN and
                owner.screen_id != screens.PORT_REBUILDING):
            return
        ship_list = owner.objects.get('ShipList')
        player_info = owner.objects.get('PlayerInfo')
        if player_info.max_ships - len(ship_list.ships) < 5:
            return {}

    def run(self):
        yield self.do_manipulator(EnhanceBestShip)
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from collections import namedtuple

import six

import onnx.checker

class DeviceType(object):
    CPU = 0
    CUDA = 1

class Device(object):
    '''
    Describes device type and device id
    syntax: device_type:device_id(optional)
    example: 'CPU', 'CUDA', 'CUDA:1'
    '''
    def __init__(self, device):
        options = device.split(':')
        self.type = getattr(DeviceType, options[0])
        self.device_id = 0
        if len(options) > 1:
            self.device_id = int(options[1])

def namedtupledict(typename, field_names, *args, **kwargs):
    field_names_map = {n: i for i, n in enumerate(field_names)}
    # Some output names are invalid python identifier, e.g. "0"
    kwargs.setdefault('rename', True)
    data = namedtuple(typename, field_names, *args, **kwargs)
    def getitem(self, key):
        if isinstance(key, six.string_types):
            key = field_names_map[key]
        return super(type(self), self).__getitem__(key)
    data.__getitem__ = getitem
    return data

class BackendRep(object):
    def run(self, inputs, **kwargs):
        pass

class Backend(object):
    @classmethod
    def prepare(cls, model, device='CPU', **kwargs):
        onnx.checker.check_model(model)

    @classmethod
    def run_model(cls, model, inputs, device='CPU', **kwargs):
        cls.prepare(model, device, **kwargs).run(inputs)

    @classmethod
    def run_node(cls, node, inputs, device='CPU'):
        onnx.checker.check_node(node)

    @classmethod
    def supports_device(device):
        """
        Checks whether the backend is compiled with particular device support.
        In particular it's used in the testing suite.
        """
        return True
<commit_msg>Fix return result in backend base class run_model<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from collections import namedtuple

import six

import onnx.checker

class DeviceType(object):
    CPU = 0
    CUDA = 1

class Device(object):
    '''
    Describes device type and device id
    syntax: device_type:device_id(optional)
    example: 'CPU', 'CUDA', 'CUDA:1'
    '''
    def __init__(self, device):
        options = device.split(':')
        self.type = getattr(DeviceType, options[0])
        self.device_id = 0
        if len(options) > 1:
            self.device_id = int(options[1])

def namedtupledict(typename, field_names, *args, **kwargs):
    field_names_map = {n: i for i, n in enumerate(field_names)}
    # Some output names are invalid python identifier, e.g. "0"
    kwargs.setdefault('rename', True)
    data = namedtuple(typename, field_names, *args, **kwargs)
    def getitem(self, key):
        if isinstance(key, six.string_types):
            key = field_names_map[key]
        return super(type(self), self).__getitem__(key)
    data.__getitem__ = getitem
    return data

class BackendRep(object):
    def run(self, inputs, **kwargs):
        pass

class Backend(object):
    @classmethod
    def prepare(cls, model, device='CPU', **kwargs):
        onnx.checker.check_model(model)

    @classmethod
    def run_model(cls, model, inputs, device='CPU', **kwargs):
        return cls.prepare(model, device, **kwargs).run(inputs)

    @classmethod
    def run_node(cls, node, inputs, device='CPU'):
        onnx.checker.check_node(node)

    @classmethod
    def supports_device(device):
        """
        Checks whether the backend is compiled with particular device support.
        In particular it's used in the testing suite.
        """
        return True
<|endoftext|>