Python tornado.gen.maybe_future() Examples

def test_after_send_error_event_called():
    tchannel = TChannel('test')
    tchannel.listen()
    with mock.patch(
        'tchannel.event.EventEmitter.fire', autospec=True,
    ) as mock_fire:
        mock_fire.return_value = maybe_future(None)
        with pytest.raises(BadRequestError):
            yield tchannel.call(
                scheme=schemes.RAW,
                service='test',
                arg1='endpoint',
                hostport=tchannel.hostport,
                timeout=0.3,
            )
        mock_fire.assert_any_call(
            mock.ANY, EventType.after_send_error, mock.ANY,
        ) 

def test_peer_incoming_connections_are_preferred(request):
    incoming = mock.MagicMock()
    incoming.closed = False

    outgoing = mock.MagicMock()
    outgoing.closed = False

    peer = tpeer.Peer(mock.MagicMock(), 'localhost:4040')
    with mock.patch(
        'tchannel.tornado.connection.StreamConnection.outgoing'
    ) as mock_outgoing:
        mock_outgoing.return_value = gen.maybe_future(outgoing)
        peer.connect()

    assert (yield peer.connect()) is outgoing

    peer.register_incoming_conn(incoming)
    assert (yield peer.connect()) is incoming 

def authorization_required(func):
    @wraps(func)
    @coroutine
    def wrap(self, *args, **kwargs):
        auth_header = self.request.headers.get('Authorization')
        if not auth_header:
            self.set_header('WWW-Authenticate', 'Basic realm="pypi"')
            self.set_status(401)
            raise Return(self.finish("Authorization required"))

        auth_type, data = auth_header.split()
        if auth_type.lower() != 'basic':
            raise Return(self.send_error(400))

        username, password = map(lambda x: unquote_plus(x.decode("utf-8")), base64.b64decode(b(data)).split(b(":")))
        try:
            self.current_user = yield check_password(username, password)
        except LookupError:
            raise HTTPError(403)

        result = yield maybe_future(func(self, *args, **kwargs))
        raise Return(result)

    return wrap 

def deprecated_build_handler(result_type, f):
    @gen.coroutine
    def handler(request, response):
        req = yield ThriftRequest._from_raw_request(request)
        res = ThriftResponse(result_type())
        try:
            # TODO: It would be nice if we could wait until write_result was
            # called or an exception was thrown instead of waiting for the
            # function to return. This would allow for use cases where the
            # implementation returns the result early but still does some work
            # after that.
            result = yield gen.maybe_future(f(req, res))
        except Exception:
            res.write_exc_info(sys.exc_info())
        else:
            if not res.finished and result is not None:
                # The user never called write_result or threw an
                # exception. The result was most likely returned by the
                # function.
                res.write_result(result)
        response.code = res.code
        response.write_header(res.headers)
        response.write_body(res.result)
    return handler 

def build_handler(result_type, f):
    @gen.coroutine
    def handler(request):

        result = ThriftResponse(result_type())
        response = Response()

        try:
            response = yield gen.maybe_future(f(request))
        except Exception:
            result.write_exc_info(sys.exc_info())
        else:
            response = response_from_mixed(response)
            result.write_result(response.body)

        response.status = result.code
        response.body = result.result

        raise gen.Return(response)
    return handler 

def process_submitZipkinBatch(self, seqid, iprot, oprot):
    args = submitZipkinBatch_args()
    args.read(iprot)
    iprot.readMessageEnd()
    result = submitZipkinBatch_result()
    result.success = yield gen.maybe_future(self._handler.submitZipkinBatch(args.spans))
    oprot.writeMessageBegin("submitZipkinBatch", TMessageType.REPLY, seqid)
    result.write(oprot)
    oprot.writeMessageEnd()
    oprot.trans.flush()


# HELPER FUNCTIONS AND STRUCTURES 

