Commit 4e3b3415 authored by Nanshu Chen's avatar Nanshu Chen Committed by Facebook Github Bot

Python async_generator bridge to folly::coro::AsyncGenerator

Summary: Given a cpp folly::coro::AsyncGenerator, a python async generator/iterator may be built and bridged to it. Users may use `async for ...` syntax to consume data from it.

Reviewed By: yfeldblum

Differential Revision: D16404979

fbshipit-source-id: dd2dfd52d00bd1a3e30e838359a27458744f6461
parent ae2d0a61
/*
* Copyright 2019-present Facebook, Inc.
*
* 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.
*/
#pragma once
#include <folly/Optional.h>
#include <folly/experimental/coro/AsyncGenerator.h>
#include <folly/experimental/coro/Task.h>
namespace folly {
namespace python {
template <typename T>
class AsyncGeneratorWrapper {
public:
AsyncGeneratorWrapper() = default;
explicit AsyncGeneratorWrapper(coro::AsyncGenerator<T>&& gen)
: gen_(std::move(gen)) {}
coro::Task<Optional<T>> getNext() {
if (!iter_) {
iter_ = co_await gen_.begin();
} else {
co_await(++*iter_);
}
if (iter_ == gen_.end()) {
co_return none;
}
co_return(**iter_);
}
private:
coro::AsyncGenerator<T> gen_;
Optional<typename coro::AsyncGenerator<T>::async_iterator> iter_;
};
} // namespace python
} // namespace folly
# Copyright 2019-present Facebook, Inc.
#
# 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.
from folly.coro cimport cFollyCoroTask
from folly.optional cimport cOptional
cdef extern from "folly/experimental/coro/AsyncGenerator.h" namespace "folly::coro":
cdef cppclass cAsyncGenerator[T]:
pass
cdef extern from "folly/python/async_generator.h" namespace "folly::python":
cdef cppclass cAsyncGeneratorWrapper "folly::python::AsyncGeneratorWrapper"[T]:
cAsyncGeneratorWrapper() except +
cAsyncGeneratorWrapper(cAsyncGenerator[T]&&) except +
cFollyCoroTask[cOptional[T]] getNext()
# Copyright 2019-present Facebook, Inc.
#
# 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.
/*
* Copyright 2019-present Facebook, Inc.
*
* 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.
*/
#include <folly/experimental/coro/BlockingWait.h>
#include <folly/experimental/coro/Task.h>
#include <folly/portability/GTest.h>
#include <folly/python/async_generator.h>
namespace folly {
namespace python {
namespace test {
coro::AsyncGenerator<int> make_generator() {
co_yield 1;
co_yield 2;
co_yield 3;
co_return;
}
using IntGeneratorWrapper = AsyncGeneratorWrapper<int>;
TEST(AsyncGeneratorTest, IterGenerator) {
coro::blockingWait([]() -> coro::Task<void> {
IntGeneratorWrapper generator{make_generator()};
auto v = co_await generator.getNext();
EXPECT_EQ(v, 1);
v = co_await generator.getNext();
EXPECT_EQ(v, 2);
v = co_await generator.getNext();
EXPECT_EQ(v, 3);
v = co_await generator.getNext();
EXPECT_EQ(v, none);
}());
}
} // namespace test
} // namespace python
} // namespace folly
#!/usr/bin/env python3
# Copyright 2019-present Facebook, Inc.
#
# 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.
import asyncio
import unittest
from .simplegenerator import SimpleGenerator
class GeneratorTest(unittest.TestCase):
def test_iter_generator(self) -> None:
loop = asyncio.get_event_loop()
async def wrapper() -> None:
gen = SimpleGenerator("normal")
expected = 1
async for v in gen:
self.assertEqual(v, expected)
expected += 1
self.assertEqual(expected, 6)
loop.run_until_complete(wrapper())
def test_iter_generator_empty(self) -> None:
loop = asyncio.get_event_loop()
async def wrapper() -> None:
gen = SimpleGenerator("empty")
async for _ in gen: # noqa: F841
self.assertFalse(True, "this should never run")
else:
self.assertTrue(
True, "this will be run when generator is empty, as expected"
)
loop.run_until_complete(wrapper())
def test_iter_generator_error(self) -> None:
loop = asyncio.get_event_loop()
async def wrapper() -> None:
gen = SimpleGenerator("error")
async for v in gen:
self.assertEqual(v, 42)
break
with self.assertRaises(RuntimeError):
async for v in gen:
pass
loop.run_until_complete(wrapper())
/*
* Copyright 2019-present Facebook, Inc.
*
* 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.
*/
#pragma once
#include <folly/experimental/coro/AsyncGenerator.h>
namespace folly {
namespace python {
namespace test {
namespace {
struct SomeError : std::exception {};
} // namespace
coro::AsyncGenerator<int> make_generator_empty() {
coro::AsyncGenerator<int> g;
return g;
}
coro::AsyncGenerator<int> make_generator() {
co_yield 1;
co_yield 2;
co_yield 3;
co_yield 4;
co_yield 5;
co_return;
}
coro::AsyncGenerator<int> make_generator_error() {
co_yield 42;
throw SomeError{};
}
} // namespace test
} // namespace python
} // namespace folly
# Copyright 2019-present Facebook, Inc.
#
# 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.
import asyncio
from cpython.ref cimport PyObject
from folly cimport cFollyTry
from folly.coro cimport bridgeCoroTaskWith
from folly.executor cimport cAsyncioExecutor, get_executor
from folly.async_generator cimport cAsyncGenerator, cAsyncGeneratorWrapper
from folly.optional cimport cOptional
cdef extern from "folly/python/test/simplegenerator.h" namespace "folly::python::test":
cAsyncGenerator[int] make_generator_empty()
cAsyncGenerator[int] make_generator()
cAsyncGenerator[int] make_generator_error()
ctypedef cAsyncGeneratorWrapper[int] cIntGenerator
cdef extern from "<utility>" namespace "std" nogil:
cdef cIntGenerator move "std::move"(cIntGenerator)
cdef class SimpleGenerator:
cdef cIntGenerator generator
cdef cAsyncioExecutor* executor
def __cinit__(self, flavor):
if flavor == "normal":
self.generator = move(cIntGenerator(make_generator()))
elif flavor == "empty":
self.generator = move(cIntGenerator(make_generator_empty()))
elif flavor == "error":
self.generator = move(cIntGenerator(make_generator_error()))
@staticmethod
cdef void callback(
cFollyTry[cOptional[int]]&& res,
PyObject* py_future,
):
cdef cOptional[int] opt_val
future = <object> py_future
if res.hasException():
try:
res.exception().throw_exception()
except Exception as ex:
future.set_exception(ex)
else:
opt_val = res.value()
if opt_val.has_value():
future.set_result(opt_val.value())
else:
future.set_exception(StopAsyncIteration())
def __aiter__(self):
return self
def __anext__(self):
loop = asyncio.get_event_loop()
future = loop.create_future()
bridgeCoroTaskWith[cOptional[int]](
get_executor(),
self.generator.getNext(),
SimpleGenerator.callback,
<PyObject *> future,
)
return future
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment