<commit_before>"""Read and write notebooks as RStudio notebook files, with .Rmd extension.

Raw and markdown cells are converted to markdown, while code cells are
converted to code chunks. The transformation is reversible and all inputs
are preserved (not outputs, though).

Authors:

* Marc Wouts
"""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------

import os
import re
from enum import Enum
from nbformat.v4.rwbase import NotebookReader, NotebookWriter
from nbformat.v4.nbbase import (
    new_code_cell, new_markdown_cell, new_raw_cell, new_notebook
)
import nbformat
import yaml

from .chunk_options import to_metadata, to_chunk_options

# -----------------------------------------------------------------------------
# Code
# -----------------------------------------------------------------------------

def _language(metadata):
    try:
        return metadata.get('main_language') or \
               metadata.language_info.name.lower()
    except AttributeError:
        return u'python'

class State(Enum):
    NONE = 0
    HEADER = 1
    MARKDOWN = 2
    CODE = 3

_header_re = re.compile(r"^---\s*")
_end_code_re = re.compile(r"^```\s*")

_jupyter_languages = ['R', 'bash', 'sh', 'python', 'python2', 'python3',
                      'javascript', 'js', 'perl']
_jupyter_languages_re = [re.compile(r"^%%{}\s*".format(lang))
                         for lang in _jupyter_languages]

class RmdReader(NotebookReader):

    def __init__(self, markdown=False):
        self.start_code_re = re.compile(r"^```(.*)\s*") if markdown \
            else re.compile(r"^```\{(.*)\}\s*")

    def reads(self, s, **kwargs):
        return self.to_notebook(s, **kwargs)

    def to_notebook(self, s, **kwargs):
        lines = s.splitlines()

        metadata = {}
        cells = []
        cell_lines = []

        def add_cell(new_cell=new_markdown_cell):
            if len(cell_lines) == 0:
                return

            if cell_lines[-1] == '':
                if len(cell_lines) > 1 or len(cells) == 0:
                    cells.append(new_cell(
                        source=u'\n'.join(cell_lines[:-1])))
                else:
                    cells[-1]['metadata']['noskipline'] = True
                    cells.append(new_cell(
                        source=u'\n'.join(cell_lines)))
            else:
                cells.append(new_cell(source=u'\n'.join(cell_lines),
                                      metadata={u'noskipline': True}))

        cell_metadata = {}
        state = State.NONE
        testblankline = False

        for line in lines:
            if state is State.NONE:
                if _header_re.match(line):
                    state = State.HEADER
                    continue
                state = State.MARKDOWN

            if state is State.HEADER:
                # Unterminated header -> treat first lines as raw
                if self.start_code_re.match(line):
                    cell_lines = ['---'] + cell_lines
                    add_cell(new_cell=new_raw_cell)
                    cell_lines = []

                    chunk_options = self.start_code_re.findall(line)[0]
                    language, cell_metadata = to_metadata(chunk_options)
                    cell_metadata['language'] = language
                    state = State.CODE
                    continue

                if _header_re.match(line):
                    header = []
                    jupyter = []
                    in_header = True
                    for l in cell_lines:
                        if l.rstrip() == 'jupyter:':
                            in_header = False
                        if in_header:
                            header.append(l)
                        else:
                            jupyter.append(l)
                    if len(header):
                        cells.append(new_raw_cell(
                            source=u'\n'.join(['---'] + header + ['---'])))
                    if len(jupyter):
                        metadata = yaml.load(u'\n'.join(jupyter))['jupyter']
                    cell_lines = []
                    state = State.MARKDOWN
                    testblankline = True
                    continue

            if testblankline:
                # Set 'noskipline' metadata if
                # no blank line is found after cell
                testblankline = False
                if line == u'':
                    continue
                else:
                    if len(cells):
                        cells[-1][u'metadata'][u'noskipline'] = True

            if state is State.MARKDOWN:
                if self.start_code_re.match(line):
                    add_cell()
                    cell_lines = []

                    chunk_options = self.start_code_re.findall(line)[0]
                    language, cell_metadata = to_metadata(chunk_options)
                    cell_metadata['language'] = language
                    state = State.CODE
                    continue

            if state is State.CODE:
                if _end_code_re.match(line):
                    cells.append(new_code_cell(source=u'\n'.join(cell_lines),
                                               metadata=cell_metadata))
                    cell_lines = []
                    cell_metadata = {}
                    state = State.MARKDOWN
                    testblankline = True
                    continue

            cell_lines.append(line)

        # Append last cell if not empty
        if state is State.MARKDOWN:
            add_cell()
        elif state is State.CODE:
            cells.append(new_code_cell(source=u'\n'.join(cell_lines),
                                       metadata=cell_metadata))
        elif state is State.HEADER:
            cell_lines = ['---'] + cell_lines
            add_cell(new_cell=new_raw_cell)

        # Determine main language
        main_language = (metadata.get('main_language') or
                         metadata.get('language_info', {}).get('name'))
        if main_language is None:
            languages = dict(python=0)
            for c in cells:
                if c['cell_type'] == 'code':
                    language = c['metadata']['language']
                    if language == 'r':
                        language = 'R'
                    languages[language] = languages.get(language, 0) + 1

            main_language = max(languages, key=languages.get)

            # save main language when not given by kernel
            if main_language is not \
                    metadata.get('language_info', {}).get('name'):
                metadata['main_language'] = main_language

        # Remove 'language' meta data and add a magic if not main language
        for c in cells:
            if c['cell_type'] == 'code':
                language = c['metadata']['language']
                del c['metadata']['language']
                if language != main_language and \
                        language in _jupyter_languages:
                    c['source'] = u'%%{}\n'.format(language) + c['source']

        nb = new_notebook(cells=cells, metadata=metadata)
        return nb

def _as_dict(metadata):
    if isinstance(metadata, nbformat.NotebookNode):
        return {k: _as_dict(metadata[k]) for k in metadata.keys()}
    return metadata

class RmdWriter(NotebookWriter):

    def __init__(self, markdown=False):
        self.markdown = markdown

    def writes(self, nb):
        default_language = _language(nb.metadata)
        metadata = _as_dict(nb.metadata)

        if 'main_language' in metadata:
            # is 'main language' redundant with kernel info?
            if metadata['main_language'] is \
                    metadata.get('language_info', {}).get('name'):
                del metadata['main_language']
            # is 'main language' redundant with cell language?
            elif metadata.get('language_info', {}).get('name') is None:
                languages = dict(python=0)
                for c in nb.cells:
                    if c.cell_type == 'code':
                        input = c.source.splitlines()
                        language = default_language
                        if len(input):
                            for lang, pattern in zip(_jupyter_languages,
                                                     _jupyter_languages_re):
                                if pattern.match(input[0]):
                                    language = lang

                        if language == 'r':
                            language = 'R'

                        languages[language] = 1 + languages.get(
                            language, 0)

                cell_main_language = max(languages, key=languages.get)
                if metadata['main_language'] == cell_main_language:
                    del metadata['main_language']

        lines = []
        header_inserted = len(metadata) == 0

        for cell in nb.cells:
            if cell.cell_type == u'raw':
                # Is this the Rmd header?
                # Starts and ends with '---',
                # and can be parsed with yaml
                if len(lines) == 0 and not header_inserted:
                    header = cell.get(u'source', '').splitlines()
                    if len(header) >= 2 and _header_re.match(header[0]) \
                            and _header_re.match(header[-1]):
                        try:
                            header = header[1:-1]
                            yaml.load(u'\n'.join(header))
                            if not self.markdown:
                                header.extend(
                                    yaml.safe_dump(
                                        {u'jupyter': metadata},
                                        default_flow_style=False).splitlines())
                            lines = [u'---'] + header + [u'---']
                            header_inserted = True
                        except yaml.ScannerError:
                            pass
                    if not header_inserted:
                        lines.append(cell.get(u'source', ''))
                else:
                    lines.append(cell.get(u'source', ''))
                if not cell.get(u'metadata', {}).get('noskipline', False):
                    lines.append(u'')
            elif cell.cell_type == u'markdown':
                lines.append(cell.get(u'source', ''))
                if not cell.get(u'metadata', {}).get('noskipline', False):
                    lines.append(u'')
            elif cell.cell_type == u'code':
                input = cell.get(u'source').splitlines()
                cell_metadata = cell.get(u'metadata', {})
                if 'noskipline' in cell_metadata:
                    noskipline = cell_metadata['noskipline']
                    del cell_metadata['noskipline']
                else:
                    noskipline = False
                language = None
                if len(input):
                    for lang, pattern in zip(_jupyter_languages,
                                             _jupyter_languages_re):
                        if pattern.match(input[0]):
                            language = lang
                            input = input[1:]
                            break
                if language is None:
                    language = default_language
                if self.markdown:
                    lines.append(
                        u'```' + to_chunk_options(language, cell_metadata))
                else:
                    lines.append(
                        u'```{' +
                        to_chunk_options(language, cell_metadata) + '}')
                if input is not None:
                    lines.extend(input)
                lines.append(u'```')
                if not noskipline:
                    lines.append(u'')

        if not self.markdown and not header_inserted and len(metadata):
            header = yaml.safe_dump({u'jupyter': metadata},
                                    default_flow_style=False).splitlines()
            lines = [u'---'] + header + [u'---', u''] + lines

        lines.append(u'')

        return u'\n'.join(lines)

_reader = RmdReader()
_writer = RmdWriter()

reads = _reader.reads
read = _reader.read
to_notebook = _reader.to_notebook
write = _writer.write
writes = _writer.writes

_md_reader = RmdReader(markdown=True)
_md_writer = RmdWriter(markdown=True)

md_reads = _md_reader.reads
md_read = _md_reader.read
md_to_notebook = _md_reader.to_notebook
md_write = _md_writer.write
md_writes = _md_writer.writes

def readf(nb_file):
    """
    Load the notebook from the desired file
    :param nb_file: file with .ipynb or .Rmd extension
    :return: the notebook
    """
    file, ext = os.path.splitext(nb_file)
    with open(nb_file) as fp:
        if ext == '.Rmd':
            return read(fp)
        elif ext == '.md':
            return md_read(fp)
        elif ext == '.ipynb':
            return nbformat.read(fp, as_version=4)
        else:
            raise TypeError(
                'File {} has incorrect extension (.Rmd or .md or '
                '.ipynb expected)'.format(nb_file))

def writef(nb, nb_file):
    """
    Save the notebook in the desired file
    :param nb: notebook
    :param nb_file: file with .ipynb or .Rmd extension
    :return:
    """

    file, ext = os.path.splitext(nb_file)
    with open(nb_file, 'w') as fp:
        if ext == '.Rmd':
            write(nb, fp)
        elif ext == '.md':
            md_write(nb, fp)
        elif ext == '.ipynb':
            nbformat.write(nb, fp)
        else:
            raise TypeError(
                'File {} has incorrect extension (.Rmd or .md or '
                '.ipynb expected)'.format(nb_file))

def readme():
    """
    Contents of README.md
    :return:
    """
    readme_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                               '..', 'README.md')
    with open(readme_path) as fh:
        return fh.read()
<commit_msg>python as prefered language<commit_after>"""Read and write notebooks as RStudio notebook files, with .Rmd extension.

Raw and markdown cells are converted to markdown, while code cells are
converted to code chunks. The transformation is reversible and all inputs
are preserved (not outputs, though).

Authors:

* Marc Wouts
"""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------

import os
import re
from enum import Enum
from nbformat.v4.rwbase import NotebookReader, NotebookWriter
from nbformat.v4.nbbase import (
    new_code_cell, new_markdown_cell, new_raw_cell, new_notebook
)
import nbformat
import yaml

from .chunk_options import to_metadata, to_chunk_options

# -----------------------------------------------------------------------------
# Code
# -----------------------------------------------------------------------------

def _language(metadata):
    try:
        return metadata.get('main_language') or \
               metadata.language_info.name.lower()
    except AttributeError:
        return u'python'

class State(Enum):
    NONE = 0
    HEADER = 1
    MARKDOWN = 2
    CODE = 3

_header_re = re.compile(r"^---\s*")
_end_code_re = re.compile(r"^```\s*")

_jupyter_languages = ['R', 'bash', 'sh', 'python', 'python2', 'python3',
                      'javascript', 'js', 'perl']
_jupyter_languages_re = [re.compile(r"^%%{}\s*".format(lang))
                         for lang in _jupyter_languages]

class RmdReader(NotebookReader):

    def __init__(self, markdown=False):
        self.start_code_re = re.compile(r"^```(.*)\s*") if markdown \
            else re.compile(r"^```\{(.*)\}\s*")

    def reads(self, s, **kwargs):
        return self.to_notebook(s, **kwargs)

    def to_notebook(self, s, **kwargs):
        lines = s.splitlines()

        metadata = {}
        cells = []
        cell_lines = []

        def add_cell(new_cell=new_markdown_cell):
            if len(cell_lines) == 0:
                return

            if cell_lines[-1] == '':
                if len(cell_lines) > 1 or len(cells) == 0:
                    cells.append(new_cell(
                        source=u'\n'.join(cell_lines[:-1])))
                else:
                    cells[-1]['metadata']['noskipline'] = True
                    cells.append(new_cell(
                        source=u'\n'.join(cell_lines)))
            else:
                cells.append(new_cell(source=u'\n'.join(cell_lines),
                                      metadata={u'noskipline': True}))

        cell_metadata = {}
        state = State.NONE
        testblankline = False

        for line in lines:
            if state is State.NONE:
                if _header_re.match(line):
                    state = State.HEADER
                    continue
                state = State.MARKDOWN

            if state is State.HEADER:
                # Unterminated header -> treat first lines as raw
                if self.start_code_re.match(line):
                    cell_lines = ['---'] + cell_lines
                    add_cell(new_cell=new_raw_cell)
                    cell_lines = []

                    chunk_options = self.start_code_re.findall(line)[0]
                    language, cell_metadata = to_metadata(chunk_options)
                    cell_metadata['language'] = language
                    state = State.CODE
                    continue

                if _header_re.match(line):
                    header = []
                    jupyter = []
                    in_header = True
                    for l in cell_lines:
                        if l.rstrip() == 'jupyter:':
                            in_header = False
                        if in_header:
                            header.append(l)
                        else:
                            jupyter.append(l)
                    if len(header):
                        cells.append(new_raw_cell(
                            source=u'\n'.join(['---'] + header + ['---'])))
                    if len(jupyter):
                        metadata = yaml.load(u'\n'.join(jupyter))['jupyter']
                    cell_lines = []
                    state = State.MARKDOWN
                    testblankline = True
                    continue

            if testblankline:
                # Set 'noskipline' metadata if
                # no blank line is found after cell
                testblankline = False
                if line == u'':
                    continue
                else:
                    if len(cells):
                        cells[-1][u'metadata'][u'noskipline'] = True

            if state is State.MARKDOWN:
                if self.start_code_re.match(line):
                    add_cell()
                    cell_lines = []

                    chunk_options = self.start_code_re.findall(line)[0]
                    language, cell_metadata = to_metadata(chunk_options)
                    cell_metadata['language'] = language
                    state = State.CODE
                    continue

            if state is State.CODE:
                if _end_code_re.match(line):
                    cells.append(new_code_cell(source=u'\n'.join(cell_lines),
                                               metadata=cell_metadata))
                    cell_lines = []
                    cell_metadata = {}
                    state = State.MARKDOWN
                    testblankline = True
                    continue

            cell_lines.append(line)

        # Append last cell if not empty
        if state is State.MARKDOWN:
            add_cell()
        elif state is State.CODE:
            cells.append(new_code_cell(source=u'\n'.join(cell_lines),
                                       metadata=cell_metadata))
        elif state is State.HEADER:
            cell_lines = ['---'] + cell_lines
            add_cell(new_cell=new_raw_cell)

        # Determine main language
        main_language = (metadata.get('main_language') or
                         metadata.get('language_info', {}).get('name'))
        if main_language is None:
            languages = dict(python=0.5)
            for c in cells:
                if c['cell_type'] == 'code':
                    language = c['metadata']['language']
                    if language == 'r':
                        language = 'R'
                    languages[language] = languages.get(language, 0.0) + 1

            main_language = max(languages, key=languages.get)

            # save main language when not given by kernel
            if main_language is not \
                    metadata.get('language_info', {}).get('name'):
                metadata['main_language'] = main_language

        # Remove 'language' meta data and add a magic if not main language
        for c in cells:
            if c['cell_type'] == 'code':
                language = c['metadata']['language']
                del c['metadata']['language']
                if language != main_language and \
                        language in _jupyter_languages:
                    c['source'] = u'%%{}\n'.format(language) + c['source']

        nb = new_notebook(cells=cells, metadata=metadata)
        return nb

def _as_dict(metadata):
    if isinstance(metadata, nbformat.NotebookNode):
        return {k: _as_dict(metadata[k]) for k in metadata.keys()}
    return metadata

class RmdWriter(NotebookWriter):

    def __init__(self, markdown=False):
        self.markdown = markdown

    def writes(self, nb):
        default_language = _language(nb.metadata)
        metadata = _as_dict(nb.metadata)

        if 'main_language' in metadata:
            # is 'main language' redundant with kernel info?
            if metadata['main_language'] is \
                    metadata.get('language_info', {}).get('name'):
                del metadata['main_language']
            # is 'main language' redundant with cell language?
            elif metadata.get('language_info', {}).get('name') is None:
                languages = dict(python=0.5)
                for c in nb.cells:
                    if c.cell_type == 'code':
                        input = c.source.splitlines()
                        language = default_language
                        if len(input):
                            for lang, pattern in zip(_jupyter_languages,
                                                     _jupyter_languages_re):
                                if pattern.match(input[0]):
                                    language = lang

                        if language == 'r':
                            language = 'R'

                        languages[language] = 1 + languages.get(
                            language, 0.0)

                cell_main_language = max(languages, key=languages.get)
                if metadata['main_language'] == cell_main_language:
                    del metadata['main_language']

        lines = []
        header_inserted = len(metadata) == 0

        for cell in nb.cells:
            if cell.cell_type == u'raw':
                # Is this the Rmd header?
                # Starts and ends with '---',
                # and can be parsed with yaml
                if len(lines) == 0 and not header_inserted:
                    header = cell.get(u'source', '').splitlines()
                    if len(header) >= 2 and _header_re.match(header[0]) \
                            and _header_re.match(header[-1]):
                        try:
                            header = header[1:-1]
                            yaml.load(u'\n'.join(header))
                            if not self.markdown:
                                header.extend(
                                    yaml.safe_dump(
                                        {u'jupyter': metadata},
                                        default_flow_style=False).splitlines())
                            lines = [u'---'] + header + [u'---']
                            header_inserted = True
                        except yaml.ScannerError:
                            pass
                    if not header_inserted:
                        lines.append(cell.get(u'source', ''))
                else:
                    lines.append(cell.get(u'source', ''))
                if not cell.get(u'metadata', {}).get('noskipline', False):
                    lines.append(u'')
            elif cell.cell_type == u'markdown':
                lines.append(cell.get(u'source', ''))
                if not cell.get(u'metadata', {}).get('noskipline', False):
                    lines.append(u'')
            elif cell.cell_type == u'code':
                input = cell.get(u'source').splitlines()
                cell_metadata = cell.get(u'metadata', {})
                if 'noskipline' in cell_metadata:
                    noskipline = cell_metadata['noskipline']
                    del cell_metadata['noskipline']
                else:
                    noskipline = False
                language = None
                if len(input):
                    for lang, pattern in zip(_jupyter_languages,
                                             _jupyter_languages_re):
                        if pattern.match(input[0]):
                            language = lang
                            input = input[1:]
                            break
                if language is None:
                    language = default_language
                if self.markdown:
                    lines.append(
                        u'```' + to_chunk_options(language, cell_metadata))
                else:
                    lines.append(
                        u'```{' +
                        to_chunk_options(language, cell_metadata) + '}')
                if input is not None:
                    lines.extend(input)
                lines.append(u'```')
                if not noskipline:
                    lines.append(u'')

        if not self.markdown and not header_inserted and len(metadata):
            header = yaml.safe_dump({u'jupyter': metadata},
                                    default_flow_style=False).splitlines()
            lines = [u'---'] + header + [u'---', u''] + lines

        lines.append(u'')

        return u'\n'.join(lines)

_reader = RmdReader()
_writer = RmdWriter()

reads = _reader.reads
read = _reader.read
to_notebook = _reader.to_notebook
write = _writer.write
writes = _writer.writes

_md_reader = RmdReader(markdown=True)
_md_writer = RmdWriter(markdown=True)

md_reads = _md_reader.reads
md_read = _md_reader.read
md_to_notebook = _md_reader.to_notebook
md_write = _md_writer.write
md_writes = _md_writer.writes

def readf(nb_file):
    """
    Load the notebook from the desired file
    :param nb_file: file with .ipynb or .Rmd extension
    :return: the notebook
    """
    file, ext = os.path.splitext(nb_file)
    with open(nb_file) as fp:
        if ext == '.Rmd':
            return read(fp)
        elif ext == '.md':
            return md_read(fp)
        elif ext == '.ipynb':
            return nbformat.read(fp, as_version=4)
        else:
            raise TypeError(
                'File {} has incorrect extension (.Rmd or .md or '
                '.ipynb expected)'.format(nb_file))

def writef(nb, nb_file):
    """
    Save the notebook in the desired file
    :param nb: notebook
    :param nb_file: file with .ipynb or .Rmd extension
    :return:
    """

    file, ext = os.path.splitext(nb_file)
    with open(nb_file, 'w') as fp:
        if ext == '.Rmd':
            write(nb, fp)
        elif ext == '.md':
            md_write(nb, fp)
        elif ext == '.ipynb':
            nbformat.write(nb, fp)
        else:
            raise TypeError(
                'File {} has incorrect extension (.Rmd or .md or '
                '.ipynb expected)'.format(nb_file))

def readme():
    """
    Contents of README.md
    :return:
    """
    readme_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                               '..', 'README.md')
    with open(readme_path) as fh:
        return fh.read()
<|endoftext|>


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

"""
Converts GFF3 representing gene models to Genbank flat-file format.

GFF3 specification:
http://www.sequenceontology.org/gff3.shtml

Genbank flat file specification:
https://www.ncbi.nlm.nih.gov/Sitemap/samplerecord.html

--molecule_type:
    http://www.ncbi.nlm.nih.gov/Sequin/sequin.hlp.html#Molecule

--genbank_division:
    http://www.ncbi.nlm.nih.gov/Sitemap/samplerecord.html#GenBankDivisionB

"""

import argparse
import biocodegenbank
import biocodegff
import os
import sys

from jinja2 import Environment, FileSystemLoader

PATH = os.path.dirname(os.path.abspath(__file__))
TEMPLATE_ENVIRONMENT = Environment(
    autoescape=False,
    loader=FileSystemLoader(os.path.join(PATH, '../data')),
    trim_blocks=False)

def main():
    parser = argparse.ArgumentParser( description='Converts GFF3 into a GenBank flat file')

    ## output file to be written
    parser.add_argument('-i', '--input_file', type=str, required=True, help='Path to an input GFF3 file to be read' )
    parser.add_argument('-o', '--output_file', type=str, required=False, help='Path to a Genbank flat file to be created. Supersedes --output_dir if both are specified.' )
    parser.add_argument('-od', '--output_dir', type=str, required=False, help='Path to an output directory. If this option is specified then each input assembly will be written to a separate GenBank output file, named with the assembly_id.' )
    parser.add_argument('-mt', '--molecule_type', type=str, required=False, default='DNA', help='Molecule type' )
    parser.add_argument('-gbd', '--genbank_division', type=str, required=False, default='.', help='GenBank Division (3-letter abbreviation)' )
    parser.add_argument('-md', '--modification_date', type=str, required=False, default='DD-MMM-YYYY', help='The modification date for header in format like 21-JUN-1999' )
    parser.add_argument('-org', '--organism', type=str, required=False, default='.', help='Full organism name (including strain)' )
    parser.add_argument('-str', '--strain', type=str, required=False, help="Only the strain designation, which is written to the FEATURES.source element" )
    parser.add_argument('-d', '--definition', type=str, required=False, default='.', help='Brief description of sequence; includes information such as source organism, gene name/protein name, or some description of the sequence\'s function.' )
    parser.add_argument('-s', '--source', type=str, required=False, default='.', help='Free-format information including an abbreviated form of the organism name, sometimes followed by a molecule type.' )
    parser.add_argument('-t', '--taxon_id', type=int, required=False, help='NCBI taxon ID, if known' )
    parser.add_argument('-l', '--lineage', type=str, required=False, default='Unknown', help='Semicolon-delimited lineage of the organism e.g., "Eukaryota; Alveolata; Apicomplexa; Aconoidasida; Piroplasmida; Theileriidae; Theileria"' )
    parser.add_argument('-seq', '--include_sequence', action='store_true', help='Include sequence (if present) in the output GenBank flat file(s).' )
    parser.add_argument('-p', '--locus_id_prefix', required=False, default='', help='Prefix to add to the GenBank LOCUS id in the output GenBank flat file(s).' )
    args = parser.parse_args()

    # check that output directory exists
    if args.output_dir is not None:
        if not os.path.isdir(args.output_dir):
            sys.stderr.write("FATAL: the specified output directory (" + args.output_dir + ") does not exist\n");
            exit(1)

    # line-wrap lineage to stay below 79 character GenBank flat file width
    lineage = biocodegenbank.line_wrap_lineage_string( args.lineage )

    (assemblies, features) = biocodegff.get_gff3_features( args.input_file )
    ofh = sys.stdout
    if args.output_file is not None:
        if args.output_dir is None:
            ofh = open(args.output_file, 'wt')
        else:
            sys.stderr.write("WARN: both -o/--output_file and -od/--output_dir were passed so the former will be ignored\n")

    for assembly_id in assemblies:
        locus_id = args.locus_id_prefix + assembly_id
        if args.output_dir is not None:
            ofn = args.output_dir + "/" + locus_id + ".gbk"
            ofh = open(ofn, 'wt')
        assembly = assemblies[assembly_id]

        context = { 'locus':locus_id, 'molecule_size':assembly.length, 'molecule_type':args.molecule_type,
                    'division':args.genbank_division, 'modification_date':args.modification_date,
                    'accession':'.', 'version':'.',
                    'source':args.source, 'definition':args.definition, 'organism':args.organism,
                    'lineage':lineage
        }
        header = TEMPLATE_ENVIRONMENT.get_template('genbank_flat_file_header.template').render(context)
        ofh.write(header)
        ofh.write("\nFEATURES             Location/Qualifiers\n")
        ofh.write("     source          1..{0}\n".format(assembly.length))
        ofh.write("                     /organism=\"{0}\"\n".format(args.organism))
        ofh.write("                     /mol_type=\"genomic DNA\"\n")

        if args.strain is not None:
            ofh.write("                     /strain=\"{0}\"\n".format(args.strain))

        if args.taxon_id is not None:
            ofh.write("                     /db_xref=\"taxon:{0}\"\n".format(args.taxon_id))

        for gene in assemblies[assembly_id].genes():
            biocodegenbank.print_biogene( gene=gene, fh=ofh, on=assembly )

        if args.include_sequence:
            ofh.write("ORIGIN\n")
            biocodegenbank.print_sequence( seq=assembly.residues, fh=ofh )

        ofh.write("//\n")
        # there may be multiple output files
        if args.output_dir is not None:
            ofh.close()

    # there is only one output file
    if args.output_dir is None:
        ofh.close()

if __name__ == '__main__':
    main()

<commit_msg>minor: added assumption documentation<commit_after>#!/usr/bin/env python3

"""
Converts GFF3 representing gene models to Genbank flat-file format.

GFF3 specification:
http://www.sequenceontology.org/gff3.shtml

Genbank flat file specification:
https://www.ncbi.nlm.nih.gov/Sitemap/samplerecord.html

--molecule_type:
    http://www.ncbi.nlm.nih.gov/Sequin/sequin.hlp.html#Molecule

--genbank_division:
    http://www.ncbi.nlm.nih.gov/Sitemap/samplerecord.html#GenBankDivisionB

MAJOR assumptions:
- Functional annotation is attached to a polypeptide feature

"""

import argparse
import biocodegenbank
import biocodegff
import os
import sys

from jinja2 import Environment, FileSystemLoader

PATH = os.path.dirname(os.path.abspath(__file__))
TEMPLATE_ENVIRONMENT = Environment(
    autoescape=False,
    loader=FileSystemLoader(os.path.join(PATH, '../data')),
    trim_blocks=False)

def main():
    parser = argparse.ArgumentParser( description='Converts GFF3 into a GenBank flat file')

    ## output file to be written
    parser.add_argument('-i', '--input_file', type=str, required=True, help='Path to an input GFF3 file to be read' )
    parser.add_argument('-o', '--output_file', type=str, required=False, help='Path to a Genbank flat file to be created. Supersedes --output_dir if both are specified.' )
    parser.add_argument('-od', '--output_dir', type=str, required=False, help='Path to an output directory. If this option is specified then each input assembly will be written to a separate GenBank output file, named with the assembly_id.' )
    parser.add_argument('-mt', '--molecule_type', type=str, required=False, default='DNA', help='Molecule type' )
    parser.add_argument('-gbd', '--genbank_division', type=str, required=False, default='.', help='GenBank Division (3-letter abbreviation)' )
    parser.add_argument('-md', '--modification_date', type=str, required=False, default='DD-MMM-YYYY', help='The modification date for header in format like 21-JUN-1999' )
    parser.add_argument('-org', '--organism', type=str, required=False, default='.', help='Full organism name (including strain)' )
    parser.add_argument('-str', '--strain', type=str, required=False, help="Only the strain designation, which is written to the FEATURES.source element" )
    parser.add_argument('-d', '--definition', type=str, required=False, default='.', help='Brief description of sequence; includes information such as source organism, gene name/protein name, or some description of the sequence\'s function.' )
    parser.add_argument('-s', '--source', type=str, required=False, default='.', help='Free-format information including an abbreviated form of the organism name, sometimes followed by a molecule type.' )
    parser.add_argument('-t', '--taxon_id', type=int, required=False, help='NCBI taxon ID, if known' )
    parser.add_argument('-l', '--lineage', type=str, required=False, default='Unknown', help='Semicolon-delimited lineage of the organism e.g., "Eukaryota; Alveolata; Apicomplexa; Aconoidasida; Piroplasmida; Theileriidae; Theileria"' )
    parser.add_argument('-seq', '--include_sequence', action='store_true', help='Include sequence (if present) in the output GenBank flat file(s).' )
    parser.add_argument('-p', '--locus_id_prefix', required=False, default='', help='Prefix to add to the GenBank LOCUS id in the output GenBank flat file(s).' )
    args = parser.parse_args()

    # check that output directory exists
    if args.output_dir is not None:
        if not os.path.isdir(args.output_dir):
            sys.stderr.write("FATAL: the specified output directory (" + args.output_dir + ") does not exist\n");
            exit(1)

    # line-wrap lineage to stay below 79 character GenBank flat file width
    lineage = biocodegenbank.line_wrap_lineage_string( args.lineage )

    (assemblies, features) = biocodegff.get_gff3_features( args.input_file )
    ofh = sys.stdout
    if args.output_file is not None:
        if args.output_dir is None:
            ofh = open(args.output_file, 'wt')
        else:
            sys.stderr.write("WARN: both -o/--output_file and -od/--output_dir were passed so the former will be ignored\n")

    for assembly_id in assemblies:
        locus_id = args.locus_id_prefix + assembly_id
        if args.output_dir is not None:
            ofn = args.output_dir + "/" + locus_id + ".gbk"
            ofh = open(ofn, 'wt')
        assembly = assemblies[assembly_id]

        context = { 'locus':locus_id, 'molecule_size':assembly.length, 'molecule_type':args.molecule_type,
                    'division':args.genbank_division, 'modification_date':args.modification_date,
                    'accession':'.', 'version':'.',
                    'source':args.source, 'definition':args.definition, 'organism':args.organism,
                    'lineage':lineage
        }
        header = TEMPLATE_ENVIRONMENT.get_template('genbank_flat_file_header.template').render(context)
        ofh.write(header)
        ofh.write("\nFEATURES             Location/Qualifiers\n")
        ofh.write("     source          1..{0}\n".format(assembly.length))
        ofh.write("                     /organism=\"{0}\"\n".format(args.organism))
        ofh.write("                     /mol_type=\"genomic DNA\"\n")

        if args.strain is not None:
            ofh.write("                     /strain=\"{0}\"\n".format(args.strain))

        if args.taxon_id is not None:
            ofh.write("                     /db_xref=\"taxon:{0}\"\n".format(args.taxon_id))

        for gene in assemblies[assembly_id].genes():
            biocodegenbank.print_biogene( gene=gene, fh=ofh, on=assembly )

        if args.include_sequence:
            ofh.write("ORIGIN\n")
            biocodegenbank.print_sequence( seq=assembly.residues, fh=ofh )

        ofh.write("//\n")
        # there may be multiple output files
        if args.output_dir is not None:
            ofh.close()

    # there is only one output file
    if args.output_dir is None:
        ofh.close()

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Fetch Tweets utility.