def handle_stream(self, stream, address):
        host, port = address
        trans = TTornadoStreamTransport(
            host=host, port=port, stream=stream,
            io_loop=self.__io_loop, read_timeout=self.transport_read_timeout)
        try:
            oprot = self._oprot_factory.get_protocol(trans)
            iprot = self._iprot_factory.get_protocol(TMemoryBuffer())

            while not trans.stream.closed():
                # TODO: maybe read multiple frames in advance for concurrency
                try:
                    frame = yield trans.read_frame()
                except TTransportException as e:
                    if e.type == TTransportException.END_OF_FILE:
                        break
                    else:
                        raise

                iprot.trans.setvalue(frame)
                api, seqid, result, call = self._processor.process_in(iprot)
                if isinstance(result, TApplicationException):
                    self._processor.send_exception(oprot, api, result, seqid)
                else:
                    try:
                        result.success = yield gen.maybe_future(call())
                    except Exception as e:
                        # raise if api don't have throws
                        if not self._processor.handle_exception(e, result):
                            raise

                    self._processor.send_result(oprot, api, result, seqid)
        except Exception:
            logger.exception('thrift exception in handle_stream')
            trans.close()

        logger.info('client disconnected %s:%d', host, port) 

def wrap_uncaught(func=None, reraise=None):
    """Catches uncaught exceptions and raises VCRServiceErrors instead.

    :param reraise:
        Collection of exception clasess that should be re-raised as-is.
    """
    reraise = reraise or ()

    def decorator(f):

        @wraps(f)
        @gen.coroutine
        def new_f(*args, **kwargs):
            try:
                result = yield gen.maybe_future(f(*args, **kwargs))
            except Exception as e:
                if any(isinstance(e, cls) for cls in reraise):
                    raise e
                raise proxy.VCRServiceError(str(e) +
                                            ' ' + traceback.format_exc())
            else:
                raise gen.Return(result)

        return new_f

    if func is not None:
        return decorator(func)
    else:
        return decorator 

def process_joinTrace(self, seqid, iprot, oprot):
    args = joinTrace_args()
    args.read(iprot)
    iprot.readMessageEnd()
    result = joinTrace_result()
    result.success = yield gen.maybe_future(self._handler.joinTrace(args.request))
    oprot.writeMessageBegin("joinTrace", TMessageType.REPLY, seqid)
    result.write(oprot)
    oprot.writeMessageEnd()
    oprot.trans.flush()


# HELPER FUNCTIONS AND STRUCTURES 

def process_startTrace(self, seqid, iprot, oprot):
    args = startTrace_args()
    args.read(iprot)
    iprot.readMessageEnd()
    result = startTrace_result()
    result.success = yield gen.maybe_future(self._handler.startTrace(args.request))
    oprot.writeMessageBegin("startTrace", TMessageType.REPLY, seqid)
    result.write(oprot)
    oprot.writeMessageEnd()
    oprot.trans.flush() 

def process_submitBatches(self, seqid, iprot, oprot):
    args = submitBatches_args()
    args.read(iprot)
    iprot.readMessageEnd()
    result = submitBatches_result()
    result.success = yield gen.maybe_future(self._handler.submitBatches(args.batches))
    oprot.writeMessageBegin("submitBatches", TMessageType.REPLY, seqid)
    result.write(oprot)
    oprot.writeMessageEnd()
    oprot.trans.flush()


# HELPER FUNCTIONS AND STRUCTURES 

def process_getSamplingStrategy(self, seqid, iprot, oprot):
    args = getSamplingStrategy_args()
    args.read(iprot)
    iprot.readMessageEnd()
    result = getSamplingStrategy_result()
    result.success = yield gen.maybe_future(self._handler.getSamplingStrategy(args.serviceName))
    oprot.writeMessageBegin("getSamplingStrategy", TMessageType.REPLY, seqid)
    result.write(oprot)
    oprot.writeMessageEnd()
    oprot.trans.flush()


# HELPER FUNCTIONS AND STRUCTURES 

def process_emitBatch(self, seqid, iprot, oprot):
    args = emitBatch_args()
    args.read(iprot)
    iprot.readMessageEnd()
    yield gen.maybe_future(self._handler.emitBatch(args.batch))


# HELPER FUNCTIONS AND STRUCTURES 

def process_emitZipkinBatch(self, seqid, iprot, oprot):
    args = emitZipkinBatch_args()
    args.read(iprot)
    iprot.readMessageEnd()
    yield gen.maybe_future(self._handler.emitZipkinBatch(args.spans)) 

def async_fetch(self, task, callback=None):
        '''Do one fetch'''
        url = task.get('url', 'data:,')
        if callback is None:
            callback = self.send_result

        type = 'None'
        start_time = time.time()
        try:
            if url.startswith('data:'):
                type = 'data'
                result = yield gen.maybe_future(self.data_fetch(url, task))
            elif task.get('fetch', {}).get('fetch_type') in ('js', 'phantomjs'):
                type = 'phantomjs'
                result = yield self.phantomjs_fetch(url, task)
            elif task.get('fetch', {}).get('fetch_type') in ('splash', ):
                type = 'splash'
                result = yield self.splash_fetch(url, task)
            elif task.get('fetch', {}).get('fetch_type') in ('puppeteer', ):
                type = 'puppeteer'
                result = yield self.puppeteer_fetch(url, task)
            else:
                type = 'http'
                result = yield self.http_fetch(url, task)
        except Exception as e:
            logger.exception(e)
            result = self.handle_error(type, url, task, start_time, e)

        callback(type, task, result)
        self.on_result(type, task, result)
        raise gen.Return(result) 

def get_future_for_response(corr_id):
        f = RequestRedirectionHandler.outstanding_responses.get(corr_id, None)
        if f is None:
            # Form a future that gets populated when a response for corr_id is seen
            f = Future() #gen.maybe_future(("some response","text/html"))
            RequestRedirectionHandler.outstanding_responses[corr_id] = f
        return f 

def wrapped_get(self, kernel_id):
    # TODO wrap in maybe_future
    yield current_get(self, kernel_id) 

def get(self, *args, **kwargs):
        if self.request.headers.get("Upgrade", "").lower() != 'websocket':
            return await self.http_get(*args, **kwargs)
        else:
            await maybe_future(super().get(*args, **kwargs)) 

def fetch(self, req_or_url, *args, **kwargs):
        """Mocked HTTP fetch

        If the request URL is in self.mocks, build a response from the cached response.
        Otherwise, run the actual request and store the response in self.records.
        """
        if isinstance(req_or_url, HTTPRequest):
            request = req_or_url
        else:
            request = HTTPRequest(req_or_url, *args, **kwargs)

        url_key = self.url_key(request.url)

        if url_key in self.mocks:
            fetch = self.fetch_mock
        else:
            fetch = super().fetch

        error = None
        try:
            response = await gen.maybe_future(fetch(request))
        except HTTPError as e:
            error = e
            response = e.response

        self._record_response(url_key, response)
        # return or raise the original result
        if error:
            raise error
        else:
            return response


# async-request utility from jupyterhub.tests.utils v0.8.1
# used under BSD license 

def _async_response(self, data):
        data = yield maybe_future(data)
        resp = yield self._to_json(data)

        if not self._finished:
            log.debug("Sending: %r", resp)
            self.finish(resp) 

def post(self):
        try:
            action = self.get_body_argument(':action')
            self.request.body_arguments.pop(':action')

            log.debug("Request to call action: %s", action)
            method = getattr(self, "action_{0}".format(action), self._action_not_found)
        except:
            raise HTTPError(400)

        log.info("Calling action: %s", action)
        yield maybe_future(method()) 

def add_slash(cls):
    def redirect(self, *args, **kwargs):
        pass

    class WrappedClass(cls):
        @coroutine
        def prepare(self, *args, **kwargs):
            if not self.request.path.endswith('/'):
                raise Return(self.redirect("{0}/".format(self.request.path)))
            else:
                raise Return((yield maybe_future(cls.prepare(self, *args, **kwargs))))

    WrappedClass.__name__ = cls.__name__

    return WrappedClass 

def load_user_options(self):
        """Load user options from self.user_options dict

        This can be set via POST to the API or via options_from_form

        Only supported argument by default is 'profile'.
        Override in subclasses to support other options.
        """

        if self._profile_list is None:
            if callable(self.profile_list):
                profile_list = yield gen.maybe_future(self.profile_list(self))
            else:
                profile_list = self.profile_list

            self._profile_list = self._init_profile_list(profile_list)

        selected_profile = self.user_options.get('profile', None)
        if self._profile_list:
            yield self._load_profile(selected_profile)
        elif selected_profile:
            self.log.warning("Profile %r requested, but profiles are not enabled", selected_profile)

        # help debugging by logging any option fields that are not recognized
        option_keys = set(self.user_options)
        unrecognized_keys = option_keys.difference(self._user_option_keys)
        if unrecognized_keys:
            self.log.warning(
                "Ignoring unrecognized KubeSpawner user_options: %s",
                ", ".join(
                    map(
                        str,
                        sorted(unrecognized_keys)
                    )
                )
            ) 

def _render_options_form_dynamically(self, current_spawner):
        profile_list = yield gen.maybe_future(self.profile_list(current_spawner))
        profile_list = self._init_profile_list(profile_list)
        return self._render_options_form(profile_list) 

def can_fetch(self, user_agent, url):
        parsed = urlsplit(url)
        domain = parsed.netloc
        if domain in self.robots_txt_cache:
            robot_txt = self.robots_txt_cache[domain]
            if time.time() - robot_txt.mtime() > self.robot_txt_age:
                robot_txt = None
        else:
            robot_txt = None

        if robot_txt is None:
            robot_txt = RobotFileParser()
            try:
                response = yield gen.maybe_future(self.http_client.fetch(
                    urljoin(url, '/robots.txt'), connect_timeout=10, request_timeout=30))
                content = response.body
            except tornado.httpclient.HTTPError as e:
                logger.error('load robots.txt from %s error: %r', domain, e)
                content = ''

            try:
                content = content.decode('utf8', 'ignore')
            except UnicodeDecodeError:
                content = ''

            robot_txt.parse(content.splitlines())
            self.robots_txt_cache[domain] = robot_txt

        raise gen.Return(robot_txt.can_fetch(user_agent, url)) 

def send(self, request_by_broker):
        """
        Sends a dict of requests keyed on broker ID and handles responses.

        Returns a dictionary of the results of
        ``handle_<response.api>_response`` method calls, keyed to the
        corresponding broker ID.

        Raises ``UnhandledResponseError`` if the client subclass does not have
        a ``handle_<response.api>_response`` method available to handle an
        incoming response object.

        If an error occurs in a response, the ``heal_cluster`` flag is set
        and the ``heal()`` method on the cluster is called after processing
        each response.

        Responses are handled in the order they come in, but this method does
        not yield a value until all responses are handled.
        """
        iterator = gen.WaitIterator(**{
            str(broker_id): self.cluster[broker_id].send(request)
            for broker_id, request in six.iteritems(request_by_broker)
        })

        results = {}
        while not iterator.done():
            try:
                response = yield iterator.next()
            except BrokerConnectionError as e:
                log.info("Connection to %s:%s lost", e.host, e.port)
                self.heal_cluster = True
                continue
            except iostream.StreamClosedError:
                log.info("Connection to broker lost.")
                continue
            except Exception:
                log.exception("Error sending request.")
                self.heal_cluster = True
                continue

            handler = getattr(self, "handle_%s_response" % response.api, None)
            if handler is None:
                raise UnhandledResponseError(response.api)

            result = yield gen.maybe_future(handler(response))
            results[int(iterator.current_index)] = result

        if self.heal_cluster:
            yield self.cluster.heal()
            self.heal_cluster = False

        raise gen.Return(results) 

def prepare(self):
        if self.request.method.upper() == 'POST' and not self.request.body.startswith(b("\r\n")):
            boundary = dict(
                filter(
                    lambda x: x[0] == 'boundary',
                    map(
                        lambda x: x.strip().split("="),
                        self.request.headers.get('Content-Type', '').split(';')
                    )
                )
            ).get('boundary')

            if not boundary:
                raise HTTPError(400)

            def normalize(chunk):
                if b('\n\n') not in chunk:
                    return b('\r\n') + chunk[1:]

                ret = b('')
                ret += b('\r\n')
                data, content = chunk.split(b('\n\n'), 1)
                ret += data[1:]
                ret += b('\r\n\r\n')
                ret += content[:-1]
                ret += b('\r\n')
                return ret

            boundary = b("--{0}".format(boundary))

            new_body = boundary.join(
                map(
                    normalize,
                    self.request.body.split(boundary)
                )
            )

            new_body = new_body[:-4]
            new_body += b('--\r\n')

            self.request.body = new_body
            self.request._parse_body()

        yield maybe_future(super(XmlRPC, self).prepare()) 

def build_handler(function, handler):
    # response_cls is a class that represents the response union for this
    # function. It accepts one parameter for each exception defined on the
    # method and another parameter 'success' for the result of the call. The
    # success kwarg is absent if the function doesn't return anything.
    response_cls = function._response_cls
    response_spec = response_cls.type_spec

    @gen.coroutine
    def handle(request):
        # kwargs for this function's response_cls constructor
        response_kwargs = {}
        status = OK

        try:
            response = yield gen.maybe_future(handler(request))
        except Exception as e:
            response = Response()

            for exc_spec in response_spec.exception_specs:
                # Each exc_spec is a thriftrw.spec.FieldSpec. The spec
                # attribute on that is the TypeSpec for the Exception class
                # and the surface on the TypeSpec is the exception class.
                exc_cls = exc_spec.spec.surface
                if isinstance(e, exc_cls):
                    status = FAILED
                    response_kwargs[exc_spec.name] = e
                    break
            else:
                raise_exc_info(sys.exc_info())
        else:
            response = response_from_mixed(response)

            if response_spec.return_spec is not None:
                assert response.body is not None, (
                    'Expected a value to be returned for %s, '
                    'but recieved None - only void procedures can '
                    'return None.' % function.endpoint
                )
                response_kwargs['success'] = response.body

        response.status = status
        response.body = response_cls(**response_kwargs)
        raise gen.Return(response)

    handle.__name__ = function.spec.name

    return handle 

def test_loop_failure(tornado_pair):
    server, client = tornado_pair
    headers = dummy_headers()

    # ... yeah
    server.tchannel = mock.MagicMock()
    server.tchannel.event_emitter.fire.return_value = gen.maybe_future(None)

    client.tchannel = mock.MagicMock()
    client.tchannel.event_emitter.fire.return_value = gen.maybe_future(None)

    handshake_future = client.initiate_handshake(headers=headers)
    yield server.expect_handshake(headers=headers)
    yield handshake_future

    assert client._handshake_performed
    assert server._handshake_performed

    # We'll put an invalid message into the reader queue. This should cause one
    # iteration of the loop to fail but the system should continue working
    # afterwards.

    yield server.reader.queue.put(gen.maybe_future(42))  # not a message

    id = client.writer.next_message_id()
    response_future = client.send_request(Request(
        id=id,
        service='server',
        endpoint='bar',
        headers={'cn': 'client'},
    ))

    call_req = yield server._await()
    assert call_req.message_type == messages.Types.CALL_REQ

    response = Response(id=id)
    response.close_argstreams(force=True)

    yield server.post_response(response)
    yield response_future

    assert client._handshake_performed
    assert server._handshake_performed

    client.close()

    # The system needs a little time to recognize that the connections were
    # closed.
    yield gen.sleep(0.15)

    assert client.closed
    assert server.closed