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

# Copyright (C) 2017-2020 The btclib developers
#
# This file is part of btclib. It is subject to the license terms in the
# LICENSE file found in the top-level directory of this distribution.
#
# No part of btclib including this file, may be copied, modified, propagated,
# or distributed except according to the terms contained in the LICENSE file.

"""Bitcoin Script.

https://en.bitcoin.it/wiki/Script

Scripts are represented by List[Token], where Token = Union[int, str, bytes]:

* int [-1, 16] are shorcuts for 'OP_1NEGATE', 'OP_0' - 'OP_16'
* str are for opcodes (e.g. 'OP_HASH160') or hexstring data
* bytes are for data (but integers are often casted to int)
"""

from io import BytesIO
from typing import List

from . import varint
from .alias import BinaryData, Octets, Token
from .utils import bytes_from_octets, bytesio_from_binarydata

SIGHASH_ALL = 1
SIGHASH_NONE = 2
SIGHASH_SINGLE = 3
SIGHASH_ANYONECANPAY = 0x80

SIGHASHES = [
    SIGHASH_ALL,
    SIGHASH_NONE,
    SIGHASH_SINGLE,
    SIGHASH_ANYONECANPAY + SIGHASH_ALL,
    SIGHASH_ANYONECANPAY + SIGHASH_NONE,
    SIGHASH_ANYONECANPAY + SIGHASH_SINGLE,
]

OP_CODES = {
    # Constants
    "OP_0": b"\x00",
    "OP_FALSE": b"\x00",
    "OP_PUSHDATA1": b"\x4c",
    "OP_PUSHDATA2": b"\x4d",
    "OP_PUSHDATA4": b"\x4e",
    "OP_1NEGATE": b"\x4f",
    "OP_1": b"\x51",
    "OP_TRUE": b"\x51",
    "OP_2": b"\x52",
    "OP_3": b"\x53",
    "OP_4": b"\x54",
    "OP_5": b"\x55",
    "OP_6": b"\x56",
    "OP_7": b"\x57",
    "OP_8": b"\x58",
    "OP_9": b"\x59",
    "OP_10": b"\x5a",
    "OP_11": b"\x5b",
    "OP_12": b"\x5c",
    "OP_13": b"\x5d",
    "OP_14": b"\x5e",
    "OP_15": b"\x5f",
    "OP_16": b"\x60",
    # Flow control
    "OP_NOP": b"\x61",
    "OP_IF": b"\x63",
    "OP_NOTIF": b"\x64",
    "OP_ELSE": b"\x67",
    "OP_ENDIF": b"\x68",
    "OP_VERIFY": b"\x69",
    "OP_RETURN": b"\x6a",
    # Stack
    "OP_TOALTSTACK": b"\x6b",
    "OP_FROMALTSTACK": b"\x6c",
    "OP_IFDUP": b"\x73",
    "OP_DEPTH": b"\x74",
    "OP_DROP": b"\x75",
    "OP_DUP": b"\x76",
    "OP_NIP": b"\x77",
    "OP_OVER": b"\x78",
    "OP_PICK": b"\x79",
    "OP_ROLL": b"\x7a",
    "OP_ROT": b"\x7b",
    "OP_SWAP": b"\x7c",
    "OP_TUCK": b"\x7d",
    "OP_2DROP": b"\x6d",
    "OP_2DUP": b"\x6e",
    "OP_3DUP": b"\x6f",
    "OP_2OVER": b"\x70",
    "OP_2ROT": b"\x71",
    "OP_2SWAP": b"\x72",
    # Splice
    "OP_SIZE": b"\x82",
    # Bitwise logic
    "OP_EQUAL": b"\x87",
    "OP_EQUALVERIFY": b"\x88",
    # Arithmetic
    "OP_1ADD": b"\x8b",  # without OP_, 1ADD would be a number
    "OP_1SUB": b"\x8c",
    "OP_NEGATE": b"\x8f",
    "OP_ABS": b"\x90",
    "OP_NOT": b"\x91",
    "OP_0NOTEQUAL": b"\x92",
    "OP_ADD": b"\x93",
    "OP_SUB": b"\x94",
    "OP_BOOLAND": b"\x9a",
    "OP_BOOLOR": b"\x9b",
    "OP_NUMEQUAL": b"\x9c",
    "OP_NUMEQUALVERIFY": b"\x9d",
    "OP_NUMNOTEQUAL": b"\x9e",
    "OP_LESSTHAN": b"\x9f",
    "OP_GREATERTHAN": b"\xa0",
    "OP_LESSTHANOREQUAL": b"\xa1",
    "OP_GREATERTHANOREQUAL": b"\xa2",
    "OP_MIN": b"\xa3",
    "OP_MAX": b"\xa4",
    "OP_WITHIN": b"\xa5",
    # Crypto
    "OP_RIPEMD160": b"\xa6",
    "OP_SHA1": b"\xa7",
    "OP_SHA256": b"\xa8",
    "OP_HASH160": b"\xa9",
    "OP_HASH256": b"\xaa",
    "OP_CODESEPARATOR": b"\xab",
    "OP_CHECKSIG": b"\xac",
    "OP_CHECKSIGVERIFY": b"\xad",
    "OP_CHECKMULTISIG": b"\xae",
    "OP_CHECKMULTISIGVERIFY": b"\xaf",
    # Locktime
    "OP_NOP2": b"\xb1",
    "OP_CHECKLOCKTIMEVERIFY": b"\xb1",
    "OP_NOP3": b"\xb2",
    "OP_CHECKSEQUENCEVERIFY": b"\xb2",
    # Reserved words
    "OP_NOP1": b"\xb0",
    "OP_NOP4": b"\xb3",
    "OP_NOP5": b"\xb4",
    "OP_NOP6": b"\xb5",
    "OP_NOP7": b"\xb6",
    "OP_NOP8": b"\xb7",
    "OP_NOP9": b"\xb8",
    "OP_NOP10": b"\xb9",
}

OP_CODE_NAMES = {
    0: "OP_0",
    76: "OP_PUSHDATA1",
    77: "OP_PUSHDATA2",
    78: "OP_PUSHDATA4",
    79: "OP_1NEGATE",
    81: "OP_1",
    82: "OP_2",
    83: "OP_3",
    84: "OP_4",
    85: "OP_5",
    86: "OP_6",
    87: "OP_7",
    88: "OP_8",
    89: "OP_9",
    90: "OP_10",
    91: "OP_11",
    92: "OP_12",
    93: "OP_13",
    94: "OP_14",
    95: "OP_15",
    96: "OP_16",
    97: "OP_NOP",
    99: "OP_IF",
    100: "OP_NOTIF",
    103: "OP_ELSE",
    104: "OP_ENDIF",
    105: "OP_VERIFY",
    106: "OP_RETURN",
    107: "OP_TOALTSTACK",
    108: "OP_FROMALTSTACK",
    109: "OP_2DROP",
    110: "OP_2DUP",
    111: "OP_3DUP",
    112: "OP_2OVER",
    113: "OP_2ROT",
    114: "OP_2SWAP",
    115: "OP_IFDUP",
    116: "OP_DEPTH",
    117: "OP_DROP",
    118: "OP_DUP",
    119: "OP_NIP",
    120: "OP_OVER",
    121: "OP_PICK",
    122: "OP_ROLL",
    123: "OP_ROT",
    124: "OP_SWAP",
    125: "OP_TUCK",
    130: "OP_SIZE",
    135: "OP_EQUAL",
    136: "OP_EQUALVERIFY",
    139: "OP_1ADD",
    140: "OP_1SUB",
    143: "OP_NEGATE",
    144: "OP_ABS",
    145: "OP_NOT",
    146: "OP_0NOTEQUAL",
    147: "OP_ADD",
    148: "OP_SUB",
    154: "OP_BOOLAND",
    155: "OP_BOOLOR",
    156: "OP_NUMEQUAL",
    157: "OP_NUMEQUALVERIFY",
    158: "OP_NUMNOTEQUAL",
    159: "OP_LESSTHAN",
    160: "OP_GREATERTHAN",
    161: "OP_LESSTHANOREQUAL",
    162: "OP_GREATERTHANOREQUAL",
    163: "OP_MIN",
    164: "OP_MAX",
    165: "OP_WITHIN",
    166: "OP_RIPEMD160",
    167: "OP_SHA1",
    168: "OP_SHA256",
    169: "OP_HASH160",
    170: "OP_HASH256",
    171: "OP_CODESEPARATOR",
    172: "OP_CHECKSIG",
    173: "OP_CHECKSIGVERIFY",
    174: "OP_CHECKMULTISIG",
    175: "OP_CHECKMULTISIGVERIFY",
    176: "OP_NOP1",
    177: "OP_CHECKLOCKTIMEVERIFY",
    178: "OP_CHECKSEQUENCEVERIFY",
    179: "OP_NOP4",
    180: "OP_NOP5",
    181: "OP_NOP6",
    182: "OP_NOP7",
    183: "OP_NOP8",
    184: "OP_NOP9",
    185: "OP_NOP10",
}

def _op_pushdata(data: Octets) -> bytes:
    """Convert to canonical push: OP_PUSHDATA (if needed) | length | data.

    According to standardness rules (BIP-62) the
    minimum possible PUSHDATA operator must be used.
    Byte vectors on the stack are not allowed to be more than 520 bytes long.
    """

    data = bytes_from_octets(data)
    r = b""
    length = len(data)
    if length < 76:  # 1-byte-length
        r += length.to_bytes(1, byteorder="little")
    elif length < 256:  # OP_PUSHDATA1 | 1-byte-length
        r += OP_CODES["OP_PUSHDATA1"]
        r += length.to_bytes(1, byteorder="little")
    elif length < 521:  # OP_PUSHDATA2 | 2-byte-length
        r += OP_CODES["OP_PUSHDATA2"]
        r += length.to_bytes(2, byteorder="little")
    else:
        # because of the 520 bytes limit
        # there is no need to use OP_PUSHDATA4
        # r += OP_CODES['OP_PUSHDATA4']
        # r += length.to_bytes(4, byteorder='little')
        raise ValueError(f"Too many bytes for OP_PUSHDATA: {length}")
    r += data
    return r

def _op_int(token: int) -> bytes:
    # Short 1-byte opcodes exist
    # to push numbers in [-1, 16]
    if token >= 0 and token <= 16:
        return OP_CODES["OP_" + str(token)]
    elif token == -1:
        return OP_CODES["OP_1NEGATE"]
    # Pushing any other number requires an
    # explicit push operation of its bytes encoding
    # FIXME: negative numbers?
    else:
        v = token
        # Convert number to bitcoin-specific little endian format
        # We need v.bit_length() bits, plus a sign bit for every nonzero number.
        n_bits = v.bit_length() + (v != 0)
        # The number of bytes for that is:
        n_bytes = (n_bits + 7) // 8
        # Convert number to absolute value + sign in top bit.
        encoded_v = 0 if v == 0 else abs(v) | ((v < 0) << (n_bytes * 8 - 1))
        # Serialize to bytes
        data = encoded_v.to_bytes(n_bytes, "little")
        return _op_pushdata(data)

def _op_str(token: str) -> bytes:
    token = token.strip()
    token = token.upper()
    if token in OP_CODES:
        return OP_CODES[token]
    try:
        data = bytes.fromhex(token)
    except Exception:
        raise ValueError(f"invalid string token: {token}")
    return _op_pushdata(data)

def encode(script: List[Token]) -> bytes:
    r = b""
    for token in script:
        if isinstance(token, int):
            r += _op_int(token)
        elif isinstance(token, str):
            r += _op_str(token)
        elif isinstance(token, bytes):
            r += _op_pushdata(token)
        else:
            raise ValueError(f"Unmanaged {type(token)} token type")
    return r

def decode(script: Octets) -> List[Token]:

    script = bytes_from_octets(script)

    # initialize the result list
    r: List[Token] = []

    length = len(script)
    s = BytesIO(script)  # TODO: avoid unnecessary streamification
    counter = 0
    while counter < length:
        # get one byte and convert it to an integer
        i = s.read(1)[0]
        counter += 1
        if i == 0:
            # numeric value 0 (OP_0)
            # r.append(OP_CODE_NAMES[i])
            r.append(i)
        elif i == 79:
            # numeric value -1 (OP_1NEGATE)
            r.append(-1)
        elif i > 80 and i < 97:
            # numeric values 1-16 (OP_1-OP_16)
            # r.append(OP_CODE_NAMES[i])
            r.append(i - 80)
        elif i < 76:
            # 1-byte-data-length | data
            data = s.read(i)
            counter += i
            r.append(data.hex().upper())
        elif i == 76:
            # OP_PUSHDATA1 | 1-byte-data-length | data
            data_length = int.from_bytes(s.read(1), byteorder="little")
            data = s.read(data_length)
            counter += 1 + data_length
            r.append(data.hex().upper())
        elif i == 77:
            # OP_PUSHDATA2 | 2-byte-data-length | data
            data_length = int.from_bytes(s.read(2), byteorder="little")
            data = s.read(data_length)
            counter += 2 + data_length
            r.append(data.hex().upper())
        elif i == 78:
            # OP_PUSHDATA4 | 4-byte-data-length | data
            data_length = int.from_bytes(s.read(4), byteorder="little")
            data = s.read(data_length)
            counter += 4 + data_length
            r.append(data.hex().upper())
        else:
            # OP_CODE
            r.append(OP_CODE_NAMES[i])

    return r

def serialize(script: List[Token]) -> bytes:
    r = encode(script)
    # prepend length as varint
    length = len(r)
    return varint.encode(length) + r

def deserialize(stream: BinaryData) -> List[Token]:

    stream = bytesio_from_binarydata(stream)

    length = varint.decode(stream)
    script = stream.read(length)
    return decode(script)
<commit_msg>cleaned up code<commit_after>#!/usr/bin/env python3

# Copyright (C) 2017-2020 The btclib developers
#
# This file is part of btclib. It is subject to the license terms in the
# LICENSE file found in the top-level directory of this distribution.
#
# No part of btclib including this file, may be copied, modified, propagated,
# or distributed except according to the terms contained in the LICENSE file.

"""Bitcoin Script.

https://en.bitcoin.it/wiki/Script

Scripts are represented by List[Token], where Token = Union[int, str, bytes]:

* int [-1, 16] are shorcuts for 'OP_1NEGATE', 'OP_0' - 'OP_16'
* str are for opcodes (e.g. 'OP_HASH160') or hexstring data
* bytes are for data (but integers are often casted to int)
"""

from typing import List

from . import varint
from .alias import BinaryData, Octets, Token
from .utils import bytes_from_octets, bytesio_from_binarydata

SIGHASH_ALL = 1
SIGHASH_NONE = 2
SIGHASH_SINGLE = 3
SIGHASH_ANYONECANPAY = 0x80

SIGHASHES = [
    SIGHASH_ALL,
    SIGHASH_NONE,
    SIGHASH_SINGLE,
    SIGHASH_ANYONECANPAY + SIGHASH_ALL,
    SIGHASH_ANYONECANPAY + SIGHASH_NONE,
    SIGHASH_ANYONECANPAY + SIGHASH_SINGLE,
]

OP_CODES = {
    # Constants
    "OP_0": b"\x00",
    "OP_FALSE": b"\x00",
    "OP_PUSHDATA1": b"\x4c",
    "OP_PUSHDATA2": b"\x4d",
    "OP_PUSHDATA4": b"\x4e",
    "OP_1NEGATE": b"\x4f",
    "OP_1": b"\x51",
    "OP_TRUE": b"\x51",
    "OP_2": b"\x52",
    "OP_3": b"\x53",
    "OP_4": b"\x54",
    "OP_5": b"\x55",
    "OP_6": b"\x56",
    "OP_7": b"\x57",
    "OP_8": b"\x58",
    "OP_9": b"\x59",
    "OP_10": b"\x5a",
    "OP_11": b"\x5b",
    "OP_12": b"\x5c",
    "OP_13": b"\x5d",
    "OP_14": b"\x5e",
    "OP_15": b"\x5f",
    "OP_16": b"\x60",
    # Flow control
    "OP_NOP": b"\x61",
    "OP_IF": b"\x63",
    "OP_NOTIF": b"\x64",
    "OP_ELSE": b"\x67",
    "OP_ENDIF": b"\x68",
    "OP_VERIFY": b"\x69",
    "OP_RETURN": b"\x6a",
    # Stack
    "OP_TOALTSTACK": b"\x6b",
    "OP_FROMALTSTACK": b"\x6c",
    "OP_IFDUP": b"\x73",
    "OP_DEPTH": b"\x74",
    "OP_DROP": b"\x75",
    "OP_DUP": b"\x76",
    "OP_NIP": b"\x77",
    "OP_OVER": b"\x78",
    "OP_PICK": b"\x79",
    "OP_ROLL": b"\x7a",
    "OP_ROT": b"\x7b",
    "OP_SWAP": b"\x7c",
    "OP_TUCK": b"\x7d",
    "OP_2DROP": b"\x6d",
    "OP_2DUP": b"\x6e",
    "OP_3DUP": b"\x6f",
    "OP_2OVER": b"\x70",
    "OP_2ROT": b"\x71",
    "OP_2SWAP": b"\x72",
    # Splice
    "OP_SIZE": b"\x82",
    # Bitwise logic
    "OP_EQUAL": b"\x87",
    "OP_EQUALVERIFY": b"\x88",
    # Arithmetic
    "OP_1ADD": b"\x8b",  # without OP_, 1ADD would be a number
    "OP_1SUB": b"\x8c",
    "OP_NEGATE": b"\x8f",
    "OP_ABS": b"\x90",
    "OP_NOT": b"\x91",
    "OP_0NOTEQUAL": b"\x92",
    "OP_ADD": b"\x93",
    "OP_SUB": b"\x94",
    "OP_BOOLAND": b"\x9a",
    "OP_BOOLOR": b"\x9b",
    "OP_NUMEQUAL": b"\x9c",
    "OP_NUMEQUALVERIFY": b"\x9d",
    "OP_NUMNOTEQUAL": b"\x9e",
    "OP_LESSTHAN": b"\x9f",
    "OP_GREATERTHAN": b"\xa0",
    "OP_LESSTHANOREQUAL": b"\xa1",
    "OP_GREATERTHANOREQUAL": b"\xa2",
    "OP_MIN": b"\xa3",
    "OP_MAX": b"\xa4",
    "OP_WITHIN": b"\xa5",
    # Crypto
    "OP_RIPEMD160": b"\xa6",
    "OP_SHA1": b"\xa7",
    "OP_SHA256": b"\xa8",
    "OP_HASH160": b"\xa9",
    "OP_HASH256": b"\xaa",
    "OP_CODESEPARATOR": b"\xab",
    "OP_CHECKSIG": b"\xac",
    "OP_CHECKSIGVERIFY": b"\xad",
    "OP_CHECKMULTISIG": b"\xae",
    "OP_CHECKMULTISIGVERIFY": b"\xaf",
    # Locktime
    "OP_NOP2": b"\xb1",
    "OP_CHECKLOCKTIMEVERIFY": b"\xb1",
    "OP_NOP3": b"\xb2",
    "OP_CHECKSEQUENCEVERIFY": b"\xb2",
    # Reserved words
    "OP_NOP1": b"\xb0",
    "OP_NOP4": b"\xb3",
    "OP_NOP5": b"\xb4",
    "OP_NOP6": b"\xb5",
    "OP_NOP7": b"\xb6",
    "OP_NOP8": b"\xb7",
    "OP_NOP9": b"\xb8",
    "OP_NOP10": b"\xb9",
}

OP_CODE_NAMES = {
    0: "OP_0",
    76: "OP_PUSHDATA1",
    77: "OP_PUSHDATA2",
    78: "OP_PUSHDATA4",
    79: "OP_1NEGATE",
    81: "OP_1",
    82: "OP_2",
    83: "OP_3",
    84: "OP_4",
    85: "OP_5",
    86: "OP_6",
    87: "OP_7",
    88: "OP_8",
    89: "OP_9",
    90: "OP_10",
    91: "OP_11",
    92: "OP_12",
    93: "OP_13",
    94: "OP_14",
    95: "OP_15",
    96: "OP_16",
    97: "OP_NOP",
    99: "OP_IF",
    100: "OP_NOTIF",
    103: "OP_ELSE",
    104: "OP_ENDIF",
    105: "OP_VERIFY",
    106: "OP_RETURN",
    107: "OP_TOALTSTACK",
    108: "OP_FROMALTSTACK",
    109: "OP_2DROP",
    110: "OP_2DUP",
    111: "OP_3DUP",
    112: "OP_2OVER",
    113: "OP_2ROT",
    114: "OP_2SWAP",
    115: "OP_IFDUP",
    116: "OP_DEPTH",
    117: "OP_DROP",
    118: "OP_DUP",
    119: "OP_NIP",
    120: "OP_OVER",
    121: "OP_PICK",
    122: "OP_ROLL",
    123: "OP_ROT",
    124: "OP_SWAP",
    125: "OP_TUCK",
    130: "OP_SIZE",
    135: "OP_EQUAL",
    136: "OP_EQUALVERIFY",
    139: "OP_1ADD",
    140: "OP_1SUB",
    143: "OP_NEGATE",
    144: "OP_ABS",
    145: "OP_NOT",
    146: "OP_0NOTEQUAL",
    147: "OP_ADD",
    148: "OP_SUB",
    154: "OP_BOOLAND",
    155: "OP_BOOLOR",
    156: "OP_NUMEQUAL",
    157: "OP_NUMEQUALVERIFY",
    158: "OP_NUMNOTEQUAL",
    159: "OP_LESSTHAN",
    160: "OP_GREATERTHAN",
    161: "OP_LESSTHANOREQUAL",
    162: "OP_GREATERTHANOREQUAL",
    163: "OP_MIN",
    164: "OP_MAX",
    165: "OP_WITHIN",
    166: "OP_RIPEMD160",
    167: "OP_SHA1",
    168: "OP_SHA256",
    169: "OP_HASH160",
    170: "OP_HASH256",
    171: "OP_CODESEPARATOR",
    172: "OP_CHECKSIG",
    173: "OP_CHECKSIGVERIFY",
    174: "OP_CHECKMULTISIG",
    175: "OP_CHECKMULTISIGVERIFY",
    176: "OP_NOP1",
    177: "OP_CHECKLOCKTIMEVERIFY",
    178: "OP_CHECKSEQUENCEVERIFY",
    179: "OP_NOP4",
    180: "OP_NOP5",
    181: "OP_NOP6",
    182: "OP_NOP7",
    183: "OP_NOP8",
    184: "OP_NOP9",
    185: "OP_NOP10",
}

def _op_pushdata(data: Octets) -> bytes:
    """Convert to canonical push: OP_PUSHDATA (if needed) | length | data.

    According to standardness rules (BIP-62) the
    minimum possible PUSHDATA operator must be used.
    Byte vectors on the stack are not allowed to be more than 520 bytes long.
    """

    data = bytes_from_octets(data)
    r = b""
    length = len(data)
    if length < 76:  # 1-byte-length
        r += length.to_bytes(1, byteorder="little")
    elif length < 256:  # OP_PUSHDATA1 | 1-byte-length
        r += OP_CODES["OP_PUSHDATA1"]
        r += length.to_bytes(1, byteorder="little")
    elif length < 521:  # OP_PUSHDATA2 | 2-byte-length
        r += OP_CODES["OP_PUSHDATA2"]
        r += length.to_bytes(2, byteorder="little")
    else:
        # because of the 520 bytes limit
        # there is no need to use OP_PUSHDATA4
        # r += OP_CODES['OP_PUSHDATA4']
        # r += length.to_bytes(4, byteorder='little')
        raise ValueError(f"Too many bytes for OP_PUSHDATA: {length}")
    r += data
    return r

def _op_int(token: int) -> bytes:
    # Short 1-byte opcodes exist
    # to push numbers in [-1, 16]
    if token >= 0 and token <= 16:
        return OP_CODES["OP_" + str(token)]
    elif token == -1:
        return OP_CODES["OP_1NEGATE"]
    # Pushing any other number requires an
    # explicit push operation of its bytes encoding
    # FIXME: negative numbers?
    else:
        v = token
        # Convert number to bitcoin-specific little endian format
        # We need v.bit_length() bits, plus a sign bit for every nonzero number.
        n_bits = v.bit_length() + (v != 0)
        # The number of bytes for that is:
        n_bytes = (n_bits + 7) // 8
        # Convert number to absolute value + sign in top bit.
        encoded_v = 0 if v == 0 else abs(v) | ((v < 0) << (n_bytes * 8 - 1))
        # Serialize to bytes
        data = encoded_v.to_bytes(n_bytes, "little")
        return _op_pushdata(data)

def _op_str(token: str) -> bytes:
    token = token.strip()
    token = token.upper()
    if token in OP_CODES:
        return OP_CODES[token]
    try:
        data = bytes.fromhex(token)
    except Exception:
        raise ValueError(f"invalid string token: {token}")
    return _op_pushdata(data)

def encode(script: List[Token]) -> bytes:
    r = b""
    for token in script:
        if isinstance(token, int):
            r += _op_int(token)
        elif isinstance(token, str):
            r += _op_str(token)
        elif isinstance(token, bytes):
            r += _op_pushdata(token)
        else:
            raise ValueError(f"Unmanaged {type(token)} token type")
    return r

def decode(stream: BinaryData) -> List[Token]:

    s = bytesio_from_binarydata(stream)
    # initialize the result list
    r: List[Token] = []
    while True:
        # get one byte
        t = s.read(1)
        if not t:
            break
        # convert it to an integer
        i = t[0]
        if i == 0:
            # numeric value 0 (OP_0)
            # r.append(OP_CODE_NAMES[i])
            r.append(i)
        elif i == 79:
            # numeric value -1 (OP_1NEGATE)
            r.append(-1)
        elif i > 80 and i < 97:
            # numeric values 1-16 (OP_1-OP_16)
            # r.append(OP_CODE_NAMES[i])
            r.append(i - 80)
        elif i < 76:
            # 1-byte-data-length | data
            data = s.read(i)
            r.append(data.hex().upper())
        elif i == 76:
            # OP_PUSHDATA1 | 1-byte-data-length | data
            data_length = int.from_bytes(s.read(1), byteorder="little")
            data = s.read(data_length)
            r.append(data.hex().upper())
        elif i == 77:
            # OP_PUSHDATA2 | 2-byte-data-length | data
            data_length = int.from_bytes(s.read(2), byteorder="little")
            data = s.read(data_length)
            r.append(data.hex().upper())
        elif i == 78:
            # OP_PUSHDATA4 | 4-byte-data-length | data
            data_length = int.from_bytes(s.read(4), byteorder="little")
            data = s.read(data_length)
            r.append(data.hex().upper())
        else:
            # OP_CODE
            r.append(OP_CODE_NAMES[i])

    return r

def serialize(script: List[Token]) -> bytes:
    r = encode(script)
    # prepend length as varint
    length = len(r)
    return varint.encode(length) + r

def deserialize(stream: BinaryData) -> List[Token]:

    stream = bytesio_from_binarydata(stream)

    length = varint.decode(stream)
    script = stream.read(length)
    return decode(script)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
__author__ = 'mdc_hk'
version = '1.0'

### Please remember to import xlwt module for the .to_excel function to work; lxml module is an important parser library too.
import logging, os, re, sys
import pandas as pd
from pandas import Series
from bs4 import BeautifulSoup
import numpy as np
from scipy import stats
logging.basicConfig(filename='Log.txt', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

#===========================================================================================================

try:
    filein=sys.argv[1]
    xmlResequencingRunStatistics=sys.argv[2]
    reference_annotated_fa = sys.argv[3]
    confidenceValue = sys.argv[4]
    proceed = True
except:
    print('\nProgram: varValidity\nVersion: {}\n'.format(version) + 'Usage: ' + sys.argv[0] + ' <input_txt_file> '
            '<ResequencingRunStatistics_xml> <reference_annotated_fa> <Confidence value>')
    print('Examples:', sys.argv[0], 'sample001_table.txt', 'ResequencingRunStatistics.xml', 'reference_annotated.fa',
          '0.9500')
    proceed = False

if proceed == True:

    programFolder = os.path.expanduser('~/FluSeq/')
    dataFromTable = pd.read_table(filein, sep='\t')
    columnID = dataFromTable.columns[7]

    try:
        # dfStat = pd.read_csv('ms2SysCtrl.db', sep=',')
        dfStat = pd.read_csv(programFolder + 'ms2SysCtrl.db', sep=',')
        dfStat = dfStat.set_index(['GeneSegment', 'POS'])

    except IOError:
        errorFile = open('Error.txt', 'a')
        errorFile.write('Please make sure the MS2SysCtrl.db is placed in the correct PATH' + '\n')
        errorFile.close()
        exit(1)

    # The statistics for contamination starts here.

    def expectedConR(x2, gene, pos):
        dfStat_tmp = dfStat.ix[gene, pos]
        x = dfStat_tmp['PFreads_inMillion']
        y = dfStat_tmp['Total_Depth']

        # Modeling with Numpy
        p, cov = np.polyfit(x,y,1,cov=True)           # parameters and covariance from of the fit
        y_model = np.polyval(p, x)                    # model using the fit parameters; NOTE: parameters here are coefficients

        # Statistics
        n = y.size                                    # number of observations
        m = p.size                                    # number of parameters
        DF = n - m                                    # degrees of freedom
        t = stats.t.ppf(float(confidenceValue), n - m)       # used for CI and PI bands

        # Estimates of Error in Data/Model
        resid = y - y_model
        s_err = np.sqrt(np.sum(resid**2)/(DF))        # standard deviation of the error

        y2 = np.polyval(p, x2)

        # Prediction Interval
        PI = t*s_err*np.sqrt(1+1/n+(x2-np.mean(x))**2/np.sum((x-np.mean(x))**2))

        uppConfLimit_ContRead = y2 + PI

        return uppConfLimit_ContRead

### To import the main table result
    dataFromTable['SAMPLE'] = columnID[:-3]
    dataFromTable = dataFromTable.set_index(['SAMPLE'])

### to create table/library to get the percentage of reads identified for each sample

    try:
        bsObj = BeautifulSoup(open(xmlResequencingRunStatistics), 'lxml-xml')
        xmlOfSamples = bsObj.findAll(re.compile(r'SummarizedSampleStatis(t)*ics'))

    except IOError:
        errorFile = open('Error.txt', 'a')
        errorFile.write('Please make sure the run-specific ResequencingRunStatistics.xml is placed in the run folder'
                        + '\n')
        errorFile.close()
        exit(1)

    listOfSamples = []
    for name in xmlOfSamples:
        listOfSamples.append(name.SampleName.get_text())

    listOfPFReads = []
    for name in xmlOfSamples:
        listOfPFReads.append(round(int(name.NumberOfClustersPF.get_text())/1000000, 2))
    if len(listOfSamples) == len(listOfPFReads):
        dictOfSamplesAndPFreads = dict(zip(listOfSamples, listOfPFReads))

    dataFromTable['PFreads_inMillion'] = dictOfSamplesAndPFreads[columnID[:-3]]

### to set the index as 'CHROM' for easy gene data separation
    dataFromTable = dataFromTable.set_index(['CHROM'], drop=False)
    geneList = []

    def validityTester(gene):
        validityOfWildtype = []
        for index, row in gene.iterrows():
            if row['refCopy'] > row['UppPredLimit_ContRead']:
                validityOfWildtype.append(1)
            else:
                validityOfWildtype.append(0)

        validityOfAlt1 = []
        for index, row in gene.iterrows():
            if row['alt1Copy'] > row['UppPredLimit_ContRead']:
                validityOfAlt1.append(2)
            else:
                validityOfAlt1.append(0)

        validityOfAlt2 = []
        for index, row in gene.iterrows():
            if row['alt2Copy'] > row['UppPredLimit_ContRead']:
                validityOfAlt2.append(4)
            else:
                validityOfAlt2.append(0)

        scoring = [x + y + z for x, y, z in zip(validityOfWildtype, validityOfAlt1, validityOfAlt2)]
        Ref = list(gene['REF'])
        Alt12 = gene['ALT'].str.split(",").apply(Series, 1)
        if len(Alt12.columns) == 1:
            Alt12[1] = ''
        Alt1 = list(Alt12[0])
        Alt2 = list(Alt12[1])
        refCopyNumber = list(gene['refCopy'])
        alt1CopyNumber = list(gene['alt1Copy'])
        alt2CopyNumber = list(gene['alt2Copy'])

        call = []
        percentage = []
        for i in range(len(scoring)):
            if scoring[i] == 0:
                call.append('Invalid')
                percentage.append(np.NaN)
            elif scoring[i] == 1:
                call.append(Ref[i])
                percentage.append(1.00)
            elif scoring[i] == 2:
                call.append(Alt1[i])
                percentage.append(1.00)
            elif scoring[i] == 3:
                call.append(Ref[i]+','+Alt1[i])
                percentage.append(alt1CopyNumber[i]/(alt1CopyNumber[i]+refCopyNumber[i]))
            elif scoring[i] == 4:
                call.append(Alt2[i])
                percentage.append(1.00)
            elif scoring[i] == 5:
                call.append(Ref[i]+','+Alt2[i])
                percentage.append(alt2CopyNumber/(alt2CopyNumber[i]+refCopyNumber[i]))
            elif scoring[i] == 6:
                call.append(Alt1[i]+','+Alt2[i])
                percentage.append(alt2CopyNumber[i]/(alt1CopyNumber[i]+alt2CopyNumber[i]))
            elif scoring[i] == 7:
                call.append(Ref[i]+','+Alt1[i]+','+Alt2[i])
                percentage.append(alt2CopyNumber[i]/(alt1CopyNumber[i]+alt2CopyNumber[i]))

        gene['Call(R1,R2)'] = call
        gene['Percentage(R2/R1+R2)'] = [ '%.3f' % elem for elem in percentage ]

        return

    def geneAnalysis(Chrom, Gene):
        GeneData = dataFromTable.ix[Chrom]
        GeneData = GeneData.reset_index(drop=True)
        UppPredLimit_ContRead = []
        for i in GeneData.index:
            if GeneData.ix[i]['POS'] <= 30:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1))
            elif 31 <= GeneData.ix[i]['POS'] <= 225:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 151))
            elif 226 <= GeneData.ix[i]['POS'] <= 450:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 301))
            elif 451 <= GeneData.ix[i]['POS'] <= 525:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 451))
            elif 526 <= GeneData.ix[i]['POS'] <= 675:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 601))
            elif 676 <= GeneData.ix[i]['POS'] <= 825:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 751))
            elif 826 <= GeneData.ix[i]['POS'] <= 975:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 901))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 751))
            elif 976 <= GeneData.ix[i]['POS'] <= 1125:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1051))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 901))
            elif 1126 <= GeneData.ix[i]['POS'] <= 1275:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1201))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1051))
            elif 1276 <= GeneData.ix[i]['POS'] <= 1425:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1351))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1201))
            elif 1426 <= GeneData.ix[i]['POS'] <= 1575:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1501))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1351))
            elif 1576 <= GeneData.ix[i]['POS'] <= 1725:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1651))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1501))
            elif 1726 <= GeneData.ix[i]['POS'] <= 1801:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1801))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1651))
            elif 1802 <= GeneData.ix[i]['POS'] <= 1951:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1951))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1801))
            elif 1952 <= GeneData.ix[i]['POS'] <= 2101:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 2101))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1951))
            else:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 2251))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 2101))

        GeneData['UppPredLimit_ContRead'] = UppPredLimit_ContRead
        GeneData_tojoin = GeneData[(columnID[:-3]+'.AD')].str.split(",").apply(Series, 1)
        if len(GeneData_tojoin.columns) == 2:
            GeneData_tojoin[2] = np.NaN
        GeneData_tojoin.columns = ['refCopy', 'alt1Copy', 'alt2Copy']
        GeneData = pd.merge(GeneData, GeneData_tojoin, left_index=True, right_index=True)
        GeneData.refCopy = GeneData.refCopy.astype(float)
        GeneData.alt1Copy = GeneData.alt1Copy.astype(float)
        GeneData.alt2Copy = GeneData.alt2Copy.astype(float)
        GeneData = GeneData.set_index('POS')
        validityTester(GeneData)
        geneList.append(GeneData)

        return

    with open(reference_annotated_fa,'r') as refFile:
        sequences = refFile.read()

    segmentRegex = re.compile(r'((Segment\d)_[\S]*)')
    segment = segmentRegex.findall(sequences)
    for seg in segment:

        try:
            geneAnalysis(seg[0], seg[1])
        except:
            logging.info(seg[1] + ' of ' + columnID[:-3] + ' was not analysed')
            errorFile = open('Error.txt', 'a')
            errorFile.write(seg[1] + ' of ' + columnID[:-3] + ' was not analysed. Suggested solution: Something wrong '
                                                              'with the formating of the AllGenes.xls file' + '\n')
            errorFile.close()

    try:
        pd.concat(geneList, axis=0).to_excel(columnID[:-3] + "_AllGenes" + ".xls")
    except ValueError:
        print('All segments of sample ' + columnID[:-3] + ' were not analysed')
        logging.info('All segments of sample ' + columnID[:-3] + ' were not analysed. '
                    'Suggested solution: please use bwa-0.6.2 but not the latest version')
        errorFile = open('Error.txt', 'a')
        errorFile.write('All gene segments of sample ' + columnID[:-3] + ' were not analysed. '
                    'Suggested solution: please use bwa-0.6.2 but not the latest version' + '\n')
        errorFile.close()<commit_msg>Update<commit_after>#!/usr/bin/env python3.4
__author__ = 'mdc_hk'
version = '1.0'

### Please remember to import xlwt module for the .to_excel function to work; lxml module is an important parser library too.
import logging, os, re, sys
import pandas as pd
from pandas import Series
from bs4 import BeautifulSoup
import numpy as np
from scipy import stats
logging.basicConfig(filename='Log.txt', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

#===========================================================================================================

try:
    filein=sys.argv[1]
    xmlResequencingRunStatistics=sys.argv[2]
    reference_annotated_fa = sys.argv[3]
    confidenceValue = sys.argv[4]
    proceed = True
except:
    print('\nProgram: varValidity\nVersion: {}\n'.format(version) + 'Usage: ' + sys.argv[0] + ' <input_txt_file> '
            '<ResequencingRunStatistics_xml> <reference_annotated_fa> <Confidence value>')
    print('Examples:', sys.argv[0], 'sample001_table.txt', 'ResequencingRunStatistics.xml', 'reference_annotated.fa',
          '0.9500')
    proceed = False

if proceed == True:

    programFolder = os.path.expanduser('~/FluSeq/')
    dataFromTable = pd.read_table(filein, sep='\t')
    columnID = dataFromTable.columns[7]

    try:
        # dfStat = pd.read_csv('ms2SysCtrl.db', sep=',')
        dfStat = pd.read_csv(programFolder + 'ms2SysCtrl.db', sep=',')
        dfStat = dfStat.set_index(['GeneSegment', 'POS'])

    except IOError:
        errorFile = open('Error.txt', 'a')
        errorFile.write('Please make sure the MS2SysCtrl.db is placed in the correct PATH' + '\n')
        errorFile.close()
        exit(1)

    # The statistics for contamination starts here.

    def expectedConR(x2, gene, pos):
        dfStat_tmp = dfStat.ix[gene, pos]
        x = dfStat_tmp['PFreads_inMillion']
        y = dfStat_tmp['Total_Depth']

        # Modeling with Numpy
        p, cov = np.polyfit(x,y,1,cov=True)           # parameters and covariance from of the fit
        y_model = np.polyval(p, x)                    # model using the fit parameters; NOTE: parameters here are coefficients

        # Statistics
        n = y.size                                    # number of observations
        m = p.size                                    # number of parameters
        DF = n - m                                    # degrees of freedom
        t = stats.t.ppf(float(confidenceValue), n - m)       # used for CI and PI bands

        # Estimates of Error in Data/Model
        resid = y - y_model
        s_err = np.sqrt(np.sum(resid**2)/(DF))        # standard deviation of the error

        y2 = np.polyval(p, x2)

        # Prediction Interval
        PI = t*s_err*np.sqrt(1+1/n+(x2-np.mean(x))**2/np.sum((x-np.mean(x))**2))

        uppConfLimit_ContRead = y2 + PI

        return uppConfLimit_ContRead

### To import the main table result
    dataFromTable['SAMPLE'] = columnID[:-3]
    dataFromTable = dataFromTable.set_index(['SAMPLE'])

### to create table/library to get the percentage of reads identified for each sample

    try:
        bsObj = BeautifulSoup(open(xmlResequencingRunStatistics), 'lxml-xml')
        xmlOfSamples = bsObj.findAll(re.compile(r'SummarizedSampleStatis(t)*ics'))

    except IOError:
        errorFile = open('Error.txt', 'a')
        errorFile.write('Please make sure the run-specific ResequencingRunStatistics.xml is placed in the run folder'
                        + '\n')
        errorFile.close()
        exit(1)

    listOfSamples = []
    for name in xmlOfSamples:
        listOfSamples.append(name.SampleName.get_text())

    listOfPFReads = []
    for name in xmlOfSamples:
        listOfPFReads.append(round(int(name.NumberOfClustersPF.get_text())/1000000, 2))
    if len(listOfSamples) == len(listOfPFReads):
        dictOfSamplesAndPFreads = dict(zip(listOfSamples, listOfPFReads))

    dataFromTable['PFreads_inMillion'] = dictOfSamplesAndPFreads[columnID[:-3]]

### to set the index as 'CHROM' for easy gene data separation
    dataFromTable = dataFromTable.set_index(['CHROM'], drop=False)
    geneList = []

    def validityTester(gene):
        validityOfWildtype = []
        for index, row in gene.iterrows():
            if row['refCopy'] > row['UppPredLimit_ContRead']:
                validityOfWildtype.append(1)
            else:
                validityOfWildtype.append(0)

        validityOfAlt1 = []
        for index, row in gene.iterrows():
            if row['alt1Copy'] > row['UppPredLimit_ContRead']:
                validityOfAlt1.append(2)
            else:
                validityOfAlt1.append(0)

        validityOfAlt2 = []
        for index, row in gene.iterrows():
            if row['alt2Copy'] > row['UppPredLimit_ContRead']:
                validityOfAlt2.append(4)
            else:
                validityOfAlt2.append(0)

        scoring = [x + y + z for x, y, z in zip(validityOfWildtype, validityOfAlt1, validityOfAlt2)]
        Ref = list(gene['REF'])
        Alt12 = gene['ALT'].str.split(",").apply(Series, 1)
        if len(Alt12.columns) == 1:
            Alt12[1] = ''
        Alt1 = list(Alt12[0])
        Alt2 = list(Alt12[1])
        refCopyNumber = list(gene['refCopy'])
        alt1CopyNumber = list(gene['alt1Copy'])
        alt2CopyNumber = list(gene['alt2Copy'])

        call = []
        percentage = []
        for i in range(len(scoring)):
            if scoring[i] == 0:
                call.append('Invalid')
                percentage.append(np.NaN)
            elif scoring[i] == 1:
                call.append(Ref[i])
                percentage.append(1.00)
            elif scoring[i] == 2:
                call.append(Alt1[i])
                percentage.append(1.00)
            elif scoring[i] == 3:
                call.append(Ref[i]+','+Alt1[i])
                percentage.append(alt1CopyNumber[i]/(alt1CopyNumber[i]+refCopyNumber[i]))
            elif scoring[i] == 4:
                call.append(Alt2[i])
                percentage.append(1.00)
            elif scoring[i] == 5:
                call.append(Ref[i]+','+Alt2[i])
                percentage.append(alt2CopyNumber/(alt2CopyNumber[i]+refCopyNumber[i]))
            elif scoring[i] == 6:
                call.append(Alt1[i]+','+Alt2[i])
                percentage.append(alt2CopyNumber[i]/(alt1CopyNumber[i]+alt2CopyNumber[i]))
            elif scoring[i] == 7:
                call.append(Ref[i]+','+Alt1[i]+','+Alt2[i])
                percentage.append(alt2CopyNumber[i]/(alt1CopyNumber[i]+alt2CopyNumber[i]))

        gene['Call(R1,R2)'] = call
        gene['Percentage(R2/R1+R2)'] = [ '%.3f' % elem for elem in percentage ]

        return

    def geneAnalysis(Chrom, Gene):
        GeneData = dataFromTable.ix[Chrom]
        GeneData = GeneData.reset_index(drop=True)
        UppPredLimit_ContRead = []
        for i in GeneData.index:
            if GeneData.ix[i]['POS'] <= 30:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1))
            elif 31 <= GeneData.ix[i]['POS'] <= 225:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 151))
            elif 226 <= GeneData.ix[i]['POS'] <= 450:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 301))
            elif 451 <= GeneData.ix[i]['POS'] <= 525:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 451))
            elif 526 <= GeneData.ix[i]['POS'] <= 675:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 601))
            elif 676 <= GeneData.ix[i]['POS'] <= 825:
                UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 751))
            elif 826 <= GeneData.ix[i]['POS'] <= 975:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 901))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 751))
            elif 976 <= GeneData.ix[i]['POS'] <= 1125:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1051))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 901))
            elif 1126 <= GeneData.ix[i]['POS'] <= 1275:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1201))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1051))
            elif 1276 <= GeneData.ix[i]['POS'] <= 1425:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1351))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1201))
            elif 1426 <= GeneData.ix[i]['POS'] <= 1575:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1501))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1351))
            elif 1576 <= GeneData.ix[i]['POS'] <= 1725:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1651))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1501))
            elif 1726 <= GeneData.ix[i]['POS'] <= 1801:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1801))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1651))
            elif 1802 <= GeneData.ix[i]['POS'] <= 1951:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1951))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1801))
            elif 1952 <= GeneData.ix[i]['POS'] <= 2101:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 2101))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 1951))
            else:
                try:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 2251))
                except:
                    UppPredLimit_ContRead.append(expectedConR(GeneData.ix[i]['PFreads_inMillion'], Gene, 2101))

        GeneData['UppPredLimit_ContRead'] = UppPredLimit_ContRead
        GeneData_tojoin = GeneData[(columnID[:-3]+'.AD')].str.split(",").apply(Series, 1)
        if len(GeneData_tojoin.columns) == 2:
            GeneData_tojoin[2] = np.NaN
        GeneData_tojoin.columns = ['refCopy', 'alt1Copy', 'alt2Copy']
        GeneData = pd.merge(GeneData, GeneData_tojoin, left_index=True, right_index=True)
        GeneData.refCopy = GeneData.refCopy.astype(float)
        GeneData.alt1Copy = GeneData.alt1Copy.astype(float)
        GeneData.alt2Copy = GeneData.alt2Copy.astype(float)
        GeneData = GeneData.set_index('POS')
        validityTester(GeneData)
        geneList.append(GeneData)

        return

    with open(reference_annotated_fa,'r') as refFile:
        sequences = refFile.read()

    segmentRegex = re.compile(r'((Segment\d)_[\S]*)')
    segment = segmentRegex.findall(sequences)
    for seg in segment:

        try:
            geneAnalysis(seg[0], seg[1])
        except:
            logging.info(seg[1] + ' of ' + columnID[:-3] + ' was not analysed')
            errorFile = open('Error.txt', 'a')
            errorFile.write(seg[1] + ' of ' + columnID[:-3] + ' was not analysed. Suggested solution: Something wrong '
                                                              'with the formating of the AllGenes.xls file' + '\n')
            errorFile.close()

    try:
        pd.concat(geneList, axis=0).to_excel(columnID[:-3] + "_AllGenes" + ".xls")
    except ValueError:
        print('All segments of sample ' + columnID[:-3] + ' were not analysed')
        logging.info('All segments of sample ' + columnID[:-3] + ' were not analysed. '
                    'Suggested solution: please use bwa-0.6.2 but not the latest version')
        errorFile = open('Error.txt', 'a')
        errorFile.write('All gene segments of sample ' + columnID[:-3] + ' were not analysed. '
                    'Suggested solution: please use bwa-0.6.2 but not the latest version' + '\n')
        errorFile.close()<|endoftext|>


<commit_before>"""
Various utilities for Varda.

.. moduleauthor:: Martijn Vermaat <martijn@vermaat.name>

.. Licensed under the MIT license, see the LICENSE file.
"""

import hashlib

from flask import current_app

from . import genome

class ReferenceMismatch(Exception):
    """
    Exception thrown mismatch with reference.
    """
    pass

def digest(data):
    """
    Given a file-like object opened for reading, calculate a digest as SHA1
    checksum and number of records.

    Calculating the number of records is done in a naive way by counting the
    number of lines in the file, and as such includes empty and header lines.
    """
    def read_chunks(data, chunksize=0xf00000):
        # Default chunksize is 16 megabytes.
        while True:
            chunk = data.read(chunksize)
            if not chunk:
                break
            yield chunk

    sha1 = hashlib.sha1()
    records = 0
    for chunk in read_chunks(data):
        sha1.update(chunk)
        records += chunk.count('\n')
    return sha1.hexdigest(), records

def normalize_chromosome(chromosome):
    """
    Try to get normalized chromosome name by reference lookup.
    """
    chromosome_aliases = [['M', 'MT', 'NC_012920.1', 'chrM', 'chrMT']]

    if not genome:
        for aliases in chromosome_aliases:
            if chromosome in aliases:
                return aliases[0]
        if chromosome.startswith('chr'):
            return chromosome[3:]
        return chromosome

    if chromosome in genome:
        return chromosome

    if chromosome.startswith('chr'):
        if chromosome[3:] in genome:
            return chromosome[3:]
    else:
        if 'chr' + chromosome in genome:
            return 'chr' + chromosome

    for aliases in chromosome_aliases:
        if chromosome in aliases:
            for alias in aliases:
                if alias in genome:
                    return alias

    raise ReferenceMismatch('Chromosome "%s" not in reference genome' %
                            chromosome)

def normalize_region(chromosome, begin, end):
    """
    Use reference to normalize chromosome name and validate location.
    """
    chromosome = normalize_chromosome(chromosome)

    if genome:
        if end > len(genome[chromosome]):
            raise ReferenceMismatch('Position %d does not exist on chromosome'
                                    ' "%s" in reference genome' %
                                    (end, chromosome))

    return chromosome, begin, end

def normalize_variant(chromosome, position, reference, observed):
    """
    Use reference to create a normalized representation of the variant.
    """
    reference = reference.upper()
    observed = observed.upper()

    chromosome = normalize_chromosome(chromosome)

    if genome:
        if position > len(genome[chromosome]):
            raise ReferenceMismatch('Position %d does not exist on chromosome'
                                    ' "%s" in reference genome' %
                                    (position, chromosome))
        if (genome[chromosome][position - 1
                               :position + len(reference) - 1].upper() !=
            reference):
            raise ReferenceMismatch('Sequence "%s" does not match reference'
                                    ' genome on "%s" at position %d' %
                                    (reference, chromosome, position))

    prefix, reference, observed, _ = trim_common(reference, observed)
    position += prefix

    if not genome:
        return chromosome, position, reference, observed

    # Insertions and deletions can be moved to the left by looking for cyclic
    # permutations.
    if reference == '':
        position, observed = move_left(genome[chromosome], position, observed)
        observed = observed.upper()
    elif observed == '':
        position, reference = move_left(genome[chromosome], position, reference)
        reference = reference.upper()

    return chromosome, position, reference, observed

def trim_common(s1, s2):
    """
    Trim two sequences by removing the longest common prefix and suffix. Also
    report the lengths of the removed parts. We start by removing the suffix.

    Standard convention with VCF is to place an indel at the left-most
    position, but some tools add additional context to the right of the
    sequences (e.g. samtools). These common suffixes are undesirable when
    comparing variants, for example in variant databases.

    Also, VCF requires to report at least one reference base, even for
    insertions.

        >>> trim_common('TATATATA', 'TATATA')
        (0, 'TA', '', 6)

        >>> trim_common('ACCCCC', 'ACCCCCCCC')
        (1, '', 'CCC', 5)

    :arg s1: First sequence.
    :type s1: str
    :arg s2: Second sequence.
    :type s2: str

    :return: Tuple (cpl, trimmed s1, trimmed s2, csl) where cpl and csl are
        the lengths of the common prefix and suffix, respectively.
    """
    suffix = 0
    while suffix < min(len(s1), len(s2)) and s1[-1 - suffix] == s2[-1 - suffix]:
        suffix += 1

    if suffix:
        s1 = s1[:-suffix]
        s2 = s2[:-suffix]

    prefix = 0
    while prefix < min(len(s1), len(s2)) and s1[prefix] == s2[prefix]:
        prefix += 1

    return prefix, s1[prefix:], s2[prefix:], suffix

def move_left(context, position, sequence):
    """
    Move ``sequence`` as far as possible to the left, starting at
    ``position`` (one-based) in ``context``, while staying in cyclic
    permutations.

    Schematic example::

                                       [=== sequence ====]
          [======================= context =======================]
                             <-- [== permutation ==]

                                       ^
                                    position

    Code examples::

        >>> move_left('abbaabbaabba', 5, 'abba')
        (1, 'abba')
        >>> move_left('abbaabbaabba', 6, 'bbaa')
        (1, 'abba')
        >>> move_left('abbaabbaabba', 6, 'bba')
        (5, 'abb')

    :arg context: Context sequence.
    :type context: str (or really a subscriptable yielding strings)
    :arg position: Start position of ``sequence`` in ``context``.
    :type position: int
    :arg sequence: Sequence to find cyclic permutations of in ``context``.
    :type sequence: str

    :return: A tuple (permutation, position) being the resulting cyclic
        permutation of ``sequence`` and its position in ``context``.
    :rtype: (str, int)
    """
    def lookup(p):
        if position <= p < position + len(sequence):
            return sequence[p - position].upper()
        return context[p - 1].upper()

    move = 0
    while (position - move > 1 and
           lookup(position - move - 1) ==
           lookup(position + len(sequence) - move - 1)):
        move += 1

    if not move:
        # Note: This case is only needed because the general case fails for
        #     move == 0 since sequence[:-0] == ''.
        return position, sequence

    return (position - move,
            context[position - move - 1
                    :min(position, position - move + len(sequence)) - 1]
            + sequence[:-move])
<commit_msg>Add NC_012920_1 as an MT alias<commit_after>"""
Various utilities for Varda.

.. moduleauthor:: Martijn Vermaat <martijn@vermaat.name>

.. Licensed under the MIT license, see the LICENSE file.
"""

import hashlib

from flask import current_app

from . import genome

class ReferenceMismatch(Exception):
    """
    Exception thrown mismatch with reference.
    """
    pass

def digest(data):
    """
    Given a file-like object opened for reading, calculate a digest as SHA1
    checksum and number of records.

    Calculating the number of records is done in a naive way by counting the
    number of lines in the file, and as such includes empty and header lines.
    """
    def read_chunks(data, chunksize=0xf00000):
        # Default chunksize is 16 megabytes.
        while True:
            chunk = data.read(chunksize)
            if not chunk:
                break
            yield chunk

    sha1 = hashlib.sha1()
    records = 0
    for chunk in read_chunks(data):
        sha1.update(chunk)
        records += chunk.count('\n')
    return sha1.hexdigest(), records

def normalize_chromosome(chromosome):
    """
    Try to get normalized chromosome name by reference lookup.
    """
    chromosome_aliases = [['M', 'MT', 'NC_012920.1', 'NC_012920_1',
                           'NC_012920', 'chrM', 'chrMT']]

    if not genome:
        for aliases in chromosome_aliases:
            if chromosome in aliases:
                return aliases[0]
        if chromosome.startswith('chr'):
            return chromosome[3:]
        return chromosome

    if chromosome in genome:
        return chromosome

    if chromosome.startswith('chr'):
        if chromosome[3:] in genome:
            return chromosome[3:]
    else:
        if 'chr' + chromosome in genome:
            return 'chr' + chromosome

    for aliases in chromosome_aliases:
        if chromosome in aliases:
            for alias in aliases:
                if alias in genome:
                    return alias

    raise ReferenceMismatch('Chromosome "%s" not in reference genome' %
                            chromosome)

def normalize_region(chromosome, begin, end):
    """
    Use reference to normalize chromosome name and validate location.
    """
    chromosome = normalize_chromosome(chromosome)

    if genome:
        if end > len(genome[chromosome]):
            raise ReferenceMismatch('Position %d does not exist on chromosome'
                                    ' "%s" in reference genome' %
                                    (end, chromosome))

    return chromosome, begin, end

def normalize_variant(chromosome, position, reference, observed):
    """
    Use reference to create a normalized representation of the variant.
    """
    reference = reference.upper()
    observed = observed.upper()

    chromosome = normalize_chromosome(chromosome)

    if genome:
        if position > len(genome[chromosome]):
            raise ReferenceMismatch('Position %d does not exist on chromosome'
                                    ' "%s" in reference genome' %
                                    (position, chromosome))
        if (genome[chromosome][position - 1
                               :position + len(reference) - 1].upper() !=
            reference):
            raise ReferenceMismatch('Sequence "%s" does not match reference'
                                    ' genome on "%s" at position %d' %
                                    (reference, chromosome, position))

    prefix, reference, observed, _ = trim_common(reference, observed)
    position += prefix

    if not genome:
        return chromosome, position, reference, observed

    # Insertions and deletions can be moved to the left by looking for cyclic
    # permutations.
    if reference == '':
        position, observed = move_left(genome[chromosome], position, observed)
        observed = observed.upper()
    elif observed == '':
        position, reference = move_left(genome[chromosome], position, reference)
        reference = reference.upper()

    return chromosome, position, reference, observed

def trim_common(s1, s2):
    """
    Trim two sequences by removing the longest common prefix and suffix. Also
    report the lengths of the removed parts. We start by removing the suffix.

    Standard convention with VCF is to place an indel at the left-most
    position, but some tools add additional context to the right of the
    sequences (e.g. samtools). These common suffixes are undesirable when
    comparing variants, for example in variant databases.

    Also, VCF requires to report at least one reference base, even for
    insertions.

        >>> trim_common('TATATATA', 'TATATA')
        (0, 'TA', '', 6)

        >>> trim_common('ACCCCC', 'ACCCCCCCC')
        (1, '', 'CCC', 5)

    :arg s1: First sequence.
    :type s1: str
    :arg s2: Second sequence.
    :type s2: str

    :return: Tuple (cpl, trimmed s1, trimmed s2, csl) where cpl and csl are
        the lengths of the common prefix and suffix, respectively.
    """
    suffix = 0
    while suffix < min(len(s1), len(s2)) and s1[-1 - suffix] == s2[-1 - suffix]:
        suffix += 1

    if suffix:
        s1 = s1[:-suffix]
        s2 = s2[:-suffix]

    prefix = 0
    while prefix < min(len(s1), len(s2)) and s1[prefix] == s2[prefix]:
        prefix += 1

    return prefix, s1[prefix:], s2[prefix:], suffix

def move_left(context, position, sequence):
    """
    Move ``sequence`` as far as possible to the left, starting at
    ``position`` (one-based) in ``context``, while staying in cyclic
    permutations.

    Schematic example::

                                       [=== sequence ====]
          [======================= context =======================]
                             <-- [== permutation ==]

                                       ^
                                    position

    Code examples::

        >>> move_left('abbaabbaabba', 5, 'abba')
        (1, 'abba')
        >>> move_left('abbaabbaabba', 6, 'bbaa')
        (1, 'abba')
        >>> move_left('abbaabbaabba', 6, 'bba')
        (5, 'abb')

    :arg context: Context sequence.
    :type context: str (or really a subscriptable yielding strings)
    :arg position: Start position of ``sequence`` in ``context``.
    :type position: int
    :arg sequence: Sequence to find cyclic permutations of in ``context``.
    :type sequence: str

    :return: A tuple (permutation, position) being the resulting cyclic
        permutation of ``sequence`` and its position in ``context``.
    :rtype: (str, int)
    """
    def lookup(p):
        if position <= p < position + len(sequence):
            return sequence[p - position].upper()
        return context[p - 1].upper()

    move = 0
    while (position - move > 1 and
           lookup(position - move - 1) ==
           lookup(position + len(sequence) - move - 1)):
        move += 1

    if not move:
        # Note: This case is only needed because the general case fails for
        #     move == 0 since sequence[:-0] == ''.
        return position, sequence

    return (position - move,
            context[position - move - 1
                    :min(position, position - move + len(sequence)) - 1]
            + sequence[:-move])
<|endoftext|>


<commit_before>from six import BytesIO, text_type
from six.moves.urllib.parse import urlparse, urlencode, urlunparse
import copy
import json
import zlib

from .util import CaseInsensitiveDict

def replace_headers(request, replacements):
    """
    Replace headers in request according to replacements. The replacements
    should be a list of (key, value) pairs where the value can be any of:
      1. A simple replacement string value.
      2. None to remove the given header.
      3. A callable which accepts (key, value, request) and returns a string
         value or None.
    """
    new_headers = request.headers.copy()
    for k, rv in replacements:
        if k in new_headers:
            ov = new_headers.pop(k)
            if callable(rv):
                rv = rv(key=k, value=ov, request=request)
            if rv is not None:
                new_headers[k] = rv
    request.headers = new_headers
    return request

def remove_headers(request, headers_to_remove):
    """
    Wrap replace_headers() for API backward compatibility.
    """
    replacements = [(k, None) for k in headers_to_remove]
    return replace_headers(request, replacements)

def replace_query_parameters(request, replacements):
    """
    Replace query parameters in request according to replacements. The
    replacements should be a list of (key, value) pairs where the value can be
    any of:
      1. A simple replacement string value.
      2. None to remove the given header.
      3. A callable which accepts (key, value, request) and returns a string
         value or None.
    """
    query = request.query
    new_query = []
    replacements = dict(replacements)
    for k, ov in query:
        if k not in replacements:
            new_query.append((k, ov))
        else:
            rv = replacements[k]
            if callable(rv):
                rv = rv(key=k, value=ov, request=request)
            if rv is not None:
                new_query.append((k, rv))
    uri_parts = list(urlparse(request.uri))
    uri_parts[4] = urlencode(new_query)
    request.uri = urlunparse(uri_parts)
    return request

def remove_query_parameters(request, query_parameters_to_remove):
    """
    Wrap replace_query_parameters() for API backward compatibility.
    """
    replacements = [(k, None) for k in query_parameters_to_remove]
    return replace_query_parameters(request, replacements)

def replace_post_data_parameters(request, replacements):
    """
    Replace post data in request--either form data or json--according to
    replacements. The replacements should be a list of (key, value) pairs where
    the value can be any of:
      1. A simple replacement string value.
      2. None to remove the given header.
      3. A callable which accepts (key, value, request) and returns a string
         value or None.
    """
    replacements = dict(replacements)
    if request.method == 'POST' and not isinstance(request.body, BytesIO):
        if (request.headers.get('Content-Type') or '').startswith('application/json'):
            json_data = json.loads(request.body.decode('utf-8'))
            for k, rv in replacements.items():
                if k in json_data:
                    ov = json_data.pop(k)
                    if callable(rv):
                        rv = rv(key=k, value=ov, request=request)
                    if rv is not None:
                        json_data[k] = rv
            request.body = json.dumps(json_data).encode('utf-8')
        else:
            if isinstance(request.body, text_type):
                request.body = request.body.encode('utf-8')
            splits = [p.partition(b'=') for p in request.body.split(b'&')]
            new_splits = []
            for k, sep, ov in splits:
                if sep is None:
                    new_splits.append((k, sep, ov))
                else:
                    rk = k.decode('utf-8')
                    if rk not in replacements:
                        new_splits.append((k, sep, ov))
                    else:
                        rv = replacements[rk]
                        if callable(rv):
                            rv = rv(key=rk, value=ov.decode('utf-8'),
                                    request=request)
                        if rv is not None:
                            new_splits.append((k, sep, rv.encode('utf-8')))
            request.body = b'&'.join(k if sep is None else b''.join([k, sep, v])
                                     for k, sep, v in new_splits)
    return request

def remove_post_data_parameters(request, post_data_parameters_to_remove):
    """
    Wrap replace_post_data_parameters() for API backward compatibility.
    """
    replacements = [(k, None) for k in post_data_parameters_to_remove]
    return replace_post_data_parameters(request, replacements)

def decode_response(response):
    """
    If the response is compressed with gzip or deflate:
      1. decompress the response body
      2. delete the content-encoding header
      3. update content-length header to decompressed length
    """
    def is_compressed(headers):
        encoding = headers.get('content-encoding', [])
        return encoding and encoding[0] in ('gzip', 'deflate')

    def decompress_body(body, encoding):
        """Returns decompressed body according to encoding using zlib.
        to (de-)compress gzip format, use wbits = zlib.MAX_WBITS | 16
        """
        if encoding == 'gzip':
            return zlib.decompress(body, zlib.MAX_WBITS | 16)
        else:  # encoding == 'deflate'
            return zlib.decompress(body)

    # Deepcopy here in case `headers` contain objects that could
    # be mutated by a shallow copy and corrupt the real response.
    response = copy.deepcopy(response)
    headers = CaseInsensitiveDict(response['headers'])
    if is_compressed(headers):
        encoding = headers['content-encoding'][0]
        headers['content-encoding'].remove(encoding)
        if not headers['content-encoding']:
            del headers['content-encoding']

        new_body = decompress_body(response['body']['string'], encoding)
        response['body']['string'] = new_body
        headers['content-length'] = [str(len(new_body))]
        response['headers'] = dict(headers)
    return response
<commit_msg>Use cgi.parse_header to get Content-Type<commit_after>from six import BytesIO, text_type
from six.moves.urllib.parse import urlparse, urlencode, urlunparse
import copy
import json
import zlib
import cgi

from .util import CaseInsensitiveDict

def replace_headers(request, replacements):
    """
    Replace headers in request according to replacements. The replacements
    should be a list of (key, value) pairs where the value can be any of:
      1. A simple replacement string value.
      2. None to remove the given header.
      3. A callable which accepts (key, value, request) and returns a string
         value or None.
    """
    new_headers = request.headers.copy()
    for k, rv in replacements:
        if k in new_headers:
            ov = new_headers.pop(k)
            if callable(rv):
                rv = rv(key=k, value=ov, request=request)
            if rv is not None:
                new_headers[k] = rv
    request.headers = new_headers
    return request

def remove_headers(request, headers_to_remove):
    """
    Wrap replace_headers() for API backward compatibility.
    """
    replacements = [(k, None) for k in headers_to_remove]
    return replace_headers(request, replacements)

def replace_query_parameters(request, replacements):
    """
    Replace query parameters in request according to replacements. The
    replacements should be a list of (key, value) pairs where the value can be
    any of:
      1. A simple replacement string value.
      2. None to remove the given header.
      3. A callable which accepts (key, value, request) and returns a string
         value or None.
    """
    query = request.query
    new_query = []
    replacements = dict(replacements)
    for k, ov in query:
        if k not in replacements:
            new_query.append((k, ov))
        else:
            rv = replacements[k]
            if callable(rv):
                rv = rv(key=k, value=ov, request=request)
            if rv is not None:
                new_query.append((k, rv))
    uri_parts = list(urlparse(request.uri))
    uri_parts[4] = urlencode(new_query)
    request.uri = urlunparse(uri_parts)
    return request

def remove_query_parameters(request, query_parameters_to_remove):
    """
    Wrap replace_query_parameters() for API backward compatibility.
    """
    replacements = [(k, None) for k in query_parameters_to_remove]
    return replace_query_parameters(request, replacements)

def replace_post_data_parameters(request, replacements):
    """
    Replace post data in request--either form data or json--according to
    replacements. The replacements should be a list of (key, value) pairs where
    the value can be any of:
      1. A simple replacement string value.
      2. None to remove the given header.
      3. A callable which accepts (key, value, request) and returns a string
         value or None.
    """
    replacements = dict(replacements)
    if request.method == 'POST' and not isinstance(request.body, BytesIO):
        # Handles headers['Content-Type'] in [application/json, application/json;charset=XXX, None]
        # Does not accept Content-Type in [application/json-not-really-json, application/custom-json]
        if 'application/json' in cgi.parse_header(request.headers.get('Content-Type') or ''):
            json_data = json.loads(request.body.decode('utf-8'))
            for k, rv in replacements.items():
                if k in json_data:
                    ov = json_data.pop(k)
                    if callable(rv):
                        rv = rv(key=k, value=ov, request=request)
                    if rv is not None:
                        json_data[k] = rv
            request.body = json.dumps(json_data).encode('utf-8')
        else:
            if isinstance(request.body, text_type):
                request.body = request.body.encode('utf-8')
            splits = [p.partition(b'=') for p in request.body.split(b'&')]
            new_splits = []
            for k, sep, ov in splits:
                if sep is None:
                    new_splits.append((k, sep, ov))
                else:
                    rk = k.decode('utf-8')
                    if rk not in replacements:
                        new_splits.append((k, sep, ov))
                    else:
                        rv = replacements[rk]
                        if callable(rv):
                            rv = rv(key=rk, value=ov.decode('utf-8'),
                                    request=request)
                        if rv is not None:
                            new_splits.append((k, sep, rv.encode('utf-8')))
            request.body = b'&'.join(k if sep is None else b''.join([k, sep, v])
                                     for k, sep, v in new_splits)
    return request

def remove_post_data_parameters(request, post_data_parameters_to_remove):
    """
    Wrap replace_post_data_parameters() for API backward compatibility.
    """
    replacements = [(k, None) for k in post_data_parameters_to_remove]
    return replace_post_data_parameters(request, replacements)

def decode_response(response):
    """
    If the response is compressed with gzip or deflate:
      1. decompress the response body
      2. delete the content-encoding header
      3. update content-length header to decompressed length
    """
    def is_compressed(headers):
        encoding = headers.get('content-encoding', [])
        return encoding and encoding[0] in ('gzip', 'deflate')

    def decompress_body(body, encoding):
        """Returns decompressed body according to encoding using zlib.
        to (de-)compress gzip format, use wbits = zlib.MAX_WBITS | 16
        """
        if encoding == 'gzip':
            return zlib.decompress(body, zlib.MAX_WBITS | 16)
        else:  # encoding == 'deflate'
            return zlib.decompress(body)

    # Deepcopy here in case `headers` contain objects that could
    # be mutated by a shallow copy and corrupt the real response.
    response = copy.deepcopy(response)
    headers = CaseInsensitiveDict(response['headers'])
    if is_compressed(headers):
        encoding = headers['content-encoding'][0]
        headers['content-encoding'].remove(encoding)
        if not headers['content-encoding']:
            del headers['content-encoding']

        new_body = decompress_body(response['body']['string'], encoding)
        response['body']['string'] = new_body
        headers['content-length'] = [str(len(new_body))]
        response['headers'] = dict(headers)
    return response
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""CLI utilities for vcspull.

vcspull.cli
~~~~~~~~~~~

"""

from __future__ import absolute_import, division, print_function, \
    with_statement

import os
import logging

import click

from .__about__ import __version__
from .util import lookup_repos
from .config import find_config_files, load_configs
from .log import DebugLogFormatter
from .repo import create_repo

log = logging.getLogger(__name__)

NO_REPOS_FOUND = """
    No repositories found.

    Check out the documentation at http://vcspull.rtfd.org for
    examples.
"""

config_dir = os.path.expanduser('~/.vcspull/')

def setup_logger(log=None, level='INFO'):
    """Setup logging for CLI use.

    :param log: instance of logger
    :type log: :py:class:`Logger`

    """
    if not log:
        log = logging.getLogger()
    if not log.handlers:
        channel = logging.StreamHandler()
        channel.setFormatter(DebugLogFormatter())

        log.setLevel(level)
        log.addHandler(channel)

class AliasedGroup(click.Group):

    def get_command(self, ctx, cmd_name):
        rv = click.Group.get_command(self, ctx, cmd_name)
        if rv is not None:
            return rv
        matches = [x for x in self.list_commands(ctx)
                   if x.startswith(cmd_name)]
        if not matches:
            return None
        elif len(matches) == 1:
            return click.Group.get_command(self, ctx, matches[0])
        ctx.fail('Too many matches: %s' % ', '.join(sorted(matches)))

@click.command(cls=AliasedGroup)
@click.option('--log_level', default='INFO',
              help='Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)')
@click.option('--version', is_flag=True, help='Print version info')
@click.version_option(version=__version__, message='%(prog)s %(version)s')
@click.pass_context
def cli(log_level, version):
    setup_logger(
        level=log_level.upper()
    )

@click.command(name='update')
@click.argument('repos', nargs=-1)
def update(repos):
    configs = load_configs(find_config_files(include_home=True))
    for repo in repos:
        dirmatch, vcsurlmatch = None, None
        if any(repo.startswith(n) for n in ['./', '/', '~', '$HOME']):
            dirmatch = repo
            repo = None
        elif any(repo.startswith(n) for n in ['http', 'git', 'svn']):
            vcsurlmatch = repo
            repo = None

        repos = lookup_repos(
            configs,
            dirmatch=dirmatch,
            vcsurlmatch=vcsurlmatch,
            namematch=repo
        )
        for repo_dict in repos:
            r = create_repo(**repo_dict)
            log.debug('%s' % r)
            r.update_repo()

cli.add_command(update)
<commit_msg>fix duplicate --version option, rm pass_context from cli<commit_after># -*- coding: utf-8 -*-
"""CLI utilities for vcspull.

vcspull.cli
~~~~~~~~~~~

"""

from __future__ import absolute_import, division, print_function, \
    with_statement

import os
import logging

import click

from .__about__ import __version__
from .util import lookup_repos
from .config import find_config_files, load_configs
from .log import DebugLogFormatter
from .repo import create_repo

log = logging.getLogger(__name__)

NO_REPOS_FOUND = """
    No repositories found.

    Check out the documentation at http://vcspull.rtfd.org for
    examples.
"""

config_dir = os.path.expanduser('~/.vcspull/')

def setup_logger(log=None, level='INFO'):
    """Setup logging for CLI use.

    :param log: instance of logger
    :type log: :py:class:`Logger`

    """
    if not log:
        log = logging.getLogger()
    if not log.handlers:
        channel = logging.StreamHandler()
        channel.setFormatter(DebugLogFormatter())

        log.setLevel(level)
        log.addHandler(channel)

class AliasedGroup(click.Group):

    def get_command(self, ctx, cmd_name):
        rv = click.Group.get_command(self, ctx, cmd_name)
        if rv is not None:
            return rv
        matches = [x for x in self.list_commands(ctx)
                   if x.startswith(cmd_name)]
        if not matches:
            return None
        elif len(matches) == 1:
            return click.Group.get_command(self, ctx, matches[0])
        ctx.fail('Too many matches: %s' % ', '.join(sorted(matches)))

@click.command(cls=AliasedGroup)
@click.option('--log_level', default='INFO',
              help='Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)')
@click.version_option(version=__version__, message='%(prog)s %(version)s')
def cli(log_level, version):
    setup_logger(
        level=log_level.upper()
    )

@click.command(name='update')
@click.argument('repos', nargs=-1)
def update(repos):
    configs = load_configs(find_config_files(include_home=True))
    for repo in repos:
        dirmatch, vcsurlmatch = None, None
        if any(repo.startswith(n) for n in ['./', '/', '~', '$HOME']):
            dirmatch = repo
            repo = None
        elif any(repo.startswith(n) for n in ['http', 'git', 'svn']):
            vcsurlmatch = repo
            repo = None

        repos = lookup_repos(
            configs,
            dirmatch=dirmatch,
            vcsurlmatch=vcsurlmatch,
            namematch=repo
        )
        for repo_dict in repos:
            r = create_repo(**repo_dict)
            log.debug('%s' % r)
            r.update_repo()

cli.add_command(update)
<|endoftext|>