Get Tweet and Profile data from the Twitter API, after filtering Profiles
by Category. The Category could be an industry name, or it could be a
compiled favourites list of Profiles which the application user wants to
routinely fetch Tweets for. Note that this utility cannot be used to lookup
a specific screen name, but rather Categories containing Profile records.

A configured campaign name is allocated to fetched Tweets, in addition to
any possible existing campaign names on Tweets which are updated. No custom
campaign is necessary to allow to the Tweets, since the fetch is done based on
screen names rather than a Twitter API search query. If a Category
was used to store up a Profile's tweets, the Tweets can always be selected
from the db later by filtering on Tweets of Profiles in a given Category.
"""
import argparse
import os
import sys

from sqlobject import SQLObjectNotFound
from sqlobject.sqlbuilder import IN, AND

# Allow imports to be done when executing this file directly.
appDir = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                      os.path.pardir, os.path.pardir))
sys.path.insert(0, appDir)

from lib import database as db
from lib.tweets import insertOrUpdateTweetBatch
from lib.query.tweets.categories import printAvailableCategories

# We fetch Tweets in this script by getting recent activity of Profiles,
# rather than through a search or fetching by Tweet GUID. Therefore assign
# this Campaign name to Tweets.
CAMPAIGN_NAME = u"_PROFILE_TIMELINE"

def main():
    """
    Command-line interface to fetch Tweet data for Profile Categories.
    """
    parser = argparse.ArgumentParser(
        description="""Fetch Tweets utility. Filter Profiles in the db using
            a Category input, update them with new data and insert or update
            the most recent Tweets for each. Tweets are assigned to the
            '{0}' Campaign.""".format(CAMPAIGN_NAME))

    viewGrp = parser.add_argument_group("View", "Print data to stdout")
    viewGrp.add_argument(
        '-a', '--available',
        action='store_true',
        help="Output available Categories in db, with Profile counts for each."
    )

    updateGrp = parser.add_argument_group("Update", "Create or update Tweet"
                                                    " records.")
    updateGrp.add_argument(
        '-c', '--categories',
        metavar='CATEGORY',
        nargs='+',
        help="""List of one or more existing Categories names in the db.
            Profiles are filtered to only those which have been assigned to at
            least one of the supplied CATEGORIES values, then their Tweets
            are fetched and stored. Values must be separated by a space and any
            multi-word values or values containing a hash symbol must be
            enclosed in single quotes.
            e.g. -c 'first cat' second 'third cat' '#fourth'"""
    )
    updateGrp.add_argument(
        '-t', '--tweets-per-profile',
        type=int,
        metavar='N',
        default=200,
        help="""Default 200. Count of Tweets to fetch and store for each
            profile. A value up to 200 takes a fixed time to query one
            page of Tweets for a Profile, while higher values require
            querying more pages and therefore will take longer per
            Profile and lead to a higher chance of hitting rate limits.
            A higher value also requires additional time to create or update
            records."""
    )
    updateGrp.add_argument(
        '-v', '--verbose',
        action='store_true',
        help="""If supplied, pretty print Tweet data fetched from the
            Twitter API. Otherwise only a count of Tweets is printed
            upon completion.""")
    updateGrp.add_argument(
        '-n', '--no-write',
        action='store_true',
        help="If supplied, do not write data to the db."
    )

    args = parser.parse_args()

    if args.available:
        printAvailableCategories()
    elif args.categories:
        inputCategories = args.categories

        categoryResult = db.Category.select(
            IN(db.Category.q.name,
               inputCategories)
        )
        dbCategoryNames = [c.name for c in list(categoryResult)]
        missing = set(inputCategories) - set(dbCategoryNames)
        assert not missing, u"Input categories not found in db: \n- {0}"\
                            .format('\n- '.join(missing))

        # Here the AND is required to include SQLObject j-magic, so that
        # Profiles are filtered by Category.
        profResults = db.Profile.select(
            AND(db.Profile.j.categories,
                IN(db.Category.q.name,
                   inputCategories))
        )
        profCount = profResults.count()
        print "Fetching Tweets for {0:,d} Profiles".format(profCount)

        try:
            campaignRec = db.Campaign.byName(CAMPAIGN_NAME)
        except SQLObjectNotFound:
            campaignRec = db.Campaign(name=CAMPAIGN_NAME, searchQuery=None)

        if profCount:
            insertOrUpdateTweetBatch(
                profResults,
                args.tweets_per_profile,
                verbose=args.verbose,
                writeToDB=not(args.no_write),
                campaignRec=campaignRec
            )

if __name__ == '__main__':
    main()
<commit_msg>feature: Add flag to fetchTweets.py to only update tweet engagements<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Fetch Tweets utility.

Get Tweet and Profile data from the Twitter API, after filtering Profiles
by Category. The Category could be an industry name, or it could be a
compiled favourites list of Profiles which the application user wants to
routinely fetch Tweets for. Note that this utility cannot be used to lookup
a specific screen name, but rather Categories containing Profile records.

A configured campaign name is allocated to fetched Tweets, in addition to
any possible existing campaign names on Tweets which are updated. No custom
campaign is necessary to allow to the Tweets, since the fetch is done based on
screen names rather than a Twitter API search query. If a Category
was used to store up a Profile's tweets, the Tweets can always be selected
from the db later by filtering on Tweets of Profiles in a given Category.
"""
import argparse
import os
import sys

from sqlobject import SQLObjectNotFound
from sqlobject.sqlbuilder import IN, AND

# Allow imports to be done when executing this file directly.
appDir = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                      os.path.pardir, os.path.pardir))
sys.path.insert(0, appDir)

from lib import database as db
from lib.tweets import insertOrUpdateTweetBatch
from lib.query.tweets.categories import printAvailableCategories

# We fetch Tweets in this script by getting recent activity of Profiles,
# rather than through a search or fetching by Tweet GUID. Therefore assign
# this Campaign name to Tweets.
CAMPAIGN_NAME = u"_PROFILE_TIMELINE"

def main():
    """
    Command-line interface to fetch Tweet data for Profile Categories.
    """
    parser = argparse.ArgumentParser(
        description="""Fetch Tweets utility. Filter Profiles in the db using
            a Category input, update them with new data and insert or update
            the most recent Tweets for each. Tweets are assigned to the
            '{0}' Campaign.""".format(CAMPAIGN_NAME))

    viewGrp = parser.add_argument_group("View", "Print data to stdout")
    viewGrp.add_argument(
        '-a', '--available',
        action='store_true',
        help="Output available Categories in db, with Profile counts for each."
    )

    updateGrp = parser.add_argument_group("Update", "Create or update Tweet"
                                                    " records.")
    updateGrp.add_argument(
        '-c', '--categories',
        metavar='CATEGORY',
        nargs='+',
        help="""List of one or more existing Categories names in the db.
            Profiles are filtered to only those which have been assigned to at
            least one of the supplied CATEGORIES values, then their Tweets
            are fetched and stored. Values must be separated by a space and any
            multi-word values or values containing a hash symbol must be
            enclosed in single quotes.
            e.g. -c 'first cat' second 'third cat' '#fourth'"""
    )
    updateGrp.add_argument(
        '-t', '--tweets-per-profile',
        type=int,
        metavar='N',
        default=200,
        help="""Default 200. Count of Tweets to fetch and store for each
            profile. A value up to 200 takes a fixed time to query one
            page of Tweets for a Profile, while higher values require
            querying more pages and therefore will take longer per
            Profile and lead to a higher chance of hitting rate limits.
            A higher value also requires additional time to create or update
            records."""
    )
    updateGrp.add_argument(
        '-v', '--verbose',
        action='store_true',
        help="""If supplied, pretty print Tweet data fetched from the
            Twitter API. Otherwise only a count of Tweets is printed
            upon completion.""")
    updateGrp.add_argument(
        '-n', '--no-write',
        action='store_true',
        help="If supplied, do not write data to the db."
    )
    updateGrp.add_argument(
        '-u', '--update-all-fields',
        action='store_true',
        help="""If supplied, update all fields when updating an existing
            local Tweet record. Otherwise, the default behavior is to
            only update the favorite and retweet counts of the record."""
    )

    args = parser.parse_args()

    if args.available:
        printAvailableCategories()
    elif args.categories:
        inputCategories = args.categories

        categoryResult = db.Category.select(
            IN(db.Category.q.name,
               inputCategories)
        )
        dbCategoryNames = [c.name for c in list(categoryResult)]
        missing = set(inputCategories) - set(dbCategoryNames)
        assert not missing, u"Input categories not found in db: \n- {0}"\
                            .format(u'\n- '.join(missing))

        # Here the AND is required to include SQLObject j-magic, so that
        # Profiles are filtered by Category.
        profResults = db.Profile.select(
            AND(db.Profile.j.categories,
                IN(db.Category.q.name, inputCategories))
        )
        profCount = profResults.count()
        print "Fetching Tweets for {0:,d} Profiles".format(profCount)

        try:
            campaignRec = db.Campaign.byName(CAMPAIGN_NAME)
        except SQLObjectNotFound:
            campaignRec = db.Campaign(name=CAMPAIGN_NAME, searchQuery=None)

        if profCount:
            insertOrUpdateTweetBatch(
                profResults,
                args.tweets_per_profile,
                verbose=args.verbose,
                writeToDB=not(args.no_write),
                campaignRec=campaignRec,
                onlyUpdateEngagements=not(args.update_all_fields)
            )

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


<commit_before>#!/usr/bin/env python3
# Copyright (c) 2015-2018 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test the txwitness for elements

1) Make sure python serialization stuff matches whatever the node/wallet creates, signed and unsigned
2) Segwit transactions have witness data still, and is validated correctly
3) Fast merkle root test? (where oh where are we going to find sha2 internals ripped open in python)
4) transaction round-trips through rpc, p2p, inside and outside of blocks
5) make sure non-segwit transactions have no witness data too
6) If we’re not touching all other tests, we’ll need to actually copy our CTransaction python stuff directly into this test, or another adjacent file, otherwise other tests will still break
7) Try to give it some bitcoin serialization transactions, make sure it fails to decode

"""

from test_framework.messages import CTransaction, CBlock, ser_uint256, FromHex, CBlockHeader, uint256_from_str, CTxOut, ToHex
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import assert_equal, bytes_to_hex_str, hex_str_to_bytes, assert_raises_rpc_error, assert_greater_than
from test_framework import util
from test_framework.blocktools import create_coinbase, create_block, get_witness_script

from io import BytesIO
import copy

class TxWitnessTest(BitcoinTestFramework):
    def set_test_params(self):
        self.setup_clean_chain = True
        self.num_nodes = 2

    def skip_test_if_missing_module(self):
        self.skip_if_no_wallet()

    def assert_tx_format_also_signed(self, utxo, segwit):
        raw = self.nodes[0].createrawtransaction(
            [{"txid": utxo["txid"], "vout": utxo["vout"]}],
            [{self.unknown_addr: "49.9"}, {"fee": "0.1"}]
        )

        unsigned_decoded = self.nodes[0].decoderawtransaction(raw)
        assert_equal(len(unsigned_decoded["vin"]), 1)
        assert('txinwitness' not in unsigned_decoded["vin"][0])

        # Cross-check python serialization
        tx = CTransaction()
        tx.deserialize(BytesIO(hex_str_to_bytes(raw)))
        assert_equal(tx.vin[0].prevout.hash, int("0x"+utxo["txid"], 0))
        assert_equal(len(tx.vin), len(unsigned_decoded["vin"]))
        assert_equal(len(tx.vout), len(unsigned_decoded["vout"]))
        # assert re-encoding
        serialized = bytes_to_hex_str(tx.serialize())
        assert_equal(serialized, raw)

        # Now sign and repeat tests
        signed_raw = self.nodes[0].signrawtransactionwithwallet(raw)["hex"]
        signed_decoded = self.nodes[0].decoderawtransaction(signed_raw)
        assert_equal(len(signed_decoded["vin"]), 1)
        assert(("txinwitness" in signed_decoded["vin"][0]) == segwit)

        # Cross-check python serialization
        tx = CTransaction()
        tx.deserialize(BytesIO(hex_str_to_bytes(signed_raw)))
        assert_equal(tx.vin[0].prevout.hash, int("0x"+utxo["txid"], 0))
        assert_equal(bytes_to_hex_str(tx.vin[0].scriptSig), signed_decoded["vin"][0]["scriptSig"]["hex"])
        # test witness
        if segwit:
            wit_decoded = signed_decoded["vin"][0]["txinwitness"]
            for i in range(len(wit_decoded)):
                assert_equal(bytes_to_hex_str(tx.wit.vtxinwit[0].scriptWitness.stack[i]), wit_decoded[i])
        # assert re-encoding
        serialized = bytes_to_hex_str(tx.serialize())
        assert_equal(serialized, signed_raw)

        txid = self.nodes[0].sendrawtransaction(serialized)
        nodetx = self.nodes[0].getrawtransaction(txid, 1)
        assert_equal(nodetx["txid"], tx.rehash())
        # cross-check wtxid report from node
        wtxid = bytes_to_hex_str(ser_uint256(tx.calc_sha256(True))[::-1])
        assert_equal(nodetx["wtxid"], wtxid)
        assert_equal(nodetx["hash"], wtxid)

        # witness hash stuff
        assert_equal(nodetx["withash"], tx.calc_witness_hash())
        return (txid, wtxid)

    def test_transaction_serialization(self):
        legacy_addr = self.nodes[0].getnewaddress("", "legacy")
        p2sh_addr = self.nodes[0].getnewaddress("", "p2sh-segwit")
        bech32_addr = self.nodes[0].getnewaddress("", "bech32")
        self.unknown_addr = self.nodes[1].getnewaddress()

        # directly seed types of utxos required
        self.nodes[0].generatetoaddress(1, legacy_addr)
        self.nodes[0].generatetoaddress(1, p2sh_addr)
        self.nodes[0].generatetoaddress(1, bech32_addr)
        self.nodes[0].generatetoaddress(101, self.unknown_addr)

        # grab utxos filtering by age
        legacy_utxo = self.nodes[0].listunspent(104, 104)[0]
        p2sh_utxo = self.nodes[0].listunspent(103, 103)[0]
        bech32_utxo = self.nodes[0].listunspent(102, 102)[0]

        submitted_txids = []
        self.log.info("Testing legacy UTXO")
        submitted_txids.append(self.assert_tx_format_also_signed(legacy_utxo, segwit=False))
        self.log.info("Testing p2sh UTXO")
        submitted_txids.append(self.assert_tx_format_also_signed(p2sh_utxo, segwit=True))
        self.log.info("Testing bech32 UTXO")
        submitted_txids.append(self.assert_tx_format_also_signed(bech32_utxo, segwit=True))

        blockhash = self.nodes[0].generate(1)[0]
        hexblock = self.nodes[0].getblock(blockhash, 0)
        block_details = self.nodes[0].getblock(blockhash, 2)
        block = CBlock()
        block.deserialize(BytesIO(hex_str_to_bytes(hexblock)))
        assert(len(block.vtx) == len(submitted_txids) + 1)
        assert_equal(len(block_details["tx"]), len(block.vtx))
        for tx1, tx2 in zip(block.vtx[1:], block_details["tx"][1:]):
            # no tuple wildcard, just re-used tx2 on first one
            assert((tx1.rehash(), tx2["wtxid"]) in submitted_txids)
            assert((tx2["txid"], tx2["hash"]) in submitted_txids)
            assert((tx2["txid"], tx2["wtxid"]) in submitted_txids)
        block.rehash()
        assert_equal(block.hash, self.nodes[0].getbestblockhash())

    def test_coinbase_witness(self):

        def WitToHex(obj):
            return bytes_to_hex_str(obj.serialize(with_witness=True))

        block = self.nodes[0].getnewblockhex()
        block_struct = FromHex(CBlock(), block)

        # Test vanilla block round-trip
        self.nodes[0].testproposedblock(WitToHex(block_struct))

        # Assert there's scriptWitness in the coinbase input that is the witness nonce and nothing else
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].scriptWitness.stack, [b'\x00'*32])
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].vchIssuanceAmountRangeproof, b'')
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].vchInflationKeysRangeproof, b'')
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].peginWitness.stack, [])

        # Add extra witness that isn't covered by witness merkle root, make sure blocks are still valid
        block_witness_stuffed = copy.deepcopy(block_struct)
        block_witness_stuffed.vtx[0].wit.vtxinwit[0].vchIssuanceAmountRangeproof = b'\x00'
        assert_raises_rpc_error(-25, "bad-cb-witness", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        block_witness_stuffed = copy.deepcopy(block_struct)
        block_witness_stuffed.vtx[0].wit.vtxinwit[0].vchInflationKeysRangeproof = b'\x00'
        assert_raises_rpc_error(-25, "bad-cb-witness", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        block_witness_stuffed = copy.deepcopy(block_struct)

        # Let's blow out block weight limit by adding 4MW here
        block_witness_stuffed.vtx[0].wit.vtxinwit[0].peginWitness.stack = [b'\x00'*4000000]
        assert_raises_rpc_error(-25, "bad-cb-witness", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))

        # Test that node isn't blinded to the block
        # Previously an over-stuffed block >4MW would have been marked permanently bad
        # as it already passes witness merkle and regular merkle root checks
        block_height = self.nodes[0].getblockcount()
        assert_equal(self.nodes[0].submitblock(WitToHex(block_witness_stuffed)), "bad-cb-witness")
        assert_equal(block_height, self.nodes[0].getblockcount())
        assert_equal(self.nodes[0].submitblock(WitToHex(block_struct)), None)
        assert_equal(block_height+1, self.nodes[0].getblockcount())

        # New block since we used the first one
        block_struct = FromHex(CBlock(), self.nodes[0].getnewblockhex())
        block_witness_stuffed = copy.deepcopy(block_struct)

        # Add extra witness data that is covered by witness merkle root, make sure invalid
        assert_equal(block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchSurjectionproof, b'')
        assert_equal(block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchRangeproof, b'')
        block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchRangeproof = b'\x00'*100000
        block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchSurjectionproof = b'\x00'*100000
        assert_raises_rpc_error(-25, "bad-witness-merkle-match", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        witness_root_hex = block_witness_stuffed.calc_witness_merkle_root()
        witness_root = uint256_from_str(hex_str_to_bytes(witness_root_hex)[::-1])
        block_witness_stuffed.vtx[0].vout[-1] = CTxOut(0, get_witness_script(witness_root, 0))
        block_witness_stuffed.vtx[0].rehash()
        block_witness_stuffed.hashMerkleRoot = block_witness_stuffed.calc_merkle_root()
        block_witness_stuffed.rehash()
        assert_raises_rpc_error(-25, "bad-cb-amount", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        assert_greater_than(len(WitToHex(block_witness_stuffed)), 100000*4) # Make sure the witness data is actually serialized

        # Test that issuance inputs in coinbase don't survive a serialization round-trip
        # (even though this can't cause issuance to occur either way due to VerifyCoinbaseAmount semantics)
        block_witness_stuffed = copy.deepcopy(block_struct)
        block_witness_stuffed.vtx[0].vin[0].assetIssuance.nAmount.setToAmount(1)
        assert(not block_witness_stuffed.vtx[0].vin[0].assetIssuance.isNull())
        stuffed_decoded = self.nodes[0].decoderawtransaction(ToHex(block_witness_stuffed.vtx[0]))
        assert("issuance" not in stuffed_decoded["vin"][0])

    def run_test(self):
        util.node_fastmerkle = self.nodes[0]
        self.test_coinbase_witness()
        self.test_transaction_serialization()

if __name__ == '__main__':
    TxWitnessTest().main()
<commit_msg>Add broken python serialization to always give issuance input, test for decode failure<commit_after>#!/usr/bin/env python3
# Copyright (c) 2015-2018 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test the txwitness for elements

1) Make sure python serialization stuff matches whatever the node/wallet creates, signed and unsigned
2) Segwit transactions have witness data still, and is validated correctly
3) Fast merkle root test? (where oh where are we going to find sha2 internals ripped open in python)
4) transaction round-trips through rpc, p2p, inside and outside of blocks
5) make sure non-segwit transactions have no witness data too
6) If we’re not touching all other tests, we’ll need to actually copy our CTransaction python stuff directly into this test, or another adjacent file, otherwise other tests will still break
7) Try to give it some bitcoin serialization transactions, make sure it fails to decode

"""

from test_framework.messages import CTransaction, CBlock, ser_uint256, FromHex, uint256_from_str, CTxOut, ToHex, CTxIn, COutPoint, OUTPOINT_ISSUANCE_FLAG, ser_string
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import assert_equal, bytes_to_hex_str, hex_str_to_bytes, assert_raises_rpc_error, assert_greater_than
from test_framework import util
from test_framework.blocktools import get_witness_script

from io import BytesIO
import copy
import struct

class TxWitnessTest(BitcoinTestFramework):
    def set_test_params(self):
        self.setup_clean_chain = True
        self.num_nodes = 2

    def skip_test_if_missing_module(self):
        self.skip_if_no_wallet()

    def assert_tx_format_also_signed(self, utxo, segwit):
        raw = self.nodes[0].createrawtransaction(
            [{"txid": utxo["txid"], "vout": utxo["vout"]}],
            [{self.unknown_addr: "49.9"}, {"fee": "0.1"}]
        )

        unsigned_decoded = self.nodes[0].decoderawtransaction(raw)
        assert_equal(len(unsigned_decoded["vin"]), 1)
        assert('txinwitness' not in unsigned_decoded["vin"][0])

        # Cross-check python serialization
        tx = CTransaction()
        tx.deserialize(BytesIO(hex_str_to_bytes(raw)))
        assert_equal(tx.vin[0].prevout.hash, int("0x"+utxo["txid"], 0))
        assert_equal(len(tx.vin), len(unsigned_decoded["vin"]))
        assert_equal(len(tx.vout), len(unsigned_decoded["vout"]))
        # assert re-encoding
        serialized = bytes_to_hex_str(tx.serialize())
        assert_equal(serialized, raw)

        # Now sign and repeat tests
        signed_raw = self.nodes[0].signrawtransactionwithwallet(raw)["hex"]
        signed_decoded = self.nodes[0].decoderawtransaction(signed_raw)
        assert_equal(len(signed_decoded["vin"]), 1)
        assert(("txinwitness" in signed_decoded["vin"][0]) == segwit)

        # Cross-check python serialization
        tx = CTransaction()
        tx.deserialize(BytesIO(hex_str_to_bytes(signed_raw)))
        assert_equal(tx.vin[0].prevout.hash, int("0x"+utxo["txid"], 0))
        assert_equal(bytes_to_hex_str(tx.vin[0].scriptSig), signed_decoded["vin"][0]["scriptSig"]["hex"])
        # test witness
        if segwit:
            wit_decoded = signed_decoded["vin"][0]["txinwitness"]
            for i in range(len(wit_decoded)):
                assert_equal(bytes_to_hex_str(tx.wit.vtxinwit[0].scriptWitness.stack[i]), wit_decoded[i])
        # assert re-encoding
        serialized = bytes_to_hex_str(tx.serialize())
        assert_equal(serialized, signed_raw)

        txid = self.nodes[0].sendrawtransaction(serialized)
        nodetx = self.nodes[0].getrawtransaction(txid, 1)
        assert_equal(nodetx["txid"], tx.rehash())
        # cross-check wtxid report from node
        wtxid = bytes_to_hex_str(ser_uint256(tx.calc_sha256(True))[::-1])
        assert_equal(nodetx["wtxid"], wtxid)
        assert_equal(nodetx["hash"], wtxid)

        # witness hash stuff
        assert_equal(nodetx["withash"], tx.calc_witness_hash())
        return (txid, wtxid)

    def test_transaction_serialization(self):
        legacy_addr = self.nodes[0].getnewaddress("", "legacy")
        p2sh_addr = self.nodes[0].getnewaddress("", "p2sh-segwit")
        bech32_addr = self.nodes[0].getnewaddress("", "bech32")
        self.unknown_addr = self.nodes[1].getnewaddress()

        # directly seed types of utxos required
        self.nodes[0].generatetoaddress(1, legacy_addr)
        self.nodes[0].generatetoaddress(1, p2sh_addr)
        self.nodes[0].generatetoaddress(1, bech32_addr)
        self.nodes[0].generatetoaddress(101, self.unknown_addr)

        # grab utxos filtering by age
        legacy_utxo = self.nodes[0].listunspent(104, 104)[0]
        p2sh_utxo = self.nodes[0].listunspent(103, 103)[0]
        bech32_utxo = self.nodes[0].listunspent(102, 102)[0]

        submitted_txids = []
        self.log.info("Testing legacy UTXO")
        submitted_txids.append(self.assert_tx_format_also_signed(legacy_utxo, segwit=False))
        self.log.info("Testing p2sh UTXO")
        submitted_txids.append(self.assert_tx_format_also_signed(p2sh_utxo, segwit=True))
        self.log.info("Testing bech32 UTXO")
        submitted_txids.append(self.assert_tx_format_also_signed(bech32_utxo, segwit=True))

        blockhash = self.nodes[0].generate(1)[0]
        hexblock = self.nodes[0].getblock(blockhash, 0)
        block_details = self.nodes[0].getblock(blockhash, 2)
        block = CBlock()
        block.deserialize(BytesIO(hex_str_to_bytes(hexblock)))
        assert(len(block.vtx) == len(submitted_txids) + 1)
        assert_equal(len(block_details["tx"]), len(block.vtx))
        for tx1, tx2 in zip(block.vtx[1:], block_details["tx"][1:]):
            # no tuple wildcard, just re-used tx2 on first one
            assert((tx1.rehash(), tx2["wtxid"]) in submitted_txids)
            assert((tx2["txid"], tx2["hash"]) in submitted_txids)
            assert((tx2["txid"], tx2["wtxid"]) in submitted_txids)
        block.rehash()
        assert_equal(block.hash, self.nodes[0].getbestblockhash())

    def test_coinbase_witness(self):

        def WitToHex(obj):
            return bytes_to_hex_str(obj.serialize(with_witness=True))

        block = self.nodes[0].getnewblockhex()
        block_struct = FromHex(CBlock(), block)

        # Test vanilla block round-trip
        self.nodes[0].testproposedblock(WitToHex(block_struct))

        # Assert there's scriptWitness in the coinbase input that is the witness nonce and nothing else
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].scriptWitness.stack, [b'\x00'*32])
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].vchIssuanceAmountRangeproof, b'')
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].vchInflationKeysRangeproof, b'')
        assert_equal(block_struct.vtx[0].wit.vtxinwit[0].peginWitness.stack, [])

        # Add extra witness that isn't covered by witness merkle root, make sure blocks are still valid
        block_witness_stuffed = copy.deepcopy(block_struct)
        block_witness_stuffed.vtx[0].wit.vtxinwit[0].vchIssuanceAmountRangeproof = b'\x00'
        assert_raises_rpc_error(-25, "bad-cb-witness", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        block_witness_stuffed = copy.deepcopy(block_struct)
        block_witness_stuffed.vtx[0].wit.vtxinwit[0].vchInflationKeysRangeproof = b'\x00'
        assert_raises_rpc_error(-25, "bad-cb-witness", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        block_witness_stuffed = copy.deepcopy(block_struct)

        # Let's blow out block weight limit by adding 4MW here
        block_witness_stuffed.vtx[0].wit.vtxinwit[0].peginWitness.stack = [b'\x00'*4000000]
        assert_raises_rpc_error(-25, "bad-cb-witness", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))

        # Test that node isn't blinded to the block
        # Previously an over-stuffed block >4MW would have been marked permanently bad
        # as it already passes witness merkle and regular merkle root checks
        block_height = self.nodes[0].getblockcount()
        assert_equal(self.nodes[0].submitblock(WitToHex(block_witness_stuffed)), "bad-cb-witness")
        assert_equal(block_height, self.nodes[0].getblockcount())
        assert_equal(self.nodes[0].submitblock(WitToHex(block_struct)), None)
        assert_equal(block_height+1, self.nodes[0].getblockcount())

        # New block since we used the first one
        block_struct = FromHex(CBlock(), self.nodes[0].getnewblockhex())
        block_witness_stuffed = copy.deepcopy(block_struct)

        # Add extra witness data that is covered by witness merkle root, make sure invalid
        assert_equal(block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchSurjectionproof, b'')
        assert_equal(block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchRangeproof, b'')
        block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchRangeproof = b'\x00'*100000
        block_witness_stuffed.vtx[0].wit.vtxoutwit[0].vchSurjectionproof = b'\x00'*100000
        assert_raises_rpc_error(-25, "bad-witness-merkle-match", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        witness_root_hex = block_witness_stuffed.calc_witness_merkle_root()
        witness_root = uint256_from_str(hex_str_to_bytes(witness_root_hex)[::-1])
        block_witness_stuffed.vtx[0].vout[-1] = CTxOut(0, get_witness_script(witness_root, 0))
        block_witness_stuffed.vtx[0].rehash()
        block_witness_stuffed.hashMerkleRoot = block_witness_stuffed.calc_merkle_root()
        block_witness_stuffed.rehash()
        assert_raises_rpc_error(-25, "bad-cb-amount", self.nodes[0].testproposedblock, WitToHex(block_witness_stuffed))
        assert_greater_than(len(WitToHex(block_witness_stuffed)), 100000*4) # Make sure the witness data is actually serialized

        # A CTxIn that always serializes the asset issuance, even for coinbases.
        class AlwaysIssuanceCTxIn(CTxIn):
            def serialize(self):
                r = b''
                outpoint = COutPoint()
                outpoint.hash = self.prevout.hash
                outpoint.n = self.prevout.n
                outpoint.n |= OUTPOINT_ISSUANCE_FLAG
                r += outpoint.serialize()
                r += ser_string(self.scriptSig)
                r += struct.pack("<I", self.nSequence)
                r += self.assetIssuance.serialize()
                return r

        # Test that issuance inputs in coinbase don't survive a serialization round-trip
        # (even though this can't cause issuance to occur either way due to VerifyCoinbaseAmount semantics)
        block_witness_stuffed = copy.deepcopy(block_struct)
        coinbase_orig = copy.deepcopy(block_witness_stuffed.vtx[0].vin[0])
        coinbase_ser_size = len(block_witness_stuffed.vtx[0].vin[0].serialize())
        block_witness_stuffed.vtx[0].vin[0] = AlwaysIssuanceCTxIn()
        block_witness_stuffed.vtx[0].vin[0].prevout = coinbase_orig.prevout
        block_witness_stuffed.vtx[0].vin[0].scriptSig = coinbase_orig.scriptSig
        block_witness_stuffed.vtx[0].vin[0].nSequence = coinbase_orig.nSequence
        block_witness_stuffed.vtx[0].vin[0].assetIssuance.nAmount.setToAmount(1)
        bad_coinbase_ser_size = len(block_witness_stuffed.vtx[0].vin[0].serialize())
        # 32+32+9+1 should be serialized for each assetIssuance field
        assert_equal(bad_coinbase_ser_size, coinbase_ser_size+32+32+9+1)
        assert(not block_witness_stuffed.vtx[0].vin[0].assetIssuance.isNull())
        assert_raises_rpc_error(-22, "TX decode failed", self.nodes[0].decoderawtransaction, ToHex(block_witness_stuffed.vtx[0]))

    def run_test(self):
        util.node_fastmerkle = self.nodes[0]
        self.test_coinbase_witness()
        self.test_transaction_serialization()

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


<commit_before>#!/usr/bin/env python
"""
Tag reads from dual barcoded experiments
"""

from Bio import SeqIO
import sys
import gzip
import time
#import re
from collections import Counter

barcodeMaxDiff = 1
primerMaxDiff = 4

barcodeFile = "MetaData/barcodeLookupTable.txt"
primerForward = "MetaData/Forward.fasta"
primerReverse = "MetaData/Reverse.fasta"

#R1_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R1_001.fastq.gz'
#R2_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R2_001.fastq.gz'
#R3_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R3_001.fastq.gz'
#R4_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R4_001.fastq.gz'

R1_File = 'Amplicon_Raw_fastq/test40k_R1_001.fastq.gz'
R2_File = 'Amplicon_Raw_fastq/test40k_R2_001.fastq.gz'
R3_File = 'Amplicon_Raw_fastq/test40k_R3_001.fastq.gz'
R4_File = 'Amplicon_Raw_fastq/test40k_R4_001.fastq.gz'

Output_prefix = 'TESTWhichmanKelsie'
Output_report = 'TESTWhichmanKelsie_report.txt'
Output_unique = 'TESTWhichmanKelsie_uniques.txt'

# ------- read in barcodes and make a dictionary for lookup ----------------
barcodes = {}
P5 = []
P7 = []
IDS = []
bcfile = open(barcodeFile,'rb')
f = bcfile.readlines()[1:]
for row in f:
    row = row.rstrip()
    ID, P5Name, P5BC, P7Name, P7BC = row.split('\t')
    if P5BC not in P5:
        P5.extend([P5BC])
    if P7BC not in P7:
        P7.extend([P7BC])
    IDS.extend([ID])
    barcodes["%s %s" % (P7BC, P5BC)] = ID

bcfile.close()

# ------- build primer dictionary --------
primersP5 = {}
primersP7 = {}

for seq in SeqIO.parse(open(primerForward, 'r'), 'fasta'):
    primersP5[seq.seq.tostring()] = seq.id

for seq in SeqIO.parse(open(primerReverse, 'r'), 'fasta'):
    primersP7[seq.seq.tostring()] = seq.id

# ---------- input files -----------------------
R1 = SeqIO.parse(gzip.open(R1_File, 'rb'), 'fastq')
R2 = SeqIO.parse(gzip.open(R2_File, 'rb'), 'fastq')
R3 = SeqIO.parse(gzip.open(R3_File, 'rb'), 'fastq')
R4 = SeqIO.parse(gzip.open(R4_File, 'rb'), 'fastq')

# ------- setup output files ------------
outf = {'identified':[gzip.open(Output_prefix + '_R1.fastq.gz', 'wb'), gzip.open(Output_prefix + '_R2.fastq.gz', 'wb')], 'unidentified':[gzip.open(Output_prefix + '_Unidentified_R1.fastq.gz', 'wb'), gzip.open(Output_prefix + '_Unidentified_R2.fastq.gz', 'wb')]}

barcodesFile = open(Output_report, 'w')
uniquesFile = open(Output_unique, 'w')

# ------- make some counters for reporting ------------
counters = {}
for bc in barcodes:
    counters[barcodes[bc]] = [0, 0, 0]  # first value is perfect matches, second is 1bp mismatch in bc1, third is 1bp mismatch in bc2
goodReadsCounter = 0
otherCounter = 0
uniqueCounter = 0
reads = 0
lasttime = time.time()

# ---------- Unique Dictionary ------------- #
uniqueDict = {}
#------------------- functions ------------------------------
def barcodeDist(b_1, b_2):
    'counts mismatches between two equal-length strings'
    if len(b_1) == len(b_2) and len(b_1) > 0:
        return sum(map(lambda x: x[0] != x[1], zip(b_1, b_2) )) # python is bad-ass
    else:
        print "ERROR lengths of barcodes and index read do not match!"
        print "Target", b_1
        print "Index read:", b_2
        sys.exit()

#------------------- functions ------------------------------
def primerDist(read, primer):
    'counts mismatches between primer and sequence'
    read = read[0:(len(primer)-1)]
    if primer == read:
        return 0
    else:
        return sum(map(lambda x: x[0] != x[1], zip(read, primer) )) # python is bad-ass

# ------- work horse function ---------------
try:
    while 1:
        read1 = R1.next()
        read2 = R2.next()
        read3 = R3.next()
        read4 = R4.next()
        reads += 1

        ### Barcode Matching ###
        bc1 = None
        bc1Mismatch = False
        if read2.seq.tostring() in P7:
            bc1 = read2.seq.tostring()
        else:
            for key in P7:
                if barcodeDist(key, read2.seq.tostring()) < barcodeMaxDiff:
                    bc1 = key
                    bc1Mismatch = True

        bc2 = None
        bc2Mismatch = False
        if read3.seq.tostring() in P5:
            bc2 = read3.seq.tostring()
        else:
            for key in P5:
                if barcodeDist(key, read3.seq.tostring()) < barcodeMaxDiff:
                    bc2 = key
                    bc2Mismatch = True

        combined_bc = None
        if "%s %s" % (bc1, bc2) in barcodes:
            combined_bc = barcodes["%s %s" % (bc1, bc2)]
            counters[combined_bc][0] += 1
            if bc1Mismatch:
                counters[combined_bc][1] += 1
            if bc2Mismatch:
                counters[combined_bc][2] += 1

        ### Primer Matching ###
        primer1 = None
        for primer in primersP5.keys():
            pmismatches = primerDist(read1.seq.tostring(), primer)
            if pmismatches < primerMaxDiff:
                primer1 = primersP5[primer]
                primer1Mismatch = pmismatches
                read1 = read1[len(primer):]

        primer2 = None
        for primer in primersP7.keys():
            pmismatches = primerDist(read4.seq.tostring(), primer)
            if pmismatches < primerMaxDiff:
                primer2 = primersP7[primer]
                primer2Mismatch = pmismatches
                read4 = read4[len(primer):]

        primer_id = None
        if primer1 == primer2:
            primer_id = primer1

        ### Output Reads ###
        if combined_bc is not None and primer_id is not None:
            read1.id = read1.name =  "%s 1:N:0:%s-%s|%s|%s|%s|%s" % (read1.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer_id, primer1Mismatch, primer2Mismatch)
            read1.description = ""
            SeqIO.write(read1, outf['identified'][0], "fastq")
            read4.id = read4.name =  "%s 2:N:0:%s-%s|%s|%s|%s|%s" % (read4.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer_id, primer1Mismatch, primer2Mismatch)
            read4.description = ""
            SeqIO.write(read4, outf['identified'][1], "fastq")
            txt = '\t'.join([read1.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer_id, str(primer1Mismatch), str(primer2Mismatch)]) + '\n'
            barcodesFile.write(txt)
            goodReadsCounter += 1
            comb_read = read1 + read4
            comb_read = "".join([comb_read.seq.tostring(), primer_id])
            if (comb_read in uniqueDict):
                uniqueDict[comb_read]['bc_counts'][combined_bc] += 1
            else:
                uniqueDict[comb_read] = {}
                uniqueDict[comb_read]['read1'] = read1.seq.tostring()
                uniqueDict[comb_read]['read2'] = read4.seq.tostring()
                uniqueDict[comb_read]['bc_counts'] = Counter()
                uniqueDict[comb_read]['bc_counts'][combined_bc] += 1
                uniqueCounter += 1

        else:
            read1.id = read1.name =  "%s 1:N:0:%s-%s|%s|%s|%s|%s|%s" % (read1.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer1, primer1Mismatch, primer2, primer2Mismatch)
            read1.description = ""
            SeqIO.write(read1, outf['unidentified'][0], "fastq")
            read4.id = read4.name =  "%s 2:N:0:%s-%s|%s|%s|%s|%s|%s" % (read4.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer1, primer1Mismatch, primer2, primer2Mismatch)
            read4.description = ""
            SeqIO.write(read4, outf['unidentified'][1], "fastq")
            otherCounter += 1

        ### Report every 100K ###
        if reads % 10000 == 0:
            print"-------------------"
            print "%s reads processed.Reads/second %s" % (reads, reads/(time.time() - lasttime))
            countersums = [0, 0, 0]
            for k in counters:
                countersums[0] += counters[k][0]
                countersums[1] += counters[k][1]
                countersums[2] += counters[k][2]
            print "identified:%s,\tmismatch bc1:%s,\tmismatch bc2:%s" % (countersums[0], countersums[1], countersums[2])
            print "good reads: %s" % goodReadsCounter
            print "other: %s" % otherCounter
            print "uniques: %s" % uniqueCounter
            print"-------------------"
except StopIteration:
    pass
finally:
    print "Finished processing "

    print"-------------------"
    print "%s reads processed in %s hours" % (reads,(time.time()-lasttime)/(60*60))
    for k in IDS:
        print "%s:\tidentified:%s,\tmismatch bc1:%s,\tmismatch bc2:%s" % (k, counters[k][0], counters[k][1], counters[k][2])
    print "good reads: %s" % goodReadsCounter
    print "other: %s" % otherCounter
    print "uniques: %s" % uniqueCounter
    print"-------------------"
    txt = '\t'.join(["read1" ,"read2" ,"primer_id"])
    for bc in IDS:
        txt = '\t'.join([txt, bc])
    txt = txt  + '\n'
    uniquesFile.write(txt)
    for k in uniqueDict:
        txt = '\t'.join([uniqueDict[k]['read1'], uniqueDict[k]['read2'], uniqueDict[k]['primer_id']])
        for bc in IDS:
            txt = '\t'.join([txt, str(uniqueDict[k]['bc_counts'][bc])])
        txt = txt  + '\n'
        uniquesFile.write(txt)
    for k in outf:
        outf[k][0].close()
        outf[k][1].close()
    print uniqueDict
    barcodesFile.close()
    uniquesFile.close()
<commit_msg>Fixed a missing key in Dual_BC_tagger.py<commit_after>#!/usr/bin/env python
"""
Tag reads from dual barcoded experiments
"""

from Bio import SeqIO
import sys
import gzip
import time
#import re
from collections import Counter

barcodeMaxDiff = 1
primerMaxDiff = 4

barcodeFile = "MetaData/barcodeLookupTable.txt"
primerForward = "MetaData/Forward.fasta"
primerReverse = "MetaData/Reverse.fasta"

#R1_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R1_001.fastq.gz'
#R2_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R2_001.fastq.gz'
#R3_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R3_001.fastq.gz'
#R4_File ='Amplicon_Raw_fastq/20130408DBCFluidigmHollyKelsie/20130408DBCFluidigmHollyKelsie_NoIndex_L001_R4_001.fastq.gz'

R1_File = 'Amplicon_Raw_fastq/test40k_R1_001.fastq.gz'
R2_File = 'Amplicon_Raw_fastq/test40k_R2_001.fastq.gz'
R3_File = 'Amplicon_Raw_fastq/test40k_R3_001.fastq.gz'
R4_File = 'Amplicon_Raw_fastq/test40k_R4_001.fastq.gz'

Output_prefix = 'TESTWhichmanKelsie'
Output_report = 'TESTWhichmanKelsie_report.txt'
Output_unique = 'TESTWhichmanKelsie_uniques.txt'

# ------- read in barcodes and make a dictionary for lookup ----------------
barcodes = {}
P5 = []
P7 = []
IDS = []
bcfile = open(barcodeFile,'rb')
f = bcfile.readlines()[1:]
for row in f:
    row = row.rstrip()
    ID, P5Name, P5BC, P7Name, P7BC = row.split('\t')
    if P5BC not in P5:
        P5.extend([P5BC])
    if P7BC not in P7:
        P7.extend([P7BC])
    IDS.extend([ID])
    barcodes["%s %s" % (P7BC, P5BC)] = ID

bcfile.close()

# ------- build primer dictionary --------
primersP5 = {}
primersP7 = {}

for seq in SeqIO.parse(open(primerForward, 'r'), 'fasta'):
    primersP5[seq.seq.tostring()] = seq.id

for seq in SeqIO.parse(open(primerReverse, 'r'), 'fasta'):
    primersP7[seq.seq.tostring()] = seq.id

# ---------- input files -----------------------
R1 = SeqIO.parse(gzip.open(R1_File, 'rb'), 'fastq')
R2 = SeqIO.parse(gzip.open(R2_File, 'rb'), 'fastq')
R3 = SeqIO.parse(gzip.open(R3_File, 'rb'), 'fastq')
R4 = SeqIO.parse(gzip.open(R4_File, 'rb'), 'fastq')

# ------- setup output files ------------
outf = {'identified':[gzip.open(Output_prefix + '_R1.fastq.gz', 'wb'), gzip.open(Output_prefix + '_R2.fastq.gz', 'wb')], 'unidentified':[gzip.open(Output_prefix + '_Unidentified_R1.fastq.gz', 'wb'), gzip.open(Output_prefix + '_Unidentified_R2.fastq.gz', 'wb')]}

barcodesFile = open(Output_report, 'w')
uniquesFile = open(Output_unique, 'w')

# ------- make some counters for reporting ------------
counters = {}
for bc in barcodes:
    counters[barcodes[bc]] = [0, 0, 0]  # first value is perfect matches, second is 1bp mismatch in bc1, third is 1bp mismatch in bc2
goodReadsCounter = 0
otherCounter = 0
uniqueCounter = 0
reads = 0
lasttime = time.time()

# ---------- Unique Dictionary ------------- #
uniqueDict = {}
#------------------- functions ------------------------------
def barcodeDist(b_1, b_2):
    'counts mismatches between two equal-length strings'
    if len(b_1) == len(b_2) and len(b_1) > 0:
        return sum(map(lambda x: x[0] != x[1], zip(b_1, b_2) )) # python is bad-ass
    else:
        print "ERROR lengths of barcodes and index read do not match!"
        print "Target", b_1
        print "Index read:", b_2
        sys.exit()

#------------------- functions ------------------------------
def primerDist(read, primer):
    'counts mismatches between primer and sequence'
    read = read[0:(len(primer)-1)]
    if primer == read:
        return 0
    else:
        return sum(map(lambda x: x[0] != x[1], zip(read, primer) )) # python is bad-ass

# ------- work horse function ---------------
try:
    while 1:
        read1 = R1.next()
        read2 = R2.next()
        read3 = R3.next()
        read4 = R4.next()
        reads += 1

        ### Barcode Matching ###
        bc1 = None
        bc1Mismatch = False
        if read2.seq.tostring() in P7:
            bc1 = read2.seq.tostring()
        else:
            for key in P7:
                if barcodeDist(key, read2.seq.tostring()) < barcodeMaxDiff:
                    bc1 = key
                    bc1Mismatch = True

        bc2 = None
        bc2Mismatch = False
        if read3.seq.tostring() in P5:
            bc2 = read3.seq.tostring()
        else:
            for key in P5:
                if barcodeDist(key, read3.seq.tostring()) < barcodeMaxDiff:
                    bc2 = key
                    bc2Mismatch = True

        combined_bc = None
        if "%s %s" % (bc1, bc2) in barcodes:
            combined_bc = barcodes["%s %s" % (bc1, bc2)]
            counters[combined_bc][0] += 1
            if bc1Mismatch:
                counters[combined_bc][1] += 1
            if bc2Mismatch:
                counters[combined_bc][2] += 1

        ### Primer Matching ###
        primer1 = None
        for primer in primersP5.keys():
            pmismatches = primerDist(read1.seq.tostring(), primer)
            if pmismatches < primerMaxDiff:
                primer1 = primersP5[primer]
                primer1Mismatch = pmismatches
                read1 = read1[len(primer):]

        primer2 = None
        for primer in primersP7.keys():
            pmismatches = primerDist(read4.seq.tostring(), primer)
            if pmismatches < primerMaxDiff:
                primer2 = primersP7[primer]
                primer2Mismatch = pmismatches
                read4 = read4[len(primer):]

        primer_id = None
        if primer1 == primer2:
            primer_id = primer1

        ### Output Reads ###
        if combined_bc is not None and primer_id is not None:
            read1.id = read1.name =  "%s 1:N:0:%s-%s|%s|%s|%s|%s" % (read1.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer_id, primer1Mismatch, primer2Mismatch)
            read1.description = ""
            SeqIO.write(read1, outf['identified'][0], "fastq")
            read4.id = read4.name =  "%s 2:N:0:%s-%s|%s|%s|%s|%s" % (read4.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer_id, primer1Mismatch, primer2Mismatch)
            read4.description = ""
            SeqIO.write(read4, outf['identified'][1], "fastq")
            txt = '\t'.join([read1.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer_id, str(primer1Mismatch), str(primer2Mismatch)]) + '\n'
            barcodesFile.write(txt)
            goodReadsCounter += 1
            comb_read = read1 + read4
            comb_read = "".join([comb_read.seq.tostring(), primer_id])
            if (comb_read in uniqueDict):
                uniqueDict[comb_read]['bc_counts'][combined_bc] += 1
            else:
                uniqueDict[comb_read] = {}
                uniqueDict[comb_read]['read1'] = read1.seq.tostring()
                uniqueDict[comb_read]['read2'] = read4.seq.tostring()
                uniqueDict[comb_read]['primer_id'] = primer_id
                uniqueDict[comb_read]['bc_counts'] = Counter()
                uniqueDict[comb_read]['bc_counts'][combined_bc] += 1
                uniqueCounter += 1

        else:
            read1.id = read1.name =  "%s 1:N:0:%s-%s|%s|%s|%s|%s|%s" % (read1.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer1, primer1Mismatch, primer2, primer2Mismatch)
            read1.description = ""
            SeqIO.write(read1, outf['unidentified'][0], "fastq")
            read4.id = read4.name =  "%s 2:N:0:%s-%s|%s|%s|%s|%s|%s" % (read4.id.split()[0], read2.seq.tostring(), read3.seq.tostring(), combined_bc, primer1, primer1Mismatch, primer2, primer2Mismatch)
            read4.description = ""
            SeqIO.write(read4, outf['unidentified'][1], "fastq")
            otherCounter += 1

        ### Report every 100K ###
        if reads % 10000 == 0:
            print"-------------------"
            print "%s reads processed.Reads/second %s" % (reads, reads/(time.time() - lasttime))
            countersums = [0, 0, 0]
            for k in counters:
                countersums[0] += counters[k][0]
                countersums[1] += counters[k][1]
                countersums[2] += counters[k][2]
            print "identified:%s,\tmismatch bc1:%s,\tmismatch bc2:%s" % (countersums[0], countersums[1], countersums[2])
            print "good reads: %s" % goodReadsCounter
            print "other: %s" % otherCounter
            print "uniques: %s" % uniqueCounter
            print"-------------------"
except StopIteration:
    pass
finally:
    print "Finished processing "

    print"-------------------"
    print "%s reads processed in %s hours" % (reads,(time.time()-lasttime)/(60*60))
    for k in IDS:
        print "%s:\tidentified:%s,\tmismatch bc1:%s,\tmismatch bc2:%s" % (k, counters[k][0], counters[k][1], counters[k][2])
    print "good reads: %s" % goodReadsCounter
    print "other: %s" % otherCounter
    print "uniques: %s" % uniqueCounter
    print"-------------------"
    txt = '\t'.join(["read1" ,"read2" ,"primer_id"])
    for bc in IDS:
        txt = '\t'.join([txt, bc])
    txt = txt  + '\n'
    uniquesFile.write(txt)
    for k in uniqueDict:
        txt = '\t'.join([uniqueDict[k]['read1'], uniqueDict[k]['read2'], uniqueDict[k]['primer_id']])
        for bc in IDS:
            txt = '\t'.join([txt, str(uniqueDict[k]['bc_counts'][bc])])
        txt = txt  + '\n'
        uniquesFile.write(txt)
    for k in outf:
        outf[k][0].close()
        outf[k][1].close()
    print uniqueDict
    barcodesFile.close()
    uniquesFile.close()
<|endoftext|>