<commit_before># -*- coding: utf-8 -*-
"""Git Repo object for libvcs.

libvcs.git
~~~~~~~~~~

From https://github.com/saltstack/salt (Apache License):

- :py:meth:`GitRepo.remote`
- :py:meth:`GitRepo.remote_get`
- :py:meth:`GitRepo.remote_set`

From pip (MIT Licnese):

- :py:meth:`GitRepo.get_url_and_revision_from_pip_url` (get_url_rev)
- :py:meth:`GitRepo.get_revision`

"""
from __future__ import absolute_import, print_function, unicode_literals

import logging
import os
import re

from . import exc
from ._compat import urlparse
from .base import BaseRepo

logger = logging.getLogger(__name__)

class GitRepo(BaseRepo):
    bin_name = 'git'
    schemes = (
        'git', 'git+http', 'git+https', 'git+ssh', 'git+git', 'git+file',
    )

    def __init__(self, url, remotes=None, **kwargs):
        """A git repository.

        :param url: URL of repo
        :type url: str

        :param remotes: list of remotes in dict format::

            [{
            "remote_name": "myremote",
            "url": "https://github.com/tony/vim-config.git"
            }]
        :type remotes: list

        :param git_remote_name: name of the remote (default "origin")
        :type git_remote_name: str

        :param git_shallow: clone with ``--depth 1`` (default False)
        :type git_shallow: bool

        :param git_submodules: Git submodules that shall be updated, all if
            empty
        :type git_submodules: list

        :param tls_verify: Should certificate for https be checked (default
            False)
        :type tls_verify: bool
        """
        if 'git_remote_name' not in kwargs:
            self.git_remote_name = "origin"
        if 'git_shallow' not in kwargs:
            self.git_shallow = False
        if 'git_submodules' not in kwargs:
            self.git_submodules = []
        if 'tls_verify' not in kwargs:
            self.tls_verify = False
        BaseRepo.__init__(self, url, **kwargs)

        self.remotes = remotes

    def get_revision(self):
        """Return current revision. Initial repositories return 'initial'."""
        try:
            return self.run(['rev-parse', '--verify', 'HEAD'])
        except exc.SubprocessError:
            return 'initial'

    @classmethod
    def get_url_and_revision_from_pip_url(cls, pip_url):
        """
        Prefixes stub URLs like 'user@hostname:user/repo.git' with 'ssh://'.
        That's required because although they use SSH they sometimes doesn't
        work with a ssh:// scheme (e.g. Github). But we need a scheme for
        parsing. Hence we remove it again afterwards and return it as a stub.
        """
        if '://' not in pip_url:
            assert 'file:' not in pip_url
            pip_url = pip_url.replace('git+', 'git+ssh://')
            url, rev = super(
                GitRepo, cls).get_url_and_revision_from_pip_url(pip_url)
            url = url.replace('ssh://', '')
        elif 'github.com:' in pip_url:
            raise exc.LibVCSException(
                "Repo %s is malformatted, please use the convention %s for"
                "ssh / private GitHub repositories." % (
                    pip_url, "git+https://github.com/username/repo.git"
                )
            )
        else:
            url, rev = super(
                GitRepo, cls).get_url_and_revision_from_pip_url(pip_url)

        return url, rev

    def obtain(self, quiet=False):
        """Retrieve the repository, clone if doesn't exist.

        :param quiet: Suppress stderr output.
        :type quiet: bool

        """
        self.check_destination()

        url = self.url

        cmd = ['clone', '--progress']
        if self.git_shallow:
            cmd.extend(['--depth', '1'])
        if self.tls_verify:
            cmd.extend(['-c', 'http.sslVerify=false'])
        cmd.extend([url, self.path])

        self.info('Cloning.')
        self.run_buffered(cmd)

        if self.remotes:
            for r in self.remotes:
                self.error('Adding remote %s <%s>' %
                           (r['remote_name'], r['url']))
                self.remote_set(
                    name=r['remote_name'],
                    url=r['url']
                )

        self.info('Initializing submodules.')
        self.run_buffered(['submodule', 'init'],)
        cmd = ['submodule', 'update', '--recursive', '--init']
        cmd.extend(self.git_submodules)
        self.run_buffered(cmd)

    def update_repo(self):
        self.check_destination()

        if not os.path.isdir(os.path.join(self.path, '.git')):
            self.obtain()
            self.update_repo()
            return

        # Get requested revision or tag
        url, git_tag = self.url, self.rev

        if not git_tag:
            self.debug("No git revision set, defaulting to origin/master")
            symref = self.run(['symbolic-ref', '--short', 'HEAD'])
            if symref:
                git_tag = symref.rstrip()
            else:
                git_tag = 'origin/master'
        self.debug("git_tag: %s" % git_tag)

        self.info("Updating to '%s'." % git_tag)

        # Get head sha
        try:
            head_sha = self.run(['rev-list', '--max-count=1', 'HEAD'])
        except exc.SubprocessError as e:
            self.error("Failed to get the hash for HEAD")
            return

        self.info("head_sha: %s" % head_sha)

        # If a remote ref is asked for, which can possibly move around,
        # we must always do a fetch and checkout.
        try:
            show_ref_output = self.run(['show-ref', git_tag])
        except exc.SubprocessError as e:
            self.error("Ref %s not found" % git_tag)
            return
        self.debug("show_ref_output: %s" % show_ref_output)
        is_remote_ref = "remotes" in show_ref_output
        self.debug("is_remote_ref: %s" % is_remote_ref)

        # Tag is in the form <remote>/<tag> (i.e. origin/master) we must
        # strip the remote from the tag.
        git_remote_name = self.git_remote_name
        if "refs/remotes/%s" % git_tag in show_ref_output:
            m = re.match(r'^(?P<git_remote_name>[^/]+)/(?P<git_tag>.+)$',
                         show_ref_output)
            git_remote_name = m.group('git_remote_name')
            git_tag = m.group('git_tag')

        # This will fail if the tag does not exist (it probably has not
        # been fetched yet).
        try:
            error_code = 0
            tag_sha = self.run(['rev-list', '--max-count=1', git_tag])
        except exc.SubprocessError as e:
            error_code = e.subprocess.returncode
        self.debug("tag_sha: %s" % tag_sha)

        # Is the hash checkout out what we want?
        somethings_up = (error_code, is_remote_ref, tag_sha != head_sha,)
        if all(not x for x in somethings_up):
            self.info("Already up-to-date.")
            return

        process = self.run_buffered(['fetch'])
        if process.returncode:
            self.error("Failed to fetch repository '%s'" % url)
            return

        if is_remote_ref:
            # Check if stash is needed
            try:
                process = self.run(['status', '--porcelain'])
            except exc.SubprocessError as e:
                self.error("Failed to get the status")
                return
            need_stash = len(process) > 0

            # If not in clean state, stash changes in order to be able
            # to be able to perform git pull --rebase
            if need_stash:
                # If Git < 1.7.6, uses --quiet --all
                git_stash_save_options = '--quiet'
                try:
                    process = self.run([
                        'stash', 'save', git_stash_save_options
                    ])
                except exc.SubprocessError as e:
                    self.error("Failed to stash changes")

            # Pull changes from the remote branch
            try:
                process = self.run_buffered([
                    'rebase', git_remote_name + '/' + git_tag
                ], print_stdout_on_progress_end=False)
            except exc.SubprocessError as e:
                # Rebase failed: Restore previous state.
                self.run(['rebase', '--abort'])
                if need_stash:
                    self.run(['stash', 'pop', '--index', '--quiet'])

                self.error(
                    "\nFailed to rebase in: '%s'.\n"
                    "You will have to resolve the conflicts manually" %
                    self.path)
                return

            if need_stash:
                try:
                    process = self.run([
                        'stash', 'pop', '--index', '--quiet'
                    ])
                except exc.SubprocessError as e:
                    # Stash pop --index failed: Try again dropping the index
                    self.run(['reset', '--hard', '--quiet'])
                    try:
                        process = self.run(['stash', 'pop', '--quiet'])
                    except exc.SubprocessError as e:
                        # Stash pop failed: Restore previous state.
                        self.run(['reset', '--hard', '--quiet', head_sha])
                        self.run(['stash', 'pop', '--index', '--quiet'])
                        self.error("\nFailed to rebase in: '%s'.\n"
                                   "You will have to resolve the "
                                   "conflicts manually" % self.path)
                        return

        else:
            try:
                process = self.run(['checkout', git_tag])
            except exc.SubprocessError as e:
                self.error("Failed to checkout tag: '%s'" % git_tag)
                return

        cmd = ['submodule', 'update', '--recursive', '--init']
        cmd.extend(self.git_submodules)
        self.run(cmd)

    @property
    def remotes_get(self):
        """Return remotes like git remote -v.

        :rtype: dict of tuples
        """
        remotes = {}

        cmd = self.run(['remote'])
        ret = filter(None, cmd.split('\n'))

        for remote_name in ret:
            remotes[remote_name] = self.remote_get(remote_name)
        return remotes

    def remote_get(self, remote='origin'):
        """Get the fetch and push URL for a specified remote name.

        :param remote: the remote name used to define the fetch and push URL
        :type remote: str
        :returns: remote name and url in tuple form
        :rtype: tuple
        """
        try:
            ret = self.run(['remote', 'show', '-n', remote])
            lines = ret.split('\n')
            remote_fetch_url = lines[1].replace('Fetch URL: ', '').strip()
            remote_push_url = lines[2].replace('Push  URL: ', '').strip()
            if remote_fetch_url != remote and remote_push_url != remote:
                res = (remote_fetch_url, remote_push_url)
                return res
            else:
                return None
        except exc.LibVCSException:
            return None

    def remote_set(self, url, name='origin'):
        """Set remote with name and URL like git remote add.

        :param url: defines the remote URL
        :type url: string
        :param name: defines the remote name.
        :type name: str
        """

        url = self.chomp_protocol(url)

        if self.remote_get(name):
            self.run(['remote', 'rm', 'name'])

        self.run(['remote', 'add', name, url])
        return self.remote_get(remote=name)

    @staticmethod
    def chomp_protocol(url):
        """Return clean VCS url from RFC-style url

        :param url: url
        :type url: string
        :return type: string
        :returns: url as VCS software would accept it
        :seealso: #14
        """
        if '+' in url:
            url = url.split('+', 1)[1]
        scheme, netloc, path, query, frag = urlparse.urlsplit(url)
        rev = None
        if '@' in path:
            path, rev = path.rsplit('@', 1)
        url = urlparse.urlunsplit((scheme, netloc, path, query, ''))
        if url.startswith('ssh://git@github.com/'):
            url = url.replace('ssh://', 'git+ssh://')
        elif '://' not in url:
            assert 'file:' not in url
            url = url.replace('git+', 'git+ssh://')
            url = url.replace('ssh://', '')
        return url
<commit_msg>only show head_sha in debug level<commit_after># -*- coding: utf-8 -*-
"""Git Repo object for libvcs.

libvcs.git
~~~~~~~~~~

From https://github.com/saltstack/salt (Apache License):

- :py:meth:`GitRepo.remote`
- :py:meth:`GitRepo.remote_get`
- :py:meth:`GitRepo.remote_set`

From pip (MIT Licnese):

- :py:meth:`GitRepo.get_url_and_revision_from_pip_url` (get_url_rev)
- :py:meth:`GitRepo.get_revision`

"""
from __future__ import absolute_import, print_function, unicode_literals

import logging
import os
import re

from . import exc
from ._compat import urlparse
from .base import BaseRepo

logger = logging.getLogger(__name__)

class GitRepo(BaseRepo):
    bin_name = 'git'
    schemes = (
        'git', 'git+http', 'git+https', 'git+ssh', 'git+git', 'git+file',
    )

    def __init__(self, url, remotes=None, **kwargs):
        """A git repository.

        :param url: URL of repo
        :type url: str

        :param remotes: list of remotes in dict format::

            [{
            "remote_name": "myremote",
            "url": "https://github.com/tony/vim-config.git"
            }]
        :type remotes: list

        :param git_remote_name: name of the remote (default "origin")
        :type git_remote_name: str

        :param git_shallow: clone with ``--depth 1`` (default False)
        :type git_shallow: bool

        :param git_submodules: Git submodules that shall be updated, all if
            empty
        :type git_submodules: list

        :param tls_verify: Should certificate for https be checked (default
            False)
        :type tls_verify: bool
        """
        if 'git_remote_name' not in kwargs:
            self.git_remote_name = "origin"
        if 'git_shallow' not in kwargs:
            self.git_shallow = False
        if 'git_submodules' not in kwargs:
            self.git_submodules = []
        if 'tls_verify' not in kwargs:
            self.tls_verify = False
        BaseRepo.__init__(self, url, **kwargs)

        self.remotes = remotes

    def get_revision(self):
        """Return current revision. Initial repositories return 'initial'."""
        try:
            return self.run(['rev-parse', '--verify', 'HEAD'])
        except exc.SubprocessError:
            return 'initial'

    @classmethod
    def get_url_and_revision_from_pip_url(cls, pip_url):
        """
        Prefixes stub URLs like 'user@hostname:user/repo.git' with 'ssh://'.
        That's required because although they use SSH they sometimes doesn't
        work with a ssh:// scheme (e.g. Github). But we need a scheme for
        parsing. Hence we remove it again afterwards and return it as a stub.
        """
        if '://' not in pip_url:
            assert 'file:' not in pip_url
            pip_url = pip_url.replace('git+', 'git+ssh://')
            url, rev = super(
                GitRepo, cls).get_url_and_revision_from_pip_url(pip_url)
            url = url.replace('ssh://', '')
        elif 'github.com:' in pip_url:
            raise exc.LibVCSException(
                "Repo %s is malformatted, please use the convention %s for"
                "ssh / private GitHub repositories." % (
                    pip_url, "git+https://github.com/username/repo.git"
                )
            )
        else:
            url, rev = super(
                GitRepo, cls).get_url_and_revision_from_pip_url(pip_url)

        return url, rev

    def obtain(self, quiet=False):
        """Retrieve the repository, clone if doesn't exist.

        :param quiet: Suppress stderr output.
        :type quiet: bool

        """
        self.check_destination()

        url = self.url

        cmd = ['clone', '--progress']
        if self.git_shallow:
            cmd.extend(['--depth', '1'])
        if self.tls_verify:
            cmd.extend(['-c', 'http.sslVerify=false'])
        cmd.extend([url, self.path])

        self.info('Cloning.')
        self.run_buffered(cmd)

        if self.remotes:
            for r in self.remotes:
                self.error('Adding remote %s <%s>' %
                           (r['remote_name'], r['url']))
                self.remote_set(
                    name=r['remote_name'],
                    url=r['url']
                )

        self.info('Initializing submodules.')
        self.run_buffered(['submodule', 'init'],)
        cmd = ['submodule', 'update', '--recursive', '--init']
        cmd.extend(self.git_submodules)
        self.run_buffered(cmd)

    def update_repo(self):
        self.check_destination()

        if not os.path.isdir(os.path.join(self.path, '.git')):
            self.obtain()
            self.update_repo()
            return

        # Get requested revision or tag
        url, git_tag = self.url, self.rev

        if not git_tag:
            self.debug("No git revision set, defaulting to origin/master")
            symref = self.run(['symbolic-ref', '--short', 'HEAD'])
            if symref:
                git_tag = symref.rstrip()
            else:
                git_tag = 'origin/master'
        self.debug("git_tag: %s" % git_tag)

        self.info("Updating to '%s'." % git_tag)

        # Get head sha
        try:
            head_sha = self.run(['rev-list', '--max-count=1', 'HEAD'])
        except exc.SubprocessError as e:
            self.error("Failed to get the hash for HEAD")
            return

        self.debug("head_sha: %s" % head_sha)

        # If a remote ref is asked for, which can possibly move around,
        # we must always do a fetch and checkout.
        try:
            show_ref_output = self.run(['show-ref', git_tag])
        except exc.SubprocessError as e:
            self.error("Ref %s not found" % git_tag)
            return
        self.debug("show_ref_output: %s" % show_ref_output)
        is_remote_ref = "remotes" in show_ref_output
        self.debug("is_remote_ref: %s" % is_remote_ref)

        # Tag is in the form <remote>/<tag> (i.e. origin/master) we must
        # strip the remote from the tag.
        git_remote_name = self.git_remote_name
        if "refs/remotes/%s" % git_tag in show_ref_output:
            m = re.match(r'^(?P<git_remote_name>[^/]+)/(?P<git_tag>.+)$',
                         show_ref_output)
            git_remote_name = m.group('git_remote_name')
            git_tag = m.group('git_tag')

        # This will fail if the tag does not exist (it probably has not
        # been fetched yet).
        try:
            error_code = 0
            tag_sha = self.run(['rev-list', '--max-count=1', git_tag])
        except exc.SubprocessError as e:
            error_code = e.subprocess.returncode
        self.debug("tag_sha: %s" % tag_sha)

        # Is the hash checkout out what we want?
        somethings_up = (error_code, is_remote_ref, tag_sha != head_sha,)
        if all(not x for x in somethings_up):
            self.info("Already up-to-date.")
            return

        process = self.run_buffered(['fetch'])
        if process.returncode:
            self.error("Failed to fetch repository '%s'" % url)
            return

        if is_remote_ref:
            # Check if stash is needed
            try:
                process = self.run(['status', '--porcelain'])
            except exc.SubprocessError as e:
                self.error("Failed to get the status")
                return
            need_stash = len(process) > 0

            # If not in clean state, stash changes in order to be able
            # to be able to perform git pull --rebase
            if need_stash:
                # If Git < 1.7.6, uses --quiet --all
                git_stash_save_options = '--quiet'
                try:
                    process = self.run([
                        'stash', 'save', git_stash_save_options
                    ])
                except exc.SubprocessError as e:
                    self.error("Failed to stash changes")

            # Pull changes from the remote branch
            try:
                process = self.run_buffered([
                    'rebase', git_remote_name + '/' + git_tag
                ], print_stdout_on_progress_end=False)
            except exc.SubprocessError as e:
                # Rebase failed: Restore previous state.
                self.run(['rebase', '--abort'])
                if need_stash:
                    self.run(['stash', 'pop', '--index', '--quiet'])

                self.error(
                    "\nFailed to rebase in: '%s'.\n"
                    "You will have to resolve the conflicts manually" %
                    self.path)
                return

            if need_stash:
                try:
                    process = self.run([
                        'stash', 'pop', '--index', '--quiet'
                    ])
                except exc.SubprocessError as e:
                    # Stash pop --index failed: Try again dropping the index
                    self.run(['reset', '--hard', '--quiet'])
                    try:
                        process = self.run(['stash', 'pop', '--quiet'])
                    except exc.SubprocessError as e:
                        # Stash pop failed: Restore previous state.
                        self.run(['reset', '--hard', '--quiet', head_sha])
                        self.run(['stash', 'pop', '--index', '--quiet'])
                        self.error("\nFailed to rebase in: '%s'.\n"
                                   "You will have to resolve the "
                                   "conflicts manually" % self.path)
                        return

        else:
            try:
                process = self.run(['checkout', git_tag])
            except exc.SubprocessError as e:
                self.error("Failed to checkout tag: '%s'" % git_tag)
                return

        cmd = ['submodule', 'update', '--recursive', '--init']
        cmd.extend(self.git_submodules)
        self.run(cmd)

    @property
    def remotes_get(self):
        """Return remotes like git remote -v.

        :rtype: dict of tuples
        """
        remotes = {}

        cmd = self.run(['remote'])
        ret = filter(None, cmd.split('\n'))

        for remote_name in ret:
            remotes[remote_name] = self.remote_get(remote_name)
        return remotes

    def remote_get(self, remote='origin'):
        """Get the fetch and push URL for a specified remote name.

        :param remote: the remote name used to define the fetch and push URL
        :type remote: str
        :returns: remote name and url in tuple form
        :rtype: tuple
        """
        try:
            ret = self.run(['remote', 'show', '-n', remote])
            lines = ret.split('\n')
            remote_fetch_url = lines[1].replace('Fetch URL: ', '').strip()
            remote_push_url = lines[2].replace('Push  URL: ', '').strip()
            if remote_fetch_url != remote and remote_push_url != remote:
                res = (remote_fetch_url, remote_push_url)
                return res
            else:
                return None
        except exc.LibVCSException:
            return None

    def remote_set(self, url, name='origin'):
        """Set remote with name and URL like git remote add.

        :param url: defines the remote URL
        :type url: string
        :param name: defines the remote name.
        :type name: str
        """

        url = self.chomp_protocol(url)

        if self.remote_get(name):
            self.run(['remote', 'rm', 'name'])

        self.run(['remote', 'add', name, url])
        return self.remote_get(remote=name)

    @staticmethod
    def chomp_protocol(url):
        """Return clean VCS url from RFC-style url

        :param url: url
        :type url: string
        :return type: string
        :returns: url as VCS software would accept it
        :seealso: #14
        """
        if '+' in url:
            url = url.split('+', 1)[1]
        scheme, netloc, path, query, frag = urlparse.urlsplit(url)
        rev = None
        if '@' in path:
            path, rev = path.rsplit('@', 1)
        url = urlparse.urlunsplit((scheme, netloc, path, query, ''))
        if url.startswith('ssh://git@github.com/'):
            url = url.replace('ssh://', 'git+ssh://')
        elif '://' not in url:
            assert 'file:' not in url
            url = url.replace('git+', 'git+ssh://')
            url = url.replace('ssh://', '')
        return url
<|endoftext|>


<commit_before>import unittest

import mock
import numpy
import six

import chainer
from chainer import cuda
from chainer import functions
from chainer.functions.normalization import batch_normalization
from chainer import gradient_check
from chainer import testing
from chainer.testing import attr
from chainer.testing import condition

def _batch_normalization(expander, gamma, beta, x, mean, var):
    mean = mean[expander]
    std = numpy.sqrt(var)[expander]
    y_expect = (gamma[expander] * (x - mean) / std + beta[expander])
    return y_expect

@testing.parameterize(*(testing.product({
    'param_shape': [(3, 4), (3, 2, 3)],
    'ndim': [0, 1, 2],
    'dtype': [numpy.float32],
}) + testing.product({
    'param_shape': [(3,)],
    'ndim': [1],
    'dtype': [numpy.float16, numpy.float32, numpy.float64],
})))
class TestBatchNormalization(unittest.TestCase):

    def setUp(self):
        self.expander = (None, Ellipsis) + (None,) * self.ndim
        self.eps = 2e-5
        self.decay = 0.9

        self.gamma = numpy.random.uniform(.5, 1,
                                          self.param_shape).astype(self.dtype)
        self.beta = numpy.random.uniform(-1, 1,
                                         self.param_shape).astype(self.dtype)
        head_ndim = self.gamma.ndim + 1
        shape = (5,) + self.param_shape + (2,) * self.ndim
        self.x = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.gy = numpy.random.uniform(-1, 1, shape).astype(self.dtype)

        self.args = [self.x, self.gamma, self.beta]
        self.aggr_axes = (0,) + tuple(six.moves.range(head_ndim, self.x.ndim))
        self.mean = self.x.mean(axis=self.aggr_axes)
        self.var = self.x.var(axis=self.aggr_axes) + self.eps
        self.train = True
        self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
        self.check_backward_options = {'dtype': numpy.float64}
        if self.dtype == numpy.float16:
            self.check_forward_options = {'atol': 1e-2, 'rtol': 1e-2}
            self.check_backward_options = {
                'dtype': numpy.float64, 'atol': 1e-2, 'rtol': 1e-2}

    def check_forward(self, args, use_cudnn=True):
        y = functions.batch_normalization(
            *[chainer.Variable(i) for i in args], running_mean=None,
            running_var=None, decay=self.decay, eps=self.eps,
            use_cudnn=use_cudnn)
        self.assertEqual(y.data.dtype, self.dtype)

        y_expect = _batch_normalization(
            self.expander, self.gamma, self.beta, self.x, self.mean, self.var)

        testing.assert_allclose(
            y_expect, y.data, **self.check_forward_options)

    @condition.retry(3)
    def test_forward_cpu(self):
        self.check_forward(self.args)

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args])

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu_no_cudnn(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args], False)

    def check_backward(self, args, y_grad):
        gradient_check.check_backward(
            batch_normalization.BatchNormalizationFunction(
                mean=None, var=None, train=self.train,
                decay=self.decay, eps=self.eps), args, y_grad,
            **self.check_backward_options)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(self.args, self.gy)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.check_backward(
            [cuda.to_gpu(i) for i in self.args], cuda.to_gpu(self.gy))

@testing.parameterize(*(testing.product({
    'param_shape': [(3, 4), (3, 2, 3)],
    'ndim': [0, 1, 2],
    'dtype': [numpy.float32],
}) + testing.product({
    'param_shape': [(3,)],
    'ndim': [1],
    'dtype': [numpy.float16, numpy.float32, numpy.float64],
})))
class TestFixedBatchNormalization(unittest.TestCase):

    def setUp(self):
        self.gamma = numpy.random.uniform(.5, 1,
                                          self.param_shape).astype(self.dtype)
        self.beta = numpy.random.uniform(-1, 1,
                                         self.param_shape).astype(self.dtype)
        self.expander = (None, Ellipsis) + (None,) * self.ndim
        shape = (5,) + self.param_shape + (2,) * self.ndim
        self.x = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.gy = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.eps = 2e-5
        self.decay = 0.0
        head_ndim = self.gamma.ndim + 1
        self.aggr_axes = (0,) + tuple(six.moves.range(head_ndim, self.x.ndim))
        self.mean = numpy.random.uniform(-1, 1,
                                         self.param_shape).astype(self.dtype)
        self.var = numpy.random.uniform(
            0.5, 1, self.param_shape).astype(self.dtype)
        self.args = [self.x, self.gamma, self.beta, self.mean, self.var]
        self.train = False
        self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
        self.check_backward_options = {'dtype': numpy.float64}
        if self.dtype == numpy.float16:
            self.check_forward_options = {'atol': 1e-2, 'rtol': 1e-2}
            self.check_backward_options = {
                'dtype': numpy.float64, 'atol': 1e-2, 'rtol': 1e-2}

    def check_forward(self, args, use_cudnn=True):
        y = functions.fixed_batch_normalization(
            *[chainer.Variable(i) for i in args],
            eps=self.eps, use_cudnn=use_cudnn)
        self.assertEqual(y.data.dtype, self.dtype)

        y_expect = _batch_normalization(
            self.expander, self.gamma, self.beta, self.x, self.mean, self.var)

        testing.assert_allclose(
            y_expect, y.data, **self.check_forward_options)

    @condition.retry(3)
    def test_forward_cpu(self):
        self.check_forward(self.args)

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args])

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu_no_cudnn(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args], False)

    def check_backward(self, args, y_grad):
        gradient_check.check_backward(
            batch_normalization.BatchNormalizationFunction(
                mean=None, var=None, train=self.train,
                decay=self.decay, eps=self.eps),
            args, y_grad,  **self.check_backward_options)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(self.args, self.gy)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.check_backward(
            [cuda.to_gpu(i) for i in self.args], cuda.to_gpu(self.gy))

@testing.parameterize(*testing.product({
    'use_cudnn': [True, False],
    # TODO(bkvogel): Check float16 support again in next cuDNN version.
    'dtype': [numpy.float32, numpy.float64],
}))
@attr.cudnn
class TestBatchNormalizationCudnnCall(unittest.TestCase):

    def setUp(self):
        ndim = 0
        param_shape = (3,)
        self.gamma = cuda.cupy.random.uniform(.5, 1,
                                              param_shape).astype(self.dtype)
        self.beta = cuda.cupy.random.uniform(-1, 1,
                                             param_shape).astype(self.dtype)
        self.eps = 2e-5
        shape = (7,) + param_shape + (2,) * ndim
        print('za shape: ', shape)
        self.x = cuda.cupy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.gy = cuda.cupy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.args = [self.x, self.gamma, self.beta]
        head_ndim = self.gamma.ndim + 1
        print('head_ndim: ', head_ndim)
        self.aggr_axes = (0,) + tuple(six.moves.range(head_ndim, self.x.ndim))
        print('self.aggr_axes: ', self.aggr_axes)
        self.mean = self.x.mean(axis=self.aggr_axes)
        print('setUp, self.mean.shape: ', self.mean.shape)
        self.var = self.x.var(axis=self.aggr_axes) + self.eps
        self.expect = self.use_cudnn and (
            cuda.cudnn.cudnn.getVersion() >= 5000)
        print('self.expect: ', self.expect)

    def forward(self):
        return functions.batch_normalization(
            *[chainer.Variable(i) for i in self.args], eps=self.eps,
            running_mean=self.mean, running_var=self.var,
            use_cudnn=self.use_cudnn)

    def test_call_cudnn_forward(self):
        with mock.patch('cupy.cudnn.cudnn.batchNormalizationForwardTraining') \
                as func:
            self.forward()
            self.assertEqual(func.called, self.expect)

    def test_call_cudnn_backward(self):
        y = self.forward()
        y.grad = self.gy
        with mock.patch('cupy.cudnn.cudnn.batchNormalizationBackward') as func:
            y.backward()
            self.assertEqual(func.called, self.expect)

testing.run_module(__name__, __file__)
<commit_msg>add tests for non-c-contiguous inputs<commit_after>import unittest

import mock
import numpy
import six

import chainer
from chainer import cuda
from chainer import functions
from chainer.functions.normalization import batch_normalization
from chainer import gradient_check
from chainer import testing
from chainer.testing import attr
from chainer.testing import condition

def _batch_normalization(expander, gamma, beta, x, mean, var):
    mean = mean[expander]
    std = numpy.sqrt(var)[expander]
    y_expect = (gamma[expander] * (x - mean) / std + beta[expander])
    return y_expect

@testing.parameterize(*(testing.product({
    'param_shape': [(3,), (3, 4), (3, 2, 3)],
    'ndim': [0, 1, 2],
    'dtype': [numpy.float32],
}) + testing.product({
    'param_shape': [(3,)],
    'ndim': [1],
    'dtype': [numpy.float16, numpy.float32, numpy.float64],
})))
class TestBatchNormalization(unittest.TestCase):

    def setUp(self):
        self.expander = (None, Ellipsis) + (None,) * self.ndim
        self.eps = 2e-5
        self.decay = 0.9

        self.gamma = numpy.random.uniform(.5, 1,
                                          self.param_shape).astype(self.dtype)
        self.beta = numpy.random.uniform(-1, 1,
                                         self.param_shape).astype(self.dtype)
        head_ndim = self.gamma.ndim + 1
        shape = (5,) + self.param_shape + (2,) * self.ndim
        self.x = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.gy = numpy.random.uniform(-1, 1, shape).astype(self.dtype)

        self.args = [self.x, self.gamma, self.beta]
        self.aggr_axes = (0,) + tuple(six.moves.range(head_ndim, self.x.ndim))
        self.mean = self.x.mean(axis=self.aggr_axes)
        self.var = self.x.var(axis=self.aggr_axes) + self.eps
        self.train = True
        self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
        self.check_backward_options = {'dtype': numpy.float64}
        if self.dtype == numpy.float16:
            self.check_forward_options = {'atol': 1e-2, 'rtol': 1e-2}
            self.check_backward_options = {
                'dtype': numpy.float64, 'atol': 1e-2, 'rtol': 1e-2}

    def check_forward(self, args, use_cudnn=True):
        y = functions.batch_normalization(
            *[chainer.Variable(i) for i in args], running_mean=None,
            running_var=None, decay=self.decay, eps=self.eps,
            use_cudnn=use_cudnn)
        self.assertEqual(y.data.dtype, self.dtype)

        y_expect = _batch_normalization(
            self.expander, self.gamma, self.beta, self.x, self.mean, self.var)

        testing.assert_allclose(
            y_expect, y.data, **self.check_forward_options)

    @condition.retry(3)
    def test_forward_cpu(self):
        self.check_forward(self.args)

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args])

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu_no_cudnn(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args], False)

    @attr.cudnn
    @condition.retry(3)
    def test_forward_gpu_non_contiguous(self):
        self.check_forward([cuda.cupy.asfortranarray(cuda.to_gpu(i))
                            for i in self.args])

    def check_backward(self, args, y_grad, use_cudnn=True):
        gradient_check.check_backward(
            batch_normalization.BatchNormalizationFunction(
                mean=None, var=None, train=self.train,
                decay=self.decay, eps=self.eps, use_cudnn=use_cudnn),
            args, y_grad, **self.check_backward_options)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(self.args, self.gy)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.check_backward(
            [cuda.to_gpu(i) for i in self.args], cuda.to_gpu(self.gy))

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu_no_cudnn(self):
        self.check_backward(
            [cuda.to_gpu(i) for i in self.args], cuda.to_gpu(self.gy), False)

    @attr.cudnn
    @condition.retry(3)
    def test_backward_gpu_non_contiguous(self):
        self.check_backward(
            [cuda.cupy.asfortranarray(cuda.to_gpu(i)) for i in self.args],
            cuda.cupy.asfortranarray(cuda.to_gpu(self.gy)))

@testing.parameterize(*(testing.product({
    'param_shape': [(3, 4), (3, 2, 3)],
    'ndim': [0, 1, 2],
    'dtype': [numpy.float32],
}) + testing.product({
    'param_shape': [(3,)],
    'ndim': [1],
    'dtype': [numpy.float16, numpy.float32, numpy.float64],
})))
class TestFixedBatchNormalization(unittest.TestCase):

    def setUp(self):
        self.gamma = numpy.random.uniform(.5, 1,
                                          self.param_shape).astype(self.dtype)
        self.beta = numpy.random.uniform(-1, 1,
                                         self.param_shape).astype(self.dtype)
        self.expander = (None, Ellipsis) + (None,) * self.ndim
        shape = (5,) + self.param_shape + (2,) * self.ndim
        self.x = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.gy = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.eps = 2e-5
        self.decay = 0.0
        head_ndim = self.gamma.ndim + 1
        self.aggr_axes = (0,) + tuple(six.moves.range(head_ndim, self.x.ndim))
        self.mean = numpy.random.uniform(-1, 1,
                                         self.param_shape).astype(self.dtype)
        self.var = numpy.random.uniform(
            0.5, 1, self.param_shape).astype(self.dtype)
        self.args = [self.x, self.gamma, self.beta, self.mean, self.var]
        self.train = False
        self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
        self.check_backward_options = {'dtype': numpy.float64}
        if self.dtype == numpy.float16:
            self.check_forward_options = {'atol': 1e-2, 'rtol': 1e-2}
            self.check_backward_options = {
                'dtype': numpy.float64, 'atol': 1e-2, 'rtol': 1e-2}

    def check_forward(self, args, use_cudnn=True):
        y = functions.fixed_batch_normalization(
            *[chainer.Variable(i) for i in args],
            eps=self.eps, use_cudnn=use_cudnn)
        self.assertEqual(y.data.dtype, self.dtype)

        y_expect = _batch_normalization(
            self.expander, self.gamma, self.beta, self.x, self.mean, self.var)

        testing.assert_allclose(
            y_expect, y.data, **self.check_forward_options)

    @condition.retry(3)
    def test_forward_cpu(self):
        self.check_forward(self.args)

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args])

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu_no_cudnn(self):
        self.check_forward([cuda.to_gpu(i) for i in self.args], False)

    @attr.cudnn
    @condition.retry(3)
    def test_forward_gpu_non_contiguous(self):
        self.check_forward([cuda.cupy.asfortranarray(cuda.to_gpu(i))
                            for i in self.args])

    def check_backward(self, args, y_grad, use_cudnn=True):
        gradient_check.check_backward(
            batch_normalization.BatchNormalizationFunction(
                mean=None, var=None, train=self.train,
                decay=self.decay, eps=self.eps, use_cudnn=use_cudnn),
            args, y_grad,  **self.check_backward_options)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(self.args, self.gy)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.check_backward(
            [cuda.to_gpu(i) for i in self.args], cuda.to_gpu(self.gy))

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu_no_cudnn(self):
        self.check_backward(
            [cuda.to_gpu(i) for i in self.args], cuda.to_gpu(self.gy), False)

    @attr.cudnn
    @condition.retry(3)
    def test_backward_gpu_no_contiguous(self):
        self.check_backward(
            [cuda.cupy.asfortranarray(cuda.to_gpu(i)) for i in self.args],
            cuda.cupy.asfortranarray(cuda.to_gpu(self.gy)))

@testing.parameterize(*testing.product({
    'use_cudnn': [True, False],
    # TODO(bkvogel): Check float16 support again in next cuDNN version.
    'dtype': [numpy.float32, numpy.float64],
}))
@attr.cudnn
class TestBatchNormalizationCudnnCall(unittest.TestCase):

    def setUp(self):
        ndim = 0
        param_shape = (3,)
        self.gamma = cuda.cupy.random.uniform(.5, 1,
                                              param_shape).astype(self.dtype)
        self.beta = cuda.cupy.random.uniform(-1, 1,
                                             param_shape).astype(self.dtype)
        self.eps = 2e-5
        shape = (7,) + param_shape + (2,) * ndim
        print('za shape: ', shape)
        self.x = cuda.cupy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.gy = cuda.cupy.random.uniform(-1, 1, shape).astype(self.dtype)
        self.args = [self.x, self.gamma, self.beta]
        head_ndim = self.gamma.ndim + 1
        print('head_ndim: ', head_ndim)
        self.aggr_axes = (0,) + tuple(six.moves.range(head_ndim, self.x.ndim))
        print('self.aggr_axes: ', self.aggr_axes)
        self.mean = self.x.mean(axis=self.aggr_axes)
        print('setUp, self.mean.shape: ', self.mean.shape)
        self.var = self.x.var(axis=self.aggr_axes) + self.eps
        self.expect = self.use_cudnn and (
            cuda.cudnn.cudnn.getVersion() >= 5000)
        print('self.expect: ', self.expect)

    def forward(self):
        return functions.batch_normalization(
            *[chainer.Variable(i) for i in self.args], eps=self.eps,
            running_mean=self.mean, running_var=self.var,
            use_cudnn=self.use_cudnn)

    def test_call_cudnn_forward(self):
        with mock.patch('cupy.cudnn.cudnn.batchNormalizationForwardTraining') \
                as func:
            self.forward()
            self.assertEqual(func.called, self.expect)

    def test_call_cudnn_backward(self):
        y = self.forward()
        y.grad = self.gy
        with mock.patch('cupy.cudnn.cudnn.batchNormalizationBackward') as func:
            y.backward()
            self.assertEqual(func.called, self.expect)

testing.run_module(__name__, __file__)
<|endoftext|>


<commit_before>"""A basic Shock(https://github.com/MG-RAST/Shock) python access class.
Uses shock-client for high performance uploads and download if it is in
the users path.

Authors:

* Jared Wilkening
* Travis Harrison
"""

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

import cStringIO
import json
import os
import requests
import subprocess

#-----------------------------------------------------------------------------
# Classes
#-----------------------------------------------------------------------------

class Client:

    shock_url = ''
    transport_method = ''
    auth_header = {}
    token = ''

    def __init__(self, shock_url, token=''):
        self.shock_url = shock_url
        if token != '':
            self.set_auth(token)
        if self._cmd_exists('shock-client'):
            self.transport_method = 'shock-client'
        else:
            self.transport_method = 'requests'

    def set_auth(self, token):
        self.auth_header = {'Authorization': 'OAuth %s'%token}
        if self.transport_method == 'shock-client':
            self._set_shockclient_auth(token)

    def _set_shockclient_auth(self, token):
        proc = subprocess.Popen("shock-client auth set-token \'{\"access_token\": \"%s\"}\'"%(token), shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        return_code = proc.wait()
        if return_code > 0:
            err = ""
            for line in proc.stderr:
                err += line
            raise Exception(u'Error setting auth token in shock-client: %s'%err)

    def get_node(self, node):
        url = self.shock_url+'/node/'+node
        try:
            rget = requests.get(url, headers=self.auth_header, allow_redirects=True)
        except Exception as e:
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, e))
        if not (rget.ok and rget.text):
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, rget.raise_for_status()))
        rj = rget.json()
        if not (rj and isinstance(rj, dict) and all([key in rj for key in ['status','data','error']])):
            raise Exception(u'Return data not valid Shock format')
        if rj['error']:
            raise Exception('Shock error: %d: %s'%(rj['status'], rj['error'][0]))
        return rj['data']

    def download_to_path(self, node, path):
        if node == '' or path == '':
            raise Exception(u'download_to_path requires non-empty node & path parameters')
        if self.transport_method == 'shock-client':
            return self._download_shockclient(node, path)
        url = '%s/node/%s?download'%(self.shock_url, node)
        try:
            rget = requests.get(url, headers=self.auth_header, allow_redirects=True, stream = True)
        except Exception as e:
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, e))
        if not (rget.ok):
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, rget.raise_for_status()))
        with open(path, 'wb') as f:
            for chunk in rget.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    f.flush()
        return path

    def _download_shockclient(self, node, path):
        proc = subprocess.Popen("shock-client pdownload -threads=4 %s %s"%(node,path), shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        return_code = proc.wait()
        if return_code > 0:
            err = ""
            for line in proc.stderr:
                err += line
            raise Exception(u'Error downloading via shock-client: %s => %s: error: %s' %(node, path, err))
        else:
            return path

    def create_node(self, data='', attr=''):
        return self.upload("", data, attr)

    def upload(self, node='', data='', attr=''):
            if self.transport_method == 'shock-client' and node == '' and os.path.exists(data):
                res = self._upload_shockclient(data)
                if attr == '':
                    return res
                else:
                    node = res['id']
                    data = ''
            method = 'post'
            (data_hdl, attr_hdl) = ('', '')
            files = {}
            url = self.shock_url+'/node'
            if node != '':
                url = '%s/%s'%(url, node)
                method = 'put'
            if data != '':
                files['upload'] = self._get_handle(data)
            if attr != '':
                files['attributes'] = self._get_handle(attr)
            try:
                req = ""
                if method == 'put':
                    req = requests.put(url, headers=self.auth_header, files=files, allow_redirects=True, stream = True)
                else:
                    req = requests.post(url, headers=self.auth_header, files=files, allow_redirects=True, stream = True)
                rj = req.json()
            except Exception as e:
                raise Exception(u'Unable to connect to Shock server %s: %s' %(url, e))
            if not (req.ok):
                raise Exception(u'Unable to connect to Shock server %s: %s' %(url, req.raise_for_status()))
            if rj['error']:
                raise Exception(u'Shock error %s : %s'%(rj['status'], rj['error'][0]))
            else:
                return rj['data']

    def _upload_shockclient(self, path):
        proc = subprocess.Popen("shock-client pcreate -threads=4 -full %s"%(path), shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        return_code = proc.wait()
        if return_code > 0:
            err = ""
            for line in proc.stderr:
                err += line
            raise Exception(u'Error uploading via shock-client: %s: error: %s' %(path, err))
        else:
            res = ""
            for line in proc.stdout:
                if 'Uploading' not in line:
                    res += line
            return json.loads(res)

    def _get_handle(self, d):
        if os.path.exists(d):
            return (os.path.basename(d), open(d))
        else:
            return ("n/a", cStringIO.StringIO(d))

    def _cmd_exists(self, cmd):
        return subprocess.call("type " + cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0

    <commit_msg>python lib: added node query function<commit_after>"""A basic Shock(https://github.com/MG-RAST/Shock) python access class.
Uses shock-client for high performance uploads and download if it is in
the users path.

Authors:

* Jared Wilkening
* Travis Harrison
"""

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

import cStringIO
import json
import os
import requests
import subprocess
import urllib

#-----------------------------------------------------------------------------
# Classes
#-----------------------------------------------------------------------------

class Client:

    shock_url = ''
    transport_method = ''
    auth_header = {}
    token = ''

    def __init__(self, shock_url, token=''):
        self.shock_url = shock_url
        if token != '':
            self.set_auth(token)
        if self._cmd_exists('shock-client'):
            self.transport_method = 'shock-client'
        else:
            self.transport_method = 'requests'

    def set_auth(self, token):
        self.auth_header = {'Authorization': 'OAuth %s'%token}
        if self.transport_method == 'shock-client':
            self._set_shockclient_auth(token)

    def _set_shockclient_auth(self, token):
        proc = subprocess.Popen("shock-client auth set-token \'{\"access_token\": \"%s\"}\'"%(token), shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        return_code = proc.wait()
        if return_code > 0:
            err = ""
            for line in proc.stderr:
                err += line
            raise Exception(u'Error setting auth token in shock-client: %s'%err)

    def get_node(self, node):
        return self._get_node_data('/'+node)

    def query_node(self, query):
        query_string = '?query&'+urllib.urlencode(query)
        return self._get_node_data(query_string)

    def _get_node_data(self, path):
        url = self.shock_url+'/node'+path
        try:
            rget = requests.get(url, headers=self.auth_header, allow_redirects=True)
        except Exception as e:
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, e))
        if not (rget.ok and rget.text):
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, rget.raise_for_status()))
        rj = rget.json()
        if not (rj and isinstance(rj, dict) and all([key in rj for key in ['status','data','error']])):
            raise Exception(u'Return data not valid Shock format')
        if rj['error']:
            raise Exception('Shock error: %d: %s'%(rj['status'], rj['error'][0]))
        return rj['data']

    def download_to_path(self, node, path):
        if node == '' or path == '':
            raise Exception(u'download_to_path requires non-empty node & path parameters')
        if self.transport_method == 'shock-client':
            return self._download_shockclient(node, path)
        url = '%s/node/%s?download'%(self.shock_url, node)
        try:
            rget = requests.get(url, headers=self.auth_header, allow_redirects=True, stream = True)
        except Exception as e:
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, e))
        if not (rget.ok):
            raise Exception(u'Unable to connect to Shock server %s: %s' %(url, rget.raise_for_status()))
        with open(path, 'wb') as f:
            for chunk in rget.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    f.flush()
        return path

    def _download_shockclient(self, node, path):
        proc = subprocess.Popen("shock-client pdownload -threads=4 %s %s"%(node,path), shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        return_code = proc.wait()
        if return_code > 0:
            err = ""
            for line in proc.stderr:
                err += line
            raise Exception(u'Error downloading via shock-client: %s => %s: error: %s' %(node, path, err))
        else:
            return path

    def create_node(self, data='', attr=''):
        return self.upload("", data, attr)

    def upload(self, node='', data='', attr=''):
            if self.transport_method == 'shock-client' and node == '' and os.path.exists(data):
                res = self._upload_shockclient(data)
                if attr == '':
                    return res
                else:
                    node = res['id']
                    data = ''
            method = 'post'
            (data_hdl, attr_hdl) = ('', '')
            files = {}
            url = self.shock_url+'/node'
            if node != '':
                url = '%s/%s'%(url, node)
                method = 'put'
            if data != '':
                files['upload'] = self._get_handle(data)
            if attr != '':
                files['attributes'] = self._get_handle(attr)
            try:
                req = ""
                if method == 'put':
                    req = requests.put(url, headers=self.auth_header, files=files, allow_redirects=True, stream = True)
                else:
                    req = requests.post(url, headers=self.auth_header, files=files, allow_redirects=True, stream = True)
                rj = req.json()
            except Exception as e:
                raise Exception(u'Unable to connect to Shock server %s: %s' %(url, e))
            if not (req.ok):
                raise Exception(u'Unable to connect to Shock server %s: %s' %(url, req.raise_for_status()))
            if rj['error']:
                raise Exception(u'Shock error %s : %s'%(rj['status'], rj['error'][0]))
            else:
                return rj['data']

    def _upload_shockclient(self, path):
        proc = subprocess.Popen("shock-client pcreate -threads=4 -full %s"%(path), shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        return_code = proc.wait()
        if return_code > 0:
            err = ""
            for line in proc.stderr:
                err += line
            raise Exception(u'Error uploading via shock-client: %s: error: %s' %(path, err))
        else:
            res = ""
            for line in proc.stdout:
                if 'Uploading' not in line:
                    res += line
            return json.loads(res)

    def _get_handle(self, d):
        if os.path.exists(d):
            return (os.path.basename(d), open(d))
        else:
            return ("n/a", cStringIO.StringIO(d))

    def _cmd_exists(self, cmd):
        return subprocess.call("type " + cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0

    <|endoftext|>


<commit_before>"""Git Repo object for libvcs.

.. todo::

    From https://github.com/saltstack/salt (Apache License):

    - [`GitRepo.remote`](libvcs.git.GitRepo.remote) (renamed to ``remote``)
    - [`GitRepo.remote`](libvcs.git.GitRepo.remote_set) (renamed to ``set_remote``)

    From pip (MIT Licnese):

    - [`GitRepo.remote`](libvcs.git.GitRepo.remote_set) (renamed to ``set_remote``)
    - [`GitRepo.convert_pip_url`](libvcs.git.GitRepo.convert_pip_url`) (``get_url_rev``)
    - [`GitRepo.get_revision`](libvcs.git.GitRepo.get_revision)
    - [`GitRepo.get_git_version`](libvcs.git.GitRepo.get_git_version)
"""  # NOQA: E501
import logging
import os
import re
from typing import Dict, NamedTuple, Optional
from urllib import parse as urlparse

from . import exc
from .base import BaseRepo, VCSLocation, convert_pip_url as base_convert_pip_url

logger = logging.getLogger(__name__)

class GitRemote(NamedTuple):
    """Structure containing git repo information.

    Supports `collections.namedtuple._asdict()`
    """

    name: str
    fetch_url: str
    push_url: str

class GitStatus(NamedTuple):
    branch_oid: Optional[str]
    branch_head: Optional[str]
    branch_upstream: Optional[str]
    branch_ab: Optional[str]
    branch_ahead: Optional[str]
    branch_behind: Optional[str]

def extract_status(value) -> GitStatus:
    """Returns ``git status -sb --porcelain=2`` extracted to a dict

    Returns
    -------
    Dictionary of git repo's status
    """
    pattern = re.compile(
        r"""[\n\r]?
        (
            #
            \W+
            branch.oid\W+
            (?P<branch_oid>
                [a-f0-9]{40}
            )
        )?
        (
            #
            \W+
            branch.head
            [\W]+
            (?P<branch_head>
                .*
            )

        )?
        (
            #
            \W+
            branch.upstream
            [\W]+
            (?P<branch_upstream>
                .*
            )
        )?
        (
            #
            \W+
            branch.ab
            [\W]+
            (?P<branch_ab>
                \+(?P<branch_ahead>\d+)
                \W{1}
                \-(?P<branch_behind>\d+)
            )
        )?
        """,
        re.VERBOSE | re.MULTILINE,
    )
    matches = pattern.search(value)
    return GitStatus(**matches.groupdict())

def convert_pip_url(pip_url: str) -> VCSLocation:
    """
    Prefixes stub URLs like 'user@hostname:user/repo.git' with 'ssh://'.
    That's required because although they use SSH they sometimes doesn't
    work with a ssh:// scheme (e.g. Github). But we need a scheme for
    parsing. Hence we remove it again afterwards and return it as a stub.
    The manpage for git-clone(1) refers to this as the "scp-like styntax".
    """
    if "://" not in pip_url:
        assert "file:" not in pip_url
        pip_url = pip_url.replace("git+", "git+ssh://")
        url, rev = base_convert_pip_url(pip_url)
        url = url.replace("ssh://", "")
    elif "github.com:" in pip_url:
        raise exc.LibVCSException(
            "Repo %s is malformatted, please use the convention %s for"
            "ssh / private GitHub repositories."
            % (pip_url, "git+https://github.com/username/repo.git")
        )
    else:
        url, rev = base_convert_pip_url(pip_url)

    return VCSLocation(url=url, rev=rev)

class GitRepo(BaseRepo):
    bin_name = "git"
    schemes = ("git", "git+http", "git+https", "git+ssh", "git+git", "git+file")

    def __init__(self, url, repo_dir, **kwargs):
        """A git repository.

        Parameters
        ----------
        url : str
            URL of repo

        tls_verify : bool
            Should certificate for https be checked (default False)
        """
        if "git_shallow" not in kwargs:
            self.git_shallow = False
        if "tls_verify" not in kwargs:
            self.tls_verify = False

        BaseRepo.__init__(self, url, repo_dir, **kwargs)

    @classmethod
    def from_pip_url(cls, pip_url, *args, **kwargs):
        url, rev = convert_pip_url(pip_url)
        self = cls(url=url, rev=rev, *args, **kwargs)

        return self

    def get_revision(self):
        """Return current revision. Initial repositories return 'initial'."""
        try:
            return self.run(["rev-parse", "--verify", "HEAD"])
        except exc.CommandError:
            return "initial"

    def obtain(self):
        """Retrieve the repository, clone if doesn't exist."""
        self.ensure_dir()

        url = self.url

        cmd = ["clone", "--progress"]
        if self.git_shallow:
            cmd.extend(["--depth", "1"])
        if self.tls_verify:
            cmd.extend(["-c", "http.sslVerify=false"])
        cmd.extend([url, self.path])

        self.info("Cloning.")
        self.run(cmd, log_in_real_time=True)

        self.info("Initializing submodules.")
        self.run(["submodule", "init"], log_in_real_time=True)
        cmd = ["submodule", "update", "--recursive", "--init"]
        self.run(cmd, log_in_real_time=True)

    def update_repo(self):
        self.ensure_dir()

        if not os.path.isdir(os.path.join(self.path, ".git")):
            self.obtain()
            self.update_repo()
            return

        # Get requested revision or tag
        url, git_tag = self.url, getattr(self, "rev", None)

        if not git_tag:
            self.debug("No git revision set, defaulting to origin/master")
            symref = self.run(["symbolic-ref", "--short", "HEAD"])
            if symref:
                git_tag = symref.rstrip()
            else:
                git_tag = "origin/master"
        self.debug("git_tag: %s" % git_tag)

        self.info("Updating to '%s'." % git_tag)

        # Get head sha
        try:
            head_sha = self.run(["rev-list", "--max-count=1", "HEAD"])
        except exc.CommandError:
            self.error("Failed to get the hash for HEAD")
            return

        self.debug("head_sha: %s" % head_sha)

        # If a remote ref is asked for, which can possibly move around,
        # we must always do a fetch and checkout.
        show_ref_output = self.run(["show-ref", git_tag], check_returncode=False)
        self.debug("show_ref_output: %s" % show_ref_output)
        is_remote_ref = "remotes" in show_ref_output
        self.debug("is_remote_ref: %s" % is_remote_ref)

        # show-ref output is in the form "<sha> refs/remotes/<remote>/<tag>"
        # we must strip the remote from the tag.
        git_remote_name = self.get_current_remote_name()

        if "refs/remotes/%s" % git_tag in show_ref_output:
            m = re.match(
                r"^[0-9a-f]{40} refs/remotes/"
                r"(?P<git_remote_name>[^/]+)/"
                r"(?P<git_tag>.+)$",
                show_ref_output,
                re.MULTILINE,
            )
            git_remote_name = m.group("git_remote_name")
            git_tag = m.group("git_tag")
        self.debug("git_remote_name: %s" % git_remote_name)
        self.debug("git_tag: %s" % git_tag)

        # This will fail if the tag does not exist (it probably has not
        # been fetched yet).
        try:
            error_code = 0
            tag_sha = self.run(
                [
                    "rev-list",
                    "--max-count=1",
                    git_remote_name + "/" + git_tag if is_remote_ref else git_tag,
                ]
            )
        except exc.CommandError as e:
            error_code = e.returncode
            tag_sha = ""
        self.debug("tag_sha: %s" % tag_sha)

        # Is the hash checkout out what we want?
        somethings_up = (error_code, is_remote_ref, tag_sha != head_sha)
        if all(not x for x in somethings_up):
            self.info("Already up-to-date.")
            return

        try:
            process = self.run(["fetch"], log_in_real_time=True)
        except exc.CommandError:
            self.error("Failed to fetch repository '%s'" % url)
            return

        if is_remote_ref:
            # Check if stash is needed
            try:
                process = self.run(["status", "--porcelain"])
            except exc.CommandError:
                self.error("Failed to get the status")
                return
            need_stash = len(process) > 0

            # If not in clean state, stash changes in order to be able
            # to be able to perform git pull --rebase
            if need_stash:
                # If Git < 1.7.6, uses --quiet --all
                git_stash_save_options = "--quiet"
                try:
                    process = self.run(["stash", "save", git_stash_save_options])
                except exc.CommandError:
                    self.error("Failed to stash changes")

            # Checkout the remote branch
            try:
                process = self.run(["checkout", git_tag])
            except exc.CommandError:
                self.error("Failed to checkout tag: '%s'" % git_tag)
                return

            # Rebase changes from the remote branch
            try:
                process = self.run(["rebase", git_remote_name + "/" + git_tag])
            except exc.CommandError as e:
                if "invalid_upstream" in str(e):
                    self.error(e)
                else:
                    # Rebase failed: Restore previous state.
                    self.run(["rebase", "--abort"])
                    if need_stash:
                        self.run(["stash", "pop", "--index", "--quiet"])

                    self.error(
                        "\nFailed to rebase in: '%s'.\n"
                        "You will have to resolve the conflicts manually" % self.path
                    )
                    return

            if need_stash:
                try:
                    process = self.run(["stash", "pop", "--index", "--quiet"])
                except exc.CommandError:
                    # Stash pop --index failed: Try again dropping the index
                    self.run(["reset", "--hard", "--quiet"])
                    try:
                        process = self.run(["stash", "pop", "--quiet"])
                    except exc.CommandError:
                        # Stash pop failed: Restore previous state.
                        self.run(["reset", "--hard", "--quiet", head_sha])
                        self.run(["stash", "pop", "--index", "--quiet"])
                        self.error(
                            "\nFailed to rebase in: '%s'.\n"
                            "You will have to resolve the "
                            "conflicts manually" % self.path
                        )
                        return

        else:
            try:
                process = self.run(["checkout", git_tag])
            except exc.CommandError:
                self.error("Failed to checkout tag: '%s'" % git_tag)
                return

        cmd = ["submodule", "update", "--recursive", "--init"]
        self.run(cmd, log_in_real_time=True)

    def remotes(self, flat=False) -> Dict:
        """Return remotes like git remote -v.

        Parameters
        ----------
        flat : bool
            Return a dict of ``tuple`` instead of ``dict``, default `False`.

        Returns
        -------
        dict of git upstream / remote URLs
        """
        remotes = {}

        cmd = self.run(["remote"])
        ret = filter(None, cmd.split("\n"))

        for remote_name in ret:
            remotes[remote_name] = (
                self.remote(remote_name) if flat else self.remote(remote_name)._asdict()
            )
        return remotes

    def remote(self, name, **kwargs) -> GitRemote:
        """Get the fetch and push URL for a specified remote name.

        Parameters
        ----------
        name : str
            The remote name used to define the fetch and push URL

        Returns
        -------
        Remote name and url in tuple form
        """

        try:
            ret = self.run(["remote", "show", "-n", name])
            lines = ret.split("\n")
            remote_fetch_url = lines[1].replace("Fetch URL: ", "").strip()
            remote_push_url = lines[2].replace("Push  URL: ", "").strip()
            if remote_fetch_url != name and remote_push_url != name:
                return GitRemote(
                    name=name, fetch_url=remote_fetch_url, push_url=remote_push_url
                )
            else:
                return None
        except exc.LibVCSException:
            return None

    def set_remote(self, name, url, overwrite=False):
        """Set remote with name and URL like git remote add.

        Parameters
        ----------
        name : str
            defines the remote name.

        url : str
            defines the remote URL
        """

        url = self.chomp_protocol(url)

        if self.remote(name) and overwrite:
            self.run(["remote", "set-url", name, url])
        else:
            self.run(["remote", "add", name, url])
        return self.remote(name=name)

    @staticmethod
    def chomp_protocol(url) -> str:
        """Return clean VCS url from RFC-style url

        Parameters
        ----------
        url : str
            PIP-style url

        Returns
        -------
        URL as VCS software would accept it
        """
        if "+" in url:
            url = url.split("+", 1)[1]
        scheme, netloc, path, query, frag = urlparse.urlsplit(url)
        rev = None
        if "@" in path:
            path, rev = path.rsplit("@", 1)
        url = urlparse.urlunsplit((scheme, netloc, path, query, ""))
        if url.startswith("ssh://git@github.com/"):
            url = url.replace("ssh://", "git+ssh://")
        elif "://" not in url:
            assert "file:" not in url
            url = url.replace("git+", "git+ssh://")
            url = url.replace("ssh://", "")
        return url

    def get_git_version(self) -> str:
        """Return current version of git binary

        Returns
        -------
        git version
        """
        VERSION_PFX = "git version "
        version = self.run(["version"])
        if version.startswith(VERSION_PFX):
            version = version[len(VERSION_PFX) :].split()[0]
        else:
            version = ""
        return ".".join(version.split(".")[:3])

    def status(self) -> dict:
        """Retrieve status of project in dict format.

        Wraps ``git status --sb --porcelain=2``. Does not include changed files, yet.

        Returns
        -------
        Status of current checked out repository

        Examples
        --------

        >>> git_repo.status()
        {
            "branch_oid": 'de6185fde0806e5c7754ca05676325a1ea4d6348',
            "branch_head": 'fix-current-remote-name',
            "branch_upstream": 'origin/fix-current-remote-name',
            "branch_ab": '+0 -0',
            "branch_ahead": '0',
            "branch_behind": '0'
        }
        """
        return extract_status(self.run(["status", "-sb", "--porcelain=2"]))

    def get_current_remote_name(self) -> str:
        """Retrieve name of the remote / upstream of currently checked out branch.

        Returns
        -------
        If upstream the same, returns ``branch_name``.
        If upstream mismatches, returns ``remote_name/branch_name``.
        """
        match = self.status()

        if match.branch_upstream is None:  # no upstream set
            return match.branch_head
        return match.branch_upstream.replace("/" + match.branch_head, "")
<commit_msg>feat(GitRepo): Add remotes<commit_after>"""Git Repo object for libvcs.

.. todo::

    From https://github.com/saltstack/salt (Apache License):

    - [`GitRepo.remote`](libvcs.git.GitRepo.remote) (renamed to ``remote``)
    - [`GitRepo.remote`](libvcs.git.GitRepo.remote_set) (renamed to ``set_remote``)

    From pip (MIT Licnese):

    - [`GitRepo.remote`](libvcs.git.GitRepo.remote_set) (renamed to ``set_remote``)
    - [`GitRepo.convert_pip_url`](libvcs.git.GitRepo.convert_pip_url`) (``get_url_rev``)
    - [`GitRepo.get_revision`](libvcs.git.GitRepo.get_revision)
    - [`GitRepo.get_git_version`](libvcs.git.GitRepo.get_git_version)
"""  # NOQA: E501
import logging
import os
import re
from typing import Dict, NamedTuple, Optional, TypedDict, Union
from urllib import parse as urlparse

from . import exc
from .base import BaseRepo, VCSLocation, convert_pip_url as base_convert_pip_url

logger = logging.getLogger(__name__)

class GitRemote(NamedTuple):
    """Structure containing git repo information.

    Supports `collections.namedtuple._asdict()`
    """

    name: str
    fetch_url: str
    push_url: str

class GitStatus(NamedTuple):
    branch_oid: Optional[str]
    branch_head: Optional[str]
    branch_upstream: Optional[str]
    branch_ab: Optional[str]
    branch_ahead: Optional[str]
    branch_behind: Optional[str]

def extract_status(value) -> GitStatus:
    """Returns ``git status -sb --porcelain=2`` extracted to a dict

    Returns
    -------
    Dictionary of git repo's status
    """
    pattern = re.compile(
        r"""[\n\r]?
        (
            #
            \W+
            branch.oid\W+
            (?P<branch_oid>
                [a-f0-9]{40}
            )
        )?
        (
            #
            \W+
            branch.head
            [\W]+
            (?P<branch_head>
                .*
            )

        )?
        (
            #
            \W+
            branch.upstream
            [\W]+
            (?P<branch_upstream>
                .*
            )
        )?
        (
            #
            \W+
            branch.ab
            [\W]+
            (?P<branch_ab>
                \+(?P<branch_ahead>\d+)
                \W{1}
                \-(?P<branch_behind>\d+)
            )
        )?
        """,
        re.VERBOSE | re.MULTILINE,
    )
    matches = pattern.search(value)
    return GitStatus(**matches.groupdict())

def convert_pip_url(pip_url: str) -> VCSLocation:
    """
    Prefixes stub URLs like 'user@hostname:user/repo.git' with 'ssh://'.
    That's required because although they use SSH they sometimes doesn't
    work with a ssh:// scheme (e.g. Github). But we need a scheme for
    parsing. Hence we remove it again afterwards and return it as a stub.
    The manpage for git-clone(1) refers to this as the "scp-like styntax".
    """
    if "://" not in pip_url:
        assert "file:" not in pip_url
        pip_url = pip_url.replace("git+", "git+ssh://")
        url, rev = base_convert_pip_url(pip_url)
        url = url.replace("ssh://", "")
    elif "github.com:" in pip_url:
        raise exc.LibVCSException(
            "Repo %s is malformatted, please use the convention %s for"
            "ssh / private GitHub repositories."
            % (pip_url, "git+https://github.com/username/repo.git")
        )
    else:
        url, rev = base_convert_pip_url(pip_url)

    return VCSLocation(url=url, rev=rev)

class RemoteDict(TypedDict):
    fetch: str
    push: str

FullRemoteDict = Dict[str, RemoteDict]
RemotesArgs = Union[None, FullRemoteDict, Dict[str, str]]

class GitRepo(BaseRepo):
    bin_name = "git"
    schemes = ("git", "git+http", "git+https", "git+ssh", "git+git", "git+file")

    def __init__(self, url, repo_dir, remotes: RemotesArgs = None, **kwargs):
        """A git repository.

        Parameters
        ----------
        url : str
            URL of repo

        tls_verify : bool
            Should certificate for https be checked (default False)
        """
        if "git_shallow" not in kwargs:
            self.git_shallow = False
        if "tls_verify" not in kwargs:
            self.tls_verify = False

        self._remotes: Union[FullRemoteDict, None]

        if remotes is None:
            self._remotes: FullRemoteDict = {"origin": url}
        elif isinstance(remotes, dict):
            self._remotes: FullRemoteDict = remotes
            for remote_name, url in remotes.items():
                if isinstance(str, dict):
                    remotes[remote_name] = {
                        "fetch": url,
                        "push": url,
                    }

        BaseRepo.__init__(self, url, repo_dir, **kwargs)

    @classmethod
    def from_pip_url(cls, pip_url, *args, **kwargs):
        url, rev = convert_pip_url(pip_url)
        self = cls(url=url, rev=rev, *args, **kwargs)

        return self

    def get_revision(self):
        """Return current revision. Initial repositories return 'initial'."""
        try:
            return self.run(["rev-parse", "--verify", "HEAD"])
        except exc.CommandError:
            return "initial"

    def set_remotes(self, overwrite: bool = False):
        remotes = self._remotes
        if isinstance(remotes, dict):
            for remote_name, url in remotes.items():
                existing_remote = self.remote(remote_name)
                if isinstance(url, dict) and "fetch" in url:
                    if not existing_remote or existing_remote.fetch_url != url:
                        self.set_remote(
                            name=remote_name, url=url["fetch"], overwrite=overwrite
                        )
                    if "push" in url:
                        if not existing_remote or existing_remote.push_url != url:
                            self.set_remote(
                                name=remote_name,
                                url=url["push"],
                                push=True,
                                overwrite=overwrite,
                            )
                else:
                    if not existing_remote or existing_remote.fetch_url != url:
                        self.set_remote(name=remote_name, url=url, overwrite=overwrite)

    def obtain(self):
        """Retrieve the repository, clone if doesn't exist."""
        self.ensure_dir()

        url = self.url

        cmd = ["clone", "--progress"]
        if self.git_shallow:
            cmd.extend(["--depth", "1"])
        if self.tls_verify:
            cmd.extend(["-c", "http.sslVerify=false"])
        cmd.extend([url, self.path])

        self.info("Cloning.")
        self.run(cmd, log_in_real_time=True)

        self.info("Initializing submodules.")
        self.run(["submodule", "init"], log_in_real_time=True)
        cmd = ["submodule", "update", "--recursive", "--init"]
        self.run(cmd, log_in_real_time=True)

        self.set_remotes()

    def update_repo(self):
        self.ensure_dir()

        if not os.path.isdir(os.path.join(self.path, ".git")):
            self.obtain()
            self.update_repo()
            return

        # Get requested revision or tag
        url, git_tag = self.url, getattr(self, "rev", None)

        if not git_tag:
            self.debug("No git revision set, defaulting to origin/master")
            symref = self.run(["symbolic-ref", "--short", "HEAD"])
            if symref:
                git_tag = symref.rstrip()
            else:
                git_tag = "origin/master"
        self.debug("git_tag: %s" % git_tag)

        self.info("Updating to '%s'." % git_tag)

        # Get head sha
        try:
            head_sha = self.run(["rev-list", "--max-count=1", "HEAD"])
        except exc.CommandError:
            self.error("Failed to get the hash for HEAD")
            return

        self.debug("head_sha: %s" % head_sha)

        # If a remote ref is asked for, which can possibly move around,
        # we must always do a fetch and checkout.
        show_ref_output = self.run(["show-ref", git_tag], check_returncode=False)
        self.debug("show_ref_output: %s" % show_ref_output)
        is_remote_ref = "remotes" in show_ref_output
        self.debug("is_remote_ref: %s" % is_remote_ref)

        # show-ref output is in the form "<sha> refs/remotes/<remote>/<tag>"
        # we must strip the remote from the tag.
        git_remote_name = self.get_current_remote_name()

        if "refs/remotes/%s" % git_tag in show_ref_output:
            m = re.match(
                r"^[0-9a-f]{40} refs/remotes/"
                r"(?P<git_remote_name>[^/]+)/"
                r"(?P<git_tag>.+)$",
                show_ref_output,
                re.MULTILINE,
            )
            git_remote_name = m.group("git_remote_name")
            git_tag = m.group("git_tag")
        self.debug("git_remote_name: %s" % git_remote_name)
        self.debug("git_tag: %s" % git_tag)

        # This will fail if the tag does not exist (it probably has not
        # been fetched yet).
        try:
            error_code = 0
            tag_sha = self.run(
                [
                    "rev-list",
                    "--max-count=1",
                    git_remote_name + "/" + git_tag if is_remote_ref else git_tag,
                ]
            )
        except exc.CommandError as e:
            error_code = e.returncode
            tag_sha = ""
        self.debug("tag_sha: %s" % tag_sha)

        # Is the hash checkout out what we want?
        somethings_up = (error_code, is_remote_ref, tag_sha != head_sha)
        if all(not x for x in somethings_up):
            self.info("Already up-to-date.")
            return

        try:
            process = self.run(["fetch"], log_in_real_time=True)
        except exc.CommandError:
            self.error("Failed to fetch repository '%s'" % url)
            return

        if is_remote_ref:
            # Check if stash is needed
            try:
                process = self.run(["status", "--porcelain"])
            except exc.CommandError:
                self.error("Failed to get the status")
                return
            need_stash = len(process) > 0

            # If not in clean state, stash changes in order to be able
            # to be able to perform git pull --rebase
            if need_stash:
                # If Git < 1.7.6, uses --quiet --all
                git_stash_save_options = "--quiet"
                try:
                    process = self.run(["stash", "save", git_stash_save_options])
                except exc.CommandError:
                    self.error("Failed to stash changes")

            # Checkout the remote branch
            try:
                process = self.run(["checkout", git_tag])
            except exc.CommandError:
                self.error("Failed to checkout tag: '%s'" % git_tag)
                return

            # Rebase changes from the remote branch
            try:
                process = self.run(["rebase", git_remote_name + "/" + git_tag])
            except exc.CommandError as e:
                if "invalid_upstream" in str(e):
                    self.error(e)
                else:
                    # Rebase failed: Restore previous state.
                    self.run(["rebase", "--abort"])
                    if need_stash:
                        self.run(["stash", "pop", "--index", "--quiet"])

                    self.error(
                        "\nFailed to rebase in: '%s'.\n"
                        "You will have to resolve the conflicts manually" % self.path
                    )
                    return

            if need_stash:
                try:
                    process = self.run(["stash", "pop", "--index", "--quiet"])
                except exc.CommandError:
                    # Stash pop --index failed: Try again dropping the index
                    self.run(["reset", "--hard", "--quiet"])
                    try:
                        process = self.run(["stash", "pop", "--quiet"])
                    except exc.CommandError:
                        # Stash pop failed: Restore previous state.
                        self.run(["reset", "--hard", "--quiet", head_sha])
                        self.run(["stash", "pop", "--index", "--quiet"])
                        self.error(
                            "\nFailed to rebase in: '%s'.\n"
                            "You will have to resolve the "
                            "conflicts manually" % self.path
                        )
                        return

        else:
            try:
                process = self.run(["checkout", git_tag])
            except exc.CommandError:
                self.error("Failed to checkout tag: '%s'" % git_tag)
                return

        cmd = ["submodule", "update", "--recursive", "--init"]
        self.run(cmd, log_in_real_time=True)

    def remotes(self, flat=False) -> Dict:
        """Return remotes like git remote -v.

        Parameters
        ----------
        flat : bool
            Return a dict of ``tuple`` instead of ``dict``, default `False`.

        Returns
        -------
        dict of git upstream / remote URLs
        """
        remotes = {}

        cmd = self.run(["remote"])
        ret = filter(None, cmd.split("\n"))

        for remote_name in ret:
            remotes[remote_name] = (
                self.remote(remote_name) if flat else self.remote(remote_name)._asdict()
            )
        return remotes

    def remote(self, name, **kwargs) -> GitRemote:
        """Get the fetch and push URL for a specified remote name.

        Parameters
        ----------
        name : str
            The remote name used to define the fetch and push URL

        Returns
        -------
        Remote name and url in tuple form
        """

        try:
            ret = self.run(["remote", "show", "-n", name])
            lines = ret.split("\n")
            remote_fetch_url = lines[1].replace("Fetch URL: ", "").strip()
            remote_push_url = lines[2].replace("Push  URL: ", "").strip()
            if remote_fetch_url != name and remote_push_url != name:
                return GitRemote(
                    name=name, fetch_url=remote_fetch_url, push_url=remote_push_url
                )
            else:
                return None
        except exc.LibVCSException:
            return None

    def set_remote(self, name, url, push: bool = False, overwrite=False):
        """Set remote with name and URL like git remote add.

        Parameters
        ----------
        name : str
            defines the remote name.

        url : str
            defines the remote URL
        """

        url = self.chomp_protocol(url)

        if self.remote(name) and overwrite:
            self.run(["remote", "set-url", name, url])
        else:
            self.run(["remote", "add", name, url])
        return self.remote(name=name)

    @staticmethod
    def chomp_protocol(url) -> str:
        """Return clean VCS url from RFC-style url

        Parameters
        ----------
        url : str
            PIP-style url

        Returns
        -------
        URL as VCS software would accept it
        """
        if "+" in url:
            url = url.split("+", 1)[1]
        scheme, netloc, path, query, frag = urlparse.urlsplit(url)
        rev = None
        if "@" in path:
            path, rev = path.rsplit("@", 1)
        url = urlparse.urlunsplit((scheme, netloc, path, query, ""))
        if url.startswith("ssh://git@github.com/"):
            url = url.replace("ssh://", "git+ssh://")
        elif "://" not in url:
            assert "file:" not in url
            url = url.replace("git+", "git+ssh://")
            url = url.replace("ssh://", "")
        return url

    def get_git_version(self) -> str:
        """Return current version of git binary

        Returns
        -------
        git version
        """
        VERSION_PFX = "git version "
        version = self.run(["version"])
        if version.startswith(VERSION_PFX):
            version = version[len(VERSION_PFX) :].split()[0]
        else:
            version = ""
        return ".".join(version.split(".")[:3])

    def status(self) -> dict:
        """Retrieve status of project in dict format.

        Wraps ``git status --sb --porcelain=2``. Does not include changed files, yet.

        Returns
        -------
        Status of current checked out repository

        Examples
        --------

        >>> git_repo.status()
        {
            "branch_oid": 'de6185fde0806e5c7754ca05676325a1ea4d6348',
            "branch_head": 'fix-current-remote-name',
            "branch_upstream": 'origin/fix-current-remote-name',
            "branch_ab": '+0 -0',
            "branch_ahead": '0',
            "branch_behind": '0'
        }
        """
        return extract_status(self.run(["status", "-sb", "--porcelain=2"]))

    def get_current_remote_name(self) -> str:
        """Retrieve name of the remote / upstream of currently checked out branch.

        Returns
        -------
        If upstream the same, returns ``branch_name``.
        If upstream mismatches, returns ``remote_name/branch_name``.
        """
        match = self.status()

        if match.branch_upstream is None:  # no upstream set
            return match.branch_head
        return match.branch_upstream.replace("/" + match.branch_head, "")
<|endoftext|>