Simplify

Author: jakubno

python/e2b_code_interpreter/models.py

@@ -425,28 +425,7 @@ def parse_output(
     on_result: Optional[OutputHandler[Result]] = None,
     on_error: Optional[OutputHandler[ExecutionError]] = None,
 ):
-    data = json.loads(output)
-    data_type = data.pop("type")
-
-    if data_type == "result":
-        result = Result(**data)
-        execution.results.append(result)
-        if on_result:
-            on_result(result)
-    elif data_type == "stdout":
-        execution.logs.stdout.append(data["text"])
-        if on_stdout:
-            on_stdout(OutputMessage(data["text"], data["timestamp"], False))
-    elif data_type == "stderr":
-        execution.logs.stderr.append(data["text"])
-        if on_stderr:
-            on_stderr(OutputMessage(data["text"], data["timestamp"], True))
-    elif data_type == "error":
-        execution.error = ExecutionError(data["name"], data["value"], data["traceback"])
-        if on_error:
-            on_error(execution.error)
-    elif data_type == "number_of_executions":
-        execution.execution_count = data["execution_count"]
+    _parse_output(execution, output, on_stdout, on_stderr, on_result, on_error)
 
 
 async def async_parse_output(
@@ -457,37 +436,46 @@ async def async_parse_output(
     on_result: Optional[OutputHandlerWithAsync[Result]] = None,
     on_error: Optional[OutputHandlerWithAsync[ExecutionError]] = None,
 ):
+    none_or_awaitable = _parse_output(
+        execution, output, on_stdout, on_stderr, on_result, on_error
+    )
+    if inspect.isawaitable(none_or_awaitable):
+        await none_or_awaitable
+
+
+def _parse_output(
+    execution: Execution,
+    output: str,
+    on_stdout: Optional[OutputHandler[OutputMessage]] = None,
+    on_stderr: Optional[OutputHandler[OutputMessage]] = None,
+    on_result: Optional[OutputHandler[Result]] = None,
+    on_error: Optional[OutputHandler[ExecutionError]] = None,
+) -> Union[None, Awaitable[Any]]:
     data = json.loads(output)
     data_type = data.pop("type")
 
     if data_type == "result":
         result = Result(**data)
         execution.results.append(result)
         if on_result:
-            cb = on_result(result)
-            if inspect.isawaitable(cb):
-                await cb
+            return on_result(result)
     elif data_type == "stdout":
         execution.logs.stdout.append(data["text"])
         if on_stdout:
-            cb = on_stdout(OutputMessage(data["text"], data["timestamp"], False))
-            if inspect.isawaitable(cb):
-                await cb
+            return on_stdout(OutputMessage(data["text"], data["timestamp"], False))
     elif data_type == "stderr":
         execution.logs.stderr.append(data["text"])
         if on_stderr:
-            cb = on_stderr(OutputMessage(data["text"], data["timestamp"], True))
-            if inspect.isawaitable(cb):
-                await cb
+            return on_stderr(OutputMessage(data["text"], data["timestamp"], True))
     elif data_type == "error":
         execution.error = ExecutionError(data["name"], data["value"], data["traceback"])
         if on_error:
-            cb = on_error(execution.error)
-            if inspect.isawaitable(cb):
-                await cb
+            return on_error(execution.error)
     elif data_type == "number_of_executions":
         execution.execution_count = data["execution_count"]
 
+    return None
+
 
 @dataclass
 class Context: