Python trio.open_nursery() Examples

def _monitor_subscription_changes(self) -> None:
        while not self.is_stopped:
            # We wait for the event to change and then immediately replace it
            # with a new event.  This **must** occur before any additional
            # `await` calls to ensure that any *new* changes to the
            # subscriptions end up operating on the *new* event and will be
            # picked up in the next iteration of the loop.
            await self._subscriptions_changed.wait()
            self._subscriptions_changed = trio.Event()

            # make a copy so that the set doesn't change while we iterate
            # over it
            subscribed_events = self.get_subscribed_events()

            async with trio.open_nursery() as nursery:
                async with self._remote_connections_changed:
                    for remote in self._connections:
                        nursery.start_soon(
                            remote.notify_subscriptions_updated,
                            subscribed_events,
                            False,
                        )
            async with self._remote_subscriptions_changed:
                self._remote_subscriptions_changed.notify_all() 

def getTypesL2(target, types, href):
    """
        取得二级分类
    """
    loger.info(colored(f'fetching {href}', 'yellow'))
    resp = await spiderSession.get(href)
    async with trio.open_nursery() as nursery:
        for item in jq(resp.text)("body > div.content-base > section > div > table > tbody > tr").items():
            name = item(
                'td:nth-child(1)>a').text().strip().replace(' ', '_').lower()
            target[name] = {}
            url = urljoin(href, item('td:nth-child(1)>a').attr('href'))
            nums = int(item('td:nth-child(2)').text().strip())
            target[name]['url'] = url
            target[name]['nums'] = nums
            target[name]['UA_list'] = []
            for page in range(1, math.ceil(nums/PERPAGE)+1):
                TASKS.add('__'.join([
                    types,
                    name,
                    f"{url}{page}"
                ])) 

def test_fetch_concurrent(asyncio_loop, nursery, autojump_clock):
    ''' If two tasks request the same robots.txt at the same time, one blocks
    while the other requests the file over the network.. '''
    db_pool = Mock()
    dl = Mock()
    dl.download = AsyncMock(side_effect=download)
    rtm = RobotsTxtManager(db_pool)
    rtm._get_robots_from_db = AsyncMock()
    rtm._save_robots_to_db = AsyncMock()
    policy = make_policy(usage='OBEY', user_agent='TestAgent1')

    async def request1():
        assert await rtm.is_allowed('https://www.example/index', policy, dl)

    async def request2():
        assert not await rtm.is_allowed('https://www.example/bar/', policy, dl)

    async with trio.open_nursery() as inner:
        inner.start_soon(request1)
        inner.start_soon(request2)

    assert dl.download.call_count == 1 

def getTypesL2(target, types, href):
        """
            取得二级分类
        """
        loger.info(colored(f'fetching {href}', 'yellow'))
        resp = await spiderSession.get(href)
        async with trio.open_nursery() as nursery:
            for item in jq(resp.text)("body > div.content-base > section > div > table > tbody > tr").items():
                name = item(
                    'td:nth-child(1)>a').text().strip().replace(' ', '_').lower()
                target[name] = {}
                url = urljoin(href, item('td:nth-child(1)>a').attr('href'))
                nums = int(item('td:nth-child(2)').text().strip())
                target[name]['url'] = url
                target[name]['nums'] = nums
                target[name]['UA_list'] = []
                for page in range(1, math.ceil(nums/PERPAGE)+1):
                    TASKS.add('__'.join([
                        types,
                        name,
                        f"{url}{page}"
                    ])) 

def _async_get_clients_report(self, clients, agent=None):
        """See :func:`burpui.misc.backend.interface.BUIbackend.get_clients_report`"""
        async def __compute_client_report(cli, queue, limit):
            async with limit:
                if not cli:
                    return
                client = await self._async_get_client(cli['name'])
                if not client or not client[-1]:
                    return
                stats = await self._async_get_backup_logs(client[-1]['number'], cli['name'])
                queue.append((cli, client, stats))

        data = []
        limiter = trio.CapacityLimiter(self.concurrency)

        async with trio.open_nursery() as nursery:
            for client in clients:
                nursery.start_soon(__compute_client_report, client, data, limiter)

        return self._do_get_clients_report(data) 

def main():
    db_config = get_config()['database']
    async with trio.open_nursery() as nursery:
        conn = await r.connect(
            host=db_config['host'],
            port=db_config['port'],
            db=db_config['db'],
            user=db_config['user'],
            password=db_config['password'],
            nursery=nursery
        )
        await clear(conn, 'captcha_solver')
        await clear(conn, 'domain_login')
        await clear(conn, 'frontier')
        await clear(conn, 'job')
        await clear(conn, 'job_schedule')
        await clear(conn, 'policy')
        await clear(conn, 'rate_limit')
        await clear(conn, 'response')
        await clear(conn, 'response_body')
        await clear(conn, 'robots_txt')
        await conn.close() 

def run_query(query, super_user=False):
    ''' Run ``query`` on RethinkDB and return result. '''
    async def async_query():
        db_config = config['database']
        kwargs = {
            'host': db_config['host'],
            'port': db_config['port'],
            'db': db_config['db'],
            'user': db_config['user'],
            'password': db_config['password'],
        }
        if super_user:
            kwargs['user'] = db_config['super_user']
            kwargs['password'] = db_config['super_password']
        async with trio.open_nursery() as nursery:
            kwargs['nursery'] = nursery
            connect_db = functools.partial(r.connect, **kwargs)
            conn = await connect_db()
            try:
                result = await query.run(conn)
            finally:
                await conn.close()
            return result

    return trio.run(async_query) 

def _async_get_clients_report(self, clients, agent=None):
        """See :func:`burpui.misc.backend.interface.BUIbackend.get_clients_report`"""
        async def __compute_client_report(cli, queue, limit):
            async with limit:
                if not cli:
                    return
                client = await self._async_get_client(cli['name'])
                if not client or not client[-1]:
                    return
                stats = await self._async_get_backup_logs(client[-1]['number'], cli['name'])
                queue.append((cli, client, stats))

        data = []
        limiter = trio.CapacityLimiter(self.concurrency)

        async with trio.open_nursery() as nursery:
            for client in clients:
                nursery.start_soon(__compute_client_report, client, data, limiter)

        return self._do_get_clients_report(data) 

def run(self):
        """
        Run the subscription.

        :returns: This function returns when the sync is complete.
        """
        logger.info("%r Starting", self)
        async with trio.open_nursery() as nursery:
            self._cancel_scope = nursery.cancel_scope
            await self._set_initial_job_status()
            nursery.start_soon(self._job_status_task)
            try:
                await self._run_sync()
            except (trio.BrokenResourceError, trio.ClosedResourceError):
                logger.info("%r Aborted", self)
            nursery.cancel_scope.cancel()
        try:
            await self._send_complete()
            logger.info("%r Finished", self)
        except (trio.BrokenResourceError, trio.ClosedResourceError):
            # If we can't send the completion message, then bail out.
            pass 

def test_client_ping_same_payload(echo_conn):
    # This test verifies that two tasks can't ping with the same payload at the
    # same time. One of them should succeed and the other should get an
    # exception.
    exc_count = 0
    async def ping_and_catch():
        nonlocal exc_count
        try:
            await echo_conn.ping(b'A')
        except ValueError:
            exc_count += 1
    async with echo_conn:
        async with trio.open_nursery() as nursery:
            nursery.start_soon(ping_and_catch)
            nursery.start_soon(ping_and_catch)
    assert exc_count == 1 

def run(self, *, task_status=trio.TASK_STATUS_IGNORED):
        '''
        Start serving incoming connections requests.

        This method supports the Trio nursery start protocol: ``server = await
        nursery.start(server.run, …)``. It will block until the server is 
        accepting connections and then return a :class:`WebSocketServer` object.

        :param task_status: Part of the Trio nursery start protocol.
        :returns: This method never returns unless cancelled.
        '''
        async with trio.open_nursery() as nursery:
            serve_listeners = partial(trio.serve_listeners,
                self._handle_connection, self._listeners,
                handler_nursery=self._handler_nursery)
            await nursery.start(serve_listeners)
            logger.debug('Listening on %s',
                ','.join([str(l) for l in self.listeners]))
            task_status.started(self)
            await trio.sleep_forever() 

def run(self):
        """
        Start the subscription stream.

        :returns: This function runs until ``cancel()`` is called.
        """
        logger.info("%r Starting", self)
        initial = True
        async with trio.open_nursery() as nursery:
            self._cancel_scope = nursery.cancel_scope

            while True:
                message = self._make_event()
                if initial or message.event.job_list.jobs:
                    await self._websocket.send_message(message.SerializeToString())
                initial = False
                await trio.sleep(self._min_interval)
        logger.info("%r Cancelled", self) 

def _async_get_clients_report(self, clients, agent=None):
        """See :func:`burpui.misc.backend.interface.BUIbackend.get_clients_report`"""
        async def __compute_client_report(cli, queue, limit):
            async with limit:
                if not cli:
                    return
                client = await self._async_get_client(cli['name'])
                if not client or not client[-1]:
                    return
                stats = await self._async_get_backup_logs(client[-1]['number'], cli['name'])
                queue.append((cli, client, stats))

        data = []
        limiter = trio.CapacityLimiter(self.concurrency)

        async with trio.open_nursery() as nursery:
            for client in clients:
                nursery.start_soon(__compute_client_report, client, data, limiter)

        return self._do_get_clients_report(data) 

def run(self, *, task_status=trio.TASK_STATUS_IGNORED):
        '''
        Run the crawl manager.

        You should call ``await nursery.start(crawl_manager.run)`` to ensure
        that the crawl manager is ready before calling any of its job methods.

        :returns: This function runs until cancelled.
        '''
        max_sequence = await self._db.get_max_sequence()
        self._sequence = itertools.count(start=max_sequence + 1)
        logger.info('%r Sequence initialized to %s', self, max_sequence + 1)

        async with trio.open_nursery() as nursery:
            self._nursery = nursery
            task_status.started()
            await trio.sleep_forever() 

def test_trio_endpoint_wait_for(endpoint_pair):
    alice, bob = endpoint_pair

    # NOTE: this test is the inverse of the broadcast test
    event = EventTest("test")

    done = trio.Event()

    async def _do_wait_for():
        result = await alice.wait_for(EventTest)

        assert isinstance(result, EventTest)
        assert result.value == "test"
        done.set()

    async with trio.open_nursery() as nursery:
        nursery.start_soon(_do_wait_for)

        await bob.wait_until_endpoint_subscribed_to(alice.name, EventTest)

        await bob.broadcast(event)
        await done.wait() 

def run(self):
        """
        Start the subscription stream.

        :returns: This method runs until cancel is called.
        """
        logger.info("%r Starting", self)
        async with trio.open_nursery() as nursery:
            self._cancel_scope = nursery.cancel_scope
            # Open a channel to the resource monitor now. We choose a channel
            # size that is big enough to hold up to 10 seconds worth of data,
            # since it may take a few seconds to send all of the requested
            # historical data.
            self._recv_channel = self._resource_monitor.get_channel(10)
            for measurement in self._resource_monitor.history(self._history):
                event = self._make_event(measurement)
                await self._websocket.send_message(event.SerializeToString())

            async for measurement in self._recv_channel:
                event = self._make_event(measurement)
                await self._websocket.send_message(event.SerializeToString())
        logger.info("%r Cancelled", self) 

def run(self, *, task_status=trio.TASK_STATUS_IGNORED):
        '''
        Run the websocket server.

        To ensure that the server is ready, call ``await
        nursery.start(server.run)``.

        :returns: Runs until cancelled.
        '''
        logger.info('Starting server on %s:%d', self._host, self._port)
        async with trio.open_nursery() as nursery:
            serve_fn = partial(serve_websocket, self._handle_connection,
                self._host, self._port, ssl_context=None,
                handler_nursery=nursery)
            server = await nursery.start(serve_fn, name='Connection Listener')
            self._port = server.port
            task_status.started()
        logger.info('Server stopped') 

def run(self):
        '''
        Run the connection: read requests and send responses.

        This opens an internal nursery in case background tasks, like
        subscriptions, need to be started.

        :returns: This runs until the connection is closed.
        '''
        try:
            async with trio.open_nursery() as nursery:
                self._nursery = nursery
                self._subscription_manager = SubscriptionManager(
                    self._subscription_db, nursery, self._ws)
                while True:
                    request_data = await self._ws.get_message()
                    nursery.start_soon(self._handle_request, request_data,
                        name='Request Handler')
        except ConnectionClosed:
            logger.info('Connection closed for %s', self._client)
        except:
            logger.exception('Connection exception')
        finally:
            await self._ws.aclose() 

def main():
    t0 = datetime.datetime.now()
    print(colorama.Fore.WHITE + "App started.", flush=True)

    """
    trio.Queue was removed in v0.11.0:
    - Replacing the call to trio.Queue() by trio.open_memory_channel()
    - Using a MemorySendChannel object in generate_data function
    - Using a MemoryReceiveChannel object in process_data function
    - Updating requirements.txt with trio v0.16.0 and trio_asyncio v0.11.0
    """

    send_channel, receive_channel = trio.open_memory_channel(max_buffer_size=10)

    with trio.move_on_after(5):
        async with trio.open_nursery() as nursery:
            nursery.start_soon(generate_data, 20, send_channel, name='Prod 1')
            nursery.start_soon(generate_data, 20, send_channel, name='Prod 2')
            nursery.start_soon(process_data, 40, receive_channel, name='Consumer')

    dt = datetime.datetime.now() - t0
    print(colorama.Fore.WHITE + "App exiting, total time: {:,.2f} sec.".format(
        dt.total_seconds()), flush=True) 

def test_multiple_contexts():
    async def recv_and_send(ctx):
        data = await ctx.arecv()
        await trio.sleep(0.05)
        await ctx.asend(data)

    with pynng.Rep0(listen=addr, recv_timeout=500) as rep, \
            pynng.Req0(dial=addr, recv_timeout=500) as req1, \
            pynng.Req0(dial=addr, recv_timeout=500) as req2:
        async with trio.open_nursery() as n:
            ctx1, ctx2 = [rep.new_context() for _ in range(2)]
            with ctx1, ctx2:
                n.start_soon(recv_and_send, ctx1)
                n.start_soon(recv_and_send, ctx2)

                await req1.asend(b'oh hi')
                await req2.asend(b'me toooo')
                assert (await req1.arecv() == b'oh hi')
                assert (await req2.arecv() == b'me toooo') 

def run(
        self, task_status: TaskStatus[None] = trio.TASK_STATUS_IGNORED
    ) -> None:
        tasks = (
            self._iterate_clock,
            self._run_validator_api,
            self._run_host,
            self._manage_peers,
            self._manage_sync_requests,
        )

        async with trio.open_nursery() as nursery:
            for task in tasks:
                await nursery.start(task)
            task_status.started() 

def _run_validator_api(
        self, task_status: TaskStatus[None] = trio.TASK_STATUS_IGNORED
    ) -> None:
        server = self._validator_api_server
        async with trio.open_nursery() as nursery:
            self.validator_api_port = await nursery.start(server.serve)
            self.logger.info(
                "validator HTTP API server listening on %d", self.validator_api_port
            )
            task_status.started() 

def XXtest_issue_241(ctx):

    inode = ctx.inode

    # Create block
    async with ctx.cache.get(inode, 0) as fh:
        fh.write(random_data(500))

    # "Fill" cache
    ctx.cache.cache.max_entries = 0

    # Mock locking to reproduce race condition
    mlock = MockMultiLock(ctx.cache.mlock)
    async with trio.open_nursery() as nursery:
        async with trio.open_nursery() as nursery2:
            with patch.object(ctx.cache, 'mlock', mlock):
                # Start expiration, will block on lock
                await nursery.start(ctx.cache.expire)

                # Remove the object while the expiration thread waits
                # for it to become available.
                await nursery2.start(ctx.cache.remove, inode, 0, 1)

                mlock.yield_to(thread2)

        # Create a new object for the same block
        async with ctx.cache.get(inode, 0) as fh:
            fh.write(random_data(500))

        # Continue first expiration run
        mlock.yield_to(thread1, block=False)
        thread1.join_and_raise(timeout=10)
        assert not thread1.is_alive() 

def test_trio_endpoint_request_and_response(endpoint_pair):
    alice, bob = endpoint_pair

    async with trio.open_nursery() as nursery:
        nursery.start_soon(_handle_double_request, alice)

        config = BroadcastConfig(alice.name)

        await bob.wait_until_endpoint_subscribed_to(alice.name, DoubleRequest)

        response = await bob.request(DoubleRequest(7), config)
        assert isinstance(response, DoubleResponse)
        assert response.result == 14 

def serve(cls, config: ConnectionConfig) -> AsyncIterator["TrioEndpoint"]:
        endpoint = cls(config.name)
        async with endpoint.run():
            async with trio.open_nursery() as nursery:
                await endpoint._start_serving(nursery, config.path)
                try:
                    yield endpoint
                finally:
                    await endpoint._stop_serving() 

def test_trio_future_exception():
    fut = Future()

    async def do_set():
        try:
            raise ForTest("testing")
        except ForTest:
            _, err, tb = sys.exc_info()
        fut.set_exception(err, tb)

    async with trio.open_nursery() as nursery:
        nursery.start_soon(do_set)
        with pytest.raises(ForTest, match="testing"):
            await fut 

def test_trio_future_result():
    fut = Future()

    async def do_set():
        fut.set_result("done")

    async with trio.open_nursery() as nursery:
        nursery.start_soon(do_set)
        result = await fut
        assert result == "done" 

def _run_server(self) -> None:
        async with trio.open_nursery() as nursery:
            # Store nursery on self so that we can access it for cancellation
            self._server_nursery = nursery

            self.logger.debug("%s: server starting", self)
            socket = trio.socket.socket(trio.socket.AF_UNIX, trio.socket.SOCK_STREAM)
            await socket.bind(self.ipc_path.__fspath__())
            socket.listen(self._connection_backlog)
            listener = trio.SocketListener(socket)

            async def set_socket_bound() -> None:
                self._socket_bound.set()

            # Use start_soon here so that we give serve_listeners() below a chance to run before
            # other endpoints start connecting to us.
            nursery.start_soon(set_socket_bound)

            try:
                await trio.serve_listeners(
                    handler=self._accept_conn,
                    listeners=(listener,),
                    handler_nursery=nursery,
                )
            finally:
                self.logger.debug("%s: server finished", self) 

def test_trio_endpoint_stream_without_limit(endpoint_pair):
    alice, bob = endpoint_pair

    async with trio.open_nursery() as nursery:
        done = trio.Event()

        async def _do_stream():
            results = []
            async for event in alice.stream(EventTest):
                results.append(event)
                if len(results) == 4:
                    break

            assert len(results) == 4
            assert all(isinstance(result, EventTest) for result in results)

            values = [result.value for result in results]
            assert values == [0, 1, 2, 3]
            done.set()

        nursery.start_soon(_do_stream)
        await bob.wait_until_endpoint_subscribed_to(alice.name, EventTest)

        await bob.broadcast(EventTest(0))
        await bob.broadcast(EventTest(1))
        await bob.broadcast(EventTest(2))
        await bob.broadcast(EventTest(3))

        await done.wait() 

def test_trio_endpoint_stream_with_limit(endpoint_pair):
    alice, bob = endpoint_pair

    async with trio.open_nursery() as nursery:
        done = trio.Event()

        async def _do_stream():
            results = []
            async for event in alice.stream(EventTest, num_events=4):
                results.append(event)
                assert len(results) <= 4

            assert len(results) == 4
            assert all(isinstance(result, EventTest) for result in results)

            values = [result.value for result in results]
            assert values == [0, 1, 2, 3]
            done.set()

        nursery.start_soon(_do_stream)

        await bob.wait_until_endpoint_subscribed_to(alice.name, EventTest)
        await bob.broadcast(EventTest(0))
        await bob.broadcast(EventTest(1))
        await bob.broadcast(EventTest(2))
        await bob.broadcast(EventTest(3))

        await done.wait()