<commit_before># -*- coding: utf-8 -*-
"""
==========================================================
Create a new coordinate class (for the Sagittarius stream)
==========================================================

This document describes in detail how to subclass and define a custom spherical
coordinate frame, as discussed in :ref:`astropy-coordinates-design` and the
docstring for `~astropy.coordinates.BaseCoordinateFrame`. In this example, we
will define a coordinate system defined by the plane of orbit of the Sagittarius
Dwarf Galaxy (hereafter Sgr; as defined in Majewski et al. 2003).  The Sgr
coordinate system is often referred to in terms of two angular coordinates,
:math:`\Lambda,B`.

To do this, wee need to define a subclass of
`~astropy.coordinates.BaseCoordinateFrame` that knows the names and units of the
coordinate system angles in each of the supported representations.  In this case
we support `~astropy.coordinates.SphericalRepresentation` with "Lambda" and
"Beta". Then we have to define the transformation from this coordinate system to
some other built-in system. Here we will use Galactic coordinates, represented
by the `~astropy.coordinates.Galactic` class.

See Also
--------

* Majewski et al. 2003, "A Two Micron All Sky Survey View of the Sagittarius
  Dwarf Galaxy. I. Morphology of the Sagittarius Core and Tidal Arms",
  http://arxiv.org/abs/astro-ph/0304198
* Law & Majewski 2010, "The Sagittarius Dwarf Galaxy: A Model for Evolution in a
  Triaxial Milky Way Halo", http://arxiv.org/abs/1003.1132
* David Law's Sgr info page http://www.stsci.edu/~dlaw/Sgr/

-------------------

*By: Adrian Price-Whelan, Erik Tollerud*

*License: BSD*

-------------------

"""

##############################################################################
# Make `print` work the same in all versions of Python, set up numpy,
# matplotlib, and use a nicer set of plot parameters:

from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
from astropy.visualization import astropy_mpl_style
plt.style.use(astropy_mpl_style)

##############################################################################
# Import the packages necessary for coordinates

from astropy.coordinates import frame_transform_graph
from astropy.coordinates.matrix_utilities import rotation_matrix
import astropy.coordinates as coord
import astropy.units as u

##############################################################################
# The first step is to create a new class, which we'll call
# ``Sagittarius`` and make it a subclass of
# `~astropy.coordinates.BaseCoordinateFrame`:

class Sagittarius(coord.BaseCoordinateFrame):
    """
    A Heliocentric spherical coordinate system defined by the orbit
    of the Sagittarius dwarf galaxy, as described in
        http://adsabs.harvard.edu/abs/2003ApJ...599.1082M
    and further explained in
        http://www.stsci.edu/~dlaw/Sgr/.

    Parameters
    ----------
    representation : `BaseRepresentation` or None
        A representation object or None to have no data (or use the other keywords)
    Lambda : `Angle`, optional, must be keyword
        The longitude-like angle corresponding to Sagittarius' orbit.
    Beta : `Angle`, optional, must be keyword
        The latitude-like angle corresponding to Sagittarius' orbit.
    distance : `Quantity`, optional, must be keyword
        The Distance for this object along the line-of-sight.

    """
    default_representation = coord.SphericalRepresentation

    frame_specific_representation_info = {
        'spherical': [coord.RepresentationMapping('lon', 'Lambda'),
                      coord.RepresentationMapping('lat', 'Beta'),
                      coord.RepresentationMapping('distance', 'distance')],
        'unitspherical': [coord.RepresentationMapping('lon', 'Lambda'),
                          coord.RepresentationMapping('lat', 'Beta')]
    }

##############################################################################
# Breaking this down line-by-line, we define the class as a subclass of
# `~astropy.coordinates.BaseCoordinateFrame`. Then we include a descriptive
# docstring.  The final lines are class-level attributes that specify the
# default representation for the data and mappings from the attribute names used
# by representation objects to the names that are to be used by ``Sagittarius``.
# In this case we override the names in the spherical representations but don't
# do anything with other representations like cartesian or cylindrical.
#
# Next we have to define the transformation from this coordinate system to some
# other built-in coordinate system; we will use Galactic coordinates. We can do
# this by defining functions that return transformation matrices, or by simply
# defining a function that accepts a coordinate and returns a new coordinate in
# the new system. We'll start by constructing the rotation matrix, using the
# ``rotation_matrix`` helper function:

SGR_PHI = np.radians(180+3.75) # Euler angles (from Law & Majewski 2010)
SGR_THETA = np.radians(90-13.46)
SGR_PSI = np.radians(180+14.111534)

# Generate the rotation matrix using the x-convention (see Goldstein)
D = rotation_matrix(SGR_PHI, "z", unit=u.radian)
C = rotation_matrix(SGR_THETA, "x", unit=u.radian)
B = rotation_matrix(SGR_PSI, "z", unit=u.radian)
SGR_MATRIX = np.array(B.dot(C).dot(D))

##############################################################################
# This is done at the module level, since it will be used by both the
# transformation from Sgr to Galactic as well as the inverse from Galactic to
# Sgr. Now we can define our first transformation function:

@frame_transform_graph.transform(coord.FunctionTransform, coord.Galactic, Sagittarius)
def galactic_to_sgr(gal_coord, sgr_frame):
    """ Compute the transformation from Galactic spherical to
        heliocentric Sgr coordinates.
    """

    l = np.atleast_1d(gal_coord.l.radian)
    b = np.atleast_1d(gal_coord.b.radian)

    X = np.cos(b)*np.cos(l)
    Y = np.cos(b)*np.sin(l)
    Z = np.sin(b)

    # Calculate X,Y,Z,distance in the Sgr system
    Xs, Ys, Zs = SGR_MATRIX.dot(np.array([X, Y, Z]))
    Zs = -Zs

    # Calculate the angular coordinates lambda,beta
    Lambda = np.arctan2(Ys,Xs)*u.radian
    Lambda[Lambda < 0] = Lambda[Lambda < 0] + 2.*np.pi*u.radian
    Beta = np.arcsin(Zs/np.sqrt(Xs*Xs+Ys*Ys+Zs*Zs))*u.radian

    return Sagittarius(Lambda=Lambda, Beta=Beta,
                       distance=gal_coord.distance)

##############################################################################
# The decorator ``@frame_transform_graph.transform(coord.FunctionTransform,
# coord.Galactic, Sagittarius)``  registers this function on the
# ``frame_transform_graph`` as a coordinate transformation. Inside the function,
# we simply follow the same procedure as detailed by David Law's `transformation
# code <http://www.astro.virginia.edu/~srm4n/Sgr/code.html>`_. Note that in this
# case, both coordinate systems are heliocentric, so we can simply copy any
# distance from the `~astropy.coordinates.Galactic` object.
#
# We then register the inverse transformation by using the transpose of the
# rotation matrix (which is faster to compute than the inverse):

@frame_transform_graph.transform(coord.FunctionTransform, Sagittarius, coord.Galactic)
def sgr_to_galactic(sgr_coord, gal_frame):
    """ Compute the transformation from heliocentric Sgr coordinates to
        spherical Galactic.
    """
    L = np.atleast_1d(sgr_coord.Lambda.radian)
    B = np.atleast_1d(sgr_coord.Beta.radian)

    Xs = np.cos(B)*np.cos(L)
    Ys = np.cos(B)*np.sin(L)
    Zs = np.sin(B)
    Zs = -Zs

    X, Y, Z = SGR_MATRIX.T.dot(np.array([Xs, Ys, Zs]))

    l = np.arctan2(Y,X)*u.radian
    b = np.arcsin(Z/np.sqrt(X*X+Y*Y+Z*Z))*u.radian

    l[l<=0] += 2*np.pi*u.radian

    return coord.Galactic(l=l, b=b, distance=sgr_coord.distance)

##############################################################################
# Now that we've registered these transformations between ``Sagittarius`` and
# `~astropy.coordinates.Galactic`, we can transform between *any* coordinate
# system and ``Sagittarius`` (as long as the other system has a path to
# transform to `~astropy.coordinates.Galactic`). For example, to transform from
# ICRS coordinates to ``Sagittarius``, we simply:

icrs = coord.ICRS(280.161732*u.degree, 11.91934*u.degree)
sgr = icrs.transform_to(Sagittarius)
print(sgr)

##############################################################################
# Or, to transform from the ``Sagittarius`` frame to ICRS coordinates:

sgr = Sagittarius(Lambda=np.linspace(0,2*np.pi,128)*u.radian,
                  Beta=np.zeros(128)*u.radian)
icrs = sgr.transform_to(coord.ICRS)

##############################################################################
# Plot the points in both coordinate systems:

fig,axes = plt.subplots(2, 1, figsize=(6,12),
                        subplot_kw={'projection': 'aitoff'})

axes[0].set_title("Sagittarius")
axes[0].scatter(sgr.Lambda.wrap_at(180*u.deg).radian, sgr.Beta.radian)

axes[1].set_title("ICRS")
axes[1].scatter(icrs.ra.wrap_at(180*u.deg).radian, icrs.dec.radian)

plt.show()
<commit_msg>switch to using matrix_product in sgr example<commit_after># -*- coding: utf-8 -*-
"""
==========================================================
Create a new coordinate class (for the Sagittarius stream)
==========================================================

This document describes in detail how to subclass and define a custom spherical
coordinate frame, as discussed in :ref:`astropy-coordinates-design` and the
docstring for `~astropy.coordinates.BaseCoordinateFrame`. In this example, we
will define a coordinate system defined by the plane of orbit of the Sagittarius
Dwarf Galaxy (hereafter Sgr; as defined in Majewski et al. 2003).  The Sgr
coordinate system is often referred to in terms of two angular coordinates,
:math:`\Lambda,B`.

To do this, wee need to define a subclass of
`~astropy.coordinates.BaseCoordinateFrame` that knows the names and units of the
coordinate system angles in each of the supported representations.  In this case
we support `~astropy.coordinates.SphericalRepresentation` with "Lambda" and
"Beta". Then we have to define the transformation from this coordinate system to
some other built-in system. Here we will use Galactic coordinates, represented
by the `~astropy.coordinates.Galactic` class.

See Also
--------

* Majewski et al. 2003, "A Two Micron All Sky Survey View of the Sagittarius
  Dwarf Galaxy. I. Morphology of the Sagittarius Core and Tidal Arms",
  http://arxiv.org/abs/astro-ph/0304198
* Law & Majewski 2010, "The Sagittarius Dwarf Galaxy: A Model for Evolution in a
  Triaxial Milky Way Halo", http://arxiv.org/abs/1003.1132
* David Law's Sgr info page http://www.stsci.edu/~dlaw/Sgr/

-------------------

*By: Adrian Price-Whelan, Erik Tollerud*

*License: BSD*

-------------------

"""

##############################################################################
# Make `print` work the same in all versions of Python, set up numpy,
# matplotlib, and use a nicer set of plot parameters:

from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
from astropy.visualization import astropy_mpl_style
plt.style.use(astropy_mpl_style)

##############################################################################
# Import the packages necessary for coordinates

from astropy.coordinates import frame_transform_graph
from astropy.coordinates.matrix_utilities import rotation_matrix, matrix_product
import astropy.coordinates as coord
import astropy.units as u

##############################################################################
# The first step is to create a new class, which we'll call
# ``Sagittarius`` and make it a subclass of
# `~astropy.coordinates.BaseCoordinateFrame`:

class Sagittarius(coord.BaseCoordinateFrame):
    """
    A Heliocentric spherical coordinate system defined by the orbit
    of the Sagittarius dwarf galaxy, as described in
        http://adsabs.harvard.edu/abs/2003ApJ...599.1082M
    and further explained in
        http://www.stsci.edu/~dlaw/Sgr/.

    Parameters
    ----------
    representation : `BaseRepresentation` or None
        A representation object or None to have no data (or use the other keywords)
    Lambda : `Angle`, optional, must be keyword
        The longitude-like angle corresponding to Sagittarius' orbit.
    Beta : `Angle`, optional, must be keyword
        The latitude-like angle corresponding to Sagittarius' orbit.
    distance : `Quantity`, optional, must be keyword
        The Distance for this object along the line-of-sight.

    """
    default_representation = coord.SphericalRepresentation

    frame_specific_representation_info = {
        'spherical': [coord.RepresentationMapping('lon', 'Lambda'),
                      coord.RepresentationMapping('lat', 'Beta'),
                      coord.RepresentationMapping('distance', 'distance')],
        'unitspherical': [coord.RepresentationMapping('lon', 'Lambda'),
                          coord.RepresentationMapping('lat', 'Beta')]
    }

##############################################################################
# Breaking this down line-by-line, we define the class as a subclass of
# `~astropy.coordinates.BaseCoordinateFrame`. Then we include a descriptive
# docstring.  The final lines are class-level attributes that specify the
# default representation for the data and mappings from the attribute names used
# by representation objects to the names that are to be used by ``Sagittarius``.
# In this case we override the names in the spherical representations but don't
# do anything with other representations like cartesian or cylindrical.
#
# Next we have to define the transformation from this coordinate system to some
# other built-in coordinate system; we will use Galactic coordinates. We can do
# this by defining functions that return transformation matrices, or by simply
# defining a function that accepts a coordinate and returns a new coordinate in
# the new system. We'll start by constructing the rotation matrix, using the
# ``rotation_matrix`` helper function:

SGR_PHI = np.radians(180+3.75) # Euler angles (from Law & Majewski 2010)
SGR_THETA = np.radians(90-13.46)
SGR_PSI = np.radians(180+14.111534)

# Generate the rotation matrix using the x-convention (see Goldstein)
D = rotation_matrix(SGR_PHI, "z", unit=u.radian)
C = rotation_matrix(SGR_THETA, "x", unit=u.radian)
B = rotation_matrix(SGR_PSI, "z", unit=u.radian)
SGR_MATRIX = matrix_product(B, C, D)

##############################################################################
# This is done at the module level, since it will be used by both the
# transformation from Sgr to Galactic as well as the inverse from Galactic to
# Sgr. Now we can define our first transformation function:

@frame_transform_graph.transform(coord.FunctionTransform, coord.Galactic, Sagittarius)
def galactic_to_sgr(gal_coord, sgr_frame):
    """ Compute the transformation from Galactic spherical to
        heliocentric Sgr coordinates.
    """

    l = np.atleast_1d(gal_coord.l.radian)
    b = np.atleast_1d(gal_coord.b.radian)

    X = np.cos(b)*np.cos(l)
    Y = np.cos(b)*np.sin(l)
    Z = np.sin(b)

    # Calculate X,Y,Z,distance in the Sgr system
    Xs, Ys, Zs = SGR_MATRIX.dot(np.array([X, Y, Z]))
    Zs = -Zs

    # Calculate the angular coordinates lambda,beta
    Lambda = np.arctan2(Ys,Xs)*u.radian
    Lambda[Lambda < 0] = Lambda[Lambda < 0] + 2.*np.pi*u.radian
    Beta = np.arcsin(Zs/np.sqrt(Xs*Xs+Ys*Ys+Zs*Zs))*u.radian

    return Sagittarius(Lambda=Lambda, Beta=Beta,
                       distance=gal_coord.distance)

##############################################################################
# The decorator ``@frame_transform_graph.transform(coord.FunctionTransform,
# coord.Galactic, Sagittarius)``  registers this function on the
# ``frame_transform_graph`` as a coordinate transformation. Inside the function,
# we simply follow the same procedure as detailed by David Law's `transformation
# code <http://www.astro.virginia.edu/~srm4n/Sgr/code.html>`_. Note that in this
# case, both coordinate systems are heliocentric, so we can simply copy any
# distance from the `~astropy.coordinates.Galactic` object.
#
# We then register the inverse transformation by using the transpose of the
# rotation matrix (which is faster to compute than the inverse):

@frame_transform_graph.transform(coord.FunctionTransform, Sagittarius, coord.Galactic)
def sgr_to_galactic(sgr_coord, gal_frame):
    """ Compute the transformation from heliocentric Sgr coordinates to
        spherical Galactic.
    """
    L = np.atleast_1d(sgr_coord.Lambda.radian)
    B = np.atleast_1d(sgr_coord.Beta.radian)

    Xs = np.cos(B)*np.cos(L)
    Ys = np.cos(B)*np.sin(L)
    Zs = np.sin(B)
    Zs = -Zs

    X, Y, Z = SGR_MATRIX.T.dot(np.array([Xs, Ys, Zs]))

    l = np.arctan2(Y,X)*u.radian
    b = np.arcsin(Z/np.sqrt(X*X+Y*Y+Z*Z))*u.radian

    l[l<=0] += 2*np.pi*u.radian

    return coord.Galactic(l=l, b=b, distance=sgr_coord.distance)

##############################################################################
# Now that we've registered these transformations between ``Sagittarius`` and
# `~astropy.coordinates.Galactic`, we can transform between *any* coordinate
# system and ``Sagittarius`` (as long as the other system has a path to
# transform to `~astropy.coordinates.Galactic`). For example, to transform from
# ICRS coordinates to ``Sagittarius``, we simply:

icrs = coord.ICRS(280.161732*u.degree, 11.91934*u.degree)
sgr = icrs.transform_to(Sagittarius)
print(sgr)

##############################################################################
# Or, to transform from the ``Sagittarius`` frame to ICRS coordinates:

sgr = Sagittarius(Lambda=np.linspace(0,2*np.pi,128)*u.radian,
                  Beta=np.zeros(128)*u.radian)
icrs = sgr.transform_to(coord.ICRS)

##############################################################################
# Plot the points in both coordinate systems:

fig,axes = plt.subplots(2, 1, figsize=(6,12),
                        subplot_kw={'projection': 'aitoff'})

axes[0].set_title("Sagittarius")
axes[0].scatter(sgr.Lambda.wrap_at(180*u.deg).radian, sgr.Beta.radian)

axes[1].set_title("ICRS")
axes[1].scatter(icrs.ra.wrap_at(180*u.deg).radian, icrs.dec.radian)

plt.show()
<|endoftext|>


<commit_before>"""
This module contains classes for representing ProblemStatistics object

For further information visit http://codeforces.com/api/help/objects#ProblemStatistics
"""

from . import BaseJsonObject

class ProblemStatistics(BaseJsonObject):
    """
    This class represents ProblemStatistics object

    For further information visit http://codeforces.com/api/help/objects#ProblemStatistics
    """

    def __init__(self, data=None):
        self._contest_id = None
        self._index = None
        self._solved_count = None

        super().__init__(data)

    def load_required_fields_from_dict(self, values):
        super().load_required_fields_from_dict(values)

        self.contest_id = values['contestId']
        self.index = values['index']
        self.solved_count = values['solvedCount']

    @property
    def contest_id(self):
        return self._contest_id

    @contest_id.setter
    def contest_id(self, value):
        assert isinstance(value, (int, str))

        if isinstance(value, str):
            value = int(value)

        self._contest_id = value

    @property
    def index(self):
        return self._index

    @index.setter
    def index(self, value):
        assert isinstance(value, str)
        self._index = value

    @property
    def solved_count(self):
        return self._solved_count

    @solved_count.setter
    def solved_count(self, value):
        assert isinstance(value, (int, str))

        if isinstance(value, str):
            value = int(value)

        self._solved_count = value
<commit_msg>Refactoring: Added doc comments Types, available for import specified explicitly Removed some redundant checks<commit_after>"""
This module contains classes for representing ProblemStatistics object

For further information visit http://codeforces.com/api/help/objects#ProblemStatistics
"""

from . import BaseJsonObject

__all__ = ['ProblemStatistics']

class ProblemStatistics(BaseJsonObject):
    """
    Represents a statistic data about a problem.

    For further information visit http://codeforces.com/api/help/objects#ProblemStatistics
    """

    def __init__(self, data=None):
        self._contest_id = None
        self._index = None
        self._solved_count = None

        super().__init__(data)

    def load_required_fields_from_dict(self, values):
        super().load_required_fields_from_dict(values)

        self.contest_id = values['contestId']
        self.index = values['index']
        self.solved_count = values['solvedCount']

    @property
    def contest_id(self):
        """
        Id of the contest, containing the problem.

        :return: Id or None if not initialized
        :rtype: int
        """
        return self._contest_id

    @contest_id.setter
    def contest_id(self, value):
        """
        Id of the contest, containing the problem.

        :param value: Id
        :type value: int or str
        """
        assert isinstance(value, (int, str))
        self._contest_id = int(value)

    @property
    def index(self):
        """
        Usually a letter or a letter, followed by a digit, that represent a problem index in a contest.

        :return: Index or None if not initialized
        :rtype: int
        """
        return self._index

    @index.setter
    def index(self, value):
        """
        Usually a letter or a letter, followed by a digit, that represent a problem index in a contest.

        :param value: Index
        :type value: str
        """
        assert isinstance(value, str)
        self._index = value

    @property
    def solved_count(self):
        """
        :return: Number of users, who solved the problem or None if not initialized
        :rtype: int
        """
        return self._solved_count

    @solved_count.setter
    def solved_count(self, value):
        """
        :param value: Number of users, who solved the problem.
        :type value: int or str
        """
        assert isinstance(value, (int, str))
        self._solved_count = int(value)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
    pyspecific.py
    ~~~~~~~~~~~~~

    Sphinx extension with Python doc-specific markup.

    :copyright: 2008, 2009 by Georg Brandl.
    :license: Python license.
"""

ISSUE_URI = 'http://bugs.python.org/issue%s'

from docutils import nodes, utils

# monkey-patch reST parser to disable alphabetic and roman enumerated lists
from docutils.parsers.rst.states import Body
Body.enum.converters['loweralpha'] = \
    Body.enum.converters['upperalpha'] = \
    Body.enum.converters['lowerroman'] = \
    Body.enum.converters['upperroman'] = lambda x: None

# monkey-patch HTML translator to give versionmodified paragraphs a class
def new_visit_versionmodified(self, node):
    self.body.append(self.starttag(node, 'p', CLASS=node['type']))
    text = versionlabels[node['type']] % node['version']
    if len(node):
        text += ': '
    else:
        text += '.'
    self.body.append('%s' % text)

from sphinx.writers.html import HTMLTranslator
from sphinx.locale import versionlabels
HTMLTranslator.visit_versionmodified = new_visit_versionmodified

# Support for marking up and linking to bugs.python.org issues

def issue_role(typ, rawtext, text, lineno, inliner, options={}, content=[]):
    issue = utils.unescape(text)
    text = 'issue ' + issue
    refnode = nodes.reference(text, text, refuri=ISSUE_URI % issue)
    return [refnode], []

# Support for marking up implementation details

from sphinx.util.compat import Directive

class ImplementationDetail(Directive):

    has_content = True
    required_arguments = 0
    optional_arguments = 1
    final_argument_whitespace = True

    def run(self):
        pnode = nodes.compound(classes=['impl-detail'])
        content = self.content
        content[0] = '**CPython implementation detail:** ' + content[0]
        self.state.nested_parse(content, self.content_offset, pnode)
        return [pnode]

# Support for building "topic help" for pydoc

pydoc_topic_labels = [
    'assert', 'assignment', 'atom-identifiers', 'atom-literals',
    'attribute-access', 'attribute-references', 'augassign', 'binary',
    'bitwise', 'bltin-code-objects', 'bltin-ellipsis-object',
    'bltin-file-objects', 'bltin-null-object', 'bltin-type-objects', 'booleans',
    'break', 'callable-types', 'calls', 'class', 'coercion-rules',
    'comparisons', 'compound', 'context-managers', 'continue', 'conversions',
    'customization', 'debugger', 'del', 'dict', 'dynamic-features', 'else',
    'exceptions', 'exec', 'execmodel', 'exprlists', 'floating', 'for',
    'formatstrings', 'function', 'global', 'id-classes', 'identifiers', 'if',
    'imaginary', 'import', 'in', 'integers', 'lambda', 'lists', 'naming',
    'numbers', 'numeric-types', 'objects', 'operator-summary', 'pass', 'power',
    'print', 'raise', 'return', 'sequence-methods', 'sequence-types',
    'shifting', 'slicings', 'specialattrs', 'specialnames',
    'string-conversions', 'string-methods', 'strings', 'subscriptions', 'truth',
    'try', 'types', 'typesfunctions', 'typesmapping', 'typesmethods',
    'typesmodules', 'typesseq', 'typesseq-mutable', 'unary', 'while', 'with',
    'yield'
]

from os import path
from time import asctime
from pprint import pformat
from docutils.io import StringOutput
from docutils.utils import new_document

from sphinx.builders import Builder
from sphinx.writers.text import TextWriter

class PydocTopicsBuilder(Builder):
    name = 'pydoc-topics'

    def init(self):
        self.topics = {}

    def get_outdated_docs(self):
        return 'all pydoc topics'

    def get_target_uri(self, docname, typ=None):
        return ''  # no URIs

    def write(self, *ignored):
        writer = TextWriter(self)
        for label in self.status_iterator(pydoc_topic_labels,
                                          'building topics... ',
                                          length=len(pydoc_topic_labels)):
            if label not in self.env.labels:
                self.warn('label %r not in documentation' % label)
                continue
            docname, labelid, sectname = self.env.labels[label]
            doctree = self.env.get_and_resolve_doctree(docname, self)
            document = new_document('
')
            document.append(doctree.ids[labelid])
            destination = StringOutput(encoding='utf-8')
            writer.write(document, destination)
            self.topics[label] = writer.output

    def finish(self):
        f = open(path.join(self.outdir, 'topics.py'), 'w')
        try:
            f.write('# Autogenerated by Sphinx on %s\n' % asctime())
            f.write('topics = ' + pformat(self.topics) + '\n')
        finally:
            f.close()

# Support for checking for suspicious markup

import suspicious

# Support for documenting Opcodes

import re
from sphinx import addnodes

opcode_sig_re = re.compile(r'(\w+(?:\+\d)?)\s*\((.*)\)')

def parse_opcode_signature(env, sig, signode):
    """Transform an opcode signature into RST nodes."""
    m = opcode_sig_re.match(sig)
    if m is None:
        raise ValueError
    opname, arglist = m.groups()
    signode += addnodes.desc_name(opname, opname)
    paramlist = addnodes.desc_parameterlist()
    signode += paramlist
    paramlist += addnodes.desc_parameter(arglist, arglist)
    return opname.strip()

def setup(app):
    app.add_role('issue', issue_role)
    app.add_directive('impl-detail', ImplementationDetail)
    app.add_builder(PydocTopicsBuilder)
    app.add_builder(suspicious.CheckSuspiciousMarkupBuilder)
    app.add_description_unit('opcode', 'opcode', '%s (opcode)',
                             parse_opcode_signature)
    app.add_description_unit('2to3fixer', '2to3fixer', '%s (2to3 fixer)')
<commit_msg>Make it more robust.<commit_after># -*- coding: utf-8 -*-
"""
    pyspecific.py
    ~~~~~~~~~~~~~

    Sphinx extension with Python doc-specific markup.

    :copyright: 2008, 2009 by Georg Brandl.
    :license: Python license.
"""

ISSUE_URI = 'http://bugs.python.org/issue%s'

from docutils import nodes, utils

# monkey-patch reST parser to disable alphabetic and roman enumerated lists
from docutils.parsers.rst.states import Body
Body.enum.converters['loweralpha'] = \
    Body.enum.converters['upperalpha'] = \
    Body.enum.converters['lowerroman'] = \
    Body.enum.converters['upperroman'] = lambda x: None

# monkey-patch HTML translator to give versionmodified paragraphs a class
def new_visit_versionmodified(self, node):
    self.body.append(self.starttag(node, 'p', CLASS=node['type']))
    text = versionlabels[node['type']] % node['version']
    if len(node):
        text += ': '
    else:
        text += '.'
    self.body.append('%s' % text)

from sphinx.writers.html import HTMLTranslator
from sphinx.locale import versionlabels
HTMLTranslator.visit_versionmodified = new_visit_versionmodified

# Support for marking up and linking to bugs.python.org issues

def issue_role(typ, rawtext, text, lineno, inliner, options={}, content=[]):
    issue = utils.unescape(text)
    text = 'issue ' + issue
    refnode = nodes.reference(text, text, refuri=ISSUE_URI % issue)
    return [refnode], []

# Support for marking up implementation details

from sphinx.util.compat import Directive

class ImplementationDetail(Directive):

    has_content = True
    required_arguments = 0
    optional_arguments = 1
    final_argument_whitespace = True

    def run(self):
        pnode = nodes.compound(classes=['impl-detail'])
        content = self.content
        add_text = nodes.strong('CPython implementation detail:',
                                'CPython implementation detail:')
        self.state.nested_parse(content, self.content_offset, pnode)
        if pnode.children and isinstance(pnode[0], nodes.paragraph):
            pnode[0].insert(0, add_text)
            pnode[0].insert(1, nodes.Text(' '))
        else:
            pnode.insert(0, nodes.paragraph('', '', add_text))
        return [pnode]

# Support for building "topic help" for pydoc

pydoc_topic_labels = [
    'assert', 'assignment', 'atom-identifiers', 'atom-literals',
    'attribute-access', 'attribute-references', 'augassign', 'binary',
    'bitwise', 'bltin-code-objects', 'bltin-ellipsis-object',
    'bltin-file-objects', 'bltin-null-object', 'bltin-type-objects', 'booleans',
    'break', 'callable-types', 'calls', 'class', 'coercion-rules',
    'comparisons', 'compound', 'context-managers', 'continue', 'conversions',
    'customization', 'debugger', 'del', 'dict', 'dynamic-features', 'else',
    'exceptions', 'exec', 'execmodel', 'exprlists', 'floating', 'for',
    'formatstrings', 'function', 'global', 'id-classes', 'identifiers', 'if',
    'imaginary', 'import', 'in', 'integers', 'lambda', 'lists', 'naming',
    'numbers', 'numeric-types', 'objects', 'operator-summary', 'pass', 'power',
    'print', 'raise', 'return', 'sequence-methods', 'sequence-types',
    'shifting', 'slicings', 'specialattrs', 'specialnames',
    'string-conversions', 'string-methods', 'strings', 'subscriptions', 'truth',
    'try', 'types', 'typesfunctions', 'typesmapping', 'typesmethods',
    'typesmodules', 'typesseq', 'typesseq-mutable', 'unary', 'while', 'with',
    'yield'
]

from os import path
from time import asctime
from pprint import pformat
from docutils.io import StringOutput
from docutils.utils import new_document

from sphinx.builders import Builder
from sphinx.writers.text import TextWriter

class PydocTopicsBuilder(Builder):
    name = 'pydoc-topics'

    def init(self):
        self.topics = {}

    def get_outdated_docs(self):
        return 'all pydoc topics'

    def get_target_uri(self, docname, typ=None):
        return ''  # no URIs

    def write(self, *ignored):
        writer = TextWriter(self)
        for label in self.status_iterator(pydoc_topic_labels,
                                          'building topics... ',
                                          length=len(pydoc_topic_labels)):
            if label not in self.env.labels:
                self.warn('label %r not in documentation' % label)
                continue
            docname, labelid, sectname = self.env.labels[label]
            doctree = self.env.get_and_resolve_doctree(docname, self)
            document = new_document('
')
            document.append(doctree.ids[labelid])
            destination = StringOutput(encoding='utf-8')
            writer.write(document, destination)
            self.topics[label] = writer.output

    def finish(self):
        f = open(path.join(self.outdir, 'topics.py'), 'w')
        try:
            f.write('# Autogenerated by Sphinx on %s\n' % asctime())
            f.write('topics = ' + pformat(self.topics) + '\n')
        finally:
            f.close()

# Support for checking for suspicious markup

import suspicious

# Support for documenting Opcodes

import re
from sphinx import addnodes

opcode_sig_re = re.compile(r'(\w+(?:\+\d)?)\s*\((.*)\)')

def parse_opcode_signature(env, sig, signode):
    """Transform an opcode signature into RST nodes."""
    m = opcode_sig_re.match(sig)
    if m is None:
        raise ValueError
    opname, arglist = m.groups()
    signode += addnodes.desc_name(opname, opname)
    paramlist = addnodes.desc_parameterlist()
    signode += paramlist
    paramlist += addnodes.desc_parameter(arglist, arglist)
    return opname.strip()

def setup(app):
    app.add_role('issue', issue_role)
    app.add_directive('impl-detail', ImplementationDetail)
    app.add_builder(PydocTopicsBuilder)
    app.add_builder(suspicious.CheckSuspiciousMarkupBuilder)
    app.add_description_unit('opcode', 'opcode', '%s (opcode)',
                             parse_opcode_signature)
    app.add_description_unit('2to3fixer', '2to3fixer', '%s (2to3 fixer)')
<|endoftext|>


<commit_before># Copyright (c) 2012, GPy authors (see AUTHORS.txt).
# Licensed under the BSD 3-clause license (see LICENSE.txt)

import numpy as np
import pylab as pb
import sys, pdb
from .. import kern
from ..core import model
from ..util.linalg import pdinv, PCA
from GP_regression import GP_regression

class GPLVM(GP_regression):
    """
    Gaussian Process Latent Variable Model

    :param Y: observed data
    :type Y: np.ndarray
    :param Q: latent dimensionality
    :type Q: int
    :param init: initialisation method for the latent space
    :type init: 'PCA'|'random'

    """
    def __init__(self, Y, Q, init='PCA', **kwargs):
        X = self.initialise_latent(init, Q, Y)
        GP_regression.__init__(self, X, Y, **kwargs)

    def initialise_latent(self, init, Q, Y):
        if init == 'PCA':
            return PCA(Y, Q)[0]
        else:
            return np.random.randn(Y.shape[0], Q)

    def get_param_names(self):
        return (sum([['X_%i_%i'%(n,q) for n in range(self.N)] for q in range(self.Q)],[])
                + self.kern.extract_param_names())

    def get_param(self):
        return np.hstack((self.X.flatten(), self.kern.extract_param()))

    def set_param(self,x):
        self.X = x[:self.X.size].reshape(self.N,self.Q).copy()
        GP_regression.set_param(self, x[self.X.size:])

    def log_likelihood_gradients(self):
        dL_dK = self.dL_dK()

        dL_dtheta = self.kern.dK_dtheta(dL_dK,self.X)
        dL_dX = 2*self.kern.dK_dX(dL_dK,self.X)

        return np.hstack((dL_dX.flatten(),dL_dtheta))

    def plot(self):
        assert self.Y.shape[1]==2
        pb.scatter(self.Y[:,0],self.Y[:,1],40,self.X[:,0].copy(),linewidth=0)
        Xnew = np.linspace(self.X.min(),self.X.max(),200)[:,None]
        mu, var = self.predict(Xnew)
        pb.plot(mu[:,0], mu[:,1],'k',linewidth=1.5)

    def plot_latent(self):
        raise NotImplementedError
<commit_msg>GPLVM accepts an initial value for X (in case you don't want to use the default random/PCA init)<commit_after># Copyright (c) 2012, GPy authors (see AUTHORS.txt).
# Licensed under the BSD 3-clause license (see LICENSE.txt)

import numpy as np
import pylab as pb
import sys, pdb
from .. import kern
from ..core import model
from ..util.linalg import pdinv, PCA
from GP_regression import GP_regression

class GPLVM(GP_regression):
    """
    Gaussian Process Latent Variable Model

    :param Y: observed data
    :type Y: np.ndarray
    :param Q: latent dimensionality
    :type Q: int
    :param init: initialisation method for the latent space
    :type init: 'PCA'|'random'

    """
    def __init__(self, Y, Q, init='PCA', X = None, **kwargs):
        if X is None:
            X = self.initialise_latent(init, Q, Y)
        GP_regression.__init__(self, X, Y, **kwargs)

    def initialise_latent(self, init, Q, Y):
        if init == 'PCA':
            return PCA(Y, Q)[0]
        else:
            return np.random.randn(Y.shape[0], Q)

    def get_param_names(self):
        return (sum([['X_%i_%i'%(n,q) for n in range(self.N)] for q in range(self.Q)],[])
                + self.kern.extract_param_names())

    def get_param(self):
        return np.hstack((self.X.flatten(), self.kern.extract_param()))

    def set_param(self,x):
        self.X = x[:self.X.size].reshape(self.N,self.Q).copy()
        GP_regression.set_param(self, x[self.X.size:])

    def log_likelihood_gradients(self):
        dL_dK = self.dL_dK()

        dL_dtheta = self.kern.dK_dtheta(dL_dK,self.X)
        dL_dX = 2*self.kern.dK_dX(dL_dK,self.X)

        return np.hstack((dL_dX.flatten(),dL_dtheta))

    def plot(self):
        assert self.Y.shape[1]==2
        pb.scatter(self.Y[:,0],self.Y[:,1],40,self.X[:,0].copy(),linewidth=0)
        Xnew = np.linspace(self.X.min(),self.X.max(),200)[:,None]
        mu, var = self.predict(Xnew)
        pb.plot(mu[:,0], mu[:,1],'k',linewidth=1.5)

    def plot_latent(self):
        raise NotImplementedError
<|endoftext|>


<commit_before>__version__ = '0.4.930'
<commit_msg>Automatic patch number increase [ci skip]<commit_after>__version__ = '0.4.931'
<|endoftext|>


<commit_before>__version__ = '0.4.770'
<commit_msg>Automatic patch number increase [ci skip]<commit_after>__version__ = '0.4.773'
<|endoftext|>


<commit_before>"""Thumber Library
Author Hannu Valtonen"""
import Image
import struct
import StringIO
try:
    import json
except:
    import simplejson as json

INDEX_VERSION = 1

class Thumber(object):
    """Thumber librarys main class, use this if you want to use everything"""
    def __init__(self, thumbnail_sizes = None, reserved_keys = None, plugins = None):
        if thumbnail_sizes:
            self.thumbnail_sizes = thumbnail_sizes
        else:
            self.thumbnail_sizes = ((128, 128), (64, 64), (32, 32)) #default sizes
        if plugins:
            self.plugins = plugins
        else:
            self.plugins = [self.create_thumbnails]
        self.thumb_indexer = ThumberIndex(reserved_keys)

    def create_thumbs_and_index(self, data_blob, extra_reserved_keys = None):
        """Create thumbnails and an index, and add possible additional keys to the file index"""
        result_dict = {}
        for plugin in self.plugins:
            result_dict.update(plugin(data_blob))
        return self.thumb_indexer.create_thumbnail_blob_with_index(result_dict, extra_reserved_keys = extra_reserved_keys)

    def create_thumbnails(self, data_blob):
        """Create required thumbnails, sizes are set in object instance creation time"""
        image = Image.open(StringIO.StringIO(data_blob))
        result_dict = {}
        for thumbnail_size in self.thumbnail_sizes:
            image.thumbnail(thumbnail_size, Image.ANTIALIAS)
            file_buffer = StringIO.StringIO()
            if image.mode != "RGB":
                image = image.convert("RGB")
            image.save(file_buffer, format = 'JPEG')
            result_dict[str(thumbnail_size[0]) + "x" + str(thumbnail_size[1]) + "xjpg"] = file_buffer.getvalue()
        return result_dict

class ThumberIndex(object):
    """Class for creating files with n thumbnails and an access index"""
    def __init__(self, reserved_keys):
        if reserved_keys:
            self.reserved_keys = reserved_keys
        else:
            self.reserved_keys = []

    def create_thumbnail_blob_with_index(self, result_dict, extra_reserved_keys = None):
        """Expects a dict like {'128x128xjpg': thumbnail_data, '64x64xjpg': thumbnail2_data, 'my_reserved_key': 1}"""
        header, data, current_offset = {'version': INDEX_VERSION}, [], 0
        for k, v in result_dict.items():
            if k not in self.reserved_keys:
                offsets = [current_offset, current_offset + len(v)]
                header[k] = offsets
                current_offset += len(v)
                data.append(v)
            else:
                header[k] = v
        json_header = json.dumps(header)
        data_list = [struct.pack("H", len(json_header)), json_header]
        data_list.extend(data)
        return ''.join(data_list)

    def read_thumbnail_blob_with_index(self, data_blob, thumbnail_key = None, extra_reserved_keys = None):
        """Expects a data blob created by create_thumbnail_file_with_index
        If thumbnail_key has been given, returns the blob in question
        otherwise return a thumbnail_dict with all available thumbnail sizes
        """
        header_length = int(struct.unpack("H", data_blob[:2])[0])
        header = json.loads(data_blob[2:2 + header_length])
        total_header_length = header_length + 2

        if thumbnail_key:
            start_offset, stop_offset = header[thumbnail_key]
            return data_blob[total_header_length + start_offset:total_header_length + stop_offset]

        return_dict = {}
        for k, v in header.items():
            if k not in self.reserved_keys:
                return_dict[k] = data_blob[total_header_length + v[0]:total_header_length + v[1]]
            else:
                return_dict[k] = v
        return return_dict

def help_msg():
    """Help msg function for thumber"""
    print "Help:"
    print "    thumber store <input_image_file> <output_thumbnail_index_file>"
    print "    thumber load <input_thumbnail_index> <size> <output_thumbnail_image.jpg>"
    print "Thumber Copyright Hannu Valtonen 2011"

def main():
    """Entry point for thumber console script"""
    import sys
    if len(sys.argv) >= 4:
        a = Thumber()
        input_data = file(sys.argv[2], "rb").read()
        if sys.argv[1] == "store":
            data_blob = a.create_thumbs_and_index(input_data)
            output_filename = sys.argv[3]
        elif sys.argv[1] == "load" and len(sys.argv) == 5:
            data_blob = a.thumb_indexer.read_thumbnail_blob_with_index(input_data, sys.argv[3] + "xjpg")
            output_filename = sys.argv[4]
        else:
            help_msg()
            sys.exit(0)
        output_file = file(output_filename, "wb")
        output_file.write(data_blob)
        output_file.close()
    else:
        help_msg()

if __name__ == "__main__":
    main()
<commit_msg>Some work to support extra user set keys in the index.<commit_after>"""Thumber Library
Author Hannu Valtonen"""
import Image
import struct
import StringIO
try:
    import json
except:
    import simplejson as json

INDEX_VERSION = 1

class Thumber(object):
    """Thumber librarys main class, use this if you want to use everything"""
    def __init__(self, thumbnail_sizes = None, reserved_keys = None, plugins = None):
        if thumbnail_sizes:
            self.thumbnail_sizes = thumbnail_sizes
        else:
            self.thumbnail_sizes = ((128, 128), (64, 64), (32, 32)) #default sizes
        if plugins:
            self.plugins = plugins
        else:
            self.plugins = [self.create_thumbnails]
        self.thumb_indexer = ThumberIndex(reserved_keys)

    def create_thumbs_and_index(self, data_blob, extra_keys_dict = None):
        """Create thumbnails and an index, and add possible additional keys to the file index"""
        result_dict = {}
        for plugin in self.plugins:
            result_dict.update(plugin(data_blob))
        return self.thumb_indexer.create_thumbnail_blob_with_index(result_dict, extra_keys_dict = extra_keys_dict)

    def create_thumbnails(self, data_blob):
        """Create required thumbnails, sizes are set in object instance creation time"""
        image = Image.open(StringIO.StringIO(data_blob))
        result_dict = {}
        for thumbnail_size in self.thumbnail_sizes:
            image.thumbnail(thumbnail_size, Image.ANTIALIAS)
            file_buffer = StringIO.StringIO()
            if image.mode != "RGB":
                image = image.convert("RGB")
            image.save(file_buffer, format = 'JPEG')
            result_dict[str(thumbnail_size[0]) + "x" + str(thumbnail_size[1]) + "xjpg"] = file_buffer.getvalue()
        return result_dict

class ThumberIndex(object):
    """Class for creating files with n thumbnails and an access index"""
    def __init__(self, reserved_keys):
        if reserved_keys:
            self.reserved_keys = reserved_keys
        else:
            self.reserved_keys = []

    def create_thumbnail_blob_with_index(self, result_dict, extra_keys_dict = None):
        """Expects a dict like {'128x128xjpg': thumbnail_data, '64x64xjpg': thumbnail2_data, 'my_reserved_key': 1}"""
        header, data, current_offset = {'version': INDEX_VERSION}, [], 0
        if extra_keys_dict:
            result_dict.update(extra_keys_dict)
        else:
            extra_keys_dict = {}

        for k, v in result_dict.items():
            if k not in self.reserved_keys + extra_keys_dict.keys():
                offsets = [current_offset, current_offset + len(v)]
                header[k] = offsets
                current_offset += len(v)
                data.append(v)
            else:
                header[k] = v
        json_header = json.dumps(header)
        data_list = [struct.pack("H", len(json_header)), json_header]
        data_list.extend(data)
        return ''.join(data_list)

    def read_thumbnail_blob_with_index(self, data_blob, thumbnail_key = None, extra_reserved_keys = None):
        """Expects a data blob created by create_thumbnail_file_with_index
        If thumbnail_key has been given, returns the blob in question
        otherwise return a thumbnail_dict with all available thumbnail sizes
        """
        if not extra_reserved_keys:
            extra_reserved_keys = []

        header_length = int(struct.unpack("H", data_blob[:2])[0])
        header = json.loads(data_blob[2:2 + header_length])
        total_header_length = header_length + 2

        if thumbnail_key:
            start_offset, stop_offset = header[thumbnail_key]
            return data_blob[total_header_length + start_offset:total_header_length + stop_offset]

        return_dict = {}
        for k, v in header.items():
            if k not in self.reserved_keys + extra_reserved_keys:
                return_dict[k] = data_blob[total_header_length + v[0]:total_header_length + v[1]]
            else:
                return_dict[k] = v
        return return_dict

def help_msg():
    """Help msg function for thumber"""
    print "Help:"
    print "    thumber store <input_image_file> <output_thumbnail_index_file>"
    print "    thumber load <input_thumbnail_index> <size> <output_thumbnail_image.jpg>"
    print "Thumber Copyright Hannu Valtonen 2011"

def main():
    """Entry point for thumber console script"""
    import sys
    if len(sys.argv) >= 4:
        a = Thumber()
        input_data = file(sys.argv[2], "rb").read()
        if sys.argv[1] == "store":
            data_blob = a.create_thumbs_and_index(input_data)
            output_filename = sys.argv[3]
        elif sys.argv[1] == "load" and len(sys.argv) == 5:
            data_blob = a.thumb_indexer.read_thumbnail_blob_with_index(input_data, sys.argv[3] + "xjpg")
            output_filename = sys.argv[4]
        else:
            help_msg()
            sys.exit(0)
        output_file = file(output_filename, "wb")
        output_file.write(data_blob)
        output_file.close()
    else:
        help_msg()

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


<commit_before>from __future__ import print_function
import json
import os
import time
from datetime import datetime

def time_string(t):
    thxtime_path = os.environ['THXTIME_PATH']
    datafile = os.path.join(thxtime_path, 'data/thxdata.json')

    with open(datafile, 'rb') as data_fh:
        data = json.load(data_fh)
        return render_time(t, data)

def render_time(t, data):
    current_time_24 = t.strftime("%H:%M")
    current_time_12 = t.strftime("%I:%M")

    if current_time_24 in data:
        return data[current_time_24]
    elif current_time_12 in data:
        return data[current_time_12]
    else:
        return current_time_24

def thx_time(t):
    return time_string(t)

def main():
    now = datetime.fromtimestamp(time.time()).time()
    print(thx_time(now))

if __name__ == '__main__':
    main()
<commit_msg>updates<commit_after>from __future__ import print_function
import json
import os
import time
from datetime import datetime

def time_string(t):
    thxtime_path = os.environ['THXTIME_PATH']
    datafile = os.path.join(thxtime_path, 'data/thxdata.json')

    with open(datafile, 'rb') as data_fh:
        data = json.load(data_fh)

    return render_time(t, data)

def render_time(t, data):
    current_time_24 = t.strftime("%H:%M")
    current_time_12 = t.strftime("%I:%M")

    if current_time_24 in data:
        return data[current_time_24]
    elif current_time_12 in data:
        return data[current_time_12]
    else:
        return current_time_24

def thx_time(t):
    return time_string(t)

def thx_time_extended(t):
    time_str = '{} {}'.format(
        datetime.utcnow().strftime('%H:%M'),
        time_string(t)
    )
    return time_str

def main():
    now = datetime.fromtimestamp(time.time()).time()
    print(thx_time_extended(now))

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


<commit_before>import os
import json
import socket
import Queue
import sys
import select
import time

ssl = False
try:
    import ssl
except ImportError:
    pass

import vim

from common import cert, msg, shared as G, utils

class AgentConnection(object):
    MAX_RETRIES = 20
    INITIAL_RECONNECT_DELAY = 500

    ''' Simple chat server using select '''
    def __init__(self, owner, workspace, host=None, port=None, secure=True, on_auth=None, Protocol=None):
        self.sock_q = Queue.Queue()
        self.sock = None
        self.net_buf = ''
        self.reconnect_delay = self.INITIAL_RECONNECT_DELAY
        self.reconnect_timeout = None
        self.username = G.USERNAME
        self.secret = G.SECRET
        self.authed = False
        G.JOINED_WORKSPACE = False
        self.host = host or G.DEFAULT_HOST
        self.port = port or G.DEFAULT_PORT
        self.secure = secure
        self.owner = owner
        self.workspace = workspace
        self.retries = self.MAX_RETRIES
        self._on_auth = on_auth
        self.empty_selects = 0
        self.workspace_info = {}
        self.protocol = Protocol(self)
        self.cert_path = os.path.join(G.BASE_DIR, 'startssl-ca.pem')

    def tick(self):
        self.protocol.push()
        self.select()

    def send_get_buf(self, buf_id):
        req = {
            'name': 'get_buf',
            'id': buf_id
        }
        self.put(req)

    def send_auth(self):
        # TODO: we shouldn't throw away all of this
        self.sock_q = Queue.Queue()
        self.put({
            'username': self.username,
            'secret': self.secret,
            'room': self.workspace,
            'room_owner': self.owner,
            'client': self.protocol.CLIENT,
            'platform': sys.platform,
            'version': G.__VERSION__
        })

    def send_saved(self, buf_id):
        self.put({'name': 'saved', 'id': buf_id})

    def send_msg(self, msg):
        self.put({'name': 'msg', 'data': msg})
        self.protocol.chat(self.username, time.time(), msg, True)

    def on_auth(self):
        self.authed = True
        G.JOINED_WORKSPACE = True
        self.retries = self.MAX_RETRIES
        msg.log('Successfully joined workspace %s/%s' % (self.owner, self.workspace))
        if self._on_auth:
            self._on_auth(self)
            self._on_auth = None

    def stop(self, log=True):
        if log:
            msg.log('Disconnecting from workspace %s/%s' % (self.owner, self.workspace))
        utils.cancel_timeout(self.reconnect_timeout)
        self.reconnect_timeout = None
        try:
            self.retries = -1
            self.sock.shutdown(2)
            self.sock.close()
        except Exception:
            return False
        if log:
            msg.log('Disconnected.')
        return True

    def is_ready(self):
        return self.authed

    def put(self, item):
        if not item:
            return
        self.sock_q.put(json.dumps(item) + '\n')
        qsize = self.sock_q.qsize()
        if qsize > 0:
            msg.debug('%s items in q' % qsize)

    def reconnect(self):
        if self.reconnect_timeout:
            return
        try:
            self.sock.close()
        except Exception:
            pass
        self.workspace_info = {}
        self.net_buf = ''
        self.sock = None
        self.authed = False
        G.JOINED_WORKSPACE = False
        self.reconnect_delay *= 1.5
        if self.reconnect_delay > 10000:
            self.reconnect_delay = 10000
        if self.retries > 0:
            msg.log('Floobits: Reconnecting in %sms' % self.reconnect_delay)
            self.reconnect_timeout = utils.set_timeout(self.connect, int(self.reconnect_delay))
        elif self.retries == 0:
            msg.error('Floobits Error! Too many reconnect failures. Giving up.')
        self.retries -= 1

    def connect(self, cb=None):
        self.stop(False)
        self.empty_selects = 0
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if self.secure:
            if ssl:
                with open(self.cert_path, 'wb') as cert_fd:
                    cert_fd.write(cert.CA_CERT.encode('utf-8'))
                self.sock = ssl.wrap_socket(self.sock, ca_certs=self.cert_path, cert_reqs=ssl.CERT_REQUIRED)
            else:
                msg.debug('No SSL module found. Connection will not be encrypted.')
                if self.port == G.DEFAULT_PORT:
                    self.port = 3148  # plaintext port
        msg.debug('Connecting to %s:%s' % (self.host, self.port))
        try:
            self.sock.connect((self.host, self.port))
            if self.secure and ssl:
                self.sock.do_handshake()
        except socket.error as e:
            msg.error('Error connecting:', e)
            self.reconnect()
            return
        self.sock.setblocking(0)
        msg.debug('Connected!')
        self.reconnect_delay = self.INITIAL_RECONNECT_DELAY
        self.send_auth()
        if cb:
            cb()

    def _get_from_queue(self):
        while True:
            try:
                yield self.sock_q.get_nowait()
            except Queue.Empty:
                break

    def handle(self, req):
        self.net_buf += req
        new_data = False
        while True:
            before, sep, after = self.net_buf.partition('\n')
            if not sep:
                break
            try:
                data = json.loads(before)
            except Exception as e:
                msg.error('Unable to parse json:', e)
                msg.error('Data:', before)
                raise e
            self.protocol.handle(data)
            new_data = True
            self.net_buf = after
        #XXX move to protocol :(
        if new_data and not G.ASYNC:
            vim.command('redraw')

    def select(self):
        if not self.sock:
            msg.error('select(): No socket.')
            return self.reconnect()

        try:
            _in, _out, _except = select.select([self.sock], [self.sock], [self.sock], 0)
        except (select.error, socket.error, Exception) as e:
            msg.error('Error in select(): %s' % str(e))
            return self.reconnect()

        if _except:
            msg.error('Socket error')
            return self.reconnect()

        if _in:
            buf = ''
            while True:
                try:
                    d = self.sock.recv(4096)
                    if not d:
                        break
                    buf += d
                except (socket.error, TypeError):
                    break
            if buf:
                self.empty_selects = 0
                self.handle(buf)
            else:
                self.empty_selects += 1
                if self.empty_selects > 10:
                    msg.error('No data from sock.recv() {0} times.'.format(self.empty_selects))
                    return self.reconnect()

        if _out:
            for p in self._get_from_queue():
                if p is None:
                    self.sock_q.task_done()
                    continue
                try:
                    self.sock.sendall(p)
                    self.sock_q.task_done()
                except Exception as e:
                    msg.error('Couldn\'t write to socket: %s' % str(e))
                    return self.reconnect()
<commit_msg>Always redraw.<commit_after>import os
import json
import socket
import Queue
import sys
import select
import time

ssl = False
try:
    import ssl
except ImportError:
    pass

import vim

from common import cert, msg, shared as G, utils

class AgentConnection(object):
    MAX_RETRIES = 20
    INITIAL_RECONNECT_DELAY = 500

    ''' Simple chat server using select '''
    def __init__(self, owner, workspace, host=None, port=None, secure=True, on_auth=None, Protocol=None):
        self.sock_q = Queue.Queue()
        self.sock = None
        self.net_buf = ''
        self.reconnect_delay = self.INITIAL_RECONNECT_DELAY
        self.reconnect_timeout = None
        self.username = G.USERNAME
        self.secret = G.SECRET
        self.authed = False
        G.JOINED_WORKSPACE = False
        self.host = host or G.DEFAULT_HOST
        self.port = port or G.DEFAULT_PORT
        self.secure = secure
        self.owner = owner
        self.workspace = workspace
        self.retries = self.MAX_RETRIES
        self._on_auth = on_auth
        self.empty_selects = 0
        self.workspace_info = {}
        self.protocol = Protocol(self)
        self.cert_path = os.path.join(G.BASE_DIR, 'startssl-ca.pem')

    def tick(self):
        self.protocol.push()
        self.select()

    def send_get_buf(self, buf_id):
        req = {
            'name': 'get_buf',
            'id': buf_id
        }
        self.put(req)

    def send_auth(self):
        # TODO: we shouldn't throw away all of this
        self.sock_q = Queue.Queue()
        self.put({
            'username': self.username,
            'secret': self.secret,
            'room': self.workspace,
            'room_owner': self.owner,
            'client': self.protocol.CLIENT,
            'platform': sys.platform,
            'version': G.__VERSION__
        })

    def send_saved(self, buf_id):
        self.put({'name': 'saved', 'id': buf_id})

    def send_msg(self, msg):
        self.put({'name': 'msg', 'data': msg})
        self.protocol.chat(self.username, time.time(), msg, True)

    def on_auth(self):
        self.authed = True
        G.JOINED_WORKSPACE = True
        self.retries = self.MAX_RETRIES
        msg.log('Successfully joined workspace %s/%s' % (self.owner, self.workspace))
        if self._on_auth:
            self._on_auth(self)
            self._on_auth = None

    def stop(self, log=True):
        if log:
            msg.log('Disconnecting from workspace %s/%s' % (self.owner, self.workspace))
        utils.cancel_timeout(self.reconnect_timeout)
        self.reconnect_timeout = None
        try:
            self.retries = -1
            self.sock.shutdown(2)
            self.sock.close()
        except Exception:
            return False
        if log:
            msg.log('Disconnected.')
        return True

    def is_ready(self):
        return self.authed

    def put(self, item):
        if not item:
            return
        self.sock_q.put(json.dumps(item) + '\n')
        qsize = self.sock_q.qsize()
        if qsize > 0:
            msg.debug('%s items in q' % qsize)

    def reconnect(self):
        if self.reconnect_timeout:
            return
        try:
            self.sock.close()
        except Exception:
            pass
        self.workspace_info = {}
        self.net_buf = ''
        self.sock = None
        self.authed = False
        G.JOINED_WORKSPACE = False
        self.reconnect_delay *= 1.5
        if self.reconnect_delay > 10000:
            self.reconnect_delay = 10000
        if self.retries > 0:
            msg.log('Floobits: Reconnecting in %sms' % self.reconnect_delay)
            self.reconnect_timeout = utils.set_timeout(self.connect, int(self.reconnect_delay))
        elif self.retries == 0:
            msg.error('Floobits Error! Too many reconnect failures. Giving up.')
        self.retries -= 1

    def connect(self, cb=None):
        self.stop(False)
        self.empty_selects = 0
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if self.secure:
            if ssl:
                with open(self.cert_path, 'wb') as cert_fd:
                    cert_fd.write(cert.CA_CERT.encode('utf-8'))
                self.sock = ssl.wrap_socket(self.sock, ca_certs=self.cert_path, cert_reqs=ssl.CERT_REQUIRED)
            else:
                msg.debug('No SSL module found. Connection will not be encrypted.')
                if self.port == G.DEFAULT_PORT:
                    self.port = 3148  # plaintext port
        msg.debug('Connecting to %s:%s' % (self.host, self.port))
        try:
            self.sock.connect((self.host, self.port))
            if self.secure and ssl:
                self.sock.do_handshake()
        except socket.error as e:
            msg.error('Error connecting:', e)
            self.reconnect()
            return
        self.sock.setblocking(0)
        msg.debug('Connected!')
        self.reconnect_delay = self.INITIAL_RECONNECT_DELAY
        self.send_auth()
        if cb:
            cb()

    def _get_from_queue(self):
        while True:
            try:
                yield self.sock_q.get_nowait()
            except Queue.Empty:
                break

    def handle(self, req):
        self.net_buf += req
        new_data = False
        while True:
            before, sep, after = self.net_buf.partition('\n')
            if not sep:
                break
            try:
                data = json.loads(before)
            except Exception as e:
                msg.error('Unable to parse json:', e)
                msg.error('Data:', before)
                raise e
            self.protocol.handle(data)
            new_data = True
            self.net_buf = after
        #XXX move to protocol :(
        if new_data:
            vim.command('redraw')

    def select(self):
        if not self.sock:
            msg.error('select(): No socket.')
            return self.reconnect()

        try:
            _in, _out, _except = select.select([self.sock], [self.sock], [self.sock], 0)
        except (select.error, socket.error, Exception) as e:
            msg.error('Error in select(): %s' % str(e))
            return self.reconnect()

        if _except:
            msg.error('Socket error')
            return self.reconnect()

        if _in:
            buf = ''
            while True:
                try:
                    d = self.sock.recv(4096)
                    if not d:
                        break
                    buf += d
                except (socket.error, TypeError):
                    break
            if buf:
                self.empty_selects = 0
                self.handle(buf)
            else:
                self.empty_selects += 1
                if self.empty_selects > 10:
                    msg.error('No data from sock.recv() {0} times.'.format(self.empty_selects))
                    return self.reconnect()

        if _out:
            for p in self._get_from_queue():
                if p is None:
                    self.sock_q.task_done()
                    continue
                try:
                    self.sock.sendall(p)
                    self.sock_q.task_done()
                except Exception as e:
                    msg.error('Couldn\'t write to socket: %s' % str(e))
                    return self.reconnect()
<|endoftext|>


<commit_before># define locations to store the rendered data

from boto import connect_s3
from boto.s3.bucket import Bucket
from builtins import range
from future.utils import raise_from
import md5
from ModestMaps.Core import Coordinate
import os
from tilequeue.metatile import metatiles_are_equal
from tilequeue.format import zip_format
import random
import threading
import time

def calc_hash(s):
    m = md5.new()
    m.update(s)
    md5_hash = m.hexdigest()
    return md5_hash[:5]

def s3_tile_key(date, path, layer, coord, extension):
    prefix = '/%s' % path if path else ''
    path_to_hash = '%(prefix)s/%(layer)s/%(z)d/%(x)d/%(y)d.%(ext)s' % dict(
        prefix=prefix,
        layer=layer,
        z=coord.zoom,
        x=coord.column,
        y=coord.row,
        ext=extension,
    )
    md5_hash = calc_hash(path_to_hash)
    s3_path = '/%(date)s/%(md5)s%(path_to_hash)s' % dict(
        date=date,
        md5=md5_hash,
        path_to_hash=path_to_hash,
    )
    return s3_path

def parse_coordinate_from_path(path, extension, layer):
    if path.endswith(extension):
        fields = path.rsplit('/', 4)
        if len(fields) == 5:
            _, tile_layer, z_str, x_str, y_fmt = fields
            if tile_layer == layer:
                y_fields = y_fmt.split('.')
                if y_fields:
                    y_str = y_fields[0]
                    try:
                        z = int(z_str)
                        x = int(x_str)
                        y = int(y_str)
                        coord = Coordinate(zoom=z, column=x, row=y)
                        return coord
                    except ValueError:
                        pass

# decorates a function to back off and retry
def _backoff_and_retry(ExceptionType, num_tries=5, retry_factor=2,
                       retry_interval=1, logger=None):
    from time import sleep
    from functools import wraps

    def decorator(f):
        @wraps(f)
        def func(*args, **kwargs):
            # do the first num_tries-1 attempts wrapped in something to catch
            # any exceptions, optionally log them, and try again.
            interval = retry_interval
            factor = retry_factor

            for _ in xrange(1, num_tries):
                try:
                    return f(*args, **kwargs)

                except ExceptionType as e:
                    if logger:
                        logger.warning("Failed. Backing off and retrying. "
                                       "Error: %s" % str(e))

                sleep(interval)
                interval *= factor

            # do final attempt without try-except, so we get the exception
            # in normal code.
            return f(*args, **kwargs)

        return func
    return decorator

class S3(object):

    def __init__(
            self, bucket, date_prefix, path, reduced_redundancy,
            delete_retry_interval, logger):
        self.bucket = bucket
        self.date_prefix = date_prefix
        self.path = path
        self.reduced_redundancy = reduced_redundancy
        self.delete_retry_interval = delete_retry_interval
        self.logger = logger

    def write_tile(self, tile_data, coord, format, layer):
        key_name = s3_tile_key(
            self.date_prefix, self.path, layer, coord, format.extension)
        key = self.bucket.new_key(key_name)

        @_backoff_and_retry(Exception, logger=self.logger)
        def write_to_s3():
            key.set_contents_from_string(
                tile_data,
                headers={'Content-Type': format.mimetype},
                policy='public-read',
                reduced_redundancy=self.reduced_redundancy,
            )

        write_to_s3()

    def read_tile(self, coord, format, layer):
        key_name = s3_tile_key(
            self.date_prefix, self.path, layer, coord, format.extension)
        key = self.bucket.get_key(key_name)
        if key is None:
            return None
        tile_data = key.get_contents_as_string()
        return tile_data

    def delete_tiles(self, coords, format, layer):
        key_names = [
            s3_tile_key(self.date_prefix, self.path, layer, coord,
                        format.extension).lstrip('/')
            for coord in coords
        ]

        num_deleted = 0
        while key_names:
            del_result = self.bucket.delete_keys(key_names)
            num_deleted += len(del_result.deleted)

            key_names = []
            for error in del_result.errors:
                # retry on internal error. documentation implies that the only
                # possible two errors are AccessDenied and InternalError.
                # retrying when access denied seems unlikely to work, but an
                # internal error might be transient.
                if error.code == 'InternalError':
                    key_names.append(error.key)

            # pause a bit to give transient errors a chance to clear.
            if key_names:
                time.sleep(self.delete_retry_interval)

        # make sure that we deleted all the tiles - this seems like the
        # expected behaviour from the calling code.
        assert num_deleted == len(coords), \
            "Failed to delete some coordinates from S3."

        return num_deleted

    def list_tiles(self, format, layer):
        ext = '.' + format.extension
        for key_obj in self.bucket.list(prefix=self.date_prefix):
            key = key_obj.key
            coord = parse_coordinate_from_path(key, ext, layer)
            if coord:
                yield coord

def make_dir_path(base_path, coord, layer):
    path = os.path.join(
        base_path, layer, str(int(coord.zoom)), str(int(coord.column)))
    return path

def make_file_path(base_path, coord, layer, extension):
    basefile_path = os.path.join(
        base_path, layer,
        str(int(coord.zoom)), str(int(coord.column)), str(int(coord.row)))
    ext_str = '.%s' % extension
    full_path = basefile_path + ext_str
    return full_path

def os_replace(src, dst):
    '''
    Simple emulation of function `os.replace(..)` from modern version
    of Python. Implementation is not fully atomic, but enough for us.
    '''

    orig_os_replace_func = getattr(os, 'replace', None)

    if orig_os_replace_func is not None:
        # not need for emulation: we using modern version of Python.
        # fully atomic for this case

        orig_os_replace_func(src, dst)
        return

    if os.name == 'posix':
        # POSIX requirement: `os.rename(..)` works as `os.replace(..)`
        # fully atomic for this case

        os.rename(src, dst)
        return

    # simple emulation for `os.name == 'nt'` and other marginal
    # operation systems.  not fully atomic implementation for this
    # case

    try:
        # trying atomic `os.rename(..)` without `os.remove(..)` or
        # other operations

        os.rename(src, dst)
        error = None
    except OSError as e:
        error = e

    if error is None:
        return

    for i in range(5):
        # some number of tries may be failed
        # because we may be in concurrent environment with other
        # processes/threads

        try:
            os.remove(dst)
        except OSError:
            # destination was not exist
            # or concurrent process/thread is removing it in parallel with us
            pass

        try:
            os.rename(src, dst)
            error = None
        except OSError as e:
            error = e
            continue

        break

    if error is not None:
        raise_from(OSError('failed to replace'), error)

class TileDirectory(object):
    '''
    Writes tiles to individual files in a local directory.
    '''

    def __init__(self, base_path):
        if os.path.exists(base_path):
            if not os.path.isdir(base_path):
                raise IOError(
                    '`{}` exists and is not a directory!'.format(base_path))
        else:
            os.makedirs(base_path)

        self.base_path = base_path

    def write_tile(self, tile_data, coord, format, layer):
        dir_path = make_dir_path(self.base_path, coord, layer)
        try:
            os.makedirs(dir_path)
        except OSError:
            pass

        file_path = make_file_path(self.base_path, coord, layer,
                                   format.extension)
        swap_file_path = '%s.swp-%s-%s-%s' % (
            file_path,
            os.getpid(),
            threading.currentThread().ident,
            random.randint(1, 1000000)
        )

        try:
            with open(swap_file_path, 'w') as tile_fp:
                tile_fp.write(tile_data)

            # write file as atomic operation
            os_replace(swap_file_path, file_path)
        except Exception as e:
            try:
                os.remove(swap_file_path)
            except OSError:
                pass
            raise e

    def read_tile(self, coord, format, layer):
        file_path = make_file_path(self.base_path, coord, layer,
                                   format.extension)
        try:
            with open(file_path, 'r') as tile_fp:
                tile_data = tile_fp.read()
            return tile_data
        except IOError:
            return None

    def delete_tiles(self, coords, format, layer):
        delete_count = 0
        for coord in coords:
            file_path = make_file_path(
                self.base_path, coord, layer, format.extension)
            if os.path.isfile(file_path):
                os.remove(file_path)
                delete_count += 1

        return delete_count

    def list_tiles(self, format, layer):
        ext = '.' + format.extension
        for root, dirs, files in os.walk(self.base_path):
            for name in files:
                tile_path = '%s/%s' % (root, name)
                coord = parse_coordinate_from_path(tile_path, ext, layer)
                if coord:
                    yield coord

def make_tile_file_store(base_path=None):
    if base_path is None:
        base_path = 'tiles'
    return TileDirectory(base_path)

class Memory(object):

    def __init__(self):
        self.data = None

    def write_tile(self, tile_data, coord, format, layer):
        self.data = tile_data, coord, format, layer

    def read_tile(self, coord, format, layer):
        if self.data is None:
            return None
        tile_data, coord, format, layer = self.data
        return tile_data

    def delete_tiles(self, coords, format, layer):
        pass

    def list_tiles(self, format, layer):
        return [self.data] if self.data else []

def make_s3_store(bucket_name,
                  aws_access_key_id=None, aws_secret_access_key=None,
                  path='osm', reduced_redundancy=False, date_prefix='',
                  delete_retry_interval=60, logger=None):
    conn = connect_s3(aws_access_key_id, aws_secret_access_key)
    bucket = Bucket(conn, bucket_name)
    s3_store = S3(bucket, date_prefix, path, reduced_redundancy,
                  delete_retry_interval, logger)
    return s3_store

def tiles_are_equal(tile_data_1, tile_data_2, fmt):
    """
    Returns True if the tile data is equal in tile_data_1 and tile_data_2. For
    most formats, this is a simple byte-wise equality check. For zipped
    metatiles, we need to check the contents, as the zip format includes
    metadata such as timestamps and doesn't control file ordering.
    """

    if fmt and fmt == zip_format:
        return metatiles_are_equal(tile_data_1, tile_data_2)

    else:
        return tile_data_1 == tile_data_2

def write_tile_if_changed(store, tile_data, coord, format, layer):
    """
    Only write tile data if different from existing.

    Try to read the tile data from the store first. If the existing
    data matches, don't write. Returns whether the tile was written.
    """

    existing_data = store.read_tile(coord, format, layer)
    if not existing_data or \
       not tiles_are_equal(existing_data, tile_data, format):
        store.write_tile(tile_data, coord, format, layer)
        return True
    else:
        return False

def make_store(yml, credentials={}, logger=None):
    store_type = yml.get('type')

    if store_type == 'directory':
        path = yml.get('path')
        name = yml.get('name')
        return make_tile_file_store(path or name)

    elif store_type == 's3':
        bucket = yml.get('name')
        path = yml.get('path')
        reduced_redundancy = yml.get('reduced-redundancy')
        date_prefix = yml.get('date-prefix')
        delete_retry_interval = yml.get('delete-retry-interval')

        assert credentials, 'S3 store configured, but no AWS credentials ' \
            'provided. AWS credentials are required to use S3.'
        aws_access_key_id = credentials.get('aws_access_key_id')
        aws_secret_access_key = credentials.get('aws_secret_access_key')

        return make_s3_store(
            bucket, aws_access_key_id, aws_secret_access_key, path=path,
            reduced_redundancy=reduced_redundancy, date_prefix=date_prefix,
            delete_retry_interval=delete_retry_interval, logger=logger)

    else:
        raise ValueError('Unrecognized store type: `{}`'.format(store_type))
<commit_msg>Rewrite store to use boto3 rather than boto2. This seems to solve some issues around accessing S3 buckets in different regions.<commit_after># define locations to store the rendered data

import boto3
from botocore.exceptions import ClientError
from builtins import range
from future.utils import raise_from
import md5
from ModestMaps.Core import Coordinate
import os
from tilequeue.metatile import metatiles_are_equal
from tilequeue.format import zip_format
import random
import threading
import time
from cStringIO import StringIO

def calc_hash(s):
    m = md5.new()
    m.update(s)
    md5_hash = m.hexdigest()
    return md5_hash[:5]

def s3_tile_key(date, path, layer, coord, extension):
    prefix = '/%s' % path if path else ''
    path_to_hash = '%(prefix)s/%(layer)s/%(z)d/%(x)d/%(y)d.%(ext)s' % dict(
        prefix=prefix,
        layer=layer,
        z=coord.zoom,
        x=coord.column,
        y=coord.row,
        ext=extension,
    )
    md5_hash = calc_hash(path_to_hash)
    s3_path = '/%(date)s/%(md5)s%(path_to_hash)s' % dict(
        date=date,
        md5=md5_hash,
        path_to_hash=path_to_hash,
    )
    return s3_path

def parse_coordinate_from_path(path, extension, layer):
    if path.endswith(extension):
        fields = path.rsplit('/', 4)
        if len(fields) == 5:
            _, tile_layer, z_str, x_str, y_fmt = fields
            if tile_layer == layer:
                y_fields = y_fmt.split('.')
                if y_fields:
                    y_str = y_fields[0]
                    try:
                        z = int(z_str)
                        x = int(x_str)
                        y = int(y_str)
                        coord = Coordinate(zoom=z, column=x, row=y)
                        return coord
                    except ValueError:
                        pass

# decorates a function to back off and retry
def _backoff_and_retry(ExceptionType, num_tries=5, retry_factor=2,
                       retry_interval=1, logger=None):
    from time import sleep
    from functools import wraps

    def decorator(f):
        @wraps(f)
        def func(*args, **kwargs):
            # do the first num_tries-1 attempts wrapped in something to catch
            # any exceptions, optionally log them, and try again.
            interval = retry_interval
            factor = retry_factor

            for _ in xrange(1, num_tries):
                try:
                    return f(*args, **kwargs)

                except ExceptionType as e:
                    if logger:
                        logger.warning("Failed. Backing off and retrying. "
                                       "Error: %s" % str(e))

                sleep(interval)
                interval *= factor

            # do final attempt without try-except, so we get the exception
            # in normal code.
            return f(*args, **kwargs)

        return func
    return decorator

class S3(object):

    def __init__(
            self, s3_client, bucket_name, date_prefix, path,
            reduced_redundancy, delete_retry_interval, logger):
        self.s3_client = s3_client
        self.bucket_name = bucket_name
        self.date_prefix = date_prefix
        self.path = path
        self.reduced_redundancy = reduced_redundancy
        self.delete_retry_interval = delete_retry_interval
        self.logger = logger

    def write_tile(self, tile_data, coord, format, layer):
        key_name = s3_tile_key(
            self.date_prefix, self.path, layer, coord, format.extension)

        storage_class = 'STANDARD'
        if self.reduced_redundancy:
            storage_class = 'REDUCED_REDUNDANCY'

        @_backoff_and_retry(Exception, logger=self.logger)
        def write_to_s3():
            self.s3_client.put_object(
                Bucket=self.bucket_name,
                Key=key_name,
                Body=tile_data,
                ContentType=format.mimetype,
                ACL='public-read',
                StorageClass=storage_class,
            )

        write_to_s3()

    def read_tile(self, coord, format, layer):
        key_name = s3_tile_key(
            self.date_prefix, self.path, layer, coord, format.extension)

        try:
            io = StringIO()
            self.s3_client.download_fileobj(self.bucket_name, key_name, io)
            return io.bytes()

        except ClientError as e:
            if e.response['Error']['Code'] != '404':
                raise

        return None

    def delete_tiles(self, coords, format, layer):
        key_names = [
            s3_tile_key(self.date_prefix, self.path, layer, coord,
                        format.extension).lstrip('/')
            for coord in coords
        ]

        num_deleted = 0
        chunk_size = 1000
        for idx in xrange(0, len(key_names), chunk_size):
            chunk = key_names[idx:idx+chunk_size]

            while chunk:
                objects = [dict(Key=k) for k in chunk]
                del_result = self.s3_client.delete_objects(
                    Bucket=self.bucket_name,
                    Delete=dict(Objects=objects),
                )
                num_deleted += len(del_result['Deleted'])

                chunk = []
                for error in del_result['Errors']:
                    # retry on internal error. documentation implies that the
                    # only possible two errors are AccessDenied and
                    # InternalError. retrying when access denied seems
                    # unlikely to work, but an internal error might be
                    # transient.
                    if error['Code'] == 'InternalError':
                        chunk.append(error['Key'])

                # pause a bit to give transient errors a chance to clear.
                if chunk:
                    time.sleep(self.delete_retry_interval)

        # make sure that we deleted all the tiles - this seems like the
        # expected behaviour from the calling code.
        assert num_deleted == len(coords), \
            "Failed to delete some coordinates from S3."

        return num_deleted

    def list_tiles(self, format, layer):
        ext = '.' + format.extension
        paginator = self.s3_client.get_paginator('list_objects_v2')
        page_iter = paginator.paginate(
            Bucket=self.bucket_name,
            Prefix=self.date_prefix
        )
        for page in page_iter:
            for key_obj in page['Contents']:
                key = key_obj['Key']
                coord = parse_coordinate_from_path(key, ext, layer)
                if coord:
                    yield coord

def make_dir_path(base_path, coord, layer):
    path = os.path.join(
        base_path, layer, str(int(coord.zoom)), str(int(coord.column)))
    return path

def make_file_path(base_path, coord, layer, extension):
    basefile_path = os.path.join(
        base_path, layer,
        str(int(coord.zoom)), str(int(coord.column)), str(int(coord.row)))
    ext_str = '.%s' % extension
    full_path = basefile_path + ext_str
    return full_path

def os_replace(src, dst):
    '''
    Simple emulation of function `os.replace(..)` from modern version
    of Python. Implementation is not fully atomic, but enough for us.
    '''

    orig_os_replace_func = getattr(os, 'replace', None)

    if orig_os_replace_func is not None:
        # not need for emulation: we using modern version of Python.
        # fully atomic for this case

        orig_os_replace_func(src, dst)
        return

    if os.name == 'posix':
        # POSIX requirement: `os.rename(..)` works as `os.replace(..)`
        # fully atomic for this case

        os.rename(src, dst)
        return

    # simple emulation for `os.name == 'nt'` and other marginal
    # operation systems.  not fully atomic implementation for this
    # case

    try:
        # trying atomic `os.rename(..)` without `os.remove(..)` or
        # other operations

        os.rename(src, dst)
        error = None
    except OSError as e:
        error = e

    if error is None:
        return

    for i in range(5):
        # some number of tries may be failed
        # because we may be in concurrent environment with other
        # processes/threads

        try:
            os.remove(dst)
        except OSError:
            # destination was not exist
            # or concurrent process/thread is removing it in parallel with us
            pass

        try:
            os.rename(src, dst)
            error = None
        except OSError as e:
            error = e
            continue

        break

    if error is not None:
        raise_from(OSError('failed to replace'), error)

class TileDirectory(object):
    '''
    Writes tiles to individual files in a local directory.
    '''

    def __init__(self, base_path):
        if os.path.exists(base_path):
            if not os.path.isdir(base_path):
                raise IOError(
                    '`{}` exists and is not a directory!'.format(base_path))
        else:
            os.makedirs(base_path)

        self.base_path = base_path

    def write_tile(self, tile_data, coord, format, layer):
        dir_path = make_dir_path(self.base_path, coord, layer)
        try:
            os.makedirs(dir_path)
        except OSError:
            pass

        file_path = make_file_path(self.base_path, coord, layer,
                                   format.extension)
        swap_file_path = '%s.swp-%s-%s-%s' % (
            file_path,
            os.getpid(),
            threading.currentThread().ident,
            random.randint(1, 1000000)
        )

        try:
            with open(swap_file_path, 'w') as tile_fp:
                tile_fp.write(tile_data)

            # write file as atomic operation
            os_replace(swap_file_path, file_path)
        except Exception as e:
            try:
                os.remove(swap_file_path)
            except OSError:
                pass
            raise e

    def read_tile(self, coord, format, layer):
        file_path = make_file_path(self.base_path, coord, layer,
                                   format.extension)
        try:
            with open(file_path, 'r') as tile_fp:
                tile_data = tile_fp.read()
            return tile_data
        except IOError:
            return None

    def delete_tiles(self, coords, format, layer):
        delete_count = 0
        for coord in coords:
            file_path = make_file_path(
                self.base_path, coord, layer, format.extension)
            if os.path.isfile(file_path):
                os.remove(file_path)
                delete_count += 1

        return delete_count

    def list_tiles(self, format, layer):
        ext = '.' + format.extension
        for root, dirs, files in os.walk(self.base_path):
            for name in files:
                tile_path = '%s/%s' % (root, name)
                coord = parse_coordinate_from_path(tile_path, ext, layer)
                if coord:
                    yield coord

def make_tile_file_store(base_path=None):
    if base_path is None:
        base_path = 'tiles'
    return TileDirectory(base_path)

class Memory(object):

    def __init__(self):
        self.data = None

    def write_tile(self, tile_data, coord, format, layer):
        self.data = tile_data, coord, format, layer

    def read_tile(self, coord, format, layer):
        if self.data is None:
            return None
        tile_data, coord, format, layer = self.data
        return tile_data

    def delete_tiles(self, coords, format, layer):
        pass

    def list_tiles(self, format, layer):
        return [self.data] if self.data else []

def make_s3_store(bucket_name,
                  path='osm', reduced_redundancy=False, date_prefix='',
                  delete_retry_interval=60, logger=None):
    s3 = boto3.resource('s3')
    s3_store = S3(s3, bucket_name, date_prefix, path, reduced_redundancy,
                  delete_retry_interval, logger)
    return s3_store

def tiles_are_equal(tile_data_1, tile_data_2, fmt):
    """
    Returns True if the tile data is equal in tile_data_1 and tile_data_2. For
    most formats, this is a simple byte-wise equality check. For zipped
    metatiles, we need to check the contents, as the zip format includes
    metadata such as timestamps and doesn't control file ordering.
    """

    if fmt and fmt == zip_format:
        return metatiles_are_equal(tile_data_1, tile_data_2)

    else:
        return tile_data_1 == tile_data_2

def write_tile_if_changed(store, tile_data, coord, format, layer):
    """
    Only write tile data if different from existing.

    Try to read the tile data from the store first. If the existing
    data matches, don't write. Returns whether the tile was written.
    """

    existing_data = store.read_tile(coord, format, layer)
    if not existing_data or \
       not tiles_are_equal(existing_data, tile_data, format):
        store.write_tile(tile_data, coord, format, layer)
        return True
    else:
        return False

def make_store(yml, credentials={}, logger=None):
    store_type = yml.get('type')

    if store_type == 'directory':
        path = yml.get('path')
        name = yml.get('name')
        return make_tile_file_store(path or name)

    elif store_type == 's3':
        bucket = yml.get('name')
        path = yml.get('path')
        reduced_redundancy = yml.get('reduced-redundancy')
        date_prefix = yml.get('date-prefix')
        delete_retry_interval = yml.get('delete-retry-interval')

        return make_s3_store(
            bucket, path=path,
            reduced_redundancy=reduced_redundancy, date_prefix=date_prefix,
            delete_retry_interval=delete_retry_interval, logger=logger)

    else:
        raise ValueError('Unrecognized store type: `{}`'.format(store_type))
<|endoftext|>