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