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

refactor AsyncGeneratorWrapper to use AsyncGenerator::NextResult

Summary: AsyncGeneratorWrapper uses folly::Optional as the return type for "getNext". Now the AsyncGenerator API has changed by adding `next()` method which returns a `NextResult` type, which works like folly::Optional while supporting both values and references, it's best to use `NextResult` for the wrapper as well.

Reviewed By: yfeldblum

Differential Revision: D16695661

fbshipit-source-id: 84e9d74e8d6b17b93bd3c22e2aaf56c7b9e4b789
parent 0c5ceda6
......@@ -15,13 +15,15 @@
*/
#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>
using NextResult = typename coro::AsyncGenerator<T>::NextResult;
template <typename T>
class AsyncGeneratorWrapper {
public:
......@@ -29,13 +31,8 @@ class AsyncGeneratorWrapper {
explicit AsyncGeneratorWrapper(coro::AsyncGenerator<T>&& gen)
: gen_(std::move(gen)) {}
coro::Task<Optional<T>> getNext() {
auto item = co_await gen_.next();
if (item) {
co_return std::move(item).value();
} else {
co_return none;
}
coro::Task<NextResult<T>> getNext() {
co_return co_await gen_.next();
}
private:
......
......@@ -13,16 +13,18 @@
# 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]:
cdef cppclass cAsyncGenerator "folly::coro::AsyncGenerator"[T]:
pass
cdef extern from "folly/python/async_generator.h" namespace "folly::python":
cdef cppclass cNextResult "folly::python::NextResult"[T]:
bint has_value()
T value()
cdef cppclass cAsyncGeneratorWrapper "folly::python::AsyncGeneratorWrapper"[T]:
cAsyncGeneratorWrapper() except +
cAsyncGeneratorWrapper(cAsyncGenerator[T]&&) except +
cFollyCoroTask[cOptional[T]] getNext()
cFollyCoroTask[cNextResult[T]] getNext()
......@@ -15,7 +15,7 @@
*/
#pragma once
#include <folly/python/executor_api.h>
#include <folly/python/executor_api.h> // @manual
namespace folly {
namespace python {
......
......@@ -23,26 +23,73 @@ namespace folly {
namespace python {
namespace test {
coro::AsyncGenerator<int> make_generator() {
coro::AsyncGenerator<int> make_generator_val() {
co_yield 1;
co_yield 2;
co_yield 3;
co_return;
}
using IntGeneratorWrapper = AsyncGeneratorWrapper<int>;
std::vector<int> values{1, 2, 3};
coro::AsyncGenerator<int&> make_generator_ref() {
co_yield values[0];
co_yield values[1];
co_yield values[2];
co_return;
}
class MoveOnly : public folly::MoveOnly {
public:
explicit MoveOnly(int val_) : val{val_} {}
int val;
};
coro::AsyncGenerator<MoveOnly&&> make_generator_rref() {
co_yield MoveOnly{1};
co_yield MoveOnly{2};
co_yield MoveOnly{3};
co_return;
}
TEST(AsyncGeneratorTest, IterGeneratorVal) {
coro::blockingWait([]() -> coro::Task<void> {
AsyncGeneratorWrapper<int> generator{make_generator_val()};
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_FALSE(v.has_value());
}());
}
TEST(AsyncGeneratorTest, IterGeneratorRef) {
coro::blockingWait([]() -> coro::Task<void> {
AsyncGeneratorWrapper<int&> generator{make_generator_ref()};
auto v = co_await generator.getNext();
EXPECT_EQ(&*v, &values[0]);
v = co_await generator.getNext();
EXPECT_EQ(&*v, &values[1]);
v = co_await generator.getNext();
EXPECT_EQ(&*v, &values[2]);
v = co_await generator.getNext();
EXPECT_FALSE(v.has_value());
}());
}
TEST(AsyncGeneratorTest, IterGenerator) {
TEST(AsyncGeneratorTest, IterGeneratorRRef) {
coro::blockingWait([]() -> coro::Task<void> {
IntGeneratorWrapper generator{make_generator()};
AsyncGeneratorWrapper<MoveOnly&&> generator{make_generator_rref()};
auto v = co_await generator.getNext();
EXPECT_EQ(v, 1);
EXPECT_EQ(v->val, 1);
v = co_await generator.getNext();
EXPECT_EQ(v, 2);
EXPECT_EQ(v->val, 2);
v = co_await generator.getNext();
EXPECT_EQ(v, 3);
EXPECT_EQ(v->val, 3);
v = co_await generator.getNext();
EXPECT_EQ(v, none);
EXPECT_FALSE(v.has_value());
}());
}
......
......@@ -17,8 +17,7 @@ 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
from folly.async_generator cimport cAsyncGenerator, cAsyncGeneratorWrapper, cNextResult
cdef extern from "folly/python/test/simplegenerator.h" namespace "folly::python::test":
......@@ -50,10 +49,9 @@ cdef class SimpleGenerator:
@staticmethod
cdef void callback(
cFollyTry[cOptional[int]]&& res,
cFollyTry[cNextResult[int]]&& res,
PyObject* py_future,
):
cdef cOptional[int] opt_val
future = <object> py_future
if res.hasException():
try:
......@@ -61,9 +59,8 @@ cdef class SimpleGenerator:
except Exception as ex:
future.set_exception(ex)
else:
opt_val = res.value()
if opt_val.has_value():
future.set_result(opt_val.value())
if res.value().has_value():
future.set_result(res.value().value())
else:
future.set_exception(StopAsyncIteration())
......@@ -73,7 +70,7 @@ cdef class SimpleGenerator:
def __anext__(self):
loop = asyncio.get_event_loop()
future = loop.create_future()
bridgeCoroTaskWith[cOptional[int]](
bridgeCoroTaskWith[cNextResult[int]](
get_executor(),
self.generator.getNext(),
SimpleGenerator.callback,
......
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