<commit_before>#!/usr/bin/env python3
# Copyright 2019 Google LLC
#
# Licensed under the the Apache License v2.0 with LLVM Exceptions (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import datetime
import json
import logging
import os
import re
import socket
import subprocess
import sys
import time
from typing import List, Optional, Tuple

from phabricator import Phabricator
from git import Repo, BadName, GitCommandError

# FIXME: maybe move to config file
"""URL of upstream LLVM repository."""
LLVM_GITHUB_URL = 'ssh://git@github.com/llvm/llvm-project'
FORK_REMOTE_URL = 'ssh://git@github.com/llvm-premerge-tests/llvm-project'

"""How far back the script searches in the git history to find Revisions that
have already landed. """
APPLIED_SCAN_LIMIT = datetime.timedelta(days=90)

class ApplyPatch:
    """Apply a diff from Phabricator on local working copy.

    This script is a rewrite of `arc patch` to accomodate for dependencies
    that have already landed, but could not be identified by `arc patch`.

    For a given diff_id, this class will get the dependencies listed on Phabricator.
    For each dependency D it will check the diff history:
    - if D has already landed, skip it.
    - If D has not landed, it will download the patch for D and try to apply it locally.
    Once this class has applied all dependencies, it will apply the diff itself.

    This script must be called from the root folder of a local checkout of
    https://github.com/llvm/llvm-project or given a path to clone into.
    """

    def __init__(self, path: str, diff_id: int, token: str, url: str, git_hash: str,
                 push_branch: bool = False):
        self.push_branch = push_branch  # type: bool
        self.conduit_token = token  # type: Optional[str]
        self.host = url  # type: Optional[str]
        self._load_arcrc()
        self.diff_id = diff_id  # type: int
        if not self.host.endswith('/api/'):
            self.host += '/api/'
        self.phab = self._create_phab()
        self.base_revision = git_hash  # type: str

        if not os.path.isdir(path):
            logging.info(f'{path} does not exist, cloning repository')
            # TODO: progress of clonning
            self.repo = Repo.clone_from(FORK_REMOTE_URL, path)
        else:
            logging.info('repository exist, will reuse')
            self.repo = Repo(path)  # type: Repo
        os.chdir(path)
        logging.info(f'working dir {os.getcwd()}')

    @property
    def branch_name(self):
        """Name used for the git branch."""
        return 'phab-diff-{}'.format(self.diff_id)

    def _load_arcrc(self):
        """Load arc configuration from file if not set."""
        if self.conduit_token is not None or self.host is not None:
            return
        logging.info('Loading configuration from ~/.arcrc file')
        with open(os.path.expanduser('~/.arcrc'), 'r') as arcrc_file:
            arcrc = json.load(arcrc_file)
        # use the first host configured in the file
        self.host = next(iter(arcrc['hosts']))
        self.conduit_token = arcrc['hosts'][self.host]['token']

    def run(self):
        """try to apply the patch from phabricator

        """

        try:
            self._refresh_master()
            revision_id, dependencies, base_revision = self._get_dependencies(self.diff_id)
            dependencies.reverse()  # Arrange deps in chronological order.
            if self.base_revision != 'auto':
                logging.info('Using base revision provided by command line\n{} instead of resolved\n{}'.format(
                    self.base_revision, base_revision))
                base_revision = self.base_revision
            self._create_branch(base_revision)
            logging.info('git reset, git cleanup...')
            self.repo.git.reset('--hard')
            self.repo.git.clean('-fdx')
            logging.info('Analyzing {}'.format(diff_to_str(revision_id)))
            if len(dependencies) > 0:
                logging.info('This diff depends on: {}'.format(diff_list_to_str(dependencies)))
                missing, landed = self._get_missing_landed_dependencies(dependencies)
                logging.info('  Already landed: {}'.format(diff_list_to_str(landed)))
                logging.info('  Will be applied: {}'.format(diff_list_to_str(missing)))
                if missing:
                    for revision in missing:
                        self._apply_revision(revision)
                    # FIXME: submit every Revision individually to get nicer history, use original user name
                    self.repo.config_writer().set_value("user", "name", "myusername").release()
                    self.repo.config_writer().set_value("user", "email", "myemail@example.com").release()
                    self.repo.git.commit('-a', '-m', 'dependencies')
                logging.info('All depended diffs are applied')
            logging.info('applying original diff')
            self._apply_diff(self.diff_id, revision_id)
            if self.push_branch:
                self._commit_and_push()
            else:
                self.repo.git.add('-u', '.')
            return 0
        except Exception as e:
            logging.error(f'exception: {e}')
            return 1

    def _refresh_master(self):
        """Update local git repo and origin.

        As origin is disjoint from upstream, it needs to be updated by this script.
        """
        if not self.push_branch:
            return

        logging.info('Syncing local, origin and upstream...')
        self.repo.git.fetch('--all')
        self.repo.git.checkout('master')
        self.repo.git.reset('--hard')
        self.repo.git.clean('-fdx')
        if 'upstream' not in self.repo.remotes:
            self.repo.create_remote('upstream', url=LLVM_GITHUB_URL)
            self.repo.remotes.upstream.fetch()
        self.repo.git.pull('origin', 'master')
        self.repo.git.pull('upstream', 'master')
        try:
            self.repo.git.push('origin', 'master')
            logging.info('refresh of master branch completed')
        except GitCommandError as e:
            logging.info('Info: Could not push to origin master.')

    def _create_branch(self, base_revision: Optional[str]):
        self.repo.git.fetch('--all')
        if self.branch_name in self.repo.heads:
            self.repo.delete_head('--force', self.branch_name)
        try:
            commit = self.repo.commit(base_revision)
        except Exception as e:
            logging.info('Cannot resolve revision {}: {}, going to use "master" instead.'.format(base_revision, e))
            commit = self.repo.heads['master'].commit
        logging.info(f'creating branch {self.branch_name} at {commit.hexsha}')
        new_branch = self.repo.create_head(self.branch_name, commit.hexsha)
        self.repo.head.reference = new_branch
        self.repo.head.reset(index=True, working_tree=True)
        logging.info('Base branch revision is {}'.format(self.repo.head.commit.hexsha))

    def _commit_and_push(self):
        """Commit the patch and push it to origin."""
        if not self.push_branch:
            return

        self.repo.git.add('-A')
        self.repo.index.commit(message='applying Diff {}'.format(
            self.diff_id))
        self.repo.git.push('--force', 'origin', self.branch_name)
        logging.info('Branch {} pushed to origin'.format(self.branch_name))
        pass

    def _create_phab(self):
        phab = Phabricator(token=self.conduit_token, host=self.host)
        try_call(lambda: phab.update_interfaces())
        return phab

    def _get_diff(self, diff_id: int):
        """Get a diff from Phabricator based on it's diff id."""
        return try_call(lambda: self.phab.differential.getdiff(diff_id=diff_id))

    def _get_revision(self, revision_id: int):
        """Get a revision from Phabricator based on its revision id."""
        return try_call(lambda: self.phab.differential.query(ids=[revision_id])[0])

    def _get_revisions(self, *, phids: List[str] = None):
        """Get a list of revisions from Phabricator based on their PH-IDs."""
        if phids is None:
            raise Exception('_get_revisions phids is None')
        if not phids:
            # Handle an empty query locally. Otherwise the connection
            # will time out.
            return []
        return try_call(lambda: self.phab.differential.query(phids=phids))

    def _get_dependencies(self, diff_id) -> Tuple[int, List[int], str]:
        """Get all dependencies for the diff.
        They are listed in reverse chronological order - from most recent to least recent."""

        logging.info('Getting dependencies of {}'.format(diff_id))
        diff = self._get_diff(diff_id)
        logging.debug(f'diff object: {diff}')
        revision_id = int(diff.revisionID)
        revision = self._get_revision(revision_id)
        logging.debug(f'revision object: {revision}')
        base_revision = diff['sourceControlBaseRevision']
        if base_revision is None or len(base_revision) == 0:
            base_revision = 'master'
        dependency_ids = revision['auxiliary']['phabricator:depends-on']
        revisions = self._get_revisions(phids=dependency_ids)
        result = []
        # Recursively resolve dependencies of those diffs.
        for r in revisions:
            _, sub, _ = self._get_dependencies(r['diffs'][0])
            result.append(r['id'])
            result.extend(sub)

        return revision_id, result, base_revision

    def _apply_diff(self, diff_id: int, revision_id: int):
        """Download and apply a diff to the local working copy."""
        logging.info('Applying diff {} for revision {}...'.format(diff_id, diff_to_str(revision_id)))
        diff = try_call(lambda: self.phab.differential.getrawdiff(diffID=str(diff_id)).response)
        logging.debug(f'diff {diff_id}:\n{diff}')
        proc = subprocess.run('git apply -', input=diff, shell=True, text=True,
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if proc.returncode != 0:
            raise Exception('Applying patch failed:\n{}'.format(proc.stdout + proc.stderr))

    def _apply_revision(self, revision_id: int):
        """Download and apply the latest  diff of a revision to the local working copy."""
        revision = self._get_revision(revision_id)
        # take the diff_id with the highest number, this should be latest one
        diff_id = max(revision['diffs'])
        logging.info(f'picking diff from revision {revision_id}: {diff_id} from {revision["diffs"]}')
        self._apply_diff(diff_id, revision_id)

    def _get_landed_revisions(self):
        """Get list of landed revisions from current git branch."""
        diff_regex = re.compile(r'^Differential Revision: https://reviews\.llvm\.org/(.*)$', re.MULTILINE)
        earliest_commit = None
        rev = self.base_revision
        age_limit = datetime.datetime.now() - APPLIED_SCAN_LIMIT
        if rev is None:
            rev = 'master'
        for commit in self.repo.iter_commits(rev):
            if datetime.datetime.fromtimestamp(commit.committed_date) < age_limit:
                break
            earliest_commit = commit
            result = diff_regex.search(commit.message)
            if result is not None:
                yield result.group(1)
        if earliest_commit is not None:
            logging.info(f'Earliest analyzed commit in history {earliest_commit.hexsha}, '
                         f'{earliest_commit.committed_datetime}')
        return

    def _get_missing_landed_dependencies(self, dependencies: List[int]) -> Tuple[List[int], List[int]]:
        """Check which of the dependencies have already landed on the current branch."""
        landed_deps = []
        missing_deps = []
        for dependency in dependencies:
            if diff_to_str(dependency) in self._get_landed_revisions():
                landed_deps.append(dependency)
            else:
                missing_deps.append(dependency)
        return missing_deps, landed_deps

def diff_to_str(diff: int) -> str:
    """Convert a diff id to a string with leading "D"."""
    return 'D{}'.format(diff)

def diff_list_to_str(diffs: List[int]) -> str:
    """Convert list of diff ids to a comma separated list, prefixed with "D"."""
    return ', '.join([diff_to_str(d) for d in diffs])

def try_call(call):
    """Tries to call function several times retrying on socked.timeout."""
    c = 0
    while True:
        try:
            return call()
        except socket.timeout as e:
            c += 1
            if c > 5:
                logging.error('Connection to Pharicator failed, giving up: {}'.format(e))
                raise
            logging.error('Connection to Pharicator failed, retrying: {}'.format(e))
            time.sleep(c * 10)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Apply Phabricator patch to working directory.')
    parser.add_argument('diff_id', type=int)
    parser.add_argument('--path', type=str, help='repository path', default=os.getcwd())
    parser.add_argument('--token', type=str, default=None, help='Conduit API token')
    parser.add_argument('--url', type=str, default=None, help='Phabricator URL')
    parser.add_argument('--commit', dest='commit', type=str, default='auto',
                        help='Use this commit as a base. For "auto" tool tries to pick the base commit itself')
    parser.add_argument('--push-branch', action='store_true', dest='push_branch',
                        help='choose if branch shall be pushed to origin')
    parser.add_argument('--log-level', type=str, default='INFO')
    args = parser.parse_args()
    logging.basicConfig(level=args.log_level, format='%(levelname)-7s %(message)s')
    patcher = ApplyPatch(args.path, args.diff_id, args.token, args.url, args.commit, args.push_branch)
    sys.exit(patcher.run())
<commit_msg>fix auto revision usage<commit_after>#!/usr/bin/env python3
# Copyright 2019 Google LLC
#
# Licensed under the the Apache License v2.0 with LLVM Exceptions (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import datetime
import json
import logging
import os
import re
import socket
import subprocess
import sys
import time
from typing import List, Optional, Tuple

from phabricator import Phabricator
from git import Repo, BadName, GitCommandError

# FIXME: maybe move to config file
"""URL of upstream LLVM repository."""
LLVM_GITHUB_URL = 'ssh://git@github.com/llvm/llvm-project'
FORK_REMOTE_URL = 'ssh://git@github.com/llvm-premerge-tests/llvm-project'

"""How far back the script searches in the git history to find Revisions that
have already landed. """
APPLIED_SCAN_LIMIT = datetime.timedelta(days=90)

class ApplyPatch:
    """Apply a diff from Phabricator on local working copy.

    This script is a rewrite of `arc patch` to accomodate for dependencies
    that have already landed, but could not be identified by `arc patch`.

    For a given diff_id, this class will get the dependencies listed on Phabricator.
    For each dependency D it will check the diff history:
    - if D has already landed, skip it.
    - If D has not landed, it will download the patch for D and try to apply it locally.
    Once this class has applied all dependencies, it will apply the diff itself.

    This script must be called from the root folder of a local checkout of
    https://github.com/llvm/llvm-project or given a path to clone into.
    """

    def __init__(self, path: str, diff_id: int, token: str, url: str, git_hash: str,
                 push_branch: bool = False):
        self.push_branch = push_branch  # type: bool
        self.conduit_token = token  # type: Optional[str]
        self.host = url  # type: Optional[str]
        self._load_arcrc()
        self.diff_id = diff_id  # type: int
        if not self.host.endswith('/api/'):
            self.host += '/api/'
        self.phab = self._create_phab()
        self.base_revision = git_hash  # type: str

        if not os.path.isdir(path):
            logging.info(f'{path} does not exist, cloning repository')
            # TODO: progress of clonning
            self.repo = Repo.clone_from(FORK_REMOTE_URL, path)
        else:
            logging.info('repository exist, will reuse')
            self.repo = Repo(path)  # type: Repo
        os.chdir(path)
        logging.info(f'working dir {os.getcwd()}')

    @property
    def branch_name(self):
        """Name used for the git branch."""
        return 'phab-diff-{}'.format(self.diff_id)

    def _load_arcrc(self):
        """Load arc configuration from file if not set."""
        if self.conduit_token is not None or self.host is not None:
            return
        logging.info('Loading configuration from ~/.arcrc file')
        with open(os.path.expanduser('~/.arcrc'), 'r') as arcrc_file:
            arcrc = json.load(arcrc_file)
        # use the first host configured in the file
        self.host = next(iter(arcrc['hosts']))
        self.conduit_token = arcrc['hosts'][self.host]['token']

    def run(self):
        """try to apply the patch from phabricator

        """

        try:
            self._refresh_master()
            revision_id, dependencies, base_revision = self._get_dependencies(self.diff_id)
            dependencies.reverse()  # Arrange deps in chronological order.
            if self.base_revision != 'auto':
                logging.info('Using base revision provided by command line\n{} instead of resolved\n{}'.format(
                    self.base_revision, base_revision))
                base_revision = self.base_revision
            self._create_branch(base_revision)
            logging.info('git reset, git cleanup...')
            self.repo.git.reset('--hard')
            self.repo.git.clean('-fdx')
            logging.info('Analyzing {}'.format(diff_to_str(revision_id)))
            if len(dependencies) > 0:
                logging.info('This diff depends on: {}'.format(diff_list_to_str(dependencies)))
                missing, landed = self._get_missing_landed_dependencies(dependencies)
                logging.info('  Already landed: {}'.format(diff_list_to_str(landed)))
                logging.info('  Will be applied: {}'.format(diff_list_to_str(missing)))
                if missing:
                    for revision in missing:
                        self._apply_revision(revision)
                    # FIXME: submit every Revision individually to get nicer history, use original user name
                    self.repo.config_writer().set_value("user", "name", "myusername").release()
                    self.repo.config_writer().set_value("user", "email", "myemail@example.com").release()
                    self.repo.git.commit('-a', '-m', 'dependencies')
                logging.info('All depended diffs are applied')
            logging.info('applying original diff')
            self._apply_diff(self.diff_id, revision_id)
            if self.push_branch:
                self._commit_and_push()
            else:
                self.repo.git.add('-u', '.')
            return 0
        except Exception as e:
            logging.error(f'exception: {e}')
            return 1

    def _refresh_master(self):
        """Update local git repo and origin.

        As origin is disjoint from upstream, it needs to be updated by this script.
        """
        if not self.push_branch:
            return

        logging.info('Syncing local, origin and upstream...')
        self.repo.git.fetch('--all')
        self.repo.git.checkout('master')
        self.repo.git.reset('--hard')
        self.repo.git.clean('-fdx')
        if 'upstream' not in self.repo.remotes:
            self.repo.create_remote('upstream', url=LLVM_GITHUB_URL)
            self.repo.remotes.upstream.fetch()
        self.repo.git.pull('origin', 'master')
        self.repo.git.pull('upstream', 'master')
        try:
            self.repo.git.push('origin', 'master')
            logging.info('refresh of master branch completed')
        except GitCommandError as e:
            logging.info('Info: Could not push to origin master.')

    def _create_branch(self, base_revision: Optional[str]):
        self.repo.git.fetch('--all')
        if self.branch_name in self.repo.heads:
            self.repo.delete_head('--force', self.branch_name)
        try:
            commit = self.repo.commit(base_revision)
        except Exception as e:
            logging.info('Cannot resolve revision {}: {}, going to use "master" instead.'.format(base_revision, e))
            commit = self.repo.heads['master'].commit
        logging.info(f'creating branch {self.branch_name} at {commit.hexsha}')
        new_branch = self.repo.create_head(self.branch_name, commit.hexsha)
        self.repo.head.reference = new_branch
        self.repo.head.reset(index=True, working_tree=True)
        logging.info('Base branch revision is {}'.format(self.repo.head.commit.hexsha))

    def _commit_and_push(self):
        """Commit the patch and push it to origin."""
        if not self.push_branch:
            return

        self.repo.git.add('-A')
        self.repo.index.commit(message='applying Diff {}'.format(
            self.diff_id))
        self.repo.git.push('--force', 'origin', self.branch_name)
        logging.info('Branch {} pushed to origin'.format(self.branch_name))
        pass

    def _create_phab(self):
        phab = Phabricator(token=self.conduit_token, host=self.host)
        try_call(lambda: phab.update_interfaces())
        return phab

    def _get_diff(self, diff_id: int):
        """Get a diff from Phabricator based on it's diff id."""
        return try_call(lambda: self.phab.differential.getdiff(diff_id=diff_id))

    def _get_revision(self, revision_id: int):
        """Get a revision from Phabricator based on its revision id."""
        return try_call(lambda: self.phab.differential.query(ids=[revision_id])[0])

    def _get_revisions(self, *, phids: List[str] = None):
        """Get a list of revisions from Phabricator based on their PH-IDs."""
        if phids is None:
            raise Exception('_get_revisions phids is None')
        if not phids:
            # Handle an empty query locally. Otherwise the connection
            # will time out.
            return []
        return try_call(lambda: self.phab.differential.query(phids=phids))

    def _get_dependencies(self, diff_id) -> Tuple[int, List[int], str]:
        """Get all dependencies for the diff.
        They are listed in reverse chronological order - from most recent to least recent."""

        logging.info('Getting dependencies of {}'.format(diff_id))
        diff = self._get_diff(diff_id)
        logging.debug(f'diff object: {diff}')
        revision_id = int(diff.revisionID)
        revision = self._get_revision(revision_id)
        logging.debug(f'revision object: {revision}')
        base_revision = diff['sourceControlBaseRevision']
        if base_revision is None or len(base_revision) == 0:
            base_revision = 'master'
        dependency_ids = revision['auxiliary']['phabricator:depends-on']
        revisions = self._get_revisions(phids=dependency_ids)
        result = []
        # Recursively resolve dependencies of those diffs.
        for r in revisions:
            _, sub, _ = self._get_dependencies(r['diffs'][0])
            result.append(r['id'])
            result.extend(sub)

        return revision_id, result, base_revision

    def _apply_diff(self, diff_id: int, revision_id: int):
        """Download and apply a diff to the local working copy."""
        logging.info('Applying diff {} for revision {}...'.format(diff_id, diff_to_str(revision_id)))
        diff = try_call(lambda: self.phab.differential.getrawdiff(diffID=str(diff_id)).response)
        logging.debug(f'diff {diff_id}:\n{diff}')
        proc = subprocess.run('git apply -', input=diff, shell=True, text=True,
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if proc.returncode != 0:
            raise Exception('Applying patch failed:\n{}'.format(proc.stdout + proc.stderr))

    def _apply_revision(self, revision_id: int):
        """Download and apply the latest  diff of a revision to the local working copy."""
        revision = self._get_revision(revision_id)
        # take the diff_id with the highest number, this should be latest one
        diff_id = max(revision['diffs'])
        logging.info(f'picking diff from revision {revision_id}: {diff_id} from {revision["diffs"]}')
        self._apply_diff(diff_id, revision_id)

    def _get_landed_revisions(self):
        """Get list of landed revisions from current git branch."""
        diff_regex = re.compile(r'^Differential Revision: https://reviews\.llvm\.org/(.*)$', re.MULTILINE)
        earliest_commit = None
        rev = self.base_revision
        age_limit = datetime.datetime.now() - APPLIED_SCAN_LIMIT
        if rev != 'auto':
            rev = 'master'
        for commit in self.repo.iter_commits(rev):
            if datetime.datetime.fromtimestamp(commit.committed_date) < age_limit:
                break
            earliest_commit = commit
            result = diff_regex.search(commit.message)
            if result is not None:
                yield result.group(1)
        if earliest_commit is not None:
            logging.info(f'Earliest analyzed commit in history {earliest_commit.hexsha}, '
                         f'{earliest_commit.committed_datetime}')
        return

    def _get_missing_landed_dependencies(self, dependencies: List[int]) -> Tuple[List[int], List[int]]:
        """Check which of the dependencies have already landed on the current branch."""
        landed_deps = []
        missing_deps = []
        for dependency in dependencies:
            if diff_to_str(dependency) in self._get_landed_revisions():
                landed_deps.append(dependency)
            else:
                missing_deps.append(dependency)
        return missing_deps, landed_deps

def diff_to_str(diff: int) -> str:
    """Convert a diff id to a string with leading "D"."""
    return 'D{}'.format(diff)

def diff_list_to_str(diffs: List[int]) -> str:
    """Convert list of diff ids to a comma separated list, prefixed with "D"."""
    return ', '.join([diff_to_str(d) for d in diffs])

def try_call(call):
    """Tries to call function several times retrying on socked.timeout."""
    c = 0
    while True:
        try:
            return call()
        except socket.timeout as e:
            c += 1
            if c > 5:
                logging.error('Connection to Pharicator failed, giving up: {}'.format(e))
                raise
            logging.error('Connection to Pharicator failed, retrying: {}'.format(e))
            time.sleep(c * 10)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Apply Phabricator patch to working directory.')
    parser.add_argument('diff_id', type=int)
    parser.add_argument('--path', type=str, help='repository path', default=os.getcwd())
    parser.add_argument('--token', type=str, default=None, help='Conduit API token')
    parser.add_argument('--url', type=str, default=None, help='Phabricator URL')
    parser.add_argument('--commit', dest='commit', type=str, default='auto',
                        help='Use this commit as a base. For "auto" tool tries to pick the base commit itself')
    parser.add_argument('--push-branch', action='store_true', dest='push_branch',
                        help='choose if branch shall be pushed to origin')
    parser.add_argument('--log-level', type=str, default='INFO')
    args = parser.parse_args()
    logging.basicConfig(level=args.log_level, format='%(levelname)-7s %(message)s')
    patcher = ApplyPatch(args.path, args.diff_id, args.token, args.url, args.commit, args.push_branch)
    sys.exit(patcher.run())
<|endoftext|>


<commit_before>"""
aia_mkmovie
============

Provides
    1. A way to download aia level1 data from the jsoc archive using drms
    2. A wrapper combining aia files into high quality movies
    3. A GUI which allows you to select a subregion of aia and set the color bar
    4. A module to create aia_images from a list
    5. A module to create a movie from a list of images

How to use the documentation
----------------------------
Documentation is available in two locations: documentation strings in the code
and on github at https://github.com/jprchlik/aia_mkmovie

The example docstring assume `aia_mkmovie` is imported as `am`:
  >>> import aia_mkmovie as am

View the function's documentation strings using the built-in ``help`` function:

  >>> help(am.aia_mkmovie)

Available subpackages
----------------------
aia_mkmovie
    Wrapper for downloading data, selecting plot regions, image creation, and movie creation
aia_select_cutout
    Only useful when called from within aia_mkmovie. It sets the color tables and location for plotting
aia_download_files
    Downloads files in a given time range with a given cadence at a given wavelength
aia_mkimage
    Makes images from a list of images
SMEARpy
    Retrieves AIA data from local SDO archives
grab_goes_xray_flux
    Downloads and processes goes x-ray fluxes and solar wind data from NOAA
make_movie
    Makes a movie with the given dimensions and cadence

"""
from aia_mkmovie import aia_mkmovie,aia_mkimage
<commit_msg>Removed space<commit_after>"""
aia_mkmovie
============

Provides
    1. A way to download aia level1 data from the jsoc archive using drms
    2. A wrapper combining aia files into high quality movies
    3. A GUI which allows you to select a subregion of aia and set the color bar
    4. A module to create aia_images from a list
    5. A module to create a movie from a list of images

How to use the documentation
----------------------------
Documentation is available in two locations: documentation strings in the code
and on github at https://github.com/jprchlik/aia_mkmovie

The example docstring assume `aia_mkmovie` is imported as `am`:
  >>> import aia_mkmovie as am

View the function's documentation strings using the built-in ``help`` function:

  >>> help(am.aia_mkmovie)

Available subpackages
----------------------
aia_mkmovie
    Wrapper for downloading data, selecting plot regions, image creation, and movie creation
aia_select_cutout
    Only useful when called from within aia_mkmovie. It sets the color tables and location for plotting
aia_download_files
    Downloads files in a given time range with a given cadence at a given wavelength
aia_mkimage
    Makes images from a list of images
SMEARpy
    Retrieves AIA data from local SDO archives
grab_goes_xray_flux
    Downloads and processes goes x-ray fluxes and solar wind data from NOAA
make_movie
    Makes a movie with the given dimensions and cadence

"""
from aia_mkmovie import aia_mkmovie,aia_mkimage
<|endoftext|>


<commit_before>from unittest.mock import MagicMock
from mpf.tests.MpfTestCase import MpfTestCase

class TestModes(MpfTestCase):

    def getConfigFile(self):
        return 'test_modes.yaml'

    def getMachinePath(self):
        return 'tests/machine_files/mode_tests/'

    def test_loading_modes(self):
        # Tests multi-case mode listing in config and checking here

        # lowercase in config, lowercase folder
        self.assertIn('mode1', self.machine.modes)
        # uppercase in config, lowercase folder
        self.assertIn('mode2', self.machine.modes)
        # lowercase config, uppercase folder
        self.assertIn('mode3', self.machine.modes)
        # lowercase config, lowercase folder, uppercase here
        self.assertIn('Mode4', self.machine.modes)

    def test_mode_start_stop(self):
        # start mode 1
        self.machine.events.post('start_mode1')
        self.advance_time_and_run()
        self.assertTrue(self.machine.mode_controller.is_active('mode1'))
        self.assertTrue(self.machine.modes.mode1.active)
        self.assertIn(self.machine.modes.mode1,
                      self.machine.mode_controller.active_modes)
        self.assertFalse(self.machine.mode_controller.is_active('mode2'))
        self.assertFalse(self.machine.modes.mode2.active)
        self.assertIn(self.machine.modes.mode1,
                      self.machine.mode_controller.active_modes)

        # test config via include
        self.assertEqual(123, self.machine.modes.mode1.config['mode_settings']['test'])

        # start a mode that's already started and make sure it doesn't explode
        self.machine.modes.mode1.start()

        # stop mode 1
        self.machine.events.post('stop_mode1')
        self.advance_time_and_run()
        self.assertFalse(self.machine.mode_controller.is_active('mode1'))
        self.assertFalse(self.machine.modes.mode1.active)

    def test_custom_mode_code(self):
        self.assertTrue(self.machine.modes.mode3.custom_code)

    def test_mode_priorities(self):

        # test the start priorities TODO there's got to be a better way?
        found_it = False

        for callback, priority, _, _ \
                in self.machine.events.registered_handlers['start_mode1']:

            if callback == self.machine.modes.mode1.start:
                found_it = True
                self.assertEqual(priority, 201)

        self.assertTrue(found_it)

        # default priorities
        self.machine.modes.mode1.start()
        self.machine.modes.mode2.start()
        self.advance_time_and_run()
        self.assertEqual(self.machine.modes.mode1.priority, 200)
        self.assertEqual(self.machine.modes.mode2.priority, 100)
        self.assertEqual(self.machine.modes.attract.priority, 10)

        # test the stop priorities
        found_it = False
        for callback, priority, _, _ \
                in self.machine.events.registered_handlers['stop_mode1']:
            if callback == self.machine.modes.mode1.stop:
                found_it = True
                # 201 because stop priorities are always +1 over mode priority
                self.assertEqual(priority, 201)
        self.assertTrue(found_it)

        found_it = False
        for callback, priority, _, _ \
                in self.machine.events.registered_handlers['stop_mode2']:
            if callback == self.machine.modes.mode2.stop:
                found_it = True
                # base priority 100, +1 default bump, +2 stop_priority config
                self.assertEqual(priority, 103)
        self.assertTrue(found_it)

        # pass a priority at start
        self.machine.modes.mode1.stop()
        self.advance_time_and_run()
        self.machine.modes.mode1.start(mode_priority=500)
        self.advance_time_and_run()
        self.assertEqual(self.machine.modes.mode1.priority, 500)
        self.advance_time_and_run()

        # test the order of the active modes list
        self.assertEqual(self.machine.modes.mode1,
                         self.machine.mode_controller.active_modes[0])
        self.assertEqual(self.machine.modes.mode2,
                         self.machine.mode_controller.active_modes[1])
        self.assertEqual(self.machine.modes.attract,
                         self.machine.mode_controller.active_modes[2])

    def test_mode_start_with_callback(self):
        self.mode_start_callback = MagicMock()

        self.machine.modes.mode1.start(callback=self.mode_start_callback)
        self.advance_time_and_run()
        self.mode_start_callback.assert_called_once_with()

    def test_use_wait_queue(self):
        self.callback = MagicMock()
        self.machine.events.post_queue('start_mode4',
                                       callback=self.callback)
        self.advance_time_and_run()

        # make sure the mode started
        self.assertTrue(self.machine.modes.mode4.active)
        self.advance_time_and_run()

        # callback should not be called since mode is set to use wait queue
        self.callback.assert_not_called()

        self.machine.modes.mode4.stop()
        self.advance_time_and_run()

        # once mode stops, verify that callback was called
        self.callback.assert_called_once_with()

    def test_ball_end(self):
        # mode1 is set to keep running on ball end
        # mode2 should stop and restart on ball end
        # mode3 should stop and stay stopped
        self.assertFalse(self.machine.modes.mode1.config['mode']['stop_on_ball_end'])
        self.assertTrue(self.machine.modes.mode2.config['mode']['stop_on_ball_end'])
        self.assertTrue(self.machine.modes.mode3.config['mode']['stop_on_ball_end'])

        # start a game
        self.machine.playfield.add_ball = MagicMock()
        self.machine.events.post('game_start')
        self.advance_time_and_run()
        self.machine.game.balls_in_play = 1
        self.assertTrue(self.machine.game)

        # start some modes
        self.machine.modes.mode1.start()
        self.machine.modes.mode2.start()
        self.machine.modes.mode3.start()
        self.advance_time_and_run()
        self.assertTrue(self.machine.modes.mode1.active)
        self.assertTrue(self.machine.modes.mode2.active)
        self.assertTrue(self.machine.modes.mode3.active)

        # change the config of modes 1 and 2 so we can verify whether they
        # actually stopped and restarted since the configs are used to start
        # the modes
        self.machine.modes.mode1.config['mode']['priority'] = 999
        self.machine.modes.mode2.config['mode']['priority'] = 999

        # end the ball
        self.machine.game.ball_ending()
        self.advance_time_and_run()

        self.assertTrue(self.machine.modes.mode1.active)
        self.assertTrue(self.machine.modes.mode2.active)
        self.assertFalse(self.machine.modes.mode3.active)

        # since mode1 should have stayed running, its config should not change
        self.assertNotEqual(self.machine.modes.mode1.priority, 999)
        # mode2 should use the new priority since it should have restarted
        self.assertEqual(self.machine.modes.mode2.priority, 999)

        # end ball 2
        self.machine.game.ball_ending()
        self.advance_time_and_run()

        # end ball 3 and end the game
        self.machine.game.ball_ending()
        self.advance_time_and_run()

        self.assertTrue(self.machine.modes.attract.active)
        self.assertFalse(self.machine.modes.game.active)
        self.assertTrue(self.machine.modes.mode1.active)
        self.assertFalse(self.machine.modes.mode2.active)
        self.assertFalse(self.machine.modes.mode3.active)
<commit_msg>fix test<commit_after>from unittest.mock import MagicMock
from mpf.tests.MpfTestCase import MpfTestCase

class TestModes(MpfTestCase):

    def getConfigFile(self):
        return 'test_modes.yaml'

    def getMachinePath(self):
        return 'tests/machine_files/mode_tests/'

    def test_loading_modes(self):
        # Tests multi-case mode listing in config and checking here

        # lowercase in config, lowercase folder
        self.assertIn('mode1', self.machine.modes)
        # uppercase in config, lowercase folder
        self.assertIn('mode2', self.machine.modes)
        # lowercase config, uppercase folder
        self.assertIn('mode3', self.machine.modes)
        # lowercase config, lowercase folder, uppercase here
        self.assertIn('Mode4', self.machine.modes)

    def test_mode_start_stop(self):
        # start mode 1
        self.machine.events.post('start_mode1')
        self.advance_time_and_run()
        self.assertTrue(self.machine.mode_controller.is_active('mode1'))
        self.assertTrue(self.machine.modes.mode1.active)
        self.assertIn(self.machine.modes.mode1,
                      self.machine.mode_controller.active_modes)
        self.assertFalse(self.machine.mode_controller.is_active('mode2'))
        self.assertFalse(self.machine.modes.mode2.active)
        self.assertIn(self.machine.modes.mode1,
                      self.machine.mode_controller.active_modes)

        # test config via include
        self.assertEqual(123, self.machine.modes.mode1.config['mode_settings']['test'])

        # start a mode that's already started and make sure it doesn't explode
        self.machine.modes.mode1.start()

        # stop mode 1
        self.machine.events.post('stop_mode1')
        self.advance_time_and_run()
        self.assertFalse(self.machine.mode_controller.is_active('mode1'))
        self.assertFalse(self.machine.modes.mode1.active)

    def test_custom_mode_code(self):
        self.assertTrue(self.machine.modes.mode3.custom_code)

    def test_mode_priorities(self):

        # test the start priorities TODO there's got to be a better way?
        found_it = False

        for handler in self.machine.events.registered_handlers['start_mode1']:

            if handler.callback == self.machine.modes.mode1.start:
                found_it = True
                self.assertEqual(handler.priority, 201)

        self.assertTrue(found_it)

        # default priorities
        self.machine.modes.mode1.start()
        self.machine.modes.mode2.start()
        self.advance_time_and_run()
        self.assertEqual(self.machine.modes.mode1.priority, 200)
        self.assertEqual(self.machine.modes.mode2.priority, 100)
        self.assertEqual(self.machine.modes.attract.priority, 10)

        # test the stop priorities
        found_it = False
        for handler in self.machine.events.registered_handlers['stop_mode1']:
            if handler.callback == self.machine.modes.mode1.stop:
                found_it = True
                # 201 because stop priorities are always +1 over mode priority
                self.assertEqual(handler.priority, 201)
        self.assertTrue(found_it)

        found_it = False
        for handler in self.machine.events.registered_handlers['stop_mode2']:
            if handler.callback == self.machine.modes.mode2.stop:
                found_it = True
                # base priority 100, +1 default bump, +2 stop_priority config
                self.assertEqual(handler.priority, 103)
        self.assertTrue(found_it)

        # pass a priority at start
        self.machine.modes.mode1.stop()
        self.advance_time_and_run()
        self.machine.modes.mode1.start(mode_priority=500)
        self.advance_time_and_run()
        self.assertEqual(self.machine.modes.mode1.priority, 500)
        self.advance_time_and_run()

        # test the order of the active modes list
        self.assertEqual(self.machine.modes.mode1,
                         self.machine.mode_controller.active_modes[0])
        self.assertEqual(self.machine.modes.mode2,
                         self.machine.mode_controller.active_modes[1])
        self.assertEqual(self.machine.modes.attract,
                         self.machine.mode_controller.active_modes[2])

    def test_mode_start_with_callback(self):
        self.mode_start_callback = MagicMock()

        self.machine.modes.mode1.start(callback=self.mode_start_callback)
        self.advance_time_and_run()
        self.mode_start_callback.assert_called_once_with()

    def test_use_wait_queue(self):
        self.callback = MagicMock()
        self.machine.events.post_queue('start_mode4',
                                       callback=self.callback)
        self.advance_time_and_run()

        # make sure the mode started
        self.assertTrue(self.machine.modes.mode4.active)
        self.advance_time_and_run()

        # callback should not be called since mode is set to use wait queue
        self.callback.assert_not_called()

        self.machine.modes.mode4.stop()
        self.advance_time_and_run()

        # once mode stops, verify that callback was called
        self.callback.assert_called_once_with()

    def test_ball_end(self):
        # mode1 is set to keep running on ball end
        # mode2 should stop and restart on ball end
        # mode3 should stop and stay stopped
        self.assertFalse(self.machine.modes.mode1.config['mode']['stop_on_ball_end'])
        self.assertTrue(self.machine.modes.mode2.config['mode']['stop_on_ball_end'])
        self.assertTrue(self.machine.modes.mode3.config['mode']['stop_on_ball_end'])

        # start a game
        self.machine.playfield.add_ball = MagicMock()
        self.machine.events.post('game_start')
        self.advance_time_and_run()
        self.machine.game.balls_in_play = 1
        self.assertTrue(self.machine.game)

        # start some modes
        self.machine.modes.mode1.start()
        self.machine.modes.mode2.start()
        self.machine.modes.mode3.start()
        self.advance_time_and_run()
        self.assertTrue(self.machine.modes.mode1.active)
        self.assertTrue(self.machine.modes.mode2.active)
        self.assertTrue(self.machine.modes.mode3.active)

        # change the config of modes 1 and 2 so we can verify whether they
        # actually stopped and restarted since the configs are used to start
        # the modes
        self.machine.modes.mode1.config['mode']['priority'] = 999
        self.machine.modes.mode2.config['mode']['priority'] = 999

        # end the ball
        self.machine.game.ball_ending()
        self.advance_time_and_run()

        self.assertTrue(self.machine.modes.mode1.active)
        self.assertTrue(self.machine.modes.mode2.active)
        self.assertFalse(self.machine.modes.mode3.active)

        # since mode1 should have stayed running, its config should not change
        self.assertNotEqual(self.machine.modes.mode1.priority, 999)
        # mode2 should use the new priority since it should have restarted
        self.assertEqual(self.machine.modes.mode2.priority, 999)

        # end ball 2
        self.machine.game.ball_ending()
        self.advance_time_and_run()

        # end ball 3 and end the game
        self.machine.game.ball_ending()
        self.advance_time_and_run()

        self.assertTrue(self.machine.modes.attract.active)
        self.assertFalse(self.machine.modes.game.active)
        self.assertTrue(self.machine.modes.mode1.active)
        self.assertFalse(self.machine.modes.mode2.active)
        self.assertFalse(self.machine.modes.mode3.active)
<|endoftext|>


<commit_before>"""
Encapsulates the SNAP aligner.

An important message from the fine folks at SNAP:

You may process more than one alignment without restarting SNAP, and if possible without reloading
the index.  In order to do this, list on the command line all of the parameters for the first
alignment, followed by a comma (separated by a space from the other parameters) followed by the
parameters for the next alignment (including single or paired).  You may have as many of these
as you please.  If two consecutive alignments use the same index, it will not be reloaded.
So, for example, you could do

'snap single hg19-20 foo.fq -o foo.sam , paired hg19-20 end1.fq end2.fq -o paired.sam'

and it would not reload the index between the single and paired alignments.

And another important message:

When specifying an input or output file, you can simply list the filename, in which case
SNAP will infer the type of the file from the file extension (.sam or .bam for example),
or you can explicitly specify the file type by preceeding the filename with one of the
 following type specifiers (which are case sensitive):
    -fastq
    -compressedFastq
    -sam
    -bam
    -pairedFastq
    -pairedInterleavedFastq
    -pairedCompressedInterleavedFastq

Input and output may also be from/to stdin/stdout. To do that, use a - for the input or output file
name and give an explicit type specifier.  So, for example,
snap single myIndex -fastq - -o -sam -
would read FASTQ from stdin and write SAM to stdout.

A couple of "paired" mode parameters to know about:

  -s   min and max spacing to allow between paired ends (default: 50 1000).
  -fs  force spacing to lie between min and max.
"""

import os
import logging
import re
import string
import sys
import operator
from subprocess import Popen, PIPE
from read import Read, Alignment
from threading import Thread
from threadutil import fh2q, q2fh
from aligner import Aligner

try:
    from Queue import Queue
except ImportError:
    from queue import Queue  # python 3.x

class SnapAligner(Aligner):

    """ Encapsulates a snap-aligner process.  The input can be a FASTQ
        file, or a Queue onto which the caller enqueues reads.
        Similarly, output can be a SAM file, or a Queue from which the
        caller dequeues SAM records.  All records are textual; parsing
        is up to the user. """

    _input_args = ['-fastq',
                   '-compressedFastq',
                   '-sam',
                   '-bam',
                   '-pairedFastq',
                   '-pairedInterleavedFastq',
                   '-pairedCompressedInterleavedFastq']

    _output_args = ['-o', '-sam', '-bam']

    def __init__(self,
                 cmd,
                 index,
                 unpaired=None,  # -fastq
                 paired=None,  # -pairedFastq
                 paired_combined=None,  # -pairedInterleavedFastq
                 pairs_only=False,
                 sam=None,
                 quiet=False,
                 input_format=None):
        """ Create new process.

            Inputs:

            'unpaired' is an iterable over unpaired input filenames.
            'paired' is an iterable over pairs of paired-end input
            filenames.  If both are None, then input reads will be
            taken over the inQ.  If either are non-None, then a call
            to inQ will raise an exception.

            Outputs:

            'sam' is a filename where output SAM records will be
            stored.  If 'sam' is none, SAM records will be added to
            the outQ.
        """

        if index is None:
            raise RuntimeError('Must specify --index when aligner is SNAP')

        cmd_toks = cmd.split()
        popen_stdin, popen_stdout, popen_stderr = None, None, None
        self.inQ, self.outQ = None, None

        #
        # Compose input arguments
        #

        for tok in self._input_args:
            assert tok not in cmd_toks

        self.input_is_queued = False
        self.output_is_queued = False

        args_single, args_paired = ['single', index], ['paired', index]

        if paired_combined is not None:
            assert paired is None
            args_paired.append('-pairedInterleavedFastq')
            args_paired.extend(paired_combined)
        elif paired is not None:
            assert paired_combined is None
            args_paired.append('-pairedFastq')
            args_paired.extend(list(reduce(operator.concat, paired)))

        if unpaired is not None:
            args_single.append('-fastq')
            args_single.extend(unpaired)

        if unpaired is None and paired is None and paired_combined is None:
            raise RuntimeError('Cannot instantiate SnapAligner without input file(s) specified')

        #
        # Compose output arguments
        #

        for tok in self._output_args:
            assert tok not in cmd_toks

        # Compose output arguments
        args_output = ['-o', '-sam']
        if sam is not None:
            args_output.append(sam)
        else:
            args_output.append('-')
            self.output_is_queued = True
            popen_stdout = PIPE

        # Put all the arguments together
        if len(args_single) > 1:
            cmd += ' '.join(args_single)
        if len(args_paired) > 1:
            if len(cmd) > 0:
                cmd += ' , '
            cmd += ' '.join(args_paired)
        cmd += ' ' + ' '.join(args_output)
        cmd += ' ' + ' '.join(cmd_toks)

        logging.info('SNAP command: ' + cmd)
        if quiet:
            popen_stderr = open(os.devnull, 'w')
        self.pipe = Popen(cmd, shell=True,
                          stdin=popen_stdin, stdout=popen_stdout, stderr=popen_stderr,
                          bufsize=-1, close_fds='posix' in sys.builtin_module_names)
        # Create queue threads, if necessary
        timeout = 0.2
        if self.input_is_queued:
            self.inQ = Queue()
            self._inThread = Thread(target=q2fh, args=(self.inQ, self.pipe.stdin, timeout))
            self._inThread.daemon = True  # thread dies with the program
            self._inThread.start()
        if self.output_is_queued:
            self.outQ = Queue()
            self._outThread = Thread(target=fh2q, args=(self.pipe.stdout, self.outQ, timeout))
            self._outThread.daemon = True  # thread dies with the program
            self._outThread.start()

    @staticmethod
    def format_read(rd1, rd2=None, truncate_name=True):
        return Read.to_interleaved_fastq(rd1, rd2, truncate_name=truncate_name) + '\n'

    def put(self, rd1, rd2=None, truncate_name=True):
        assert self.input_is_queued
        self.inQ.put(self.format_read(rd1, rd2, truncate_name=truncate_name))

    @staticmethod
    def preferred_unpaired_format():
        return 'fastq'

    @staticmethod
    def preferred_paired_format():
        return 'interleaved_fastq'

    def done(self):
        assert self.input_is_queued
        self.inQ.put(None)

    def supports_mix(self):
        """
        Note: return value of true doesn't just mean it can take some unpaired
        and some paired in a given invocation; it also means that can be
        interleaved in the input.
        """
        return False

    def supports_concurrency(self):
        """ Can take input reads on a queue and write output alignments
            to a queue?  Otherwise, . """
        return False  # might be True

    def writes_bam(self):
        """ Writes BAM directly to a file (like MOSAIK)?  Otherwise, we
            assume it writes SAM to stdout. """
        return False

class AlignmentSnap(Alignment):
    """ Encapsulates a SNAP SAM alignment record.  Parses certain
        important SAM extra fields output by snap-aligner. """

    __asRe = re.compile('AS:i:([-]?[0-9]+)')  # best score
    __xsRe = re.compile('XS:i:([-]?[0-9]+)')  # second-best score
    __mdRe = re.compile('MD:Z:([^\s]+)')  # MD:Z string
    __zupRe = re.compile('ZP:i:([-]?[0-9]+)')  # best concordant
    __zlpRe = re.compile('Zp:i:([-]?[0-9]+)')  # 2nd best concordant
    __ztRe = re.compile('ZT:Z:([^\s]*)')  # extra features

    def __init__(self):
        super(AlignmentSnap, self).__init__()
        self.name = None
        self.flags = None
        self.refid = None
        self.pos = None
        self.mapq = None
        self.cigar = None
        self.rnext = None
        self.pnext = None
        self.tlen = None
        self.seq = None
        self.qual = None
        self.extra = None
        self.fw = None
        self.mate1 = None
        self.mate2 = None
        self.paired = None
        self.concordant = None
        self.discordant = None
        self.bestScore = None
        self.secondBestScore = None
        self.bestConcordantScore = None
        self.secondBestConcordantScore = None
        self.ztzs = None
        self.mdz = None

    def parse(self, ln):
        """ Parse ln, which is a line of SAM output from SNAP.  The line
            must correspond to an aligned read. """
        self.name, self.flags, self.refid, self.pos, self.mapq, self.cigar, \
            self.rnext, self.pnext, self.tlen, self.seq, self.qual, self.extra = \
            string.split(ln, '\t', 11)
        assert self.flags != "*"
        assert self.pos != "*"
        assert self.mapq != "*"
        assert self.tlen != "*"
        self.flags = flags = int(self.flags)
        self.pos = int(self.pos) - 1
        self.mapq = int(self.mapq)
        self.tlen = int(self.tlen)
        self.pnext = int(self.pnext)
        self.fw = (flags & 16) == 0
        self.mate1 = (flags & 64) != 0
        self.mate2 = (flags & 128) != 0
        self.paired = self.mate1 or self.mate2
        assert self.paired == ((flags & 1) != 0)
        self.concordant = ((flags & 2) != 0)
        self.discordant = ((flags & 2) == 0) and ((flags & 4) == 0) and ((flags & 8) == 0)
        # Parse AS:i
        se = self.__asRe.search(self.extra)
        self.bestScore = None
        if se is not None:
            self.bestScore = int(se.group(1))
        # Parse XS:i
        se = self.__xsRe.search(self.extra)
        self.secondBestScore = None
        if se is not None:
            self.secondBestScore = int(se.group(1))
        # Parse ZP:i
        se = self.__zupRe.search(self.extra)
        self.bestConcordantScore = None
        if se is not None:
            self.bestConcordantScore = int(se.group(1))
        # Parse Zp:i
        se = self.__zlpRe.search(self.extra)
        self.secondBestConcordantScore = None
        if se is not None:
            self.secondBestConcordantScore = int(se.group(1))
        # Parse ZT.Z
        self.ztzs = None
        se = self.__ztRe.search(self.extra)
        if se is not None:
            self.ztzs = se.group(1).split(',')
        # Parse MD:Z
        self.mdz = None
        se = self.__mdRe.search(self.extra)
        if se is not None:
            self.mdz = se.group(1)
        assert self.rep_ok()

    def rep_ok(self):
        # Call parent's repOk
        assert super(AlignmentSnap, self).rep_ok()
        return True
<commit_msg>fixing issue with paired args being included spruiously<commit_after>"""
Encapsulates the SNAP aligner.

An important message from the fine folks at SNAP:

You may process more than one alignment without restarting SNAP, and if possible without reloading
the index.  In order to do this, list on the command line all of the parameters for the first
alignment, followed by a comma (separated by a space from the other parameters) followed by the
parameters for the next alignment (including single or paired).  You may have as many of these
as you please.  If two consecutive alignments use the same index, it will not be reloaded.
So, for example, you could do

'snap single hg19-20 foo.fq -o foo.sam , paired hg19-20 end1.fq end2.fq -o paired.sam'

and it would not reload the index between the single and paired alignments.

And another important message:

When specifying an input or output file, you can simply list the filename, in which case
SNAP will infer the type of the file from the file extension (.sam or .bam for example),
or you can explicitly specify the file type by preceeding the filename with one of the
 following type specifiers (which are case sensitive):
    -fastq
    -compressedFastq
    -sam
    -bam
    -pairedFastq
    -pairedInterleavedFastq
    -pairedCompressedInterleavedFastq

Input and output may also be from/to stdin/stdout. To do that, use a - for the input or output file
name and give an explicit type specifier.  So, for example,
snap single myIndex -fastq - -o -sam -
would read FASTQ from stdin and write SAM to stdout.

A couple of "paired" mode parameters to know about:

  -s   min and max spacing to allow between paired ends (default: 50 1000).
  -fs  force spacing to lie between min and max.
"""

import os
import logging
import re
import string
import sys
import operator
from subprocess import Popen, PIPE
from read import Read, Alignment
from threading import Thread
from threadutil import fh2q, q2fh
from aligner import Aligner

try:
    from Queue import Queue
except ImportError:
    from queue import Queue  # python 3.x

class SnapAligner(Aligner):

    """ Encapsulates a snap-aligner process.  The input can be a FASTQ
        file, or a Queue onto which the caller enqueues reads.
        Similarly, output can be a SAM file, or a Queue from which the
        caller dequeues SAM records.  All records are textual; parsing
        is up to the user. """

    _input_args = ['-fastq',
                   '-compressedFastq',
                   '-sam',
                   '-bam',
                   '-pairedFastq',
                   '-pairedInterleavedFastq',
                   '-pairedCompressedInterleavedFastq']

    _output_args = ['-o', '-sam', '-bam']

    def __init__(self,
                 cmd,
                 index,
                 unpaired=None,  # -fastq
                 paired=None,  # -pairedFastq
                 paired_combined=None,  # -pairedInterleavedFastq
                 pairs_only=False,
                 sam=None,
                 quiet=False,
                 input_format=None):
        """ Create new process.

            Inputs:

            'unpaired' is an iterable over unpaired input filenames.
            'paired' is an iterable over pairs of paired-end input
            filenames.  If both are None, then input reads will be
            taken over the inQ.  If either are non-None, then a call
            to inQ will raise an exception.

            Outputs:

            'sam' is a filename where output SAM records will be
            stored.  If 'sam' is none, SAM records will be added to
            the outQ.
        """

        if index is None:
            raise RuntimeError('Must specify --index when aligner is SNAP')

        cmd_toks = cmd.split()
        popen_stdin, popen_stdout, popen_stderr = None, None, None
        self.inQ, self.outQ = None, None

        #
        # Compose input arguments
        #

        for tok in self._input_args:
            assert tok not in cmd_toks

        self.input_is_queued = False
        self.output_is_queued = False

        args_single, args_paired = ['single', index], ['paired', index]

        if paired_combined is not None:
            assert paired is None
            args_paired.append('-pairedInterleavedFastq')
            args_paired.extend(paired_combined)
        elif paired is not None:
            assert paired_combined is None
            args_paired.append('-pairedFastq')
            args_paired.extend(list(reduce(operator.concat, paired)))

        if unpaired is not None:
            args_single.append('-fastq')
            args_single.extend(unpaired)

        if unpaired is None and paired is None and paired_combined is None:
            raise RuntimeError('Cannot instantiate SnapAligner without input file(s) specified')

        #
        # Compose output arguments
        #

        for tok in self._output_args:
            assert tok not in cmd_toks

        # Compose output arguments
        args_output = ['-o', '-sam']
        if sam is not None:
            args_output.append(sam)
        else:
            args_output.append('-')
            self.output_is_queued = True
            popen_stdout = PIPE

        # Put all the arguments together
        if len(args_single) > 2:
            cmd += ' '.join(args_single)
        if len(args_paired) > 2:
            if len(cmd) > 0:
                cmd += ' , '
            cmd += ' '.join(args_paired)
        cmd += ' ' + ' '.join(args_output)
        cmd += ' ' + ' '.join(cmd_toks)

        logging.info('SNAP command: ' + cmd)
        if quiet:
            popen_stderr = open(os.devnull, 'w')
        self.pipe = Popen(cmd, shell=True,
                          stdin=popen_stdin, stdout=popen_stdout, stderr=popen_stderr,
                          bufsize=-1, close_fds='posix' in sys.builtin_module_names)
        # Create queue threads, if necessary
        timeout = 0.2
        if self.input_is_queued:
            self.inQ = Queue()
            self._inThread = Thread(target=q2fh, args=(self.inQ, self.pipe.stdin, timeout))
            self._inThread.daemon = True  # thread dies with the program
            self._inThread.start()
        if self.output_is_queued:
            self.outQ = Queue()
            self._outThread = Thread(target=fh2q, args=(self.pipe.stdout, self.outQ, timeout))
            self._outThread.daemon = True  # thread dies with the program
            self._outThread.start()

    @staticmethod
    def format_read(rd1, rd2=None, truncate_name=True):
        return Read.to_interleaved_fastq(rd1, rd2, truncate_name=truncate_name) + '\n'

    def put(self, rd1, rd2=None, truncate_name=True):
        assert self.input_is_queued
        self.inQ.put(self.format_read(rd1, rd2, truncate_name=truncate_name))

    @staticmethod
    def preferred_unpaired_format():
        return 'fastq'

    @staticmethod
    def preferred_paired_format():
        return 'interleaved_fastq'

    def done(self):
        assert self.input_is_queued
        self.inQ.put(None)

    def supports_mix(self):
        """
        Note: return value of true doesn't just mean it can take some unpaired
        and some paired in a given invocation; it also means that can be
        interleaved in the input.
        """
        return False

    def supports_concurrency(self):
        """ Can take input reads on a queue and write output alignments
            to a queue?  Otherwise, . """
        return False  # might be True

    def writes_bam(self):
        """ Writes BAM directly to a file (like MOSAIK)?  Otherwise, we
            assume it writes SAM to stdout. """
        return False

class AlignmentSnap(Alignment):
    """ Encapsulates a SNAP SAM alignment record.  Parses certain
        important SAM extra fields output by snap-aligner. """

    __asRe = re.compile('AS:i:([-]?[0-9]+)')  # best score
    __xsRe = re.compile('XS:i:([-]?[0-9]+)')  # second-best score
    __mdRe = re.compile('MD:Z:([^\s]+)')  # MD:Z string
    __zupRe = re.compile('ZP:i:([-]?[0-9]+)')  # best concordant
    __zlpRe = re.compile('Zp:i:([-]?[0-9]+)')  # 2nd best concordant
    __ztRe = re.compile('ZT:Z:([^\s]*)')  # extra features

    def __init__(self):
        super(AlignmentSnap, self).__init__()
        self.name = None
        self.flags = None
        self.refid = None
        self.pos = None
        self.mapq = None
        self.cigar = None
        self.rnext = None
        self.pnext = None
        self.tlen = None
        self.seq = None
        self.qual = None
        self.extra = None
        self.fw = None
        self.mate1 = None
        self.mate2 = None
        self.paired = None
        self.concordant = None
        self.discordant = None
        self.bestScore = None
        self.secondBestScore = None
        self.bestConcordantScore = None
        self.secondBestConcordantScore = None
        self.ztzs = None
        self.mdz = None

    def parse(self, ln):
        """ Parse ln, which is a line of SAM output from SNAP.  The line
            must correspond to an aligned read. """
        self.name, self.flags, self.refid, self.pos, self.mapq, self.cigar, \
            self.rnext, self.pnext, self.tlen, self.seq, self.qual, self.extra = \
            string.split(ln, '\t', 11)
        assert self.flags != "*"
        assert self.pos != "*"
        assert self.mapq != "*"
        assert self.tlen != "*"
        self.flags = flags = int(self.flags)
        self.pos = int(self.pos) - 1
        self.mapq = int(self.mapq)
        self.tlen = int(self.tlen)
        self.pnext = int(self.pnext)
        self.fw = (flags & 16) == 0
        self.mate1 = (flags & 64) != 0
        self.mate2 = (flags & 128) != 0
        self.paired = self.mate1 or self.mate2
        assert self.paired == ((flags & 1) != 0)
        self.concordant = ((flags & 2) != 0)
        self.discordant = ((flags & 2) == 0) and ((flags & 4) == 0) and ((flags & 8) == 0)
        # Parse AS:i
        se = self.__asRe.search(self.extra)
        self.bestScore = None
        if se is not None:
            self.bestScore = int(se.group(1))
        # Parse XS:i
        se = self.__xsRe.search(self.extra)
        self.secondBestScore = None
        if se is not None:
            self.secondBestScore = int(se.group(1))
        # Parse ZP:i
        se = self.__zupRe.search(self.extra)
        self.bestConcordantScore = None
        if se is not None:
            self.bestConcordantScore = int(se.group(1))
        # Parse Zp:i
        se = self.__zlpRe.search(self.extra)
        self.secondBestConcordantScore = None
        if se is not None:
            self.secondBestConcordantScore = int(se.group(1))
        # Parse ZT.Z
        self.ztzs = None
        se = self.__ztRe.search(self.extra)
        if se is not None:
            self.ztzs = se.group(1).split(',')
        # Parse MD:Z
        self.mdz = None
        se = self.__mdRe.search(self.extra)
        if se is not None:
            self.mdz = se.group(1)
        assert self.rep_ok()

    def rep_ok(self):
        # Call parent's repOk
        assert super(AlignmentSnap, self).rep_ok()
        return True
<|endoftext|>


<commit_before>from flask import Flask, Response, Request
import flask

from werkzeug.test import EnvironBuilder

from nose.tools import raises

from .resolver import Pushrod, pushrod_view
from .renderers.base import renderer, UnrenderedResponse, RendererNotFound
from .renderers.json import json_renderer
from .renderers.jinja2 import jinja2_renderer

from unittest import TestCase
import json

import logging

test_response_str = """
{
    "spam": "eggs",
    "yes": [
        "maybe",
        "no",
    ],
    "nonsense": (
        True,
        False,
        "FILE_NOT_FOUND",
        -5.4,
    ),
}
"""

test_response = eval(test_response_str)

@renderer('repr')
def repr_renderer(unrendered, **kwargs):
    return unrendered.rendered(repr(unrendered.response), "text/plain")

def test_logging_no_app():
    pushrod = Pushrod()
    assert pushrod.logger == logging

class PushrodTestCase(TestCase):
    def setUp(self):
        app = Flask(__name__,
            template_folder='test_templates',
        )
        pushrod = Pushrod(app)

        pushrod.register_renderer(repr_renderer)
        pushrod.default_renderer = repr_renderer

        self.app = app
        self.client = app.test_client()

        if not hasattr(self, 'app_contexts'):
            self.app_contexts = []
        app_context = app.app_context()
        app_context.push()
        self.app_contexts.append(app_context)

    def tearDown(self):
        self.app_contexts.pop().pop()

    def setup_method(self, method):
        # Setup for py.test
        return self.setUp()

    def teardown_method(self, method):
        # Setup for py.test
        return self.tearDown()

class PushrodResolverTestCase(PushrodTestCase):
    def test_raw_dict_response_data(self):
        @self.app.route("/raw_dict")
        @pushrod_view()
        def test_view_raw_dict():
            return test_response

        test_view_raw_dict_response = self.client.get("/raw_dict")
        assert test_view_raw_dict_response.status_code == 200
        assert test_view_raw_dict_response.data == repr(test_response)

    def test_response_data(self):
        @self.app.route("/200_response")
        @pushrod_view()
        def test_view_200_view():
            return UnrenderedResponse(test_response)

        @self.app.route("/404_response")
        @pushrod_view()
        def test_view_404_view():
            return UnrenderedResponse(test_response, 404)

        test_view_200_response = self.client.get("/200_response")
        assert test_view_200_response.status_code == 200
        assert test_view_200_response.data == repr(test_response)

        test_view_404_response = self.client.get("/404_response")
        assert test_view_404_response.status_code == 404
        assert test_view_404_response.data == repr(test_response)

    def test_response_bypass(self):
        @self.app.route("/raw_string")
        @pushrod_view()
        def test_raw_string_view():
            return "test"

        @self.app.route("/regular_response")
        @pushrod_view()
        def test_regular_response_view():
            return Response("test")

        test_raw_string_response = self.client.get("/raw_string")
        assert test_raw_string_response.status_code == 200
        assert test_raw_string_response.data == "test"

        test_regular_response_response = self.client.get("/regular_response")
        assert test_regular_response_response.status_code == 200
        assert test_regular_response_response.data == "test"

    @raises(RendererNotFound)
    def test_renderer_not_found(self):
        @self.app.route("/")
        @pushrod_view()
        def test_renderer_not_found_view():
            return {}

        response = self.client.get("/?format=none")
        assert response.status_code == 406, \
            "Status code returned from an unexistant renderer should be 406 (not acceptable)"

        with self.app.test_request_context("/?format=none"):
            self.app.pushrod.render_response(test_renderer_not_found_view(), "none")

    def test_constructor_renderer_registration(self):
        get_request_obj = lambda *args, **kwargs: Request(EnvironBuilder(*args, **kwargs).get_environ())

        pushrod = Pushrod(renderers=['json'])

        assert pushrod.default_renderer == None

        assert 'html' not in pushrod.named_renderers
        assert 'text/html' not in pushrod.mime_type_renderers

        assert 'json' in pushrod.named_renderers
        assert 'application/json' in pushrod.mime_type_renderers

        assert pushrod.named_renderers['json'] == json_renderer
        assert pushrod.mime_type_renderers['application/json'] == json_renderer

        assert pushrod.get_renderers_for_request(get_request_obj("/")) == []
        assert pushrod.get_renderers_for_request(get_request_obj("/?format=json")) == [json_renderer]

        pushrod = Pushrod(renderers=[repr_renderer])

        assert 'html' not in pushrod.named_renderers
        assert 'text/html' not in pushrod.named_renderers

        assert 'json' not in pushrod.named_renderers
        assert 'application/json' not in pushrod.mime_type_renderers

        assert 'repr' in pushrod.named_renderers

        assert pushrod.named_renderers['repr'] == repr_renderer

        assert pushrod.get_renderers_for_request(get_request_obj("/")) == []
        assert pushrod.get_renderers_for_request(get_request_obj("/?format=repr")) == [repr_renderer]

        pushrod = Pushrod(renderers=[repr_renderer], default_renderer='repr')

        assert pushrod.get_renderers_for_request(get_request_obj("/")) == [repr_renderer]

        with self.app.app_context():
            Pushrod(renderers=[])

        Pushrod(self.app, renderers=[])

        Pushrod(self.app, default_renderer=repr_renderer)

    @raises(TypeError)
    def test_register_invalid_renderer(self):
        def dummy():  # pragma: no cover
            pass

        self.app.pushrod.register_renderer(dummy)

    def test_render_tuple_response(self):
        response = self.app.pushrod.render_response((
            {u'aaa': u"hi"},  # Response value
            978,            # Status code
            {               # Headers
                'X-Aaa': "Hi!",
            },
        ), json_renderer)

        assert response.status_code == 978
        assert 'X-Aaa' in response.headers
        assert response.headers['X-Aaa'] == "Hi!"

        response_json = json.loads(response.data)

        assert u'aaa' in response_json
        assert response_json[u'aaa'] == u"hi"

class PushrodNormalizerTestCase(PushrodTestCase):
    def test_basestring_normalizer(self):
        from_unicode = self.app.pushrod.normalize(u"testing string")
        from_str = self.app.pushrod.normalize("testing string")

        assert from_unicode == u"testing string"
        assert from_str == u"testing string"

    def test_number_normalizer(self):
        assert self.app.pushrod.normalize(10) == 10
        assert self.app.pushrod.normalize(500L) == 500L
        assert self.app.pushrod.normalize(20.5) == 20.5

    def test_iterable_normalizer(self):
        in_list = ["test string", 57]

        from_list = self.app.pushrod.normalize(in_list)
        from_tuple = self.app.pushrod.normalize(tuple(in_list))

        assert from_list == in_list
        assert from_tuple == in_list

    def test_dict_normalizer(self):
        in_dict = {
            "hi": 54,
            "bleh": "testable",
            "no": ["12"],
        }

        assert self.app.pushrod.normalize(in_dict) == in_dict

    def test_delegate_normalizer(self):
        class MyClass(object):
            __pushrod_field__ = "value"

            def __init__(self):
                self.value = 53.21

        my_instance = MyClass()

        assert self.app.pushrod.normalize(my_instance) == self.app.pushrod.normalize(my_instance.value)

class PushrodRendererTestCase(PushrodTestCase):
    def test_json_renderer(self):
        regular = json.dumps(test_response)
        rendered = self.app.pushrod.render_response(
            test_response, json_renderer)

        assert regular == rendered.data

        json.loads(rendered.data)

    @raises(RendererNotFound)
    def test_jinja_renderer_no_template(self):
        self.app.pushrod.render_response(
            test_response, jinja2_renderer)

    def test_jinja_renderer_no_template_fallback(self):
        self.app.pushrod.render_response(
            test_response, [jinja2_renderer, json_renderer])

    def test_jinja_renderer(self):
        regular = test_response_str

        with self.app.test_request_context():
            flask.current_app
            rendered = self.app.pushrod.render_response(
                test_response, jinja2_renderer, {'jinja_template': 'jinja.txt'})

        assert regular == rendered.data, \
            "'%s' does not equal '%s'" % (rendered.data, regular)

        assert eval(rendered.data) == test_response
<commit_msg>More normalizer tests<commit_after>from flask import Flask, Response, Request
import flask

from werkzeug.test import EnvironBuilder

from nose.tools import raises

from .resolver import Pushrod, pushrod_view
from .renderers.base import renderer, UnrenderedResponse, RendererNotFound
from .renderers.json import json_renderer
from .renderers.jinja2 import jinja2_renderer

from unittest import TestCase
import json

import logging

test_response_str = """
{
    "spam": "eggs",
    "yes": [
        "maybe",
        "no",
    ],
    "nonsense": (
        True,
        False,
        "FILE_NOT_FOUND",
        -5.4,
    ),
}
"""

test_response = eval(test_response_str)

@renderer('repr')
def repr_renderer(unrendered, **kwargs):
    return unrendered.rendered(repr(unrendered.response), "text/plain")

def test_logging_no_app():
    pushrod = Pushrod()
    assert pushrod.logger == logging

class PushrodTestCase(TestCase):
    def setUp(self):
        app = Flask(__name__,
            template_folder='test_templates',
        )
        pushrod = Pushrod(app)

        pushrod.register_renderer(repr_renderer)
        pushrod.default_renderer = repr_renderer

        self.app = app
        self.client = app.test_client()

        if not hasattr(self, 'app_contexts'):
            self.app_contexts = []
        app_context = app.app_context()
        app_context.push()
        self.app_contexts.append(app_context)

    def tearDown(self):
        self.app_contexts.pop().pop()

    def setup_method(self, method):
        # Setup for py.test
        return self.setUp()

    def teardown_method(self, method):
        # Setup for py.test
        return self.tearDown()

class PushrodResolverTestCase(PushrodTestCase):
    def test_raw_dict_response_data(self):
        @self.app.route("/raw_dict")
        @pushrod_view()
        def test_view_raw_dict():
            return test_response

        test_view_raw_dict_response = self.client.get("/raw_dict")
        assert test_view_raw_dict_response.status_code == 200
        assert test_view_raw_dict_response.data == repr(test_response)

    def test_response_data(self):
        @self.app.route("/200_response")
        @pushrod_view()
        def test_view_200_view():
            return UnrenderedResponse(test_response)

        @self.app.route("/404_response")
        @pushrod_view()
        def test_view_404_view():
            return UnrenderedResponse(test_response, 404)

        test_view_200_response = self.client.get("/200_response")
        assert test_view_200_response.status_code == 200
        assert test_view_200_response.data == repr(test_response)

        test_view_404_response = self.client.get("/404_response")
        assert test_view_404_response.status_code == 404
        assert test_view_404_response.data == repr(test_response)

    def test_response_bypass(self):
        @self.app.route("/raw_string")
        @pushrod_view()
        def test_raw_string_view():
            return "test"

        @self.app.route("/regular_response")
        @pushrod_view()
        def test_regular_response_view():
            return Response("test")

        test_raw_string_response = self.client.get("/raw_string")
        assert test_raw_string_response.status_code == 200
        assert test_raw_string_response.data == "test"

        test_regular_response_response = self.client.get("/regular_response")
        assert test_regular_response_response.status_code == 200
        assert test_regular_response_response.data == "test"

    @raises(RendererNotFound)
    def test_renderer_not_found(self):
        @self.app.route("/")
        @pushrod_view()
        def test_renderer_not_found_view():
            return {}

        response = self.client.get("/?format=none")
        assert response.status_code == 406, \
            "Status code returned from an unexistant renderer should be 406 (not acceptable)"

        with self.app.test_request_context("/?format=none"):
            self.app.pushrod.render_response(test_renderer_not_found_view(), "none")

    def test_constructor_renderer_registration(self):
        get_request_obj = lambda *args, **kwargs: Request(EnvironBuilder(*args, **kwargs).get_environ())

        pushrod = Pushrod(renderers=['json'])

        assert pushrod.default_renderer == None

        assert 'html' not in pushrod.named_renderers
        assert 'text/html' not in pushrod.mime_type_renderers

        assert 'json' in pushrod.named_renderers
        assert 'application/json' in pushrod.mime_type_renderers

        assert pushrod.named_renderers['json'] == json_renderer
        assert pushrod.mime_type_renderers['application/json'] == json_renderer

        assert pushrod.get_renderers_for_request(get_request_obj("/")) == []
        assert pushrod.get_renderers_for_request(get_request_obj("/?format=json")) == [json_renderer]

        pushrod = Pushrod(renderers=[repr_renderer])

        assert 'html' not in pushrod.named_renderers
        assert 'text/html' not in pushrod.named_renderers

        assert 'json' not in pushrod.named_renderers
        assert 'application/json' not in pushrod.mime_type_renderers

        assert 'repr' in pushrod.named_renderers

        assert pushrod.named_renderers['repr'] == repr_renderer

        assert pushrod.get_renderers_for_request(get_request_obj("/")) == []
        assert pushrod.get_renderers_for_request(get_request_obj("/?format=repr")) == [repr_renderer]

        pushrod = Pushrod(renderers=[repr_renderer], default_renderer='repr')

        assert pushrod.get_renderers_for_request(get_request_obj("/")) == [repr_renderer]

        with self.app.app_context():
            Pushrod(renderers=[])

        Pushrod(self.app, renderers=[])

        Pushrod(self.app, default_renderer=repr_renderer)

    @raises(TypeError)
    def test_register_invalid_renderer(self):
        def dummy():  # pragma: no cover
            pass

        self.app.pushrod.register_renderer(dummy)

    def test_render_tuple_response(self):
        response = self.app.pushrod.render_response((
            {u'aaa': u"hi"},  # Response value
            978,            # Status code
            {               # Headers
                'X-Aaa': "Hi!",
            },
        ), json_renderer)

        assert response.status_code == 978
        assert 'X-Aaa' in response.headers
        assert response.headers['X-Aaa'] == "Hi!"

        response_json = json.loads(response.data)

        assert u'aaa' in response_json
        assert response_json[u'aaa'] == u"hi"

class PushrodNormalizerTestCase(PushrodTestCase):
    def test_basestring_normalizer(self):
        from_unicode = self.app.pushrod.normalize(u"testing string")
        from_str = self.app.pushrod.normalize("testing string")

        assert from_unicode == u"testing string"
        assert from_str == u"testing string"

    def test_number_normalizer(self):
        assert self.app.pushrod.normalize(10) == 10
        assert self.app.pushrod.normalize(500L) == 500L
        assert self.app.pushrod.normalize(20.5) == 20.5

    def test_iterable_normalizer(self):
        in_list = ["test string", 57]

        from_list = self.app.pushrod.normalize(in_list)
        from_tuple = self.app.pushrod.normalize(tuple(in_list))

        assert from_list == in_list
        assert from_tuple == in_list

    def test_dict_normalizer(self):
        in_dict = {
            "hi": 54,
            "bleh": "testable",
            "no": ["12"],
        }

        assert self.app.pushrod.normalize(in_dict) == in_dict

    def test_delegate_normalizer(self):
        class MyClass(object):
            __pushrod_field__ = "value"

            def __init__(self):
                self.value = 53.21

        my_instance = MyClass()

        assert self.app.pushrod.normalize(my_instance) == self.app.pushrod.normalize(my_instance.value)

    def test_no_normalizer(self):
        class MyClass(object):
            pass

        assert self.app.pushrod.normalize(MyClass()) == NotImplemented

        self.app.pushrod.normalizer_fallbacks[MyClass] = lambda x, pushrod: NotImplemented

        assert self.app.pushrod.normalize(MyClass()) == NotImplemented

    def test_normalizer_method(self):
        class MyClass(object):
            def __pushrod_normalize__(self, pushrod):
                return u"Normalized!"

        assert self.app.pushrod.normalize(MyClass()) == u"Normalized!"

    def test_normalizer_generated_dict(self):
        class MyClass(object):
            __pushrod_fields__ = ["one", "three"]

            def __init__(self, one, two, three):
                self.one = one
                self.two = two
                self.three = three

        assert self.app.pushrod.normalize(MyClass('first', 'second', 'third')) == {u'one': u'first', u'three': u'third'}

    def test_normalizer_override(self):
        self.app.pushrod.normalizer_overrides[int].append(lambda x, pushrod: unicode(x) if x > 0 else NotImplemented)

        assert self.app.pushrod.normalize(0) == 0
        assert self.app.pushrod.normalize(1) == u"1"

class PushrodRendererTestCase(PushrodTestCase):
    def test_json_renderer(self):
        regular = json.dumps(test_response)
        rendered = self.app.pushrod.render_response(
            test_response, json_renderer)

        assert regular == rendered.data

        json.loads(rendered.data)

    @raises(RendererNotFound)
    def test_jinja_renderer_no_template(self):
        self.app.pushrod.render_response(
            test_response, jinja2_renderer)

    def test_jinja_renderer_no_template_fallback(self):
        self.app.pushrod.render_response(
            test_response, [jinja2_renderer, json_renderer])

    def test_jinja_renderer(self):
        regular = test_response_str

        with self.app.test_request_context():
            flask.current_app
            rendered = self.app.pushrod.render_response(
                test_response, jinja2_renderer, {'jinja_template': 'jinja.txt'})

        assert regular == rendered.data, \
            "'%s' does not equal '%s'" % (rendered.data, regular)

        assert eval(rendered.data) == test_response
<|endoftext|>