Unverified Commit c0454b32 authored by Christopher Chou's avatar Christopher Chou Committed by GitHub
Browse files

Add option to return metadata in async streaming (#18)

parent 8024fc5e
...@@ -365,11 +365,10 @@ class StreamExecutor: ...@@ -365,11 +365,10 @@ class StreamExecutor:
for comp, meta_info in generator: for comp, meta_info in generator:
self.text_ += comp self.text_ += comp
self.variables[name] += comp self.variables[name] += comp
self.meta_info[name] = meta_info
self.stream_var_event[name].set() self.stream_var_event[name].set()
self.stream_text_event.set() self.stream_text_event.set()
self.meta_info[name] = meta_info
self.variable_event[name].set() self.variable_event[name].set()
self.stream_var_event[name].set() self.stream_var_event[name].set()
...@@ -583,7 +582,7 @@ class ProgramState: ...@@ -583,7 +582,7 @@ class ProgramState:
else: else:
yield self.get_var(name) yield self.get_var(name)
async def text_async_iter(self, var_name=None): async def text_async_iter(self, var_name=None, return_meta_data=False):
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
if self.stream_executor.stream: if self.stream_executor.stream:
...@@ -607,7 +606,10 @@ class ProgramState: ...@@ -607,7 +606,10 @@ class ProgramState:
out = str(self.stream_executor.variables[var_name][prev:]) out = str(self.stream_executor.variables[var_name][prev:])
prev += len(out) prev += len(out)
if out: if out:
yield out if return_meta_data:
yield out, self.stream_executor.meta_info[var_name]
else:
yield out
if self.stream_executor.variable_event[var_name].is_set(): if self.stream_executor.variable_event[var_name].is_set():
break break
else: else:
......
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