<commit_before>import re
import datetime
from astrometry.util.starutil_numpy import *
from numpy import array

## FIXME -- requires at least one digit before the decimal place.
floatre = r'[+-]?[\d]+(.[\d]*)?([eE][+-]?[\d]*)?'

def floatvar(x):
	return '(?P<%s>%s)' % (x, floatre)

# For parsing orbital elements type output.
elemrexstr = (r'^' + floatvar('jd')
			  + '.*' + ' EC= ' + floatvar('e')
			  + '.*' + ' IN= ' + floatvar('i')
			  + '.*' + ' OM= ' + floatvar('Omega')
			  + ' W = ' + floatvar('pomega')
			  + '.*' + 'MA= ' + floatvar('M')
			  + '.*' + 'A = ' + floatvar('a')
			  )
elemrex = re.compile(elemrexstr, re.MULTILINE | re.DOTALL)

# For finding "System GM" -- only in orbital elements type
sysgmrexstr = '^System GM *: ' + floatvar('gm') + r' '#AU\^3/d\^2$'
sysgmrex = re.compile(sysgmrexstr)

# For parsing X,V type output
xvrexstr = ('^' + floatvar('jd') + r' = (?P<ad>A\.D\. .*?)'
			+ '^ +' + floatvar('x0')
			+ ' +' + floatvar('x1')
			+ ' +' + floatvar('x2')
			+ '.*?'
			+ '^ +' + floatvar('v0')
			+ ' +' + floatvar('v1')
			+ ' +' + floatvar('v2')
			)
xvrex = re.compile(xvrexstr, re.MULTILINE | re.DOTALL)

# For parsing "observer" type output (RA,Dec)
# With QUANTITIES=1; angle format=DEG
radecrexstr = ('^ ' + '(?P<datetime>[\d]{4}-[\w]{3}-[\d]{2} [\d]{2}:[\d]{2}) .*?'
			   + floatvar('ra') + ' *?' + floatvar('dec'))
radecrex = re.compile(radecrexstr, re.MULTILINE | re.DOTALL)

# Returns a list of lists of elements, plus a list of the JDs.
#   ([jd1, jd2, ...], [   [a1, e1, i1, Omega1, pomega1, M1, GM1], ... ])
# Where  i, Omega, pomega, M   are in radians
def parse_orbital_elements(s, needSystemGM=True):
	if needSystemGM:
		m = sysgmrex.search(s)
		if not m:
			print 'Did not find "System GM" entry'
			return None
		gm = float(m.group('gm'))
	else:
		gm = 1.
	allE = []
	alljd = []
	for m in elemrex.finditer(s):
		d = m.groupdict()
		E = [deg2rad(x) if rad else x
			 for (x,rad) in zip([float(d[x]) for x in ['a','e','i','Omega','pomega','M']],
								[False, False, True, True, True, True])]
		E.append(gm)
		allE.append(E)
		alljd.append(float(d['jd']))
	return alljd,allE

# Returns (x, v, jd), each as numpy arrays.
#     x in AU
#     v in AU/yr
def parse_phase_space(s):
	all_x = []
	all_v = []
	all_jd = []
	for m in xvrex.finditer(s):
		d = m.groupdict()
		x = array([float(d[k]) for k in ['x0','x1','x2']])
		v = array([float(d[k]) for k in ['v0','v1','v2']])
		all_x.append(x)
		all_v.append(v)
		all_jd.append(float(d['jd']))
	return (array(all_x), array(all_v), array(all_jd))

# Returns (ra,dec,jd), each as numpy arrays.
#   RA,Dec in J2000 deg
def parse_radec(s):
	all_ra = []
	all_dec = []
	all_jd = []
	for m in radecrex.finditer(s):
		d = m.groupdict()
		all_ra.append(float(d['ra']))
		all_dec.append(float(d['dec']))
		t = datetime.datetime.strptime(d['datetime'], '%Y-%b-%d %H:%M')
		# 2000-Jan-01 12:00
		all_jd.append(datetojd(t))
	return (array(all_ra), array(all_dec), array(all_jd))
<commit_msg><commit_after>import re
import datetime
from astrometry.util.starutil_numpy import *
from numpy import array

## FIXME -- requires at least one digit before the decimal place.
floatre = r'[+-]?[\d]+(.[\d]*)?([eE][+-]?[\d]*)?'

def floatvar(x):
	return '(?P<%s>%s)' % (x, floatre)

# For parsing orbital elements type output.
elemrexstr = (r'^' + floatvar('jd')
			  + '.*' + ' EC= ' + floatvar('e')
			  + '.*' + ' IN= ' + floatvar('i')
			  + '.*' + ' OM= ' + floatvar('Omega')
			  + ' W = ' + floatvar('pomega')
			  + '.*' + 'MA= ' + floatvar('M')
			  + '.*' + 'A = ' + floatvar('a')
			  )
elemrex = re.compile(elemrexstr, re.MULTILINE | re.DOTALL)

# For finding "System GM" -- only in orbital elements type
sysgmrexstr = '^System GM *: ' + floatvar('gm') + r' '#AU\^3/d\^2$'
sysgmrex = re.compile(sysgmrexstr)

# For parsing X,V type output
xvrexstr = ('^' + floatvar('jd') + r' = (?P<ad>A\.D\. .*?)'
			+ '^ +' + floatvar('x0')
			+ ' +' + floatvar('x1')
			+ ' +' + floatvar('x2')
			+ '.*?'
			+ '^ +' + floatvar('v0')
			+ ' +' + floatvar('v1')
			+ ' +' + floatvar('v2')
			)
xvrex = re.compile(xvrexstr, re.MULTILINE | re.DOTALL)

# For parsing "observer" type output (RA,Dec)
# With QUANTITIES=1; angle format=DEG
radecrexstr = ('^ ' + '(?P<datetime>[\d]{4}-[\w]{3}-[\d]{2} [\d]{2}:[\d]{2}) .*?'
			   + floatvar('ra') + ' *?' + floatvar('dec'))
radecrex = re.compile(radecrexstr, re.MULTILINE | re.DOTALL)

# Returns a list of lists of elements, plus a list of the JDs.
#   ([jd1, jd2, ...], [   [a1, e1, i1, Omega1, pomega1, M1, GM1], ... ])
# Where  i, Omega, pomega, M   are in radians
def parse_orbital_elements(s, needSystemGM=True):
	if needSystemGM:
		m = sysgmrex.search(s)
		if not m:
			print 'Did not find "System GM" entry'
			return None
		gm = float(m.group('gm'))
	else:
		gm = 1.
	allE = []
	alljd = []
	for m in elemrex.finditer(s):
		d = m.groupdict()
		E = [deg2rad(x) if rad else x
			 for (x,rad) in zip([float(d[x]) for x in ['a','e','i','Omega','pomega','M']],
								[False, False, True, True, True, True])]
		E.append(gm)
		allE.append(E)
		alljd.append(float(d['jd']))
	return alljd,allE

# Returns (x, v, jd), each as numpy arrays.
#     x in AU
#     v in AU/day
def parse_phase_space(s):
	all_x = []
	all_v = []
	all_jd = []
	for m in xvrex.finditer(s):
		d = m.groupdict()
		x = array([float(d[k]) for k in ['x0','x1','x2']])
		v = array([float(d[k]) for k in ['v0','v1','v2']])
		all_x.append(x)
		all_v.append(v)
		all_jd.append(float(d['jd']))
	return (array(all_x), array(all_v), array(all_jd))

# Returns (ra,dec,jd), each as numpy arrays.
#   RA,Dec in J2000 deg
def parse_radec(s):
	all_ra = []
	all_dec = []
	all_jd = []
	for m in radecrex.finditer(s):
		d = m.groupdict()
		all_ra.append(float(d['ra']))
		all_dec.append(float(d['dec']))
		t = datetime.datetime.strptime(d['datetime'], '%Y-%b-%d %H:%M')
		# 2000-Jan-01 12:00
		all_jd.append(datetojd(t))
	return (array(all_ra), array(all_dec), array(all_jd))
<|endoftext|>


<commit_before># No shebang line, this module is meant to be imported
#
# Copyright 2013 Oliver Palmer
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Resource
--------

Base resources which can be used to build top leve
documents, pages, or other types of data for the web.
"""

from json import loads

try:
    from httplib import (
        responses, NOT_FOUND, BAD_REQUEST, UNSUPPORTED_MEDIA_TYPE,
        METHOD_NOT_ALLOWED, INTERNAL_SERVER_ERROR, OK, NOT_ACCEPTABLE)
except ImportError:  # pragma: no cover
    from http.client import (
        responses, NOT_FOUND, BAD_REQUEST, UNSUPPORTED_MEDIA_TYPE,
        METHOD_NOT_ALLOWED, INTERNAL_SERVER_ERROR, OK, NOT_ACCEPTABLE)

from twisted.internet.defer import Deferred, inlineCallbacks
from twisted.web.server import NOT_DONE_YET
from twisted.web.resource import Resource as _Resource
from twisted.web.static import File
from voluptuous import Invalid, Schema

from pyfarm.core.enums import STRING_TYPES
from pyfarm.agent.http.core import template
from pyfarm.agent.logger import getLogger
from pyfarm.agent.utility import dumps

logger = getLogger("agent.http.resource")

class Resource(_Resource):
    """
    Basic subclass of :class:`._Resource` for passing requests to
    specific methods.  Unlike :class:`._Resource` however this will
    will also handle:

        * Templates
        * Content type discovery and validation
        * Handling of deferred responses
        * Validation of POST/PUT data against a schema

    :cvar string TEMPLATE:
        The name of the template this class will use when rendering
        an html view.

    :type SCHEMAS: dict
    :cvar SCHEMAS:
        A dictionary of schemas to validate the data of an incoming request
        against.  The structure of this dictionary is::

            {http method: <instance of voluptuous.Schema>}

        If the schema validation fails the request will be rejected with
        ``400 BAD REQUEST``.

    :type ALLOWED_CONTENT_TYPE: frozenset
    :cvar ALLOWED_CONTENT_TYPE:
        An instance of :class:`frozenset` which describes what this resource
        is going to allow in the ``Content-Type`` header.  The request
        and this instance must share at least on entry in common.  If not,
        the request will be rejected with ``415 UNSUPPORTED MEDIA TYPE``.

        **This must be defined in subclass**

    :type ALLOWED_ACCEPT: frozenset
    :cvar ALLOWED_ACCEPT:
        An instance of :class:`frozenset` which describes what this
        resource is going to allow in the ``Accept`` header.  The request
        and this instance must share at least one entry in common.  If not,
        the request will be rejected with ``406 NOT ACCEPTABLE``.

        **This must be defined in subclass**

    :type DEFAULT_ACCEPT: frozenset
    :cvar DEFAULT_ACCEPT:
        If ``Accept`` header is not present in the request, use this as the
        value instead.  This defaults to ``frozenset(["*/*"])``

    :type DEFAULT_CONTENT_TYPE: frozenset
    :cvar DEFAULT_CONTENT_TYPE:
        If ``Content-Type`` header is not present in the request, use this as
        the value instead.  This defaults to ``frozenset([""])``
    """
    TEMPLATE = NotImplemented
    SCHEMAS = {}

    # These must be set in a subclass and
    # should contain the full range of headers
    # allowed for Accept and Content-Type.
    ALLOWED_ACCEPT = NotImplemented
    ALLOWED_CONTENT_TYPE = NotImplemented

    # Default values if certain headers
    # are not present.
    DEFAULT_ACCEPT = frozenset(["*/*"])
    DEFAULT_CONTENT_TYPE = frozenset([None])

    @property
    def template(self):
        """
        Loads the template provided but the partial path in ``TEMPLATE`` on
        the class.
        """
        if self.TEMPLATE is NotImplemented:
            raise NotImplementedError("You must set `TEMPLATE` first")
        return template.load(self.TEMPLATE)

    def methods(self):
        """
        Returns a tuple of methods which an instance of this class implements
        """
        methods = []
        for method_name in ("get", "put", "post", "delete", "head"):
            method = getattr(self, method_name, None)
            if method is not None and callable(method):
                methods.append(method_name)

        return tuple(methods)

    def get_content_type(self, request):
        """
        Return the ``Content-Type`` header(s) in the request or
        ``DEFAULT_CONTENT_TYPE`` if the header is not set.
        """
        header = request.requestHeaders.getRawHeaders("Content-Type")
        if not header:
            return self.DEFAULT_CONTENT_TYPE

        content_type = set()
        for value in header:
            # Split out the various parts of the header and return them.  We
            # ignore the q parameter here for the moment.
            content_type.update(
                entry.split(";")[0] for entry in value.split(","))

        return frozenset(content_type)

    def get_accept(self, request):
        """
        Return the ``Accept`` header(s) in the request or
        ``DEFAULT_ACCEPT`` if the header is not set.
        """
        header = request.requestHeaders.getRawHeaders("Accept")
        if not header:
            return self.DEFAULT_ACCEPT

        accept = set()
        for value in header:
            # Split out the various parts of the header and return them.  We
            # ignore the q parameter here for the moment.
            accept.update(entry.split(";")[0] for entry in value.split(","))

        return frozenset(accept)

    def putChild(self, path, child):
        """
        Overrides the builtin putChild() so we can return the results for
        each call and use them externally.
        """
        assert isinstance(path, STRING_TYPES)
        assert isinstance(child, (Resource, File))
        _Resource.putChild(self, path, child)
        return child

    def error(self, request, code, message):
        """
        Writes the proper out an error response message depending on the
        content type in the request
        """
        response_types = self.get_accept(request)
        logger.error(message)

        if "text/html" in response_types:
            request.setResponseCode(code)
            html_error = template.load("error.html")
            result = html_error.render(
                code=code, code_msg=responses[code], message=message)
            request.write(result)

        elif "application/json" in response_types:
            request.setResponseCode(code)
            request.write(dumps({"error": message}))

        else:
            request.setResponseCode(UNSUPPORTED_MEDIA_TYPE)
            error = dumps(
                {"error": "Can only handle text/html or application/json here"})
            request.write(error)

        request.finish()

    def render_tuple(self, request, response):
        """
        Takes a response tuple of ``(body, code, headers)`` or
        ``(body, code)`` and renders the resulting data onto
        the request.
        """
        assert isinstance(response, (list, tuple)), type(response)
        if len(response) == 3:
            body, code, headers = response

            if isinstance(headers, dict):
                for header, value in headers.items():
                    if isinstance(value, STRING_TYPES):
                        value = [value]
                    request.responseHeaders.setRawHeaders(header, value)

            request.setResponseCode(code)

            # Cast to str, otherwise Twisted responds
            #   TypeError: Data must not be unicode
            request.write(str(body))
            request.finish()

        elif len(response) == 2:
            body, code = response
            request.setResponseCode(code)

            # Cast to str, otherwise Twisted responds
            #   TypeError: Data must not be unicode
            request.write(str(body))
            request.finish()

        else:
            self.error(
                request, INTERNAL_SERVER_ERROR,
                "Expected two or three length tuple for response"
            )

    @inlineCallbacks
    def render_deferred(self, request, deferred):
        """
        An inline callback used to unpack a deferred
        response object.
        """
        assert isinstance(deferred, Deferred)
        response = yield deferred
        self.render_tuple(request, response)

    def render(self, request):
        try:
            handler_method = getattr(self, request.method.lower())
        except AttributeError:
            self.error(
                request, METHOD_NOT_ALLOWED,
                "Method %s is not supported" % request.method)
            return NOT_DONE_YET

        assert isinstance(self.ALLOWED_CONTENT_TYPE, (set, frozenset))
        content = request.content.read().strip()
        shared_content_types = \
            self.get_content_type(request) & self.ALLOWED_CONTENT_TYPE

        # Ensure we can handle the content of the request
        if content and not shared_content_types:
            self.error(
                request, UNSUPPORTED_MEDIA_TYPE,
                "Can only support content "
                "type(s) %s" % self.ALLOWED_CONTENT_TYPE)
            return NOT_DONE_YET

        # Determine if we'll be able to produce a response for the request
        assert isinstance(self.ALLOWED_ACCEPT, (set, frozenset))
        response_types = self.get_accept(request) & self.ALLOWED_ACCEPT
        if not response_types:
            self.error(
                request, NOT_ACCEPTABLE,
                "Can only respond with %s" % self.ALLOWED_ACCEPT)
            return NOT_DONE_YET

        # Keywords to pass into `handler_method` below
        kwargs = dict(request=request)

        # Attempt to load the data for the incoming request if appropriate
        if content and "application/json" in shared_content_types:
            try:
                data = loads(content)
            except ValueError as e:
                self.error(
                    request, BAD_REQUEST,
                    "Failed to decode json data: %r" % e)
                return NOT_DONE_YET

            # We have data, check to see if we have a schema
            # and if we do does it validate.
            schema = self.SCHEMAS.get(request.method)
            if isinstance(schema, Schema):
                try:
                    schema(data)
                except Invalid as e:
                    self.error(
                        request, BAD_REQUEST,
                        "Failed to validate the request data "
                        "against the schema: %s" % e)
                    return NOT_DONE_YET

            kwargs.update(data=data)

        try:
            response = handler_method(**kwargs)
        except Exception as error:
            self.error(
                request, INTERNAL_SERVER_ERROR,
                "Unhandled error while rendering response: %s" % error
            )
            return NOT_DONE_YET

        # The handler_method is going to handle everything
        if response == NOT_DONE_YET:
            return NOT_DONE_YET

        # Flask style response
        elif isinstance(response, tuple):
            self.render_tuple(request, response)
            return NOT_DONE_YET

        # handler_method() is returns a Deferred which means
        # we have to handle writing the response ourselves
        elif isinstance(response, Deferred):
            self.render_deferred(request, response)
            return NOT_DONE_YET

        elif isinstance(response, STRING_TYPES):
            request.setResponseCode(OK)
            request.write(response)
            request.finish()
            return NOT_DONE_YET

        else:
            self.error(
                request, INTERNAL_SERVER_ERROR,
                "Unhandled %r in response" % response
            )
            return NOT_DONE_YET
<commit_msg>minor bugfix, this should return ALLOWED_ACCEPT<commit_after># No shebang line, this module is meant to be imported
#
# Copyright 2013 Oliver Palmer
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Resource
--------

Base resources which can be used to build top leve
documents, pages, or other types of data for the web.
"""

from json import loads

try:
    from httplib import (
        responses, NOT_FOUND, BAD_REQUEST, UNSUPPORTED_MEDIA_TYPE,
        METHOD_NOT_ALLOWED, INTERNAL_SERVER_ERROR, OK, NOT_ACCEPTABLE)
except ImportError:  # pragma: no cover
    from http.client import (
        responses, NOT_FOUND, BAD_REQUEST, UNSUPPORTED_MEDIA_TYPE,
        METHOD_NOT_ALLOWED, INTERNAL_SERVER_ERROR, OK, NOT_ACCEPTABLE)

from twisted.internet.defer import Deferred, inlineCallbacks
from twisted.web.server import NOT_DONE_YET
from twisted.web.resource import Resource as _Resource
from twisted.web.static import File
from voluptuous import Invalid, Schema

from pyfarm.core.enums import STRING_TYPES
from pyfarm.agent.http.core import template
from pyfarm.agent.logger import getLogger
from pyfarm.agent.utility import dumps

logger = getLogger("agent.http.resource")

class Resource(_Resource):
    """
    Basic subclass of :class:`._Resource` for passing requests to
    specific methods.  Unlike :class:`._Resource` however this will
    will also handle:

        * Templates
        * Content type discovery and validation
        * Handling of deferred responses
        * Validation of POST/PUT data against a schema

    :cvar string TEMPLATE:
        The name of the template this class will use when rendering
        an html view.

    :type SCHEMAS: dict
    :cvar SCHEMAS:
        A dictionary of schemas to validate the data of an incoming request
        against.  The structure of this dictionary is::

            {http method: <instance of voluptuous.Schema>}

        If the schema validation fails the request will be rejected with
        ``400 BAD REQUEST``.

    :type ALLOWED_CONTENT_TYPE: frozenset
    :cvar ALLOWED_CONTENT_TYPE:
        An instance of :class:`frozenset` which describes what this resource
        is going to allow in the ``Content-Type`` header.  The request
        and this instance must share at least on entry in common.  If not,
        the request will be rejected with ``415 UNSUPPORTED MEDIA TYPE``.

        **This must be defined in subclass**

    :type ALLOWED_ACCEPT: frozenset
    :cvar ALLOWED_ACCEPT:
        An instance of :class:`frozenset` which describes what this
        resource is going to allow in the ``Accept`` header.  The request
        and this instance must share at least one entry in common.  If not,
        the request will be rejected with ``406 NOT ACCEPTABLE``.

        **This must be defined in subclass**

    :type DEFAULT_ACCEPT: frozenset
    :cvar DEFAULT_ACCEPT:
        If ``Accept`` header is not present in the request, use this as the
        value instead.  This defaults to ``frozenset(["*/*"])``

    :type DEFAULT_CONTENT_TYPE: frozenset
    :cvar DEFAULT_CONTENT_TYPE:
        If ``Content-Type`` header is not present in the request, use this as
        the value instead.  This defaults to ``frozenset([""])``
    """
    TEMPLATE = NotImplemented
    SCHEMAS = {}

    # These must be set in a subclass and
    # should contain the full range of headers
    # allowed for Accept and Content-Type.
    ALLOWED_ACCEPT = NotImplemented
    ALLOWED_CONTENT_TYPE = NotImplemented

    # Default values if certain headers
    # are not present.
    DEFAULT_ACCEPT = frozenset(["*/*"])
    DEFAULT_CONTENT_TYPE = frozenset([None])

    @property
    def template(self):
        """
        Loads the template provided but the partial path in ``TEMPLATE`` on
        the class.
        """
        if self.TEMPLATE is NotImplemented:
            raise NotImplementedError("You must set `TEMPLATE` first")
        return template.load(self.TEMPLATE)

    def methods(self):
        """
        Returns a tuple of methods which an instance of this class implements
        """
        methods = []
        for method_name in ("get", "put", "post", "delete", "head"):
            method = getattr(self, method_name, None)
            if method is not None and callable(method):
                methods.append(method_name)

        return tuple(methods)

    def get_content_type(self, request):
        """
        Return the ``Content-Type`` header(s) in the request or
        ``DEFAULT_CONTENT_TYPE`` if the header is not set.
        """
        header = request.requestHeaders.getRawHeaders("Content-Type")
        if not header:
            return self.DEFAULT_CONTENT_TYPE

        content_type = set()
        for value in header:
            # Split out the various parts of the header and return them.  We
            # ignore the q parameter here for the moment.
            content_type.update(
                entry.split(";")[0] for entry in value.split(","))

        return frozenset(content_type)

    def get_accept(self, request):
        """
        Return the ``Accept`` header(s) in the request or
        ``DEFAULT_ACCEPT`` if the header is not set.
        """
        header = request.requestHeaders.getRawHeaders("Accept")
        if not header:
            return self.DEFAULT_ACCEPT

        accept = set()
        for value in header:
            # Split out the various parts of the header and return them.  We
            # ignore the q parameter here for the moment.
            accept.update(entry.split(";")[0] for entry in value.split(","))

        return frozenset(accept)

    def putChild(self, path, child):
        """
        Overrides the builtin putChild() so we can return the results for
        each call and use them externally.
        """
        assert isinstance(path, STRING_TYPES)
        assert isinstance(child, (Resource, File))
        _Resource.putChild(self, path, child)
        return child

    def error(self, request, code, message):
        """
        Writes the proper out an error response message depending on the
        content type in the request
        """
        response_types = self.get_accept(request)
        logger.error(message)

        if "text/html" in response_types:
            request.setResponseCode(code)
            html_error = template.load("error.html")
            result = html_error.render(
                code=code, code_msg=responses[code], message=message)
            request.write(result)

        elif "application/json" in response_types:
            request.setResponseCode(code)
            request.write(dumps({"error": message}))

        else:
            request.setResponseCode(UNSUPPORTED_MEDIA_TYPE)
            error = dumps(
                {"error":
                     "Can only handle one of %s here" % self.ALLOWED_ACCEPT})
            request.write(error)

        request.finish()

    def render_tuple(self, request, response):
        """
        Takes a response tuple of ``(body, code, headers)`` or
        ``(body, code)`` and renders the resulting data onto
        the request.
        """
        assert isinstance(response, (list, tuple)), type(response)
        if len(response) == 3:
            body, code, headers = response

            if isinstance(headers, dict):
                for header, value in headers.items():
                    if isinstance(value, STRING_TYPES):
                        value = [value]
                    request.responseHeaders.setRawHeaders(header, value)

            request.setResponseCode(code)

            # Cast to str, otherwise Twisted responds
            #   TypeError: Data must not be unicode
            request.write(str(body))
            request.finish()

        elif len(response) == 2:
            body, code = response
            request.setResponseCode(code)

            # Cast to str, otherwise Twisted responds
            #   TypeError: Data must not be unicode
            request.write(str(body))
            request.finish()

        else:
            self.error(
                request, INTERNAL_SERVER_ERROR,
                "Expected two or three length tuple for response"
            )

    @inlineCallbacks
    def render_deferred(self, request, deferred):
        """
        An inline callback used to unpack a deferred
        response object.
        """
        assert isinstance(deferred, Deferred)
        response = yield deferred
        self.render_tuple(request, response)

    def render(self, request):
        try:
            handler_method = getattr(self, request.method.lower())
        except AttributeError:
            self.error(
                request, METHOD_NOT_ALLOWED,
                "Method %s is not supported" % request.method)
            return NOT_DONE_YET

        assert isinstance(self.ALLOWED_CONTENT_TYPE, (set, frozenset))
        content = request.content.read().strip()
        shared_content_types = \
            self.get_content_type(request) & self.ALLOWED_CONTENT_TYPE

        # Ensure we can handle the content of the request
        if content and not shared_content_types:
            self.error(
                request, UNSUPPORTED_MEDIA_TYPE,
                "Can only support content "
                "type(s) %s" % self.ALLOWED_CONTENT_TYPE)
            return NOT_DONE_YET

        # Determine if we'll be able to produce a response for the request
        assert isinstance(self.ALLOWED_ACCEPT, (set, frozenset))
        response_types = self.get_accept(request) & self.ALLOWED_ACCEPT
        if not response_types:
            self.error(
                request, NOT_ACCEPTABLE,
                "Can only respond with %s" % self.ALLOWED_ACCEPT)
            return NOT_DONE_YET

        # Keywords to pass into `handler_method` below
        kwargs = dict(request=request)

        # Attempt to load the data for the incoming request if appropriate
        if content and "application/json" in shared_content_types:
            try:
                data = loads(content)
            except ValueError as e:
                self.error(
                    request, BAD_REQUEST,
                    "Failed to decode json data: %r" % e)
                return NOT_DONE_YET

            # We have data, check to see if we have a schema
            # and if we do does it validate.
            schema = self.SCHEMAS.get(request.method)
            if isinstance(schema, Schema):
                try:
                    schema(data)
                except Invalid as e:
                    self.error(
                        request, BAD_REQUEST,
                        "Failed to validate the request data "
                        "against the schema: %s" % e)
                    return NOT_DONE_YET

            kwargs.update(data=data)

        try:
            response = handler_method(**kwargs)
        except Exception as error:
            self.error(
                request, INTERNAL_SERVER_ERROR,
                "Unhandled error while rendering response: %s" % error
            )
            return NOT_DONE_YET

        # The handler_method is going to handle everything
        if response == NOT_DONE_YET:
            return NOT_DONE_YET

        # Flask style response
        elif isinstance(response, tuple):
            self.render_tuple(request, response)
            return NOT_DONE_YET

        # handler_method() is returns a Deferred which means
        # we have to handle writing the response ourselves
        elif isinstance(response, Deferred):
            self.render_deferred(request, response)
            return NOT_DONE_YET

        elif isinstance(response, STRING_TYPES):
            request.setResponseCode(OK)
            request.write(response)
            request.finish()
            return NOT_DONE_YET

        else:
            self.error(
                request, INTERNAL_SERVER_ERROR,
                "Unhandled %r in response" % response
            )
            return NOT_DONE_YET
<|endoftext|>


<commit_before>import os
from unittest.mock import patch, PropertyMock
import urllib.request

import requests

import api as api_test
from memsource import api, constants, models

ANY_ID = 1
ANY_LCS = ['en', 'ja']

class TestApiTranslationMemory(api_test.ApiTestCase):
    def setUp(self):
        self.url_base = 'https://cloud.memsource.com/web/api/v4/transMemory'
        self.translation_memory = api.TranslationMemory(None)
        self.test_tmx_file_path = '/tmp/test.tmx'
        self.test_export_file_path = '/tmp/export.tmx'

        self.setCleanUpFiles([self.test_tmx_file_path, self.test_export_file_path])

        with open(self.test_tmx_file_path, 'w+') as f:
            f.write('This is test tmx file.')

    @patch.object(requests.Session, 'request')
    def test_create(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        returning_id = self.gen_random_int()
        mock_request().json.return_value = {'id': returning_id}

        name = 'test translation memory'
        source_lang = 'en'
        target_langs = ['ja']

        returned_id = self.translation_memory.create(name, source_lang, target_langs)
        self.assertEqual(returning_id, returned_id)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/create'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'name': name,
                'sourceLang': source_lang,
                'targetLang': target_langs,
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_list(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = [{
            'id': 1,
            'targetLangs': [
                'ja'
            ],
            'sourceLang': 'en',
            'name': 'transMem'
        }]

        for translation_memory in self.translation_memory.list():
            self.assertIsInstance(translation_memory, models.TranslationMemory)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/list'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'page': 0
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_list_with_page_id(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = [{
            'id': 1,
            'targetLangs': [
                'ja'
            ],
            'sourceLang': 'en',
            'name': 'transMem'
        }]
        EXISTING_PAGE_ID = 1
        for translation_memory in self.translation_memory.list(EXISTING_PAGE_ID):
            self.assertIsInstance(translation_memory, models.TranslationMemory)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/list'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'page': EXISTING_PAGE_ID
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_list_with_EMPTY_PAGE_ID(self, mock_request):

        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = []
        EMPTY_PAGE_ID = 1
        self.assertEqual(0, len(self.translation_memory(EMPTY_PAGE_ID)))

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/list'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'page': EMPTY_PAGE_ID
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_upload(self, mock_request):
        accepted_segments_count = self.gen_random_int()
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = {'acceptedSegmentsCount': accepted_segments_count}

        translation_memory_id = self.gen_random_int()
        returned_value = self.translation_memory.upload(
            translation_memory_id, self.test_tmx_file_path)

        # Don't use assert_called_with because files has file object. It is difficult to test.
        self.assertTrue(mock_request.called)
        (called_args, called_kwargs) = mock_request.call_args

        # hard to test
        del called_kwargs['files']
        self.assertEqual(
            (constants.HttpMethod.post.value, '{}/import'.format(self.url_base)), called_args)

        self.assertEqual({
            'params': {
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
            },
            'timeout': constants.Base.timeout.value,
        }, called_kwargs)

        self.assertEqual(accepted_segments_count, returned_value)

    @patch.object(requests.Session, 'request')
    def test_upload_from_text(self, mock_request):
        accepted_segments_count = self.gen_random_int()
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = {'acceptedSegmentsCount': accepted_segments_count}

        translation_memory_id = self.gen_random_int()
        returned_value = self.translation_memory.uploadFromText(translation_memory_id, '<xml/>')

        # Don't use assert_called_with because files has file object. It is difficult to test.
        self.assertTrue(mock_request.called)
        (called_args, called_kwargs) = mock_request.call_args

        # hard to test
        del called_kwargs['files']
        self.assertEqual(
            (constants.HttpMethod.post.value, '{}/import'.format(self.url_base)), called_args)

        self.assertEqual({
            'params': {
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
            },
            'timeout': constants.Base.timeout.value,
        }, called_kwargs)

        self.assertEqual(accepted_segments_count, returned_value)

    @patch.object(requests.Session, 'request')
    def test_search_segment_by_task(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)

        segment = "Hello"
        next_segment = "Next segment"
        previous_segment = "Previous Segment"
        task = 'test task'

        mock_request().json.return_value = [{
            'score': 1.01,
            'grossScore': 1.01,
            'segmentId': '5023cd08e4b015e0656c4a8f',
            'source': {
                'fileName': None,
                'modifiedAt': 1418972812301,
                'project': None,
                'rtl': False,
                'previousSegment': None,
                'createdAt': 1418972812301,
                'createdBy': None,
                'modifiedBy': None,
                'domain': None,
                'id': None,
                'nextSegment': None,
                'text': segment,
                'tagMetadata': [],
                'subDomain': None,
                'lang': 'en',
                'client': None
            },
            'subSegment': False,
            'translations': [{
                'fileName': None,
                'modifiedAt': 1340641766000,
                'project': None,
                'rtl': False,
                'previousSegment': None,
                'createdAt': 1336380000000,
                'createdBy': {
                    'role': None,
                    'email': None,
                    # Because I got null as string.
                    'userName': 'null',
                    'lastName': None,
                    'active': False,
                    'firstName': None,
                    'id': 0
                },
                'modifiedBy': {
                    'role': None,
                    'email': None,
                    'userName': 'null',
                    'lastName': None,
                    'active': False,
                    'firstName': None,
                    'id': 0
                },
                'domain': None,
                'id': None,
                'nextSegment': None,
                'text': 'こんにちは',
                'tagMetadata': [],
                'subDomain': None,
                'lang': 'ja',
                'client': None
            }],
            'transMemory': {
                'name': 'Test translation memory',
                'id': '5023cb2ee4b015e0656c4a8e',
                'reverse': False
            }
        }]

        returned_value = self.translation_memory.searchSegmentByTask(
            task, segment, next_segment=next_segment, previous_segment=previous_segment
        )

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            'https://cloud.memsource.com/web/api/v4/transMemory/searchSegmentByTask',
            params={
                'token': self.translation_memory.token,
                'task': task,
                'segment': segment,
                'nextSegment': next_segment,
                'previousSegment': previous_segment,
                'scoreThreshold': 0.6,
            },
            files={},
            timeout=constants.Base.timeout.value
        )

        for segment_search_result in returned_value:
            self.assertIsInstance(segment_search_result, models.SegmentSearchResult)

    @patch.object(requests.Session, 'request')
    def test_insert(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)

        translation_memory_id = self.gen_random_int()
        target_lang = 'ja'
        source_segment = 'this is source segment'
        target_segment = 'this is target segment'
        previous_source_segment = 'this is previous source segment'
        next_source_segment = 'this is next source segment'

        self.translation_memory.insert(
            translation_memory_id,
            target_lang=target_lang,
            source_segment=source_segment,
            target_segment=target_segment,
            previous_source_segment=previous_source_segment,
            next_source_segment=next_source_segment
        )

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            'https://cloud.memsource.com/web/api/v4/transMemory/insert',
            params={
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
                'targetLang': target_lang,
                'sourceSegment': source_segment,
                'targetSegment': target_segment,
                'previousSourceSegment': previous_source_segment,
                'nextSourceSegment': next_source_segment
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_export(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)

        translation_memory_id = ANY_ID
        target_langs = ANY_LCS
        file_format = 'TMX'

        tmx_contents = ['test first chunk', 'second chunk']
        mock_request().iter_content.return_value = [
            bytes(content, 'utf-8') for content in tmx_contents]

        self.assertFalse(os.path.isfile(self.test_export_file_path))

        returned_value = self.translation_memory.export(
            translation_memory_id=translation_memory_id,
            file_format=file_format,
            target_langs=target_langs,
            file_path=self.test_export_file_path
        )
        self.assertTrue(os.path.isfile(self.test_export_file_path))

        self.assertIsNone(returned_value)

        with open(self.test_export_file_path) as f:
            self.assertEqual(''.join(tmx_contents), f.read())

        mock_request.assert_called_with(
            constants.HttpMethod.get.value,
            "{}/export".format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
                'format': file_format,
                'targetLang': target_langs,
            },
            files={},
            timeout=constants.Base.timeout.value * 5,
            stream=True
        )

    @patch.object(requests.Session, 'request')
    def test_delete_source_and_translations(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        translation_memory_id = 1
        segment_ids = ['2', '3']

        self.translation_memory.deleteSourceAndTranslations(translation_memory_id, segment_ids)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            urllib.request.urljoin(self.url_base, "transMemory/deleteSourceAndTranslations"),
            params={
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
                'segmentId': segment_ids
            },
            files={},
            timeout=constants.Base.timeout.value
        )
<commit_msg>add call to 'list()' in test for empty page<commit_after>import os
from unittest.mock import patch, PropertyMock
import urllib.request

import requests

import api as api_test
from memsource import api, constants, models

ANY_ID = 1
ANY_LCS = ['en', 'ja']

class TestApiTranslationMemory(api_test.ApiTestCase):
    def setUp(self):
        self.url_base = 'https://cloud.memsource.com/web/api/v4/transMemory'
        self.translation_memory = api.TranslationMemory(None)
        self.test_tmx_file_path = '/tmp/test.tmx'
        self.test_export_file_path = '/tmp/export.tmx'

        self.setCleanUpFiles([self.test_tmx_file_path, self.test_export_file_path])

        with open(self.test_tmx_file_path, 'w+') as f:
            f.write('This is test tmx file.')

    @patch.object(requests.Session, 'request')
    def test_create(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        returning_id = self.gen_random_int()
        mock_request().json.return_value = {'id': returning_id}

        name = 'test translation memory'
        source_lang = 'en'
        target_langs = ['ja']

        returned_id = self.translation_memory.create(name, source_lang, target_langs)
        self.assertEqual(returning_id, returned_id)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/create'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'name': name,
                'sourceLang': source_lang,
                'targetLang': target_langs,
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_list(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = [{
            'id': 1,
            'targetLangs': [
                'ja'
            ],
            'sourceLang': 'en',
            'name': 'transMem'
        }]

        for translation_memory in self.translation_memory.list():
            self.assertIsInstance(translation_memory, models.TranslationMemory)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/list'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'page': 0
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_list_with_page_id(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = [{
            'id': 1,
            'targetLangs': [
                'ja'
            ],
            'sourceLang': 'en',
            'name': 'transMem'
        }]
        EXISTING_PAGE_ID = 1
        for translation_memory in self.translation_memory.list(EXISTING_PAGE_ID):
            self.assertIsInstance(translation_memory, models.TranslationMemory)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/list'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'page': EXISTING_PAGE_ID
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_list_with_EMPTY_PAGE_ID(self, mock_request):

        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = []
        EMPTY_PAGE_ID = 1
        self.assertEqual(0, len(self.translation_memory.list(EMPTY_PAGE_ID)))

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            '{}/list'.format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'page': EMPTY_PAGE_ID
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_upload(self, mock_request):
        accepted_segments_count = self.gen_random_int()
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = {'acceptedSegmentsCount': accepted_segments_count}

        translation_memory_id = self.gen_random_int()
        returned_value = self.translation_memory.upload(
            translation_memory_id, self.test_tmx_file_path)

        # Don't use assert_called_with because files has file object. It is difficult to test.
        self.assertTrue(mock_request.called)
        (called_args, called_kwargs) = mock_request.call_args

        # hard to test
        del called_kwargs['files']
        self.assertEqual(
            (constants.HttpMethod.post.value, '{}/import'.format(self.url_base)), called_args)

        self.assertEqual({
            'params': {
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
            },
            'timeout': constants.Base.timeout.value,
        }, called_kwargs)

        self.assertEqual(accepted_segments_count, returned_value)

    @patch.object(requests.Session, 'request')
    def test_upload_from_text(self, mock_request):
        accepted_segments_count = self.gen_random_int()
        type(mock_request()).status_code = PropertyMock(return_value=200)
        mock_request().json.return_value = {'acceptedSegmentsCount': accepted_segments_count}

        translation_memory_id = self.gen_random_int()
        returned_value = self.translation_memory.uploadFromText(translation_memory_id, '<xml/>')

        # Don't use assert_called_with because files has file object. It is difficult to test.
        self.assertTrue(mock_request.called)
        (called_args, called_kwargs) = mock_request.call_args

        # hard to test
        del called_kwargs['files']
        self.assertEqual(
            (constants.HttpMethod.post.value, '{}/import'.format(self.url_base)), called_args)

        self.assertEqual({
            'params': {
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
            },
            'timeout': constants.Base.timeout.value,
        }, called_kwargs)

        self.assertEqual(accepted_segments_count, returned_value)

    @patch.object(requests.Session, 'request')
    def test_search_segment_by_task(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)

        segment = "Hello"
        next_segment = "Next segment"
        previous_segment = "Previous Segment"
        task = 'test task'

        mock_request().json.return_value = [{
            'score': 1.01,
            'grossScore': 1.01,
            'segmentId': '5023cd08e4b015e0656c4a8f',
            'source': {
                'fileName': None,
                'modifiedAt': 1418972812301,
                'project': None,
                'rtl': False,
                'previousSegment': None,
                'createdAt': 1418972812301,
                'createdBy': None,
                'modifiedBy': None,
                'domain': None,
                'id': None,
                'nextSegment': None,
                'text': segment,
                'tagMetadata': [],
                'subDomain': None,
                'lang': 'en',
                'client': None
            },
            'subSegment': False,
            'translations': [{
                'fileName': None,
                'modifiedAt': 1340641766000,
                'project': None,
                'rtl': False,
                'previousSegment': None,
                'createdAt': 1336380000000,
                'createdBy': {
                    'role': None,
                    'email': None,
                    # Because I got null as string.
                    'userName': 'null',
                    'lastName': None,
                    'active': False,
                    'firstName': None,
                    'id': 0
                },
                'modifiedBy': {
                    'role': None,
                    'email': None,
                    'userName': 'null',
                    'lastName': None,
                    'active': False,
                    'firstName': None,
                    'id': 0
                },
                'domain': None,
                'id': None,
                'nextSegment': None,
                'text': 'こんにちは',
                'tagMetadata': [],
                'subDomain': None,
                'lang': 'ja',
                'client': None
            }],
            'transMemory': {
                'name': 'Test translation memory',
                'id': '5023cb2ee4b015e0656c4a8e',
                'reverse': False
            }
        }]

        returned_value = self.translation_memory.searchSegmentByTask(
            task, segment, next_segment=next_segment, previous_segment=previous_segment
        )

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            'https://cloud.memsource.com/web/api/v4/transMemory/searchSegmentByTask',
            params={
                'token': self.translation_memory.token,
                'task': task,
                'segment': segment,
                'nextSegment': next_segment,
                'previousSegment': previous_segment,
                'scoreThreshold': 0.6,
            },
            files={},
            timeout=constants.Base.timeout.value
        )

        for segment_search_result in returned_value:
            self.assertIsInstance(segment_search_result, models.SegmentSearchResult)

    @patch.object(requests.Session, 'request')
    def test_insert(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)

        translation_memory_id = self.gen_random_int()
        target_lang = 'ja'
        source_segment = 'this is source segment'
        target_segment = 'this is target segment'
        previous_source_segment = 'this is previous source segment'
        next_source_segment = 'this is next source segment'

        self.translation_memory.insert(
            translation_memory_id,
            target_lang=target_lang,
            source_segment=source_segment,
            target_segment=target_segment,
            previous_source_segment=previous_source_segment,
            next_source_segment=next_source_segment
        )

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            'https://cloud.memsource.com/web/api/v4/transMemory/insert',
            params={
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
                'targetLang': target_lang,
                'sourceSegment': source_segment,
                'targetSegment': target_segment,
                'previousSourceSegment': previous_source_segment,
                'nextSourceSegment': next_source_segment
            },
            files={},
            timeout=constants.Base.timeout.value
        )

    @patch.object(requests.Session, 'request')
    def test_export(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)

        translation_memory_id = ANY_ID
        target_langs = ANY_LCS
        file_format = 'TMX'

        tmx_contents = ['test first chunk', 'second chunk']
        mock_request().iter_content.return_value = [
            bytes(content, 'utf-8') for content in tmx_contents]

        self.assertFalse(os.path.isfile(self.test_export_file_path))

        returned_value = self.translation_memory.export(
            translation_memory_id=translation_memory_id,
            file_format=file_format,
            target_langs=target_langs,
            file_path=self.test_export_file_path
        )
        self.assertTrue(os.path.isfile(self.test_export_file_path))

        self.assertIsNone(returned_value)

        with open(self.test_export_file_path) as f:
            self.assertEqual(''.join(tmx_contents), f.read())

        mock_request.assert_called_with(
            constants.HttpMethod.get.value,
            "{}/export".format(self.url_base),
            params={
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
                'format': file_format,
                'targetLang': target_langs,
            },
            files={},
            timeout=constants.Base.timeout.value * 5,
            stream=True
        )

    @patch.object(requests.Session, 'request')
    def test_delete_source_and_translations(self, mock_request):
        type(mock_request()).status_code = PropertyMock(return_value=200)
        translation_memory_id = 1
        segment_ids = ['2', '3']

        self.translation_memory.deleteSourceAndTranslations(translation_memory_id, segment_ids)

        mock_request.assert_called_with(
            constants.HttpMethod.post.value,
            urllib.request.urljoin(self.url_base, "transMemory/deleteSourceAndTranslations"),
            params={
                'token': self.translation_memory.token,
                'transMemory': translation_memory_id,
                'segmentId': segment_ids
            },
            files={},
            timeout=constants.Base.timeout.value
        )
<|endoftext|>


<commit_before>#!/usr/bin/python2
import os,sys
import binascii
import struct
import StringIO

# chunk names are 4 bytes long
# not quite sure how long the header is but I guess 8 bytes for now :>
dwheader=binascii.unhexlify("464F524DE6EF5A03")
datawin=open("data.win","r+")
datawin.seek(0,2)
dwsize=datawin.tell()
datawin.seek(0,0)
header=datawin.read(8)
chunks={}

if header==dwheader:
    print "Doogie doogie Dodger"
else:
    print "Wrong magic bytes"
    quit()

def read_chunk():
    chunkname=datawin.read(4)
    print "Reading "+chunkname
    chunksize=datawin.read(4)
    # TODO: _THIS_ is stupid! ... let's correct this with a nice struct or somethin
    # later...
    foo=binascii.hexlify(chunksize)
    chunksize=foo[-2:]+foo[4:6]+foo[2:4]+foo[:2]
    chunksize=int(chunksize,16)
    chunk=datawin.read(chunksize)
    return chunkname, chunk

while dwsize!=datawin.tell():
    chunk=read_chunk()
    chunks[chunk[0]]=chunk[1]
<commit_msg>A few hours of fruitless hacking and failures<commit_after>#!/usr/bin/python2
# -*- coding: utf-8 -*-
import os,sys
import binascii
import struct
import StringIO

# chunk names are 4 bytes long
# not quite sure how long the header is but I guess 8 bytes for now :>
# actually it's not really a header but the first chunk
# 4 bytes for the name (FORM) and 4 bytes for the length (filesize-8)
# so this is not needed anymore

#dwheader=binascii.unhexlify("464F524DE6EF5A03")
datawin=open("data.win","r+")
#datawin.seek(0,2)
#dwsize=datawin.tell()
#datawin.seek(0,0)
#header=datawin.read(8)
chunks={}

#if header==dwheader:
#    print "Doogie doogie Dodger"
#else:
#    print "Wrong magic bytes"
#    quit()

def read_chunk(data):
    # data has to be a StringIO/file... not!
    if type(data)!=type("foo"):
        print "Somehow the data is not a string"
        quit()
    else:
        dsize=len(binascii.hexlify(data))/2
        data=StringIO.StringIO(data)
#    data.seek(0,2)
#    dsize=datawin.tell()
#    if dsize!=realdsize:
#        print dsize
#        print realdsize
#        print "FUCK THIS SHIT I'M GOING HOME"
#        quit()
    data.seek(0,0)
    chunkname=data.read(4)
    if chunkname.isupper():
        print "Reading "+chunkname
    else:
        print "Reading "+binascii.hexlify(chunkname)
    chunksize=data.read(4)
    if len(chunksize)!=4:
        data.seek(0,0)
        return [data.read()]
    # TODO: _THIS_ is stupid! ... let's correct this with a nice struct or somethin
    # later...
    foo=binascii.hexlify(chunksize)
    chunksize=foo[-2:]+foo[4:6]+foo[2:4]+foo[:2]
    chunksize=int(chunksize,16)
    if chunksize+8==dsize:
        # print "One big chunk"
        chunk=data.read(chunksize)
        return [chunkname, chunk]
    elif chunksize+8 > dsize:
        # print "No chunk for you" # pure data
        data.seek(0,0)
        return [data.read()]
    else:
        # print "Multiple CHUNKS!!!!! xD ohh I'm sooo tired"
        chunk=data.read(chunksize)
        rest=data.read()
        # print chunksize
        # print "Actual chunk size "+str(len(chunk))
        # print dsize
        # print "REST:"
        # print len(rest)
        if len(rest)==0:
            print "Something went terrible, terrible WRONG :( WTF IS HAPPENING?????"
            return [chunkname, chunk]
        else:
            return [chunkname, chunk, rest]

def extract_chunks(data):
    if type(data)==type("Foo"):
        realdatasize=len(data)
        data=StringIO.StringIO(data)
    else:
        realdatasize=len(data.read())
    data.seek(0,2)
    datasize=data.tell()
    if datasize!=realdatasize:
        print "OK WHY ISN'T THIS WORKING??"
        quit()
    data.seek(0,0)
    while data.tell()!=datasize:
        chunk=read_chunk(data.read())
        if len(chunk)==1:
            return chunk[0]
        elif len(chunk)==2:
            print "One big chunk you chump"
            if type(chunk[1])==type("foo"):
                return {chunk[0]:extract_chunks(chunk[1])}
            else:
                print "OMG LIKE THAT'S TOTALLY LIKE A DICTIONARY THAT SOMEHOW GOT LIKE RETURNED! WHAT THE EFF"
                return {chunk[0]:chunk[1]}
        elif len(chunk)==3:
            if type(chunk[1])==type("foo"):
                newchunk={chunk[0]:extract_chunks(chunk[1])}
            else:
                newchunk={chunk[0]:chunk[1]}
            if type(chunk[2])==type("foo"):
                newdict=extract_chunks(chunk[2])
                if type(newdict)==type({}):
                    newchunk.update(newdict)
                else:
                    newchunk.update({"DEFECT":newdict})
            else:
                newchunk.update(chunk[2])
            return newchunk

final_dict=extract_chunks(datawin)
#print type(final_dict)
print len(final_dict["FORM"]["STRG"])
print type(final_dict["FORM"]["STRG"])
#for foo in final_dict["FORM"]:
#    print foo
#while datawin.tell()!=dsize:
    #extract chunk from data
    #übergib chunk to while like this

# chunk=read_chunk(datawin)
#
# chunks[chunk[0]]=StringIO.StringIO(chunk[1])
#
# chunk=read_chunk(chunks["FORM"])
#
#
# if check_chunk(StringIO.StringIO(chunk[1]))==0:
#     chunk=read_chunk(StringIO.StringIO(chunk[1]))

#while dwsize!=datawin.tell():
#    chunk=read_chunk()
#    chunks[chunk[0]]=chunk[1]

# check if chunk
# add chunk to dictionary
# check if chunk
# add to dictionary
<|endoftext|>


<commit_before>import requests
import dataset
import os
import re
import datetime as dt
import time
from tqdm import tqdm
from selenium import webdriver
from selenium.common.exceptions import NoSuchElementException, ElementNotVisibleException

# Number of seconds to sleep between crawls
CRAWL_TIMEOUT = 10

THIRTY_DAY_REVIEW_REGEX = re.compile(r'^([0-9]+)% of the ([,0-9]+) user reviews in the last 30 days')
ALL_TIME_REVIEW_REGEX = re.compile(r'^([0-9]+)% of the ([,0-9]+) user reviews for this game')
DETAILS_BOX_REGEX = re.compile(r'^Title: ([^\n]+)\nGenre: ([^\n]+)\nDeveloper: ([^\n]+)(?:\nPublisher: ([^\n]+))?')
NUM_ACHIEVEMENTS_REGEX = re.compile(r'Includes ([,[0-9]+) Steam Achievements')

def upsert_all_apps(db):
    '''
    Get the full list of steam apps and upsert them in our database
    on the basis of steam's app ID.
    '''
    json = requests.get('http://api.steampowered.com/ISteamApps/GetAppList/v0001/').json()

    apps = json['applist']['apps']['app']

    db.begin()

    for app in tqdm(apps):
        db['game'].upsert({
            'steam_app_id': app['appid'],
            'game_name': app['name'],
        }, keys=['app_id'])

    db.commit()

def scrape_store_page(app_id):
    '''
    Extract all the information we can from the store page for a given app ID.
    '''
    results = {'steam_app_id': app_id}
    driver = webdriver.Chrome()
    store_base_url = "http://store.steampowered.com"
    app_url = "{}/app/{}".format(store_base_url, app_id)
    driver.get(app_url)

    if driver.current_url in (store_base_url, '{}/'.format(store_base_url)):
        # We were redirected; the app doesn't have a store page.
        driver.close()
        return results

    try:
        # If this succeeds, we need to pass through the age gate.
        driver.find_element_by_id('agegate_box')

        select_element = driver.find_element_by_id('ageYear')
        # open year dialog
        select_element.click()
        # select correct year
        select_element.find_element_by_css_selector('option[value="1993"]').click()
        # close year dialog
        select_element.click()
        # submit the form
        driver.find_element_by_id('agecheck_form').submit()
    except NoSuchElementException:
        # No age gate; we're good to continue
        pass

    # Get the description first, since it tells us whether the app is streaming video
    # (which means we don't care about it)
    descriptions = driver.find_elements_by_class_name('game_area_description')
    for description in descriptions:
        if description.text.startswith('ABOUT THIS GAME'):
            results['is_dlc'] = False
            results['long_description'] = description.text
        elif description.text.startswith('ABOUT THIS CONTENT'):
            results['is_dlc'] = True
            results['long_description'] = description.text
        elif description.text.startswith('ABOUT THIS SERIES'):
            # This is streaming video; we don't care about it, so just return
            driver.close()
            return results
    if 'long_description' not in results:
        raise RuntimeError('Unable to parse description for app_id {}'.format(app_id))

    results['game_name'] = (driver
                            .find_element_by_class_name('apphub_AppName')
                            .text)

    try:
        results['short_description'] = (driver
                                        .find_element_by_class_name('game_description_snippet')
                                        .text)
    except NoSuchElementException:
        # DLC doesn't have a short description
        pass

    reviews_texts = [element.get_attribute('data-store-tooltip')
                     for element in (driver
                                     .find_elements_by_class_name('user_reviews_summary_row')
                                     )]

    for text in reviews_texts:
        thirty_day_match = THIRTY_DAY_REVIEW_REGEX.match(text)

        if thirty_day_match:
            results['pct_positive_reviews_last_30_days'] = int(thirty_day_match.group(1))
            results['reviews_last_30_days'] = int(thirty_day_match.group(2).replace(',', ''))
        else:
            all_time_match = ALL_TIME_REVIEW_REGEX.match(text)

            if all_time_match:
                results['pct_positive_reviews_all_time'] = int(all_time_match.group(1))
                results['reviews_all_time'] = int(all_time_match.group(2).replace(',', ''))

    results['release_date'] = driver.find_element_by_css_selector('.release_date .date').text

    # There's additional detail about VR stuff here, but we're not worried about that for now
    details_text = driver.find_elements_by_css_selector('.details_block:not(.vrsupport)')[0].text
    details_match = DETAILS_BOX_REGEX.match(details_text)
    results['title'] = details_match.group(1)
    raw_genre = details_match.group(2)
    results['genres'] = raw_genre.split(', ')
    results['developer'] = details_match.group(3)
    results['publisher'] = details_match.group(4)

    block_titles_texts = [element.text
                          for element in (driver
                                          .find_elements_by_class_name('block_title')
                                          )]

    for text in block_titles_texts:
        num_achievements_match = NUM_ACHIEVEMENTS_REGEX.match(text)

        if num_achievements_match:
            results['num_achievements'] = int(num_achievements_match.group(1))

    try:
        results['metacritic_score'] = int(driver.find_element_by_class_name('score').text)
    except NoSuchElementException:
        # Some games don't have metascores
        pass

    try:
        # Not on sale or Free to Play
        raw_price = driver.find_element_by_class_name('game_purchase_price').text
        if raw_price == 'Free to Play':
            price = 0
        else:
            price = float(raw_price.replace('$', ''))
        results['full_price'] = price
    except NoSuchElementException:
        # On sale
        raw_price = driver.find_element_by_class_name('discount_original_price').text
        results['full_price'] = float(raw_price.replace('$', ''))

    results['game_details'] = []
    game_details_elements = driver.find_elements_by_class_name('game_area_details_specs')
    for element in game_details_elements:
        results['game_details'].append(element.find_element_by_css_selector('a.name').text)

    results['tags'] = []
    try:
        # Try to get the big list of tags if it's there
        driver.find_element_by_css_selector('.app_tag.add_button').click()
        tag_elements = driver.find_elements_by_css_selector('#app_tagging_modal a.app_tag')
    except ElementNotVisibleException:
        # Settle for the short list if not
        tag_elements = driver.find_elements_by_css_selector('a.app_tag')
    for element in tag_elements:
        results['tags'].append(element.text)

    driver.close()

    return results

def insert_with_mapping(*, descrs, entity_table, pk_name, join_table, mapping, app_id, crawl_time):
    '''
    Given some list of description data, a table of entities, the name of the PK,
    a many-to-many join table for the entities to game crawls,
    a mapping from descrs to entity table IDs, and an app_id/crawl_time,
    update the entity table/mapping if necessary and insert
    the entity in the many-to-many join table.

    Mutates the mapping.
    '''
    for descr in descrs:
        try:
            entity_id = mapping[descr]
        except KeyError:
            entity_id = db[entity_table]
            mapping[descr] = entity_id

        db[join_table].insert({
            'steam_app_id': app_id,
            'crawl_time': crawl_time,
            pk_name: entity_id,
        })

def do_crawl(app_ids, db):
    '''
    Given a list of steam app IDs and a db connection, do a crawl for the app IDs
    and append the results to our list of crawls in the database.
    '''
    tag_mapping = {r['descr']: r['tag_id'] for r in db['steam_tag'].find()}
    detail_mapping = {r['descr']: r['detail_id'] for r in db['steam_game_detail'].find()}
    genre_mapping = {r['descr']: r['genre_id'] for r in db['steam_genre'].find()}

    db.begin()
    for app_id in tqdm(app_ids):
        try:
            results = scrape_store_page(app_id)
        except:
            db.commit()
            print("Failed crawl for ID {}".format(app_id))
            raise

        crawl_time = dt.datetime.now()
        results['crawl_time'] = crawl_time

        # Pull the lists off before we insert the main crawl record.
        tags = results['tags']
        del results['tags']
        details = results['game_details']
        del results['game_details']
        genres = results['genres']
        del results['genres']

        db['game_crawl'].insert(results)

        insert_with_mapping(
            descrs=tags,
            entity_table='steam_tag',
            pk_name='tag_id',
            join_table='game_crawl_tag',
            mapping=tag_mapping,
            app_id=app_id,
            crawl_time=crawl_time
        )

        insert_with_mapping(
            descrs=details,
            entity_table='steam_game_detail',
            pk_name='detail_id',
            join_table='game_crawl_detail',
            mapping=detail_mapping,
            app_id=app_id,
            crawl_time=crawl_time
        )

        insert_with_mapping(
            descrs=genres,
            entity_table='steam_genre',
            pk_name='genre_id',
            join_table='game_crawl_genre',
            mapping=genre_mapping,
            app_id=app_id,
            crawl_time=crawl_time
        )

        time.sleep(CRAWL_TIMEOUT)
    db.commit()

def run():
    db = dataset.connect(os.environ['POSTGRES_URI'], ensure_schema=False)

    # Since we're not super worried about having an up-to-date list of apps,
    # run this only if the table is empty
    if db['game'].count() == 0:
        upsert_all_apps(db)

    # For now, just crawl the apps we don't already have
    missing_crawl_query = '''
    SELECT steam_app_id
    FROM game g
      LEFT JOIN game_crawl gc
        USING (steam_app_id)
    WHERE gc.steam_app_id IS NULL;
    '''

    missing_app_ids = [r['steam_app_id'] for r in db.query(missing_crawl_query)]

    # TODO remove after we're done testing the crawl function
    missing_app_ids = missing_app_ids[:1]

    do_crawl(missing_app_ids, db)

if __name__ == '__main__':
    run()

<commit_msg>Fix insert_with_mapping missing DB connection; handle Ctrl+C<commit_after>import requests
import dataset
import os
import re
import datetime as dt
import time
from tqdm import tqdm
from selenium import webdriver
from selenium.common.exceptions import NoSuchElementException, ElementNotVisibleException

# Number of seconds to sleep between crawls
CRAWL_TIMEOUT = 10

THIRTY_DAY_REVIEW_REGEX = re.compile(r'^([0-9]+)% of the ([,0-9]+) user reviews in the last 30 days')
ALL_TIME_REVIEW_REGEX = re.compile(r'^([0-9]+)% of the ([,0-9]+) user reviews for this game')
DETAILS_BOX_REGEX = re.compile(r'^Title: ([^\n]+)\nGenre: ([^\n]+)\nDeveloper: ([^\n]+)(?:\nPublisher: ([^\n]+))?')
NUM_ACHIEVEMENTS_REGEX = re.compile(r'Includes ([,[0-9]+) Steam Achievements')

def upsert_all_apps(db):
    '''
    Get the full list of steam apps and upsert them in our database
    on the basis of steam's app ID.
    '''
    json = requests.get('http://api.steampowered.com/ISteamApps/GetAppList/v0001/').json()

    apps = json['applist']['apps']['app']

    db.begin()

    for app in tqdm(apps):
        db['game'].upsert({
            'steam_app_id': app['appid'],
            'game_name': app['name'],
        }, keys=['app_id'])

    db.commit()

def scrape_store_page(app_id):
    '''
    Extract all the information we can from the store page for a given app ID.
    '''
    results = {'steam_app_id': app_id}
    driver = webdriver.Chrome()
    store_base_url = "http://store.steampowered.com"
    app_url = "{}/app/{}".format(store_base_url, app_id)
    driver.get(app_url)

    if driver.current_url in (store_base_url, '{}/'.format(store_base_url)):
        # We were redirected; the app doesn't have a store page.
        driver.close()
        return results

    try:
        # If this succeeds, we need to pass through the age gate.
        driver.find_element_by_id('agegate_box')

        select_element = driver.find_element_by_id('ageYear')
        # open year dialog
        select_element.click()
        # select correct year
        select_element.find_element_by_css_selector('option[value="1993"]').click()
        # close year dialog
        select_element.click()
        # submit the form
        driver.find_element_by_id('agecheck_form').submit()
    except NoSuchElementException:
        # No age gate; we're good to continue
        pass

    # Get the description first, since it tells us whether the app is streaming video
    # (which means we don't care about it)
    descriptions = driver.find_elements_by_class_name('game_area_description')
    for description in descriptions:
        if description.text.startswith('ABOUT THIS GAME'):
            results['is_dlc'] = False
            results['long_description'] = description.text
        elif description.text.startswith('ABOUT THIS CONTENT'):
            results['is_dlc'] = True
            results['long_description'] = description.text
        elif description.text.startswith('ABOUT THIS SERIES'):
            # This is streaming video; we don't care about it, so just return
            driver.close()
            return results
    if 'long_description' not in results:
        raise RuntimeError('Unable to parse description for app_id {}'.format(app_id))

    results['game_name'] = (driver
                            .find_element_by_class_name('apphub_AppName')
                            .text)

    try:
        results['short_description'] = (driver
                                        .find_element_by_class_name('game_description_snippet')
                                        .text)
    except NoSuchElementException:
        # DLC doesn't have a short description
        pass

    reviews_texts = [element.get_attribute('data-store-tooltip')
                     for element in (driver
                                     .find_elements_by_class_name('user_reviews_summary_row')
                                     )]

    for text in reviews_texts:
        thirty_day_match = THIRTY_DAY_REVIEW_REGEX.match(text)

        if thirty_day_match:
            results['pct_positive_reviews_last_30_days'] = int(thirty_day_match.group(1))
            results['reviews_last_30_days'] = int(thirty_day_match.group(2).replace(',', ''))
        else:
            all_time_match = ALL_TIME_REVIEW_REGEX.match(text)

            if all_time_match:
                results['pct_positive_reviews_all_time'] = int(all_time_match.group(1))
                results['reviews_all_time'] = int(all_time_match.group(2).replace(',', ''))

    results['release_date'] = driver.find_element_by_css_selector('.release_date .date').text

    # There's additional detail about VR stuff here, but we're not worried about that for now
    details_text = driver.find_elements_by_css_selector('.details_block:not(.vrsupport)')[0].text
    details_match = DETAILS_BOX_REGEX.match(details_text)
    results['title'] = details_match.group(1)
    raw_genre = details_match.group(2)
    results['genres'] = raw_genre.split(', ')
    results['developer'] = details_match.group(3)
    results['publisher'] = details_match.group(4)

    block_titles_texts = [element.text
                          for element in (driver
                                          .find_elements_by_class_name('block_title')
                                          )]

    for text in block_titles_texts:
        num_achievements_match = NUM_ACHIEVEMENTS_REGEX.match(text)

        if num_achievements_match:
            results['num_achievements'] = int(num_achievements_match.group(1))

    try:
        results['metacritic_score'] = int(driver.find_element_by_class_name('score').text)
    except NoSuchElementException:
        # Some games don't have metascores
        pass

    try:
        # Not on sale or Free to Play
        raw_price = driver.find_element_by_class_name('game_purchase_price').text
        if raw_price == 'Free to Play':
            price = 0
        else:
            price = float(raw_price.replace('$', ''))
        results['full_price'] = price
    except NoSuchElementException:
        # On sale
        raw_price = driver.find_element_by_class_name('discount_original_price').text
        results['full_price'] = float(raw_price.replace('$', ''))

    results['game_details'] = []
    game_details_elements = driver.find_elements_by_class_name('game_area_details_specs')
    for element in game_details_elements:
        results['game_details'].append(element.find_element_by_css_selector('a.name').text)

    results['tags'] = []
    try:
        # Try to get the big list of tags if it's there
        driver.find_element_by_css_selector('.app_tag.add_button').click()
        tag_elements = driver.find_elements_by_css_selector('#app_tagging_modal a.app_tag')
    except ElementNotVisibleException:
        # Settle for the short list if not
        tag_elements = driver.find_elements_by_css_selector('a.app_tag')
    for element in tag_elements:
        results['tags'].append(element.text)

    driver.close()

    return results

def insert_with_mapping(*, db, descrs, entity_table, pk_name, join_table, mapping, app_id, crawl_time):
    '''
    Given a db connection, some list of description data, a table of entities, the name of the PK,
    a many-to-many join table for the entities to game crawls,
    a mapping from descrs to entity table IDs, and an app_id/crawl_time,
    update the entity table/mapping if necessary and insert
    the entity in the many-to-many join table.

    Mutates the mapping.
    '''
    for descr in descrs:
        try:
            entity_id = mapping[descr]
        except KeyError:
            entity_id = db[entity_table].insert({'descr': descr})
            mapping[descr] = entity_id

        db[join_table].insert({
            'steam_app_id': app_id,
            'crawl_time': crawl_time,
            pk_name: entity_id,
        })

def do_crawl(app_ids, db):
    '''
    Given a list of steam app IDs and a db connection, do a crawl for the app IDs
    and append the results to our list of crawls in the database.
    '''
    tag_mapping = {r['descr']: r['tag_id'] for r in db['steam_tag'].find()}
    detail_mapping = {r['descr']: r['detail_id'] for r in db['steam_game_detail'].find()}
    genre_mapping = {r['descr']: r['genre_id'] for r in db['steam_genre'].find()}

    # Add a handler here to allow us to gracefully save our work and quit
    # if the user halts a crawl early via Ctrl + C.
    # NOTE: since background jobs ignore SIGINT, this has no effect if
    # the script is launched in the background; see
    # http://stackoverflow.com/questions/1112343/how-do-i-capture-sigint-in-python#comment68802096_1112357
    should_quit = False

    db.begin()
    for app_id in tqdm(app_ids):
        try:
            if should_quit:
                break

            try:
                results = scrape_store_page(app_id)
            except:
                db.commit()
                print("Failed crawl for ID {}".format(app_id))
                raise

            crawl_time = dt.datetime.now()
            results['crawl_time'] = crawl_time

            # Pull the lists off before we insert the main crawl record.
            tags = results['tags']
            del results['tags']
            details = results['game_details']
            del results['game_details']
            genres = results['genres']
            del results['genres']

            db['game_crawl'].insert(results)

            insert_with_mapping(
                db=db,
                descrs=tags,
                entity_table='steam_tag',
                pk_name='tag_id',
                join_table='game_crawl_tag',
                mapping=tag_mapping,
                app_id=app_id,
                crawl_time=crawl_time
            )

            insert_with_mapping(
                db=db,
                descrs=details,
                entity_table='steam_game_detail',
                pk_name='detail_id',
                join_table='game_crawl_detail',
                mapping=detail_mapping,
                app_id=app_id,
                crawl_time=crawl_time
            )

            insert_with_mapping(
                db=db,
                descrs=genres,
                entity_table='steam_genre',
                pk_name='genre_id',
                join_table='game_crawl_genre',
                mapping=genre_mapping,
                app_id=app_id,
                crawl_time=crawl_time
            )

            time.sleep(CRAWL_TIMEOUT)
        except KeyboardInterrupt:
            # Finish this iteration before we exit
            should_quit = True
    db.commit()

def run():
    db = dataset.connect(os.environ['POSTGRES_URI'], ensure_schema=False)

    # Since we're not super worried about having an up-to-date list of apps,
    # run this only if the table is empty
    if db['game'].count() == 0:
        upsert_all_apps(db)

    # For now, just crawl the apps we don't already have
    missing_crawl_query = '''
    SELECT g.steam_app_id
    FROM game g
      LEFT JOIN game_crawl gc
        USING (steam_app_id)
    WHERE gc.steam_app_id IS NULL;
    '''

    missing_app_ids = [r['steam_app_id'] for r in db.query(missing_crawl_query)]

    # TODO remove after we're done testing the crawl function
    missing_app_ids = missing_app_ids[:1]

    do_crawl(missing_app_ids, db)

if __name__ == '__main__':
    run()

<|endoftext|>


<commit_before>"""Classes which manage BCP transports."""
import asyncio

from typing import Union

from mpf.core.bcp.bcp_client import BaseBcpClient

class BcpTransportManager:

    """Manages BCP transports."""

    __slots__ = ["_machine", "_transports", "_readers", "_handlers"]

    def __init__(self, machine):
        """Initialise BCP transport manager."""
        self._machine = machine
        self._transports = []
        self._readers = {}
        self._handlers = {}
        self._machine.events.add_handler("shutdown", self.shutdown)

    def add_handler_to_transport(self, handler, transport: BaseBcpClient):
        """Register client as handler."""
        if handler not in self._handlers:
            self._handlers[handler] = []

        if transport is None:
            raise AssertionError("Cannot register None transport.")

        self._handlers[handler].append(transport)

    def remove_transport_from_handle(self, handler, transport: BaseBcpClient):
        """Remove client from a certain handler."""
        if transport in self._handlers[handler]:
            self._handlers[handler].remove(transport)

    def get_transports_for_handler(self, handler):
        """Get clients which registered for a certain handler."""
        return self._handlers.get(handler, [])

    def register_transport(self, transport, future=None, **kwargs):
        """Register a client."""
        del future
        del kwargs
        self._transports.append(transport)
        self._readers[transport] = self._machine.clock.loop.create_task(self._receive_loop(transport))
        self._readers[transport].add_done_callback(self._done)

    @staticmethod
    def _done(future):
        """Evaluate result of task.

        Will raise exceptions from within task.
        """
        try:
            future.result()
        except asyncio.CancelledError:
            pass

    @asyncio.coroutine
    def _receive_loop(self, transport: BaseBcpClient):
        while True:
            try:
                cmd, kwargs = yield from transport.read_message()
            except IOError:
                self.unregister_transport(transport)
                return

            yield from self._machine.bcp.interface.process_bcp_message(cmd, kwargs, transport)

    def unregister_transport(self, transport: BaseBcpClient):
        """Unregister client."""
        if transport in self._transports:
            self._transports.remove(transport)

        # remove transport from all handlers
        for handler in self._handlers:
            if transport in self._handlers[handler]:
                self._handlers[handler].remove(transport)

        if transport in self._readers:
            self._readers[transport].cancel()
            del self._readers[transport]

        if transport.exit_on_close:
            self._machine.stop("BCP client {} disconnected and exit_on_close is set".format(transport.name))

    def get_all_clients(self):
        """Get a list of all clients."""
        return self._transports

    def get_named_client(self, client_name) -> Union[BaseBcpClient, bool]:
        """Get a client by name."""
        for client in self._transports:
            if client.name == client_name:
                return client
        return False

    def send_to_clients(self, clients, bcp_command, **kwargs):
        """Send command to a list of clients."""
        for client in set(clients):
            self.send_to_client(client, bcp_command, **kwargs)

    def send_to_clients_with_handler(self, handler, bcp_command, **kwargs):
        """Send command to clients which registered for a specific handler."""
        clients = self.get_transports_for_handler(handler)
        self.send_to_clients(clients, bcp_command, **kwargs)

    def send_to_client(self, client: BaseBcpClient, bcp_command, **kwargs):
        """Send command to a specific bcp client."""
        try:
            client.send(bcp_command, kwargs)
        except IOError:
            client.stop()
            self.unregister_transport(client)

    def send_to_all_clients(self, bcp_command, **kwargs):
        """Send command to all bcp clients."""
        for client in self._transports:
            self.send_to_client(client, bcp_command, **kwargs)

    def shutdown(self, **kwargs):
        """Prepare the BCP clients for MPF shutdown."""
        del kwargs
        for client in list(self._transports):
            client.stop()
            self.unregister_transport(client)
<commit_msg>add type annotation<commit_after>"""Classes which manage BCP transports."""
import asyncio

from typing import Union

from mpf.core.bcp.bcp_client import BaseBcpClient

MYPY = False  # noqa
if MYPY:
    from mpf.core.machine import MachineController

class BcpTransportManager:

    """Manages BCP transports."""

    __slots__ = ["_machine", "_transports", "_readers", "_handlers"]

    def __init__(self, machine):
        """Initialise BCP transport manager."""
        self._machine = machine     # type: MachineController
        self._transports = []
        self._readers = {}
        self._handlers = {}
        self._machine.events.add_handler("shutdown", self.shutdown)

    def add_handler_to_transport(self, handler, transport: BaseBcpClient):
        """Register client as handler."""
        if handler not in self._handlers:
            self._handlers[handler] = []

        if transport is None:
            raise AssertionError("Cannot register None transport.")

        self._handlers[handler].append(transport)

    def remove_transport_from_handle(self, handler, transport: BaseBcpClient):
        """Remove client from a certain handler."""
        if transport in self._handlers[handler]:
            self._handlers[handler].remove(transport)

    def get_transports_for_handler(self, handler):
        """Get clients which registered for a certain handler."""
        return self._handlers.get(handler, [])

    def register_transport(self, transport, future=None, **kwargs):
        """Register a client."""
        del future
        del kwargs
        self._transports.append(transport)
        self._readers[transport] = self._machine.clock.loop.create_task(self._receive_loop(transport))
        self._readers[transport].add_done_callback(self._done)

    @staticmethod
    def _done(future):
        """Evaluate result of task.

        Will raise exceptions from within task.
        """
        try:
            future.result()
        except asyncio.CancelledError:
            pass

    @asyncio.coroutine
    def _receive_loop(self, transport: BaseBcpClient):
        while True:
            try:
                cmd, kwargs = yield from transport.read_message()
            except IOError:
                self.unregister_transport(transport)
                return

            yield from self._machine.bcp.interface.process_bcp_message(cmd, kwargs, transport)

    def unregister_transport(self, transport: BaseBcpClient):
        """Unregister client."""
        if transport in self._transports:
            self._transports.remove(transport)

        # remove transport from all handlers
        for handler in self._handlers:
            if transport in self._handlers[handler]:
                self._handlers[handler].remove(transport)

        if transport in self._readers:
            self._readers[transport].cancel()
            del self._readers[transport]

        if transport.exit_on_close:
            self._machine.stop("BCP client {} disconnected and exit_on_close is set".format(transport.name))

    def get_all_clients(self):
        """Get a list of all clients."""
        return self._transports

    def get_named_client(self, client_name) -> Union[BaseBcpClient, bool]:
        """Get a client by name."""
        for client in self._transports:
            if client.name == client_name:
                return client
        return False

    def send_to_clients(self, clients, bcp_command, **kwargs):
        """Send command to a list of clients."""
        for client in set(clients):
            self.send_to_client(client, bcp_command, **kwargs)

    def send_to_clients_with_handler(self, handler, bcp_command, **kwargs):
        """Send command to clients which registered for a specific handler."""
        clients = self.get_transports_for_handler(handler)
        self.send_to_clients(clients, bcp_command, **kwargs)

    def send_to_client(self, client: BaseBcpClient, bcp_command, **kwargs):
        """Send command to a specific bcp client."""
        try:
            client.send(bcp_command, kwargs)
        except IOError:
            client.stop()
            self.unregister_transport(client)

    def send_to_all_clients(self, bcp_command, **kwargs):
        """Send command to all bcp clients."""
        for client in self._transports:
            self.send_to_client(client, bcp_command, **kwargs)

    def shutdown(self, **kwargs):
        """Prepare the BCP clients for MPF shutdown."""
        del kwargs
        for client in list(self._transports):
            client.stop()
            self.unregister_transport(client)
<|endoftext|>