<commit_before>from __future__ import print_function, division

import numpy as np
from numpy import dot

from sisl.messages import warn
from sisl.utils.ranges import array_arange
import sisl._array as _a
import sisl.linalg as lin

__all__ = ['SelfEnergy', 'SemiInfinite']
__all__ += ['RecursiveSI']

class SelfEnergy(object):
    """ Self-energy object able to calculate the dense self-energy for a given sparse matrix

    The self-energy object contains a `SparseGeometry` object which, in it-self
    contains the geometry.

    This is the base class for self-energies.
    """

    def __init__(self, *args, **kwargs):
        """ Self-energy class for constructing a self-energy. """
        pass

    def _setup(self, *args, **kwargs):
        """ Class specific setup routine """
        pass

    def self_energy(self, E, k=(0, 0, 0)):
        raise NotImplementedError

    def __getattr__(self, attr):
        """ Overload attributes from the hosting object """
        pass

class SemiInfinite(SelfEnergy):
    """ Self-energy object able to calculate the dense self-energy for a given `SparseGeometry` in a semi-infinite chain.

    Parameters
    ----------
    spgeom : SparseGeometry
       any sparse geometry matrix which may return matrices
    infinite : str
       axis specification for the semi-infinite direction (`+A`/`-A`/`+B`/`-B`/`+C`/`-C`)
    eta : float, optional
       the default imaginary part of the self-energy calculation
    """

    def __init__(self, spgeom, infinite, eta=1e-6):
        """ Create a `SelfEnergy` object from any `SparseGeometry` """
        self.eta = eta

        # Determine whether we are in plus/minus direction
        if infinite.startswith('+'):
            self.semi_inf_dir = 1
        elif infinite.startswith('-'):
            self.semi_inf_dir = -1
        else:
            raise ValueError(self.__class__.__name__ + ": infinite keyword does not start with `+` or `-`.")

        # Determine the direction
        INF = infinite.upper()
        if INF.endswith('A'):
            self.semi_inf = 0
        elif INF.endswith('B'):
            self.semi_inf = 1
        elif INF.endswith('C'):
            self.semi_inf = 2

        # Check that the Hamiltonian does have a non-zero V along the semi-infinite direction
        if spgeom.geometry.sc.nsc[self.semi_inf] == 1:
            warn('Creating a semi-infinite self-energy with no couplings along the semi-infinite direction')

        # Finalize the setup by calling the class specific routine
        self._setup(spgeom)

    def __str__(self):
        """ String representation of the SemiInfinite """
        direction = {-1: '-', 1: '+'}
        axis = {0: 'A', 1: 'B', 2: 'C'}
        return self.__class__.__name__ + '{{direction: {0}{1}\n}}'.format(
            direction[self.semi_inf_dir], axis[self.semi_inf])

    def _correct_k(self, k=None):
        """ Return a corrected k-point

        Notes
        -----
        This is strictly not required because any `k` along the semi-infinite direction
        is *integrated* out and thus the self-energy is the same for all k along the
        semi-infinite direction.
        """
        if k is None:
            k = _a.zerosd([3])
        else:
            k = self._fill(k, np.float64)
            k[self.semi_inf] = 0.
        return k

class RecursiveSI(SemiInfinite):
    """ Self-energy object using the Lopez-Sancho Lopez-Sancho algorithm """

    def __getattr__(self, attr):
        """ Overload attributes from the hosting object """
        return getattr(self.spgeom0, attr)

    def __str__(self):
        """ Representation of the RecursiveSI model """
        direction = {-1: '-', 1: '+'}
        axis = {0: 'A', 1: 'B', 2: 'C'}
        return self.__class__.__name__ + '{{direction: {0}{1},\n {2}\n}}'.format(
            direction[self.semi_inf_dir], axis[self.semi_inf],
            str(self.spgeom0).replace('\n', '\n '),
        )

    def _setup(self, spgeom):
        """ Setup the Lopez-Sancho internals for easy axes """

        # Create spgeom0 and spgeom1
        self.spgeom0 = spgeom.copy()
        nsc = np.copy(spgeom.geometry.sc.nsc)
        nsc[self.semi_inf] = 1
        self.spgeom0.set_nsc(nsc)

        # For spgeom1 we have to do it slightly differently
        old_nnz = spgeom.nnz
        self.spgeom1 = spgeom.copy()
        nsc[self.semi_inf] = 3

        # Already now limit the sparse matrices
        self.spgeom1.set_nsc(nsc)
        if self.spgeom1.nnz < old_nnz:
            warn("RecursiveSI: SparseGeometry has connections across the first neighbouring cell. "
                 "These values will be forced to 0 as the principal cell-interaction is a requirement")

        # I.e. we will delete all interactions that are un-important
        n_s = self.spgeom1.geometry.sc.n_s
        n = self.spgeom1.shape[0]
        # Figure out the matrix columns we should set to zero
        nsc = [None] * 3
        nsc[self.semi_inf] = self.semi_inf_dir
        # Get all supercell indices that we should delete
        idx = np.delete(_a.arangei(n_s),
                        _a.arrayi(spgeom.geometry.sc.sc_index(nsc)))

        cols = array_arange(idx * n, (idx + 1) * n)
        # Delete all values in columns, but keep them to retain the supercell information
        self.spgeom1._csr.delete_columns(cols, keep_shape=True)

    def self_energy(self, E, k=None, eta=None, dtype=None, eps=1e-14, bulk=False):
        r""" Return a dense matrix with the self-energy at energy `E` and k-point `k` (default Gamma).

        Parameters
        ----------
        E : float
          energy at which the calculation will take place (should *not* be complex)
        k : array_like, optional
          k-point at which the self-energy should be evaluated.
          the k-point should be in units of the reciprocal lattice vectors, and
          the semi-infinite component will be automatically set to zero.
        eta : float, optional
          the imaginary value to evaluate the self-energy with. Defaults to the
          value with which the object was created
        dtype : numpy.dtype
          the resulting data type
        eps : float, optional
          convergence criteria for the recursion
        bulk : bool, optional
          if true, :math:`E\cdot \mathbf S - \mathbf H -\boldsymbol\Sigma` is returned, else
          :math:`\boldsymbol\Sigma` is returned (default).

        Returns
        -------
        self-energy : the self-energy corresponding to the semi-infinite direction
        """
        if eta is None:
            eta = self.eta
        try:
            Z = E.real + 1j * eta
        except:
            Z = E + 1j * eta

        # Get k-point
        k = self._correct_k(k)

        if dtype is None:
            dtype = np.complex128

        sp0 = self.spgeom0
        sp1 = self.spgeom1

        # As the SparseGeometry inherently works for
        # orthogonal and non-orthogonal basis, there is no
        # need to have two algorithms.
        GB = (sp0.Sk(k, dtype=dtype) * Z - sp0.Pk(k, dtype=dtype)).toarray()

        if sp1.orthogonal:
            alpha = sp1.Pk(k, dtype=dtype, format='array')
            beta  = np.conjugate(np.transpose(alpha))
        else:
            M = sp1.Pk(k, dtype=dtype)
            S = sp1.Sk(k, dtype=dtype)
            alpha = (M - S * Z).toarray()
            beta  = (M.getH() - S.getH() * Z).toarray()
            del M, S

        # Surface Green function (self-energy)
        if bulk:
            GS = np.copy(GB)
        else:
            GS = np.zeros_like(GB)

        solve = lin.solve

        i = 0
        while True:
            i += 1

            tA = solve(GB, alpha)
            tB = solve(GB, beta)

            tmp = dot(alpha, tB)
            # Update bulk Green function
            GB -= tmp + dot(beta, tA)
            # Update surface self-energy
            GS -= tmp

            # Update forward/backward
            alpha = dot(alpha, tA)
            beta = dot(beta, tB)

            # Convergence criteria, it could be stricter
            if np.amax(np.abs(tmp)) < eps:
                # Return the pristine Green function
                del tA, tB, alpha, beta, GB
                if bulk:
                    return GS
                return - GS

        raise ValueError(self.__class__.__name__+': could not converge self-energy calculation')

    def self_energy_lr(self, E, k=None, eta=None, dtype=None, eps=1e-14, bulk=False):
        r""" Return two dense matrices with the left/right self-energy at energy `E` and k-point `k` (default Gamma).

        Note calculating the LR self-energies simultaneously requires that their chemical potentials are the same.
        I.e. only when the reference energy is equivalent in the left/right schemes does this make sense.

        Parameters
        ----------
        E : float
          energy at which the calculation will take place (should *not* be complex)
        k : array_like, optional
          k-point at which the self-energy should be evaluated.
          the k-point should be in units of the reciprocal lattice vectors, and
          the semi-infinite component will be automatically set to zero.
        eta : float, optional
          the imaginary value to evaluate the self-energy with. Defaults to the
          value with which the object was created
        dtype : numpy.dtype
          the resulting data type
        eps : float, optional
          convergence criteria for the recursion
        bulk : bool, optional
          if true, :math:`E\cdot \mathbf S - \mathbf H -\boldsymbol\Sigma` is returned, else
          :math:`\boldsymbol\Sigma` is returned (default).

        Returns
        -------
        left : the left self-energy
        right : the right self-energy
        """
        if eta is None:
            eta = self.eta
        try:
            Z = E.real + 1j * eta
        except:
            Z = E + 1j * eta

        # Get k-point
        k = self._correct_k(k)

        if dtype is None:
            dtype = np.complex128

        sp0 = self.spgeom0
        sp1 = self.spgeom1

        # As the SparseGeometry inherently works for
        # orthogonal and non-orthogonal basis, there is no
        # need to have two algorithms.
        SmH0 = sp0.Sk(k, dtype=dtype) * Z - sp0.Pk(k, dtype=dtype)
        GB = SmH0.toarray()

        if sp1.orthogonal:
            alpha = sp1.Pk(k, dtype=dtype, format='array')
            beta  = np.conjugate(np.transpose(alpha))
        else:
            M = sp1.Pk(k, dtype=dtype)
            S = sp1.Sk(k, dtype=dtype)
            alpha = (M - S * Z).toarray()
            beta  = (M.getH() - S.getH() * Z).toarray()
            del M, S

        # Surface Green function (self-energy)
        if bulk:
            GS = np.copy(GB)
        else:
            GS = np.zeros_like(GB)

        solve = lin.solve

        i = 0
        while True:
            i += 1

            tA = solve(GB, alpha)
            tB = solve(GB, beta)

            tmp = dot(alpha, tB)
            # Update bulk Green function
            GB -= tmp + dot(beta, tA)
            # Update surface self-energy
            GS -= tmp

            # Update forward/backward
            alpha = dot(alpha, tA)
            beta = dot(beta, tB)

            # Convergence criteria, it could be stricter
            if np.amax(np.abs(tmp)) < eps:
                # Return the pristine Green function
                del tA, tB, alpha, beta
                if self.semi_inf_dir == 1:
                    # GS is the "right" self-energy
                    if bulk:
                        return GB - GS + SmH0.toarray(), GS
                    return GS - GB + SmH0.toarray(), - GS
                # GS is the "left" self-energy
                if bulk:
                    return GS, GB - GS + SmH0.toarray()
                return - GS, GS - GB + SmH0.toarray()

        raise ValueError(self.__class__.__name__+': could not converge self-energy (LR) calculation')
<commit_msg>enh: removed a never reached statement<commit_after>from __future__ import print_function, division

import numpy as np
from numpy import dot

from sisl.messages import warn
from sisl.utils.ranges import array_arange
import sisl._array as _a
import sisl.linalg as lin

__all__ = ['SelfEnergy', 'SemiInfinite']
__all__ += ['RecursiveSI']

class SelfEnergy(object):
    """ Self-energy object able to calculate the dense self-energy for a given sparse matrix

    The self-energy object contains a `SparseGeometry` object which, in it-self
    contains the geometry.

    This is the base class for self-energies.
    """

    def __init__(self, *args, **kwargs):
        """ Self-energy class for constructing a self-energy. """
        pass

    def _setup(self, *args, **kwargs):
        """ Class specific setup routine """
        pass

    def self_energy(self, E, k=(0, 0, 0)):
        raise NotImplementedError

    def __getattr__(self, attr):
        """ Overload attributes from the hosting object """
        pass

class SemiInfinite(SelfEnergy):
    """ Self-energy object able to calculate the dense self-energy for a given `SparseGeometry` in a semi-infinite chain.

    Parameters
    ----------
    spgeom : SparseGeometry
       any sparse geometry matrix which may return matrices
    infinite : str
       axis specification for the semi-infinite direction (`+A`/`-A`/`+B`/`-B`/`+C`/`-C`)
    eta : float, optional
       the default imaginary part of the self-energy calculation
    """

    def __init__(self, spgeom, infinite, eta=1e-6):
        """ Create a `SelfEnergy` object from any `SparseGeometry` """
        self.eta = eta

        # Determine whether we are in plus/minus direction
        if infinite.startswith('+'):
            self.semi_inf_dir = 1
        elif infinite.startswith('-'):
            self.semi_inf_dir = -1
        else:
            raise ValueError(self.__class__.__name__ + ": infinite keyword does not start with `+` or `-`.")

        # Determine the direction
        INF = infinite.upper()
        if INF.endswith('A'):
            self.semi_inf = 0
        elif INF.endswith('B'):
            self.semi_inf = 1
        elif INF.endswith('C'):
            self.semi_inf = 2

        # Check that the Hamiltonian does have a non-zero V along the semi-infinite direction
        if spgeom.geometry.sc.nsc[self.semi_inf] == 1:
            warn('Creating a semi-infinite self-energy with no couplings along the semi-infinite direction')

        # Finalize the setup by calling the class specific routine
        self._setup(spgeom)

    def __str__(self):
        """ String representation of the SemiInfinite """
        direction = {-1: '-', 1: '+'}
        axis = {0: 'A', 1: 'B', 2: 'C'}
        return self.__class__.__name__ + '{{direction: {0}{1}\n}}'.format(
            direction[self.semi_inf_dir], axis[self.semi_inf])

    def _correct_k(self, k=None):
        """ Return a corrected k-point

        Notes
        -----
        This is strictly not required because any `k` along the semi-infinite direction
        is *integrated* out and thus the self-energy is the same for all k along the
        semi-infinite direction.
        """
        if k is None:
            k = _a.zerosd([3])
        else:
            k = self._fill(k, np.float64)
            k[self.semi_inf] = 0.
        return k

class RecursiveSI(SemiInfinite):
    """ Self-energy object using the Lopez-Sancho Lopez-Sancho algorithm """

    def __getattr__(self, attr):
        """ Overload attributes from the hosting object """
        return getattr(self.spgeom0, attr)

    def __str__(self):
        """ Representation of the RecursiveSI model """
        direction = {-1: '-', 1: '+'}
        axis = {0: 'A', 1: 'B', 2: 'C'}
        return self.__class__.__name__ + '{{direction: {0}{1},\n {2}\n}}'.format(
            direction[self.semi_inf_dir], axis[self.semi_inf],
            str(self.spgeom0).replace('\n', '\n '),
        )

    def _setup(self, spgeom):
        """ Setup the Lopez-Sancho internals for easy axes """

        # Create spgeom0 and spgeom1
        self.spgeom0 = spgeom.copy()
        nsc = np.copy(spgeom.geometry.sc.nsc)
        nsc[self.semi_inf] = 1
        self.spgeom0.set_nsc(nsc)

        # For spgeom1 we have to do it slightly differently
        old_nnz = spgeom.nnz
        self.spgeom1 = spgeom.copy()
        nsc[self.semi_inf] = 3

        # Already now limit the sparse matrices
        self.spgeom1.set_nsc(nsc)
        if self.spgeom1.nnz < old_nnz:
            warn("RecursiveSI: SparseGeometry has connections across the first neighbouring cell. "
                 "These values will be forced to 0 as the principal cell-interaction is a requirement")

        # I.e. we will delete all interactions that are un-important
        n_s = self.spgeom1.geometry.sc.n_s
        n = self.spgeom1.shape[0]
        # Figure out the matrix columns we should set to zero
        nsc = [None] * 3
        nsc[self.semi_inf] = self.semi_inf_dir
        # Get all supercell indices that we should delete
        idx = np.delete(_a.arangei(n_s),
                        _a.arrayi(spgeom.geometry.sc.sc_index(nsc)))

        cols = array_arange(idx * n, (idx + 1) * n)
        # Delete all values in columns, but keep them to retain the supercell information
        self.spgeom1._csr.delete_columns(cols, keep_shape=True)

    def self_energy(self, E, k=None, eta=None, dtype=None, eps=1e-14, bulk=False):
        r""" Return a dense matrix with the self-energy at energy `E` and k-point `k` (default Gamma).

        Parameters
        ----------
        E : float
          energy at which the calculation will take place (should *not* be complex)
        k : array_like, optional
          k-point at which the self-energy should be evaluated.
          the k-point should be in units of the reciprocal lattice vectors, and
          the semi-infinite component will be automatically set to zero.
        eta : float, optional
          the imaginary value to evaluate the self-energy with. Defaults to the
          value with which the object was created
        dtype : numpy.dtype
          the resulting data type
        eps : float, optional
          convergence criteria for the recursion
        bulk : bool, optional
          if true, :math:`E\cdot \mathbf S - \mathbf H -\boldsymbol\Sigma` is returned, else
          :math:`\boldsymbol\Sigma` is returned (default).

        Returns
        -------
        self-energy : the self-energy corresponding to the semi-infinite direction
        """
        if eta is None:
            eta = self.eta
        Z = E.real + 1j * eta

        # Get k-point
        k = self._correct_k(k)

        if dtype is None:
            dtype = np.complex128

        sp0 = self.spgeom0
        sp1 = self.spgeom1

        # As the SparseGeometry inherently works for
        # orthogonal and non-orthogonal basis, there is no
        # need to have two algorithms.
        GB = (sp0.Sk(k, dtype=dtype) * Z - sp0.Pk(k, dtype=dtype)).toarray()

        if sp1.orthogonal:
            alpha = sp1.Pk(k, dtype=dtype, format='array')
            beta  = np.conjugate(np.transpose(alpha))
        else:
            M = sp1.Pk(k, dtype=dtype)
            S = sp1.Sk(k, dtype=dtype)
            alpha = (M - S * Z).toarray()
            beta  = (M.getH() - S.getH() * Z).toarray()
            del M, S

        # Surface Green function (self-energy)
        if bulk:
            GS = np.copy(GB)
        else:
            GS = np.zeros_like(GB)

        solve = lin.solve

        i = 0
        while True:
            i += 1

            tA = solve(GB, alpha)
            tB = solve(GB, beta)

            tmp = dot(alpha, tB)
            # Update bulk Green function
            GB -= tmp + dot(beta, tA)
            # Update surface self-energy
            GS -= tmp

            # Update forward/backward
            alpha = dot(alpha, tA)
            beta = dot(beta, tB)

            # Convergence criteria, it could be stricter
            if np.amax(np.abs(tmp)) < eps:
                # Return the pristine Green function
                del tA, tB, alpha, beta, GB
                if bulk:
                    return GS
                return - GS

        raise ValueError(self.__class__.__name__+': could not converge self-energy calculation')

    def self_energy_lr(self, E, k=None, eta=None, dtype=None, eps=1e-14, bulk=False):
        r""" Return two dense matrices with the left/right self-energy at energy `E` and k-point `k` (default Gamma).

        Note calculating the LR self-energies simultaneously requires that their chemical potentials are the same.
        I.e. only when the reference energy is equivalent in the left/right schemes does this make sense.

        Parameters
        ----------
        E : float
          energy at which the calculation will take place (should *not* be complex)
        k : array_like, optional
          k-point at which the self-energy should be evaluated.
          the k-point should be in units of the reciprocal lattice vectors, and
          the semi-infinite component will be automatically set to zero.
        eta : float, optional
          the imaginary value to evaluate the self-energy with. Defaults to the
          value with which the object was created
        dtype : numpy.dtype
          the resulting data type
        eps : float, optional
          convergence criteria for the recursion
        bulk : bool, optional
          if true, :math:`E\cdot \mathbf S - \mathbf H -\boldsymbol\Sigma` is returned, else
          :math:`\boldsymbol\Sigma` is returned (default).

        Returns
        -------
        left : the left self-energy
        right : the right self-energy
        """
        if eta is None:
            eta = self.eta
        try:
            Z = E.real + 1j * eta
        except:
            Z = E + 1j * eta

        # Get k-point
        k = self._correct_k(k)

        if dtype is None:
            dtype = np.complex128

        sp0 = self.spgeom0
        sp1 = self.spgeom1

        # As the SparseGeometry inherently works for
        # orthogonal and non-orthogonal basis, there is no
        # need to have two algorithms.
        SmH0 = sp0.Sk(k, dtype=dtype) * Z - sp0.Pk(k, dtype=dtype)
        GB = SmH0.toarray()

        if sp1.orthogonal:
            alpha = sp1.Pk(k, dtype=dtype, format='array')
            beta  = np.conjugate(np.transpose(alpha))
        else:
            M = sp1.Pk(k, dtype=dtype)
            S = sp1.Sk(k, dtype=dtype)
            alpha = (M - S * Z).toarray()
            beta  = (M.getH() - S.getH() * Z).toarray()
            del M, S

        # Surface Green function (self-energy)
        if bulk:
            GS = np.copy(GB)
        else:
            GS = np.zeros_like(GB)

        solve = lin.solve

        i = 0
        while True:
            i += 1

            tA = solve(GB, alpha)
            tB = solve(GB, beta)

            tmp = dot(alpha, tB)
            # Update bulk Green function
            GB -= tmp + dot(beta, tA)
            # Update surface self-energy
            GS -= tmp

            # Update forward/backward
            alpha = dot(alpha, tA)
            beta = dot(beta, tB)

            # Convergence criteria, it could be stricter
            if np.amax(np.abs(tmp)) < eps:
                # Return the pristine Green function
                del tA, tB, alpha, beta
                if self.semi_inf_dir == 1:
                    # GS is the "right" self-energy
                    if bulk:
                        return GB - GS + SmH0.toarray(), GS
                    return GS - GB + SmH0.toarray(), - GS
                # GS is the "left" self-energy
                if bulk:
                    return GS, GB - GS + SmH0.toarray()
                return - GS, GS - GB + SmH0.toarray()

        raise ValueError(self.__class__.__name__+': could not converge self-energy (LR) calculation')
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from subprocess import Popen, PIPE

output = Popen(['git','status'], stdout=PIPE).communicate()[0]
lines = output.splitlines()

import re
behead_re = re.compile(r"^# Your branch is (ahead of|behind) '(.*)' by (\d+) commit")
diverge_re = re.compile(r"^# and have (\d+) and (\d+) different")
symbols = {'ahead of': '↑', 'behind': '↓', 'staged':'♦', 'changed':'‣', 'untracked':'…', 'clean':'⚡', 'sha1':':'}

status = ''
staged = re.compile(r'^# Changes to be committed:$', re.MULTILINE)
changed = re.compile(r'^# Changed but not updated:$', re.MULTILINE)
untracked = re.compile(r'^# Untracked files:$', re.MULTILINE)

if staged.search(output):
	nb = len(Popen(['git','diff','--staged','--name-only'], stdout=PIPE).communicate()[0].splitlines())
	status += '%s%d' % (symbols['staged'], nb)
if changed.search(output):
	nb = len(Popen(['git','ls-files','-m'], stdout=PIPE).communicate()[0].splitlines())
	status += '%s%d' % (symbols['changed'], nb)
if untracked.search(output):
## 		nb = len(Popen(['git','ls-files','--others','--exclude-standard'],stdout=PIPE).communicate()[0].splitlines())
## 		status += "%s" % (symbols['untracked']*(nb//3 + 1), )
	status += symbols['untracked']
if status == '':
	status = symbols['clean']

remote = ''

bline = lines[0]
if bline.find('Not currently on any branch') != -1:
	branch = symbols['sha1']+ Popen(['git','rev-parse','--short','HEAD'], stdout=PIPE).communicate()[0][:-1]
else:
	branch = bline.split(' ')[3]
	bstatusline = lines[1]
	match = behead_re.match(bstatusline)
	if match:
		remote = symbols[match.groups()[0]]
		remote += match.groups()[2]
	elif lines[2:]:
		div_match = diverge_re.match(lines[2])
	 	if div_match:
			remote = "{behind}{1}{ahead of}{0}".format(*div_match.groups(), **symbols)

print '(%s%s%s)' % (branch,remote,status)

<commit_msg>moved symbols in beginning of file<commit_after>#!/usr/bin/env python
# -*- coding: UTF-8 -*-

# change those symbols to whatever you prefer
symbols = {'ahead of': '↑', 'behind': '↓', 'staged':'♦', 'changed':'‣', 'untracked':'…', 'clean':'⚡', 'sha1':':'}

from subprocess import Popen, PIPE

output = Popen(['git','status'], stdout=PIPE).communicate()[0]
lines = output.splitlines()

import re
behead_re = re.compile(r"^# Your branch is (ahead of|behind) '(.*)' by (\d+) commit")
diverge_re = re.compile(r"^# and have (\d+) and (\d+) different")

status = ''
staged = re.compile(r'^# Changes to be committed:$', re.MULTILINE)
changed = re.compile(r'^# Changed but not updated:$', re.MULTILINE)
untracked = re.compile(r'^# Untracked files:$', re.MULTILINE)

if staged.search(output):
	nb = len(Popen(['git','diff','--staged','--name-only'], stdout=PIPE).communicate()[0].splitlines())
	status += '%s%d' % (symbols['staged'], nb)
if changed.search(output):
	nb = len(Popen(['git','ls-files','-m'], stdout=PIPE).communicate()[0].splitlines())
	status += '%s%d' % (symbols['changed'], nb)
if untracked.search(output):
## 		nb = len(Popen(['git','ls-files','--others','--exclude-standard'],stdout=PIPE).communicate()[0].splitlines())
## 		status += "%s" % (symbols['untracked']*(nb//3 + 1), )
	status += symbols['untracked']
if status == '':
	status = symbols['clean']

remote = ''

bline = lines[0]
if bline.find('Not currently on any branch') != -1:
	branch = symbols['sha1']+ Popen(['git','rev-parse','--short','HEAD'], stdout=PIPE).communicate()[0][:-1]
else:
	branch = bline.split(' ')[3]
	bstatusline = lines[1]
	match = behead_re.match(bstatusline)
	if match:
		remote = symbols[match.groups()[0]]
		remote += match.groups()[2]
	elif lines[2:]:
		div_match = diverge_re.match(lines[2])
	 	if div_match:
			remote = "{behind}{1}{ahead of}{0}".format(*div_match.groups(), **symbols)

print '(%s%s%s)' % (branch,remote,status)

<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function
import os, sys
import fuse
from fuse import Fuse
from gitviewfs_objects import get_gitviewfs_object
import errno

if not hasattr(fuse, '__version__'):
	raise RuntimeError, \
		"your fuse-py doesn't know of fuse.__version__, probably it's too old."

fuse.fuse_python_api = (0, 2)

fuse.feature_assert('stateful_files', 'has_init')

class GitViewFS(Fuse):

	def __init__(self, *args, **kw):
		Fuse.__init__(self, *args, **kw)

		# do stuff to set up your filesystem here, if you want
		#import thread
		#thread.start_new_thread(self.mythread, ())
		self.repo = '.'

#	def mythread(self):
#
#		"""
#		The beauty of the FUSE python implementation is that with the python interp
#		running in foreground, you can have threads
#		"""
#		print "mythread: started"
#		while 1:
#			time.sleep(120)
#			print "mythread: ticking"

	def getattr(self, path):
		obj = get_gitviewfs_object(path)
		return obj.get_stat()

	def readlink(self, path):
		symlink = get_gitviewfs_object(path)
		target_path = symlink.get_target_path()
		parent_path, _ = os.path.split(path)
		relative_target_path = os.path.relpath(target_path, parent_path)
		return relative_target_path

	def readdir(self, path, offset):
		obj = get_gitviewfs_object(path)
		for item_name in obj.get_items_names():
			yield fuse.Direntry(item_name)

	def unlink(self, path):
		os.unlink("." + path)

	def rmdir(self, path):
		os.rmdir("." + path)

	def symlink(self, path, path1):
		os.symlink(path, "." + path1)

	def rename(self, path, path1):
		os.rename("." + path, "." + path1)

	def link(self, path, path1):
		os.link("." + path, "." + path1)

	def chmod(self, path, mode):
		os.chmod("." + path, mode)

	def chown(self, path, user, group):
		os.chown("." + path, user, group)

	def truncate(self, path, length):
		f = open("." + path, "a")
		f.truncate(length)
		f.close()

	def mknod(self, path, mode, dev):
		os.mknod("." + path, mode, dev)

	def mkdir(self, path, mode):
		os.mkdir("." + path, mode)

	def utime(self, path, times):
		os.utime("." + path, times)

#	The following utimens method would do the same as the above utime method.
#	We can't make it better though as the Python stdlib doesn't know of
#	subsecond preciseness in acces/modify times.
#
#	def utimens(self, path, ts_acc, ts_mod):
#	  os.utime("." + path, (ts_acc.tv_sec, ts_mod.tv_sec))

	def access(self, path, mode):
		if not os.access("." + path, mode):
			return -errno.EACCES

#	This is how we could add stub extended attribute handlers...
#	(We can't have ones which aptly delegate requests to the underlying fs
#	because Python lacks a standard xattr interface.)
#
#	def getxattr(self, path, name, size):
#		val = name.swapcase() + '@' + path
#		if size == 0:
#			# We are asked for size of the value.
#			return len(val)
#		return val
#
#	def listxattr(self, path, size):
#		# We use the "user" namespace to please XFS utils
#		aa = ["user." + a for a in ("foo", "bar")]
#		if size == 0:
#			# We are asked for size of the attr list, ie. joint size of attrs
#			# plus null separators.
#			return len("".join(aa)) + len(aa)
#		return aa

	def statfs(self):
		"""
		Should return an object with statvfs attributes (f_bsize, f_frsize...).
		Eg., the return value of os.statvfs() is such a thing (since py 2.2).
		If you are not reusing an existing statvfs object, start with
		fuse.StatVFS(), and define the attributes.

		To provide usable information (ie., you want sensible df(1)
		output, you are suggested to specify the following attributes:

			- f_bsize - preferred size of file blocks, in bytes
			- f_frsize - fundamental size of file blcoks, in bytes
				[if you have no idea, use the same as blocksize]
			- f_blocks - total number of blocks in the filesystem
			- f_bfree - number of free blocks
			- f_files - total number of file inodes
			- f_ffree - nunber of free file inodes
		"""
		return os.statvfs(".")

	def fsinit(self):
		os.chdir(self.repo)

	class GitViewFSFile(object):

		def __init__(self, path, flags, *mode):
			self.file = get_gitviewfs_object(path)

		def read(self, length, offset):
			content = self.file.get_content()
			return content[offset : offset+length]

		def write(self, buf, offset):
			self.file.seek(offset)
			self.file.write(buf)
			return len(buf)

		def release(self, flags):
			pass

		def _fflush(self):
			pass

		def fsync(self, isfsyncfile):
			self._fflush()
			if isfsyncfile and hasattr(os, 'fdatasync'):
				os.fdatasync(self.fd)
			else:
				os.fsync(self.fd)

		def flush(self):
			pass

		def fgetattr(self):
			return self.file.get_stat()

		def ftruncate(self, length):
			self.file.truncate(length)

		def lock(self, cmd, owner, **kw):
			# The code here is much rather just a demonstration of the locking
			# API than something which actually was seen to be useful.

			# Advisory file locking is pretty messy in Unix, and the Python
			# interface to this doesn't make it better.
			# We can't do fcntl(2)/F_GETLK from Python in a platfrom independent
			# way. The following implementation *might* work under Linux.
			#
			# if cmd == fcntl.F_GETLK:
			#	 import struct
			#
			#	 lockdata = struct.pack('hhQQi', kw['l_type'], os.SEEK_SET,
			#							kw['l_start'], kw['l_len'], kw['l_pid'])
			#	 ld2 = fcntl.fcntl(self.fd, fcntl.F_GETLK, lockdata)
			#	 flockfields = ('l_type', 'l_whence', 'l_start', 'l_len', 'l_pid')
			#	 uld2 = struct.unpack('hhQQi', ld2)
			#	 res = {}
			#	 for i in xrange(len(uld2)):
			#		  res[flockfields[i]] = uld2[i]
			#
			#	 return fuse.Flock(**res)

			# Convert fcntl-ish lock parameters to Python's weird
			# lockf(3)/flock(2) medley locking API...
			pass

	def main(self, *a, **kw):
		self.file_class = self.GitViewFSFile
		return Fuse.main(self, *a, **kw)

def main():

	server = GitViewFS(version="%prog 0.1", dash_s_do='setsingle')

	# Disable multithreading: if you want to use it, protect all method of
	# XmlFile class with locks, in order to prevent race conditions
	server.multithreaded = False

	server.parser.add_option(mountopt="repo", metavar="PATH", default='.',
							 help="browse Git repository from under PATH [default: %default]")
	server.parse(values=server, errex=1)

	try:
		if server.fuse_args.mount_expected():
			os.chdir(server.repo)
	except OSError:
		print("can't enter Git repo directory", file=sys.stderr)
		sys.exit(1)

	server.main()

if __name__ == '__main__':
	main()
<commit_msg>Remove functions intended to make changes<commit_after>#!/usr/bin/env python
from __future__ import print_function
import os, sys
import fuse
from fuse import Fuse
from gitviewfs_objects import get_gitviewfs_object
import errno

if not hasattr(fuse, '__version__'):
	raise RuntimeError, \
		"your fuse-py doesn't know of fuse.__version__, probably it's too old."

fuse.fuse_python_api = (0, 2)

fuse.feature_assert('stateful_files', 'has_init')

class GitViewFS(Fuse):

	def __init__(self, *args, **kw):
		Fuse.__init__(self, *args, **kw)

		# do stuff to set up your filesystem here, if you want
		#import thread
		#thread.start_new_thread(self.mythread, ())
		self.repo = '.'

#	def mythread(self):
#
#		"""
#		The beauty of the FUSE python implementation is that with the python interp
#		running in foreground, you can have threads
#		"""
#		print "mythread: started"
#		while 1:
#			time.sleep(120)
#			print "mythread: ticking"

	def getattr(self, path):
		obj = get_gitviewfs_object(path)
		return obj.get_stat()

	def readlink(self, path):
		symlink = get_gitviewfs_object(path)
		target_path = symlink.get_target_path()
		parent_path, _ = os.path.split(path)
		relative_target_path = os.path.relpath(target_path, parent_path)
		return relative_target_path

	def readdir(self, path, offset):
		obj = get_gitviewfs_object(path)
		for item_name in obj.get_items_names():
			yield fuse.Direntry(item_name)

	'''
	def access(self, path, mode):
		if not os.access("." + path, mode):
			return -errno.EACCES
	'''

#	This is how we could add stub extended attribute handlers...
#	(We can't have ones which aptly delegate requests to the underlying fs
#	because Python lacks a standard xattr interface.)
#
#	def getxattr(self, path, name, size):
#		val = name.swapcase() + '@' + path
#		if size == 0:
#			# We are asked for size of the value.
#			return len(val)
#		return val
#
#	def listxattr(self, path, size):
#		# We use the "user" namespace to please XFS utils
#		aa = ["user." + a for a in ("foo", "bar")]
#		if size == 0:
#			# We are asked for size of the attr list, ie. joint size of attrs
#			# plus null separators.
#			return len("".join(aa)) + len(aa)
#		return aa

	def statfs(self):
		"""
		Should return an object with statvfs attributes (f_bsize, f_frsize...).
		Eg., the return value of os.statvfs() is such a thing (since py 2.2).
		If you are not reusing an existing statvfs object, start with
		fuse.StatVFS(), and define the attributes.

		To provide usable information (ie., you want sensible df(1)
		output, you are suggested to specify the following attributes:

			- f_bsize - preferred size of file blocks, in bytes
			- f_frsize - fundamental size of file blcoks, in bytes
				[if you have no idea, use the same as blocksize]
			- f_blocks - total number of blocks in the filesystem
			- f_bfree - number of free blocks
			- f_files - total number of file inodes
			- f_ffree - nunber of free file inodes
		"""
		return os.statvfs(".")

	def fsinit(self):
		os.chdir(self.repo)

	class GitViewFSFile(object):

		def __init__(self, path, flags, *mode):
			self.file = get_gitviewfs_object(path)

		def read(self, length, offset):
			content = self.file.get_content()
			return content[offset : offset+length]

		def fgetattr(self):
			return self.file.get_stat()

	def main(self, *a, **kw):
		self.file_class = self.GitViewFSFile
		return Fuse.main(self, *a, **kw)

def main():

	server = GitViewFS(version="%prog 0.1", dash_s_do='setsingle')

	# Disable multithreading: if you want to use it, protect all method of
	# XmlFile class with locks, in order to prevent race conditions
	server.multithreaded = False

	server.parser.add_option(mountopt="repo", metavar="PATH", default='.',
							 help="browse Git repository from under PATH [default: %default]")
	server.parse(values=server, errex=1)

	try:
		if server.fuse_args.mount_expected():
			os.chdir(server.repo)
	except OSError:
		print("can't enter Git repo directory", file=sys.stderr)
		sys.exit(1)

	server.main()

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


<commit_before>"""
        _ _
   __ _(_) |_ _   _
  / _` | | __| | | |
 | (_| | | |_| |_| |
  \__, |_|\__|\__, |
  |___/       |___/

gity : The Git magician for Python

"""

import subprocess
import os,sys

class gity:
	def __init__(self,user='',paswd='',repo=''):
		""" (gity object, user, paswd, repo) -> NoneType

		Initalize with username , password and path to repository
		"""

		self.user = user
		self.paswd = paswd

		#check if repositry is a string
		if repo=='':
			self.repo = os.getcwd()
		elif type(repo) == str:
			self.repo = repo
			os.chdir(repo)
		else:
			raise Exception("Expected String")

	def setUser(self,user):
		"""(gity object,str) -> NoneType

		set user value to entered String
		"""
		try:
			assert type(user) == str
			self.user = user
		except:
			raise Exception("Expected String")

	def setPassword(self,paswd):
		"""(gity object,str) -> NoneType

		Set password for github account
		"""
		try:
			assert type(paswd) == str
			self.paswd = paswd
		except:
			raise Exception("Expected String")

	def setRepository(self,repo):
		"""(gity object,str) -> NoneType

		Set path to repository
		"""
		try:
			assert type(repo) == str
			self.repo = repo
			os.chdir(repo)
		except:
			raise Exception("Expected String")

	def gitInit(self):
		"""(gity object) -> bool

		Initalizes git repository at specified path
		path taken from self.repo
		"""
		try:
			args = ["git","init"]
			state = subprocess.call(args)
			status = True
		except:
			status = False

		return status

	def gitClone(self,url,arg=''):
		"""(gity object, str, str) -> bool
		Clones a repository into a newly created directory (path taken from self.repo ),
		creates remote-tracking branches for each branch in the cloned repository
		"""
		if arg!='':
			args.append(arg)
		try:
			args = ["git","init"]
			state = subprocess.call(args)
			status = True
		except:
			status = False

		return status

	def gitStatus(self):
		"""(gity object) -> str

		Executes command "git status" on command line
		Prints status of the repository
		Returns True if status is printed
				False if an exception occurs
		"""
		try:
			args = ["git","status"]
			state = subprocess.call(args)
			status = True
		except:
			status = False

		return status

	def gitAdd(self,files=[],deletion=False):
		"""(gity object, list, bool) -> bool

		Add multiple files in directory.
		Files is the list of files that you want to add.
		If the list is empty, all files in the directory will be added

		Default:
			All files in the directory will be added
			Deleted files will not be tracked/added

		gitAdd can :
			-> Add multiple files while ignoring deletion
			-> Add/Track all files including deleted
		"""

		try:
		# Check if arguments match the exoected type
			assert type(files) == list
			assert type(deletion) == bool
		except:
			raise Exception("Files expected as list\ndeletion expected as bool")

		if deletion == True:
		# If deletion is True track all files
			try:
				args = ["git","add","-A"]
				state = subprocess.call(args)
				status = True
				print "Added all the files [ including deleted ]"
			except:
				status = False
		else:
		# Deletion is false then
		#	-> Add all files while ignoring deletion
		#	-> Add only specified files while ignoring deletion
			if files==[]:
				try:
					args = ["git","add","--no-all"]
					state = subprocess.call(args)
					status = True
					print "Added all the files except deletions"
				except:
					status = False
			else:
				try:
					args = ["git","add"]

					# Add all the file name to list of args
					for fi in files:
						args.append(fi)

					state = subprocess.call(args)
					status = True
					print "Added all the files mentioned"
				except:
					status = False

			return status

	def gitCommit(self,message):
		"""(gity object,str) -> bool

		Returns True if commit is succesful
		else returns False.
		"""
		if type(message) != str:
			raise Exception("Expected message as string")
		try:
			args = ["git","commit","-m",message]
			state = subprocess.call(args)
			status = True
			print "Commit succesful"
		except:
			status = False

		return status

	def gitPush(self,remoteName='origin',branchName='master'):
		""" (gity object, str, str) -> bool

		Default:
			pushes origin
		Returns True if the push is succesful
		else returns False
		"""

		try:
			args = ["git","push",remoteName,branchName]
			state = subprocess.call(args)
			status = True
			print "Push succesful"
		except:
			status = False

		return status

	def gitPull(self):
		"""(gity object) -> bool

		Executes git pull on the mentioned repository
		"""
		try:
			args = ["git","pull"]
			state = subprocess.call(args)
			status = True
			print "Pull succesful"
		except:
			status = False

		return status

	def gitCheckout(self,branch):
		"""

		Executes git checkout to switch to given branch
		"""
		try:
			args = ["git","checkout","branch"]
			state = subprocess.call(args)
			status = True
			print "Checkout succesful"
		except:
			status = False

		return status

if __name__=="__main__":
	print """
	        _ _
	   __ _(_) |_ _   _
	  / _` | | __| | | |
	 | (_| | | |_| |_| |
	  \__, |_|\__|\__, |
	  |___/       |___/

	gity : Git Magician for Python

"""
<commit_msg>bug with gitAdd fixed<commit_after>"""
        _ _
   __ _(_) |_ _   _
  / _` | | __| | | |
 | (_| | | |_| |_| |
  \__, |_|\__|\__, |
  |___/       |___/

gity : The Git magician for Python

"""

import subprocess
import os,sys

class gity:
	def __init__(self,user='',paswd='',repo=''):
		""" (gity object, user, paswd, repo) -> NoneType

		Initalize with username , password and path to repository
		"""

		self.user = user
		self.paswd = paswd

		#check if repositry is a string
		if repo=='':
			self.repo = os.getcwd()
		elif type(repo) == str:
			self.repo = repo
			os.chdir(repo)
		else:
			raise Exception("Expected String")

	def setUser(self,user):
		"""(gity object,str) -> NoneType

		set user value to entered String
		"""
		try:
			assert type(user) == str
			self.user = user
		except:
			raise Exception("Expected String")

	def setPassword(self,paswd):
		"""(gity object,str) -> NoneType

		Set password for github account
		"""
		try:
			assert type(paswd) == str
			self.paswd = paswd
		except:
			raise Exception("Expected String")

	def setRepository(self,repo):
		"""(gity object,str) -> NoneType

		Set path to repository
		"""
		try:
			assert type(repo) == str
			self.repo = repo
			os.chdir(repo)
		except:
			raise Exception("Expected String")

	def gitInit(self):
		"""(gity object) -> bool

		Initalizes git repository at specified path
		path taken from self.repo
		"""
		try:
			args = ["git","init"]
			state = subprocess.call(args)
			status = True
		except:
			status = False

		return status

	def gitClone(self,url,arg=''):
		"""(gity object, str, str) -> bool
		Clones a repository into a newly created directory (path taken from self.repo ),
		creates remote-tracking branches for each branch in the cloned repository
		"""
		if arg!='':
			args.append(arg)
		try:
			args = ["git","init"]
			state = subprocess.call(args)
			status = True
		except:
			status = False

		return status

	def gitStatus(self):
		"""(gity object) -> str

		Executes command "git status" on command line
		Prints status of the repository
		Returns True if status is printed
				False if an exception occurs
		"""
		try:
			args = ["git","status"]
			state = subprocess.call(args)
			status = True
		except:
			status = False

		return status

	def gitAdd(self,files=[],deletion=False):
		"""(gity object, list, bool) -> bool

		Add multiple files in directory.
		Files is the list of files that you want to add.
		If the list is empty, all files in the directory will be added

		Default:
			All files in the directory will be added
			Deleted files will not be tracked/added

		gitAdd can :
			-> Add multiple files while ignoring deletion
			-> Add/Track all files including deleted
		"""

		try:
		# Check if arguments match the exoected type
			assert type(files) == list
			assert type(deletion) == bool
		except:
			raise Exception("Files expected as list\ndeletion expected as bool")

		if deletion == True:
		# If deletion is True track all files
			try:
				args = ["git","add","-A"]
				state = subprocess.call(args)
				status = True
				print "Added all the files [ including deleted ]"
			except:
				status = False
		else:
		# Deletion is false then
		#	-> Add all files while ignoring deletion
		#	-> Add only specified files while ignoring deletion
			if files==[]:
				try:
					args = ["git","add",".","--ignore-removal"]
					state = subprocess.call(args)
					status = True
					print "Added all the files except deletions"
				except:
					status = False
			else:
				try:
					args = ["git","add"]

					# Add all the file name to list of args
					for fi in files:
						args.append(fi)

					state = subprocess.call(args)
					status = True
					print "Added all the files mentioned"
				except:
					status = False

			return status

	def gitCommit(self,message):
		"""(gity object,str) -> bool

		Returns True if commit is succesful
		else returns False.
		"""
		if type(message) != str:
			raise Exception("Expected message as string")
		try:
			args = ["git","commit","-m",message]
			state = subprocess.call(args)
			status = True
			print "Commit succesful"
		except:
			status = False

		return status

	def gitPush(self,remoteName='origin',branchName='master'):
		""" (gity object, str, str) -> bool

		Default:
			pushes origin
		Returns True if the push is succesful
		else returns False
		"""

		try:
			args = ["git","push",remoteName,branchName]
			state = subprocess.call(args)
			status = True
			print "Push succesful"
		except:
			status = False

		return status

	def gitPull(self):
		"""(gity object) -> bool

		Executes git pull on the mentioned repository
		"""
		try:
			args = ["git","pull"]
			state = subprocess.call(args)
			status = True
			print "Pull succesful"
		except:
			status = False

		return status

	def gitCheckout(self,branch):
		"""

		Executes git checkout to switch to given branch
		"""
		try:
			args = ["git","checkout","branch"]
			state = subprocess.call(args)
			status = True
			print "Checkout succesful"
		except:
			status = False

		return status

if __name__=="__main__":
	print """
	        _ _
	   __ _(_) |_ _   _
	  / _` | | __| | | |
	 | (_| | | |_| |_| |
	  \__, |_|\__|\__, |
	  |___/       |___/

	gity : Git Magician for Python

"""
<|endoftext|>


<commit_before>
"""
Usage::

    import install_visual_studio_wrapper
    install_visual_studio_wrapper.install()

"""

import os
import shutil
import sys

try:
    from urllib.request import urlopen
except ImportError:
    from urllib2 import urlopen

SCRIPT_URL = \
    "https://raw.githubusercontent.com/ogrisel/python-appveyor-demo/" \
    "f54ec3593bcea682098a59b560c1850c19746e10/appveyor/run_with_env.cmd"

LOCAL_PATH_DEFAULT = \
    os.path.join(os.path.dirname(__file__), "run-with-visual-studio.cmd")

def _log(*args):
    script_name = os.path.basename(__file__)
    print("[appveyor:%s] " % script_name + " ".join(args))
    sys.stdout.flush()

def install(local_path=LOCAL_PATH_DEFAULT):
    """Download ``run_with_env.cmd`` and copy it to ``local_path``"""

    _log("Downloading %s" % SCRIPT_URL)
    remote_script = urlopen(SCRIPT_URL)
    with open(local_path, "wb") as local_script:
        _log("Renaming to %s" % local_script)
        shutil.copyfileobj(remote_script, local_script)

if __name__ == '__main__':
    install(sys.argv[1])
<commit_msg>ci/appveyor: Ensure install_visual_studio_wrapper can be used within args<commit_after>
"""
Usage::

    import install_visual_studio_wrapper
    install_visual_studio_wrapper.install()

"""

import os
import shutil
import sys

try:
    from urllib.request import urlopen
except ImportError:
    from urllib2 import urlopen

SCRIPT_URL = \
    "https://raw.githubusercontent.com/ogrisel/python-appveyor-demo/" \
    "f54ec3593bcea682098a59b560c1850c19746e10/appveyor/run_with_env.cmd"

LOCAL_PATH_DEFAULT = \
    os.path.join(os.path.dirname(__file__), "run-with-visual-studio.cmd")

def _log(*args):
    script_name = os.path.basename(__file__)
    print("[appveyor:%s] " % script_name + " ".join(args))
    sys.stdout.flush()

def install(local_path=LOCAL_PATH_DEFAULT):
    """Download ``run_with_env.cmd`` and copy it to ``local_path``"""

    _log("Downloading %s" % SCRIPT_URL)
    remote_script = urlopen(SCRIPT_URL)
    with open(local_path, "wb") as local_script:
        _log("Renaming to %s" % local_script)
        shutil.copyfileobj(remote_script, local_script)

if __name__ == '__main__':
    install(sys.argv[1] if len(sys.argv) > 1 else LOCAL_PATH_DEFAULT)
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import division

import cluster_labeling.local_django_models as local
import cluster_labeling.optemo_django_models as optemo
import cluster_labeling.text_handling as th

import cluster_labeling.words as words

from django.db import models

class Usecase(local.LocalModel):
    class Meta:
        db_table = '%s_%s' % (optemo.product_type_tablename_prefix,
                              'usecases')

    label = models.CharField(max_length=255, unique=True)
    direct_indicator_words = \
        models.ManyToManyField\
        (words.Word, related_name='directly_indicated_usecases')
    indicator_words = \
        models.ManyToManyField\
        (words.Word, related_name='indicated_usecases')

class UsecaseClusterScore(local.LocalModel):
    class Meta:
        db_table = '%s_%s' % (optemo.product_type_tablename_prefix,
                              'usecase_cluster_scores')

    usecase = models.ForeignKey(Usecase)
    score = models.FloatField()

    cluster_id = models.BigIntegerField()
    version = models.IntegerField()

# This list also includes meta-features, for now.
known_usecases = {
    optemo.Printer : {},
    optemo.Camera :
    {"Travel" : [],
     "Fun" : [],
     "Family photos" : ["Family"],
     "Landscape/scenery" : [],
     "Low light" : ["Dim"],
     "Outdoors" : [],
     "Art" : [],
     "Sports/action" : ["Sports", "Action"],
     "Video" : [],
     "Wildlife" : [],
     "Weddings" : [],
     "Home" : [],
     "Street" : [],
     "Build quality" : [],
     "Compact" : [],
     "Clunky" : [],
     "Image stabilization" : ["Stabilization"],
     "Viewfinder" : [],
     "Manual control" : [],
     "LCD" : []}}

def populate_usecases():
    Usecase.drop_tables_if_exist()
    Usecase.create_tables()

    for (usecase_label, direct_indicator_words) in \
            known_usecases[optemo.product_type].iteritems():
        if len(direct_indicator_words) == 0:
            direct_indicator_words = th.get_words_from_string(usecase_label)

        usecase = Usecase(label=usecase_label)
        usecase.save()

        existing_words_qs, dne_word_entries = \
            words.Word.create_multiple_if_dne_and_return(direct_indicator_words)

        for word in existing_words_qs:
            usecase.direct_indicator_words.add(word)
        for word in dne_word_entries:
            word.save()
            usecase.direct_indicator_words.add(word)

        usecase.save()

def populate_indicator_words_for_usecases():
    for usecase in Usecase.get_manager():
        diws = usecase.direct_indicator_words

        for diw in diws.all():
            usecase.indicator_words.add(diw)

            for synonym in diw.get_all_synonyms():
                usecase.indicator_words.add(synonym)

        usecase.save()

import cluster_labeling.nh_chi_scorer as chi

def score_usecases_for_all_clusters\
        (version=optemo.product_cluster_type.get_latest_version()):
    clusters = optemo.product_cluster_type.get_manager()\
               .filter(version=version)

    for cluster in clusters:
        score_usecases_for_cluster(cluster)

def score_usecases_for_cluster(cluster):
    usecases = Usecase.get_manager()

    for usecase in usecases:
        score = score_usecase_for_cluster(cluster, usecase)

        if score == 0:
            continue

        kwargs = {'usecase':usecase, 'version':cluster.version,
                  'cluster_id':cluster.id}
        qs = UsecaseClusterScore.get_manager().filter(**kwargs)

        assert(qs.count() == 0)

        kwargs['score'] = score
        usecase_cluster_score = UsecaseClusterScore(**kwargs)
        usecase_cluster_score.save()

def score_usecase_for_cluster(cluster, usecase):
    # Get indicator words for usecase
    iwords = usecase.indicator_words.all()

    # Compute chi-squared score for all indicator words
    iword_scores = \
        map(lambda iw:
            chi.get_chi_squared_score\
            (cluster.id, cluster.version, iw.word),
            iwords)

    # Combine all chi-squared scores in a way that is not influenced
    # by the number of indicator words associated with the cluster
    # i.e. take the average:
    iword_scores = \
        filter(lambda s: s > 0,
               map(lambda s: 0 if s is None else s, iword_scores))

    if len(iword_scores) == 0:
        return 0

    score = sum(iword_scores) / len(iword_scores)
    return score
<commit_msg>Adding a check to make sure that all direct indicator words are lower-case<commit_after>#!/usr/bin/env python
from __future__ import division

import cluster_labeling.local_django_models as local
import cluster_labeling.optemo_django_models as optemo
import cluster_labeling.text_handling as th

import cluster_labeling.words as words

from django.db import models

class Usecase(local.LocalModel):
    class Meta:
        db_table = '%s_%s' % (optemo.product_type_tablename_prefix,
                              'usecases')

    label = models.CharField(max_length=255, unique=True)
    direct_indicator_words = \
        models.ManyToManyField\
        (words.Word, related_name='directly_indicated_usecases')
    indicator_words = \
        models.ManyToManyField\
        (words.Word, related_name='indicated_usecases')

class UsecaseClusterScore(local.LocalModel):
    class Meta:
        db_table = '%s_%s' % (optemo.product_type_tablename_prefix,
                              'usecase_cluster_scores')

    usecase = models.ForeignKey(Usecase)
    score = models.FloatField()

    cluster_id = models.BigIntegerField()
    version = models.IntegerField()

# This list also includes meta-features, for now.
known_usecases = {
    optemo.Printer : {},
    optemo.Camera :
    {"Travel" : [],
     "Fun" : [],
     "Family photos" : ["Family"],
     "Landscape/scenery" : [],
     "Low light" : ["Dim"],
     "Outdoors" : [],
     "Art" : [],
     "Sports/action" : ["Sports", "Action"],
     "Video" : [],
     "Wildlife" : [],
     "Weddings" : [],
     "Home" : [],
     "Street" : [],
     "Build quality" : [],
     "Compact" : [],
     "Clunky" : [],
     "Image stabilization" : ["Stabilization"],
     "Viewfinder" : [],
     "Manual control" : [],
     "LCD" : []}}

def populate_usecases():
    Usecase.drop_tables_if_exist()
    Usecase.create_tables()

    for (usecase_label, direct_indicator_words) in \
            known_usecases[optemo.product_type].iteritems():
        if len(direct_indicator_words) == 0:
            direct_indicator_words = th.get_words_from_string(usecase_label)

        usecase = Usecase(label=usecase_label)
        usecase.save()

        existing_words_qs, dne_word_entries = \
            words.Word.create_multiple_if_dne_and_return\
            (map(lambda s: s.lower(), direct_indicator_words))

        for word in existing_words_qs:
            usecase.direct_indicator_words.add(word)
        for word in dne_word_entries:
            word.save()
            usecase.direct_indicator_words.add(word)

        usecase.save()

def populate_indicator_words_for_usecases():
    for usecase in Usecase.get_manager():
        diws = usecase.direct_indicator_words

        for diw in diws.all():
            usecase.indicator_words.add(diw)

            for synonym in diw.get_all_synonyms():
                usecase.indicator_words.add(synonym)

        usecase.save()

import cluster_labeling.nh_chi_scorer as chi

def score_usecases_for_all_clusters\
        (version=optemo.product_cluster_type.get_latest_version()):
    clusters = optemo.product_cluster_type.get_manager()\
               .filter(version=version)

    for cluster in clusters:
        score_usecases_for_cluster(cluster)

def score_usecases_for_cluster(cluster):
    usecases = Usecase.get_manager()

    for usecase in usecases:
        score = score_usecase_for_cluster(cluster, usecase)

        if score == 0:
            continue

        kwargs = {'usecase':usecase, 'version':cluster.version,
                  'cluster_id':cluster.id}
        qs = UsecaseClusterScore.get_manager().filter(**kwargs)

        assert(qs.count() == 0)

        kwargs['score'] = score
        usecase_cluster_score = UsecaseClusterScore(**kwargs)
        usecase_cluster_score.save()

def score_usecase_for_cluster(cluster, usecase):
    # Get indicator words for usecase
    iwords = usecase.indicator_words.all()

    # Compute chi-squared score for all indicator words
    iword_scores = \
        map(lambda iw:
            chi.get_chi_squared_score\
            (cluster.id, cluster.version, iw.word),
            iwords)

    # Combine all chi-squared scores in a way that is not influenced
    # by the number of indicator words associated with the cluster
    # i.e. take the average:
    iword_scores = \
        filter(lambda s: s > 0,
               map(lambda s: 0 if s is None else s, iword_scores))

    if len(iword_scores) == 0:
        return 0

    score = sum(iword_scores) / len(iword_scores)
    return score
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from math import sqrt, exp
import random
import scipy.stats

class Estimator(object):
    def __init__(self):
        self._n = 0
        self._mean = 0
        self._variance = 0

    def __str__(self):
        return '\t{:.4f}\t{:.4f}'.format(self._mean, self.confidence())

    def _observation(self, x):
        self._n += 1
        mean = self._mean + (x - self._mean) / float(self.n)
        variance = self._variance + self._mean ** 2 - (mean) ** 2
        variance += (x ** 2 - self._variance - self._mean ** 2) /float(self._n)
        self._mean = mean
        self._variance = variance

    @property
    def n(self):
        return self._n

    @property
    def mean(self):
        return self._mean

    @property
    def std_dev(self):
        return sqrt(self._n * self._variance / float(self._n - 1))

    def confidence(self, level=0.95):
        return scipy.stats.t.ppf((1 + level) / 2.0, self._n - 1) * self.std_dev / sqrt(self._n)

class MonteCarlo(Estimator):
    def observation(self, graph, budget):
        """Sample sample_budget vertices uniformly at random and return maximum degree"""
        max_degree = max(graph.degree())
        estimate = max(random.choice(graph.vs).degree() for _ in range(budget))
        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)

class Greedy(Estimator):
    def observation(self, graph, budget):
        """Greedy walk maximum degree estimator.

        Select a random vertex, look at the degrees of all adjacent vertices, select (if possible) a new neighbor that
        has maximum degree among all neighbors, move to that vertex, and repeat at most n times.
        """
        max_degree = max(graph.degree())
        estimate = 0
        v = random.choice(graph.vs)
        for i in range(budget):
            estimate = max(estimate, v.degree())
            try:
                v_prime = max(v.neighbors(), key=lambda v: v.degree())
            except ValueError:
                break
            if v_prime.degree() < v.degree():
                break
            v = v_prime

        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)

class GreedyRestart(Estimator):
    def observation(self, graph, budget):
        """Greedy walk maximum degree estimator.

        Select a random vertex, look at the degrees of all adjacent vertices, select
        (if possible) a new neighbor that has maximum degree among all neighbors,
        move to that vertex, and repeat at most n times.
        """
        max_degree = max(graph.degree())
        estimate = 0
        v = random.choice(graph.vs)
        for i in range(budget):
            estimate = max(estimate, v.degree())
            try:
                v_prime = max(v.neighbors(), key=lambda v: v.degree())
            except ValueError:
                v_prime = random.choice(graph.vs)
            if v_prime.degree() < v.degree():
                v_prime = random.choice(graph.vs)
            v = v_prime

        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)

class Annealing(Estimator):
    def observation(self, graph, budget):
        max_degree = max(graph.degree())
        estimate = 0
        cool = 1e-5 ** (1. / budget)
        temperature = 10000
        v = random.choice(graph.vs)
        while temperature > 0.1:
            temperature *= cool
            d = v.degree()
            estimate = max(estimate, d)
            # select next vertex
            try:
                v_prime = random.choice(v.neighbors())
            except IndexError:
                v_prime = random.choice(graph.vs)
            d_prime = v_prime.degree()
            p = exp(min(0, d_prime - d) / temperature)
            uniform_rv = random.random()
            if d_prime >= d or uniform_rv <= p:
                v = v_prime

        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)
<commit_msg>Welford method of iterative variance computation<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from math import sqrt, exp
import random
import scipy.stats

class Estimator(object):
    def __init__(self):
        self._n = 0
        self._M = 0
        self._S = 0

    def __str__(self):
        return '\t{:.4f}\t{:.4f}'.format(self._mean, self.confidence())

    def _observation(self, x):
        delta = x - self._M
        self._M += delta / float(self.n + 1)
        self._S += delta * (x - self._M)
        self._n += 1

    @property
    def n(self):
        return self._n

    @property
    def mean(self):
        return self._M

    @property
    def variance(self):
        return self._S / float(self._n - 1)

    def confidence(self, level=0.95):
        t = scipy.stats.t.ppf((1 + level) / 2.0, self._n - 1)
        return t * sqrt(self.variance / self._n)

class MonteCarlo(Estimator):
    def observation(self, graph, budget):
        """Sample sample_budget vertices uniformly at random and return maximum degree"""
        max_degree = max(graph.degree())
        estimate = max(random.choice(graph.vs).degree() for _ in range(budget))
        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)

class Greedy(Estimator):
    def observation(self, graph, budget):
        """Greedy walk maximum degree estimator.

        Select a random vertex, look at the degrees of all adjacent vertices, select (if possible) a new neighbor that
        has maximum degree among all neighbors, move to that vertex, and repeat at most n times.
        """
        max_degree = max(graph.degree())
        estimate = 0
        v = random.choice(graph.vs)
        for i in range(budget):
            estimate = max(estimate, v.degree())
            try:
                v_prime = max(v.neighbors(), key=lambda v: v.degree())
            except ValueError:
                break
            if v_prime.degree() < v.degree():
                break
            v = v_prime

        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)

class GreedyRestart(Estimator):
    def observation(self, graph, budget):
        """Greedy walk maximum degree estimator.

        Select a random vertex, look at the degrees of all adjacent vertices, select
        (if possible) a new neighbor that has maximum degree among all neighbors,
        move to that vertex, and repeat at most n times.
        """
        max_degree = max(graph.degree())
        estimate = 0
        v = random.choice(graph.vs)
        for i in range(budget):
            estimate = max(estimate, v.degree())
            try:
                v_prime = max(v.neighbors(), key=lambda v: v.degree())
            except ValueError:
                v_prime = random.choice(graph.vs)
            if v_prime.degree() < v.degree():
                v_prime = random.choice(graph.vs)
            v = v_prime

        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)

class Annealing(Estimator):
    def observation(self, graph, budget):
        max_degree = max(graph.degree())
        estimate = 0
        cool = 1e-5 ** (1. / budget)
        temperature = 10000
        v = random.choice(graph.vs)
        while temperature > 0.1:
            temperature *= cool
            d = v.degree()
            estimate = max(estimate, d)
            # select next vertex
            try:
                v_prime = random.choice(v.neighbors())
            except IndexError:
                v_prime = random.choice(graph.vs)
            d_prime = v_prime.degree()
            p = exp(min(0, d_prime - d) / temperature)
            uniform_rv = random.random()
            if d_prime >= d or uniform_rv <= p:
                v = v_prime

        relative_error = float(max_degree - estimate) / float(max_degree)
        self._observation(relative_error)
<|endoftext|>


<commit_before># Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

"""A base user-interface workflow, to be inherited by UI modules."""

import tempfile

import benchmark
import better_webbrowser
import config
import data_sources
import geoip
import nameserver_list
import reporter
import site_connector
import util

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

class BaseUI(object):
  """Common methods for all UI implementations."""

  def __init__(self):
    self.SetupDataStructures()

  def SetupDataStructures(self):
    """Instead of requiring users to inherit __init__(), this sets up structures."""
    self.reporter = None
    self.nameservers = None
    self.bmark = None
    self.report_path = None
    self.csv_path = None
    self.geodata = None
    self.country = None
    self.sources = {}
    self.url = None
    self.share_state = None
    self.test_records = []

  def UpdateStatus(self, msg, **kwargs):
    """Update the little status message on the bottom of the window."""
    if hasattr(self, 'status_callback') and self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print msg

  def DebugMsg(self, message):
    self.UpdateStatus(message, debug=True)

  def LoadDataSources(self):
    self.data_src = data_sources.DataSources(status_callback=self.UpdateStatus)

  def PrepareTestRecords(self):
    """Figure out what data source a user wants, and create test_records."""
    if self.options.input_source:
      src_type = self.options.input_source
    else:
      src_type = self.data_src.GetBestSourceDetails()[0]
      self.options.input_source = src_type

    self.test_records = self.data_src.GetTestsFromSource(
        src_type,
        self.options.query_count,
        select_mode=self.options.select_mode
    )

  def PrepareNameServers(self):
    """Setup self.nameservers to have a list of healthy fast servers."""
    self.nameservers = nameserver_list.NameServers(
        self.supplied_ns,
        global_servers=self.global_ns,
        regional_servers=self.regional_ns,
        include_internal=self.include_internal,
        num_servers=self.options.num_servers,
        timeout=self.options.timeout,
        ping_timeout=self.options.ping_timeout,
        health_timeout=self.options.health_timeout,
        ipv6_only=self.options.ipv6_only,
        status_callback=self.UpdateStatus
    )
    if self.options.invalidate_cache:
      self.nameservers.InvalidateSecondaryCache()

    self.nameservers.cache_dir = tempfile.gettempdir()

    # Don't waste time checking the health of the only nameserver in the list.
    if len(self.nameservers) > 1:
      self.nameservers.thread_count = int(self.options.health_thread_count)
      self.nameservers.cache_dir = tempfile.gettempdir()

    self.UpdateStatus('Checking latest sanity reference')
    (primary_checks, secondary_checks, censor_tests) = config.GetLatestSanityChecks()
    if not self.options.enable_censorship_checks:
      censor_tests = []
    else:
      self.UpdateStatus('Censorship checks enabled: %s found.' % len(censor_tests))

    self.UpdateStatus('Checking nameserver health (%s servers)' % len(self.nameservers))
    self.nameservers.CheckHealth(primary_checks, secondary_checks, censor_tests=censor_tests)

  def PrepareBenchmark(self):
    """Setup the benchmark object with the appropriate dataset."""
    if len(self.nameservers) == 1:
      thread_count = 1
    else:
      thread_count = self.options.benchmark_thread_count

    self.bmark = benchmark.Benchmark(self.nameservers,
                                     query_count=self.options.query_count,
                                     run_count=self.options.run_count,
                                     thread_count=thread_count,
                                     status_callback=self.UpdateStatus)

  def RunBenchmark(self):
    """Run the benchmark."""
    results = self.bmark.Run(self.test_records)
    index = []

    if self.options.upload_results:
      connector = site_connector.SiteConnector(self.options, status_callback=self.UpdateStatus)
      index_hosts = connector.GetIndexHosts()
      if index_hosts:
        index = self.bmark.RunIndex(index_hosts)
      else:
        index = []
      self.DiscoverLocation()
      if len(self.nameservers) > 1:
        self.nameservers.RunPortBehaviorThreads()

    self.reporter = reporter.ReportGenerator(self.options, self.nameservers,
                                             results, index=index, geodata=self.geodata)

  def DiscoverLocation(self):
    if not getattr(self, 'geodata', None):
      self.geodata = geoip.GetGeoData()
      self.country = self.geodata.get('country_name', None)
    return self.geodata

  def RunAndOpenReports(self):
    """Run the benchmark and open up the report on completion."""
    self.RunBenchmark()
    best = self.reporter.BestOverallNameServer()
    self.CreateReports()
    if self.options.template == 'html':
      self.DisplayHtmlReport()
    if self.url:
      self.UpdateStatus('Complete! Your results: %s' % self.url)
    else:
      self.UpdateStatus('Complete! %s [%s] is the best.' % (best.name, best.ip))

  def CreateReports(self):
    """Create CSV & HTML reports for the latest run."""

    if self.options.output_file:
      self.report_path = self.options.output_file
    else:
      self.report_path = util.GenerateOutputFilename(self.options.template)

    if self.options.csv_file:
      self.csv_path = self.options_csv_file
    else:
      self.csv_path = util.GenerateOutputFilename('csv')

    if self.options.upload_results:
      # This is for debugging and transparency only.
      self.json_path = util.GenerateOutputFilename('js')
      self.UpdateStatus('Saving anonymized JSON to %s' % self.json_path)
      json_data = self.reporter.CreateJsonData()
      f = open(self.json_path, 'w')
      f.write(json_data)
      f.close()

      self.UpdateStatus('Uploading results to %s' % self.options.site_url)
      connector = site_connector.SiteConnector(self.options, status_callback=self.UpdateStatus)
      self.url, self.share_state = connector.UploadJsonResults(
          json_data,
          hide_results=self.options.hide_results
      )

      if self.url:
        self.UpdateStatus('Your sharing URL: %s (%s)' % (self.url, self.share_state))

    self.UpdateStatus('Saving report to %s' % self.report_path)
    f = open(self.report_path, 'w')
    self.reporter.CreateReport(format=self.options.template,
                               output_fp=f,
                               csv_path=self.csv_path,
                               sharing_url=self.url,
                               sharing_state=self.share_state)
    f.close()

    self.UpdateStatus('Saving detailed results to %s' % self.csv_path)
    self.reporter.SaveResultsToCsv(self.csv_path)

  def DisplayHtmlReport(self):
    self.UpdateStatus('Opening %s' % self.report_path)
    better_webbrowser.output = self.DebugMsg
    better_webbrowser.open(self.report_path)

<commit_msg>Remove redundant check message<commit_after># Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

"""A base user-interface workflow, to be inherited by UI modules."""

import tempfile

import benchmark
import better_webbrowser
import config
import data_sources
import geoip
import nameserver_list
import reporter
import site_connector
import util

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

class BaseUI(object):
  """Common methods for all UI implementations."""

  def __init__(self):
    self.SetupDataStructures()

  def SetupDataStructures(self):
    """Instead of requiring users to inherit __init__(), this sets up structures."""
    self.reporter = None
    self.nameservers = None
    self.bmark = None
    self.report_path = None
    self.csv_path = None
    self.geodata = None
    self.country = None
    self.sources = {}
    self.url = None
    self.share_state = None
    self.test_records = []

  def UpdateStatus(self, msg, **kwargs):
    """Update the little status message on the bottom of the window."""
    if hasattr(self, 'status_callback') and self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print msg

  def DebugMsg(self, message):
    self.UpdateStatus(message, debug=True)

  def LoadDataSources(self):
    self.data_src = data_sources.DataSources(status_callback=self.UpdateStatus)

  def PrepareTestRecords(self):
    """Figure out what data source a user wants, and create test_records."""
    if self.options.input_source:
      src_type = self.options.input_source
    else:
      src_type = self.data_src.GetBestSourceDetails()[0]
      self.options.input_source = src_type

    self.test_records = self.data_src.GetTestsFromSource(
        src_type,
        self.options.query_count,
        select_mode=self.options.select_mode
    )

  def PrepareNameServers(self):
    """Setup self.nameservers to have a list of healthy fast servers."""
    self.nameservers = nameserver_list.NameServers(
        self.supplied_ns,
        global_servers=self.global_ns,
        regional_servers=self.regional_ns,
        include_internal=self.include_internal,
        num_servers=self.options.num_servers,
        timeout=self.options.timeout,
        ping_timeout=self.options.ping_timeout,
        health_timeout=self.options.health_timeout,
        ipv6_only=self.options.ipv6_only,
        status_callback=self.UpdateStatus
    )
    if self.options.invalidate_cache:
      self.nameservers.InvalidateSecondaryCache()

    self.nameservers.cache_dir = tempfile.gettempdir()

    # Don't waste time checking the health of the only nameserver in the list.
    if len(self.nameservers) > 1:
      self.nameservers.thread_count = int(self.options.health_thread_count)
      self.nameservers.cache_dir = tempfile.gettempdir()

    self.UpdateStatus('Checking latest sanity reference')
    (primary_checks, secondary_checks, censor_tests) = config.GetLatestSanityChecks()
    if not self.options.enable_censorship_checks:
      censor_tests = []
    else:
      self.UpdateStatus('Censorship checks enabled: %s found.' % len(censor_tests))

    self.nameservers.CheckHealth(primary_checks, secondary_checks, censor_tests=censor_tests)

  def PrepareBenchmark(self):
    """Setup the benchmark object with the appropriate dataset."""
    if len(self.nameservers) == 1:
      thread_count = 1
    else:
      thread_count = self.options.benchmark_thread_count

    self.bmark = benchmark.Benchmark(self.nameservers,
                                     query_count=self.options.query_count,
                                     run_count=self.options.run_count,
                                     thread_count=thread_count,
                                     status_callback=self.UpdateStatus)

  def RunBenchmark(self):
    """Run the benchmark."""
    results = self.bmark.Run(self.test_records)
    index = []

    if self.options.upload_results:
      connector = site_connector.SiteConnector(self.options, status_callback=self.UpdateStatus)
      index_hosts = connector.GetIndexHosts()
      if index_hosts:
        index = self.bmark.RunIndex(index_hosts)
      else:
        index = []
      self.DiscoverLocation()
      if len(self.nameservers) > 1:
        self.nameservers.RunPortBehaviorThreads()

    self.reporter = reporter.ReportGenerator(self.options, self.nameservers,
                                             results, index=index, geodata=self.geodata)

  def DiscoverLocation(self):
    if not getattr(self, 'geodata', None):
      self.geodata = geoip.GetGeoData()
      self.country = self.geodata.get('country_name', None)
    return self.geodata

  def RunAndOpenReports(self):
    """Run the benchmark and open up the report on completion."""
    self.RunBenchmark()
    best = self.reporter.BestOverallNameServer()
    self.CreateReports()
    if self.options.template == 'html':
      self.DisplayHtmlReport()
    if self.url:
      self.UpdateStatus('Complete! Your results: %s' % self.url)
    else:
      self.UpdateStatus('Complete! %s [%s] is the best.' % (best.name, best.ip))

  def CreateReports(self):
    """Create CSV & HTML reports for the latest run."""

    if self.options.output_file:
      self.report_path = self.options.output_file
    else:
      self.report_path = util.GenerateOutputFilename(self.options.template)

    if self.options.csv_file:
      self.csv_path = self.options_csv_file
    else:
      self.csv_path = util.GenerateOutputFilename('csv')

    if self.options.upload_results:
      # This is for debugging and transparency only.
      self.json_path = util.GenerateOutputFilename('js')
      self.UpdateStatus('Saving anonymized JSON to %s' % self.json_path)
      json_data = self.reporter.CreateJsonData()
      f = open(self.json_path, 'w')
      f.write(json_data)
      f.close()

      self.UpdateStatus('Uploading results to %s' % self.options.site_url)
      connector = site_connector.SiteConnector(self.options, status_callback=self.UpdateStatus)
      self.url, self.share_state = connector.UploadJsonResults(
          json_data,
          hide_results=self.options.hide_results
      )

      if self.url:
        self.UpdateStatus('Your sharing URL: %s (%s)' % (self.url, self.share_state))

    self.UpdateStatus('Saving report to %s' % self.report_path)
    f = open(self.report_path, 'w')
    self.reporter.CreateReport(format=self.options.template,
                               output_fp=f,
                               csv_path=self.csv_path,
                               sharing_url=self.url,
                               sharing_state=self.share_state)
    f.close()

    self.UpdateStatus('Saving detailed results to %s' % self.csv_path)
    self.reporter.SaveResultsToCsv(self.csv_path)

  def DisplayHtmlReport(self):
    self.UpdateStatus('Opening %s' % self.report_path)
    better_webbrowser.output = self.DebugMsg
    better_webbrowser.open(self.report_path)

<|endoftext|>