langfun 0.0.2.dev20240330__py3-none-any.whl → 0.1.2.dev202501140804__py3-none-any.whl

langfun/core/coding/python/execution.py

@@ -13,50 +13,24 @@
 # limitations under the License.
 """Python code execution."""
 
-import ast
-import contextlib
-import io
-import multiprocessing
-from typing import Any, Callable
+from typing import Any
 
-from langfun.core.coding.python import errors
 from langfun.core.coding.python import parsing
-from langfun.core.coding.python import permissions
 import pyglove as pg
 
 
-# Key in returned dict that captures stdout.
-STDOUT_KEY = '__stdout__'
-
-# Key in the returned dict that represents the final result.
-RESULT_KEY = '__result__'
-_TLS_CODE_RUN_CONTEXT = '__code_run_context__'
-
-
-@contextlib.contextmanager
-def context(**kwargs):
-  """Context manager to inject symbols for code execution."""
-  ctx = get_context()
-  ctx.update(kwargs)
-  pg.object_utils.thread_local_push(_TLS_CODE_RUN_CONTEXT, ctx)
-
-  try:
-    yield ctx
-  finally:
-    pg.object_utils.thread_local_pop(_TLS_CODE_RUN_CONTEXT)
-
-
-def get_context() -> dict[str, Any]:
-  """Gets the current context for code execution."""
-  context_stack = pg.object_utils.thread_local_get(_TLS_CODE_RUN_CONTEXT, None)
-  return dict(context_stack[-1]) if context_stack else {}
+context = pg.coding.context
+CodeError = pg.coding.CodeError
+CodePermission = pg.coding.CodePermission
+permission = pg.coding.permission
 
 
 def evaluate(
     code: str,
     *,
     global_vars: dict[str, Any] | None = None,
-    permission: permissions.CodePermission | None = None,
+    permission: CodePermission | None = None,   # pylint: disable=redefined-outer-name
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
 ) -> Any | dict[str, Any]:
   """Executes Python code.
@@ -71,191 +45,33 @@ def evaluate(
     global_vars: An optional dict as the globals that could be referenced by the
       code.
     permission: Permission for the Python code to run.
-    outputs_intermediate: If True, intermediate output will be outputted as a
-      dict, with the last line's value accessible by key '__result__'. Otherwise
-      the value of the last line will be returned.
-
-  Returns:
-    The value of the last line of the code. Or a dict of variable name to
-    their values if `outputs_intermediate` is set to True, with the final result
-    accessible by key '__result__'.
-  """
-  # Set up the permission and context.
-  permission = permission or permissions.get_permission()
-  ctx = dict(get_context())
-  if global_vars:
-    ctx.update(global_vars)
-
-  # Parse the code str.
-  code, code_block = parsing.PythonCodeParser().parse(code, permission)
-  global_vars, orig_global_vars = ctx, ctx.copy()
-
-  # No code.
-  if not code_block.body:
-    return {} if outputs_intermediate else None
-
-  stdout = io.StringIO()
-  with contextlib.redirect_stdout(stdout):
-    if hasattr(code_block.body[-1], 'value'):
-      last_expr = code_block.body.pop()  # pytype: disable=attribute-error
-      result_vars = [RESULT_KEY]
-
-      if isinstance(last_expr, ast.Assign):
-        for name_node in last_expr.targets:
-          result_vars.append(name_node.id)
-
-      last_expr = ast.Expression(last_expr.value)  # pytype: disable=attribute-error
-
-      try:
-        # Execute the lines before the last expression.
-        # NOTE(daiyip): Only a `globals` dict is specified here, which will also
-        # be used to output intermediate values by `exec`. We do not specify a
-        # separate `locals` dict here, for - "If exec gets two separate objects
-        # as globals and locals, the code will be executed as if it were
-        # embedded in a class definition." - as the Python document explains.
-        # The outcome is that new functions defined in the code block could not
-        # be called by other newly defined functions.
-        # Refer to https://stackoverflow.com/questions/
-        # 73940751/why-cant-i-call-a-function-from-another-function-using-exec
-        # for more details.
-        exec(compile(code_block, '', mode='exec'), global_vars)  # pylint: disable=exec-used
-
-        # Evaluate the last expression.
-        result = eval(  # pylint: disable=eval-used
-            compile(last_expr, '', mode='eval'), global_vars
-        )
-      except Exception as e:
-        raise errors.CodeError(code, e) from e
-
-      for result_var in result_vars:
-        global_vars[result_var] = result
-    else:
-      try:
-        exec(compile(code_block, '', mode='exec'), global_vars)  # pylint: disable=exec-used
-      except Exception as e:
-        raise errors.CodeError(code, e) from e
-      global_vars[RESULT_KEY] = list(global_vars.values())[-1]
-
-  if outputs_intermediate:
-    outputs = {}
-    for k, v in global_vars.items():
-      if k == '__builtins__':
-        continue
-      if k not in orig_global_vars or v is not orig_global_vars[k]:
-        outputs[k] = v
-    # Add stdout to outputs.
-    outputs[STDOUT_KEY] = stdout.getvalue()
-    return outputs
-  return global_vars[RESULT_KEY]
-
-
-def sandbox_call(
-    func: Callable[..., Any],
-    *args,
-    timeout: float | None = None,
-    **kwargs) -> Any:
-  """Calls a function with sandboxing.
-
-  Args:
-    func: Function to call.
-    *args: Positional arguments for `func`
-    timeout: Execution timeout in seconds. If None, wait `func` to complete.
-    **kwargs: Keyword arguments for `func`.
-
-  Returns:
-    Return value from `func`.
-
-  Raises:
-    TimeoutError: If the execution time exceeds the timeout.
-    Exception: Exception raised from `func`.
-  """
-  def _call(q, *args, **kwargs):
-    # NOTE(daiyip): if `q` is closed by the main process when `q.put` is called
-    # on a subprocess, ValueError will be raised. This is okay since the main
-    # process is no longer waiting for the result, and the subprocess could
-    # recycled with non-zero error code, which does not affect the main
-    # process.
-    def _run():
-      r = func(*args, **kwargs)
-      try:
-        return pg.to_json_str(r)
-      except Exception as e:
-        raise errors.SerializationError(
-            f'Cannot serialize sandbox result: {r}', e
-        ) from e
-
-    try:
-      q.put(_run())
-    except Exception as e:  # pylint: disable=broad-exception-caught
-      q.put(e)
-
-  q = multiprocessing.Queue()
-  try:
-    p = multiprocessing.Process(
-        target=_call, args=tuple([q] + list(args)), kwargs=kwargs)
-    p.start()
-    p.join(timeout=timeout)
-    if p.is_alive():
-      p.terminate()
-      raise TimeoutError(f'Execution time exceed {timeout} seconds.')
-    x = q.get()
-    if isinstance(x, Exception):
-      raise x
-    try:
-      return pg.from_json_str(x)
-    except Exception as e:
-      raise errors.SerializationError(
-          'Cannot deserialize the output from sandbox.', e
-      ) from e
-  finally:
-    q.close()
-
-
-def call(
-    func: Callable[..., Any],
-    *args,
-    sandbox: bool | None = None,
-    timeout: float | None = None,
-    **kwargs
-) -> Any:
-  """Calls a function with sandbox support.
-
-  Args:
-    func: Function to call.
-    *args: Postional args that will be passed to `func`.
-    sandbox: If True, run code in sandbox; If False, run code in current
-      process. If None, run in sandbox first, if the output could not be
-      serialized and pass to current process, run the code again in current
-      process.
-    timeout: Execution timeout in seconds. If None, wait the code the complete.
-    **kwargs: Keyword args that will be passed to `func`.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
 
   Returns:
-    The return value of `func`.
-
-  Raises:
-    TimeoutError: If the execution time exceeds the timeout.
-    Exception: Exception  that are raised from `func`.
+    The value of the last line of the code block. Or a dict of variable
+    names of all locals to their evaluated values as the output of the code to
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
   """
-  if sandbox is None:
-    try:
-      return sandbox_call(func, *args, timeout=timeout, **kwargs)
-    # NOTE(daiyip): output could be serialized across processes, giving it
-    # already finishes on sandbox, so it should be much safer to run under
-    # current process.
-    except errors.SerializationError:
-      return func(*args, **kwargs)
-  elif sandbox:
-    return sandbox_call(func, *args, timeout=timeout, **kwargs)
-  else:
-    return func(*args, **kwargs)
+  return pg.coding.evaluate(
+      parsing.clean(code),
+      global_vars=global_vars,
+      permission=permission,
+      returns_stdout=returns_stdout,
+      outputs_intermediate=outputs_intermediate,
+  )
 
 
 def run(
     code: str,
     *,
     global_vars: dict[str, Any] | None = None,
-    permission: permissions.CodePermission | None = None,
+    permission: CodePermission | None = None,  # pylint: disable=redefined-outer-name
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
     sandbox: bool | None = None,
     timeout: float | None = None,
@@ -271,9 +87,11 @@ def run(
     code: Python code to run.
     global_vars: An optional dict of
     permission: Permission for the Python code to run.
-    outputs_intermediate: If True, all variables created as locals will be
-      returned, with the final result accessible by key '__result__'. Otherwise
-      only the final result will be returned.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
     sandbox: If True, run code in sandbox; If False, run code in current
       process. If None, run in sandbox first, if the output could not be
       serialized and pass to current process, run the code again in current
@@ -283,13 +101,19 @@ def run(
   Returns:
     The value of the last line of the code block. Or a dict of variable
     names of all locals to their evaluated values as the output of the code to
-    run. The value for the last line can be accessed by key '__result__'.
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
 
   Raises:
     TimeoutError: If the execution time exceeds the timeout.
     Exception: Exception  that are raised from the code.
   """
-  return call(
-      evaluate, code=code, global_vars=global_vars, permission=permission,
+  return pg.coding.run(
+      parsing.clean(code),
+      global_vars=global_vars,
+      permission=permission,
+      returns_stdout=returns_stdout,
       outputs_intermediate=outputs_intermediate,
-      sandbox=sandbox, timeout=timeout)
+      sandbox=sandbox,
+      timeout=timeout,
+  )

langfun/core/coding/python/execution_test.py

@@ -14,11 +14,8 @@
 """Tests for Python code execution."""
 
 import inspect
-import time
 import unittest
-from langfun.core.coding.python import errors
 from langfun.core.coding.python import execution
-from langfun.core.coding.python import permissions
 import pyglove as pg
 
 
@@ -63,6 +60,15 @@ class EvaluateTest(unittest.TestCase):
         ),
         3,
     )
+    with self.assertRaisesRegex(execution.CodeError, 'ValueError'):
+      execution.evaluate(
+          """
+          def foo():
+            raise ValueError("intentional error")
+          foo()
+          """,
+          permission=execution.CodePermission.ALL
+      )
 
   def test_class_def(self):
     ret = execution.evaluate(
@@ -73,7 +79,7 @@ class EvaluateTest(unittest.TestCase):
           def __call__(self):
             return self.x + self.y
         """,
-        permission=permissions.CodePermission.ALL,
+        permission=execution.CodePermission.ALL,
         global_vars=dict(pg=pg),
         outputs_intermediate=True,
     )
@@ -91,7 +97,7 @@ class EvaluateTest(unittest.TestCase):
         def bar(z):
           return z + foo(z, z)
         """,
-        permission=permissions.CodePermission.ALL,
+        permission=execution.CodePermission.ALL,
         outputs_intermediate=True,
     )
     self.assertEqual(
@@ -102,23 +108,33 @@ class EvaluateTest(unittest.TestCase):
     self.assertIs(ret['__result__'], ret['bar'])
 
   def test_function_def_and_call(self):
-    ret = execution.evaluate(
+    code = (
         """
         def foo(x, y):
           return x + y
 
         def bar(z):
+          print(f'z is {z}')
           return z + foo(z, z)
 
         bar(1)
-        """,
-        permission=permissions.CodePermission.ALL,
+        """
+    )
+    ret = execution.evaluate(
+        code,
+        permission=execution.CodePermission.ALL,
         outputs_intermediate=True,
     )
     self.assertEqual(
         list(ret.keys()), ['foo', 'bar', '__result__', '__stdout__']
     )
     self.assertEqual(ret['__result__'], 3)
+    ret = execution.evaluate(
+        code,
+        permission=execution.CodePermission.ALL,
+        returns_stdout=True,
+    )
+    self.assertEqual(ret, 'z is 1\n')
 
   def test_complex(self):
     ret = execution.evaluate(
@@ -134,7 +150,7 @@ class EvaluateTest(unittest.TestCase):
         k = A(1, 2)
         k(foo(3, 4))
         """,
-        permission=permissions.CodePermission.ALL,
+        permission=execution.CodePermission.ALL,
         global_vars=dict(pg=pg),
         outputs_intermediate=True,
     )
@@ -148,18 +164,18 @@ class EvaluateTest(unittest.TestCase):
 
   def test_run_with_error(self):
     with self.assertRaisesRegex(
-        errors.CodeError, 'NameError: name .* is not defined'
+        execution.CodeError, 'NameError: name .* is not defined'
     ):
       execution.evaluate(
           """
           x = 1
           y = x + z
           """,
-          permission=permissions.CodePermission.ALL,
+          permission=execution.CodePermission.ALL,
       )
-    with self.assertRaisesRegex(errors.CodeError, 'ValueError'):
+    with self.assertRaisesRegex(execution.CodeError, 'ValueError'):
       execution.evaluate(
-          'raise ValueError()', permission=permissions.CodePermission.ALL
+          'raise ValueError()', permission=execution.CodePermission.ALL
       )
 
 
@@ -168,82 +184,6 @@ class Foo(pg.Object):
   y: int
 
 
-class SandboxCallTest(unittest.TestCase):
-
-  def test_basics(self):
-    def f(x, y):
-      return x + y
-    self.assertEqual(execution.sandbox_call(f, 1, y=2), 3)
-
-  def test_complex_type(self):
-    def f(x, y):
-      return Foo(x, y)
-
-    self.assertEqual(execution.sandbox_call(f, 1, 2), Foo(1, 2))
-
-  def test_timeout(self):
-    def f(x):
-      time.sleep(x)
-
-    self.assertIsNone(execution.sandbox_call(f, 0, timeout=1))
-    with self.assertRaises(TimeoutError):
-      execution.sandbox_call(f, 2, timeout=1)
-
-  def test_raise(self):
-    def f(x):
-      if x == 0:
-        raise ValueError()
-
-    self.assertIsNone(execution.sandbox_call(f, 1))
-    with self.assertRaises(ValueError):
-      execution.sandbox_call(f, 0)
-
-
-class CallTest(unittest.TestCase):
-
-  def test_call_without_sandboxing(self):
-    def foo(x, y):
-      return x + y
-
-    self.assertEqual(
-        execution.call(foo, 1, y=2, sandbox=False),
-        3
-    )
-
-  def test_call_with_sandboxing(self):
-    def foo(x, y):
-      return x + y
-
-    self.assertEqual(
-        execution.call(foo, 1, y=2, sandbox=True),
-        3
-    )
-
-    def make_cls():
-      class A(pg.Object):
-        x: str
-      return A
-
-    with self.assertRaises(errors.SerializationError):
-      execution.call(make_cls, sandbox=True)
-
-  def test_call_with_automatic_sandboxing(self):
-    def foo(x, y):
-      return x + y
-
-    self.assertEqual(
-        execution.call(foo, 1, y=2),
-        3
-    )
-
-    def make_cls():
-      class A(pg.Object):
-        x: str
-      return A
-
-    self.assertTrue(inspect.isclass(execution.call(make_cls)))
-
-
 class RunTest(unittest.TestCase):
 
   def test_run_without_sandboxing(self):

langfun/core/coding/python/generation.py

@@ -88,6 +88,8 @@ class PythonCode(pg.Object):
       sandbox: bool | None = None,
       timeout: int | None = 5,
       global_vars: dict[str, Any] | None = None,
+      returns_stdout: bool = False,
+      outputs_intermediate: bool = False,
       autofix: int = 3,
       autofix_lm: lf.LanguageModel | None = None,
   ) -> Any:
@@ -101,13 +103,22 @@ class PythonCode(pg.Object):
       timeout: Timeout in seconds. If None, there is no timeout. Applicable when
         sandbox is set to True.
       global_vars: Global variables that could be accessed from the source code.
+      returns_stdout: If True, the stdout (a str) will be returned.
+      outputs_intermediate: Applicable when returns_stdout is False. If True,
+        intermediate output will be outputted as a dict, with the last line's
+        value accessible by key '__result__' and the std output accessible by 
+        key '__stdout__'. Otherwise the value of the last line will be returned.
       autofix: Number of attempts to auto fix the generated code. If 0, autofix
         is disabled.
       autofix_lm: Language model to be used. If not specified, it will try to
         use the `lm` under `lf.context`.
 
     Returns:
-      The value of the last expression in the source code.
+      The value of the last expression in the source code. Or a dict of local
+      variable names defined in the source code to their values if
+      `outputs_intermediate` is set to True. The value for the last line can be
+      accessed by key '__result__'. Or the stdout as a str if `returns_stdout`
+      is set to True.
 
     Raises:
       TimeoutError: If `sandbox` is True and timeout has reached.
@@ -121,6 +132,8 @@ class PythonCode(pg.Object):
         max_attempts=autofix,
         lm=autofix_lm,
         returns_code=True,
+        returns_stdout=returns_stdout,
+        outputs_intermediate=outputs_intermediate,
     )
     self.rebind(source=updated_code)
     return result
@@ -158,18 +171,14 @@ class PythonCode(pg.Object):
       TimeoutError: If `sandbox` is True and timeout has reached.
       Exception: Any errors that the source code has raised.
     """
-    result, updated_code = correction.run_with_correction(
-        self.source,
-        global_vars=global_vars,
+    return self(
         sandbox=sandbox,
         timeout=timeout,
+        global_vars=global_vars,
+        autofix=autofix,
+        autofix_lm=autofix_lm,
         outputs_intermediate=True,
-        max_attempts=autofix,
-        lm=autofix_lm,
-        returns_code=True,
     )
-    self.rebind(source=updated_code)
-    return result
 
 
 class PythonFunction(pg.Object):
@@ -218,5 +227,6 @@ class PythonFunction(pg.Object):
       TimeoutError: If `sandbox` is True and timeout has reached.
       Exception: Any errors that the source code has raised.
     """
-    return execution.call(
-        self.implementation, *args, sandbox=sandbox, timeout=timeout, **kwargs)
+    return pg.coding.maybe_sandbox_call(
+        self.implementation, *args, sandbox=sandbox, timeout=timeout, **kwargs
+    )

langfun/core/coding/python/generation_test.py

@@ -16,7 +16,7 @@
 import inspect
 import unittest
 from langfun.core.coding.python import generation
-from langfun.core.coding.python import permissions
+import pyglove as pg
 
 
 class PythonCodeTest(unittest.TestCase):
@@ -58,7 +58,7 @@ class PythonCodeTest(unittest.TestCase):
     )
 
   def test_call_class_def(self):
-    with permissions.permission(permissions.CodePermission.CLASS_DEFINITION):
+    with pg.coding.permission(pg.coding.CodePermission.CLASS_DEFINITION):
       v = generation.PythonCode("""
           class A:
             pass