langfun 0.1.2.dev202501080804__py3-none-any.whl → 0.1.2.dev202501240804__py3-none-any.whl

langfun/core/__init__.py

@@ -72,16 +72,11 @@ from langfun.core.sampling import random_sample
 
 # Concurrent execute a function with parallel inputs with inheriting current
 # context's defaults and overrides.
+from langfun.core.concurrent import RetryEntry
 from langfun.core.concurrent import concurrent_execute
 from langfun.core.concurrent import concurrent_map
-from langfun.core.concurrent import with_context_access
 from langfun.core.concurrent import with_retry
 
-# Utility libraries for text formatting.
-from langfun.core.text_formatting import colored
-from langfun.core.text_formatting import colored_print as print  # pylint: disable=redefined-builtin
-from langfun.core.text_formatting import colored_template
-
 # Interface for natural language formattable.
 from langfun.core.natural_language import NaturalLanguageFormattable
 

langfun/core/coding/python/__init__.py

@@ -16,19 +16,13 @@
 # pylint: disable=g-bad-import-order
 # pylint: disable=g-importing-member
 
-from langfun.core.coding.python.errors import CodeError
-
-from langfun.core.coding.python.permissions import CodePermission
-from langfun.core.coding.python.permissions import permission
-from langfun.core.coding.python.permissions import get_permission
-
-from langfun.core.coding.python.parsing import PythonCodeParser
-
+# Expose from `lf.coding` as aliases for `pg.coding` for backward compatibility.
+from langfun.core.coding.python.execution import CodeError
+from langfun.core.coding.python.execution import CodePermission
 from langfun.core.coding.python.execution import context
-from langfun.core.coding.python.execution import get_context
+
+from langfun.core.coding.python.parsing import clean
 from langfun.core.coding.python.execution import evaluate
-from langfun.core.coding.python.execution import sandbox_call
-from langfun.core.coding.python.execution import call
 from langfun.core.coding.python.execution import run
 
 from langfun.core.coding.python.generation import PythonCode

langfun/core/coding/python/correction.py

@@ -14,7 +14,6 @@
 """Python code error correction."""
 from typing import Any
 import langfun.core as lf
-from langfun.core.coding.python import errors
 from langfun.core.coding.python import execution
 import pyglove as pg
 
@@ -125,7 +124,7 @@ def run_with_correction(
       correction = querying.query(
           CodeWithError(code=code, error=error), CorrectedCode, lm=lm, autofix=0
       )
-    except errors.CodeError:
+    except pg.coding.CodeError:
       break
 
     code = correction.corrected_code
@@ -133,7 +132,7 @@ def run_with_correction(
     if error is None:
       return (result, code) if returns_code else result
 
-  raise errors.CodeError(
+  raise pg.coding.CodeError(
       code,
       RuntimeError(
           f"Cannot correct code after {num_attempts} attempts. "
@@ -191,10 +190,8 @@ def correct(
 
 def _error_feedback_str(error: Exception) -> str:
   """Returns the error str for feedback."""
-  if isinstance(error, errors.CodeError):
-    return lf.text_formatting.decolored(
-        error.format(include_complete_code=False)
-    )
+  if isinstance(error, pg.coding.CodeError):
+    return pg.decolor(error.format(include_complete_code=False))
   else:
     return f"Encountered {error.__class__.__name__}: {error}"
 

langfun/core/coding/python/correction_test.py

@@ -17,7 +17,6 @@ import inspect
 import unittest
 
 from langfun.core.coding.python import correction
-from langfun.core.coding.python import errors
 from langfun.core.llms import fake
 import pyglove as pg
 
@@ -82,7 +81,7 @@ class RunWithCorrectionTest(unittest.TestCase):
         max_attempts=0,
     )
     self.assertEqual(result, 4)
-    with self.assertRaises(errors.CodeError):
+    with self.assertRaises(pg.coding.CodeError):
       correction.run_with_correction(
           inspect.cleandoc("""
             x = 1,
@@ -125,7 +124,7 @@ class CorrectTest(unittest.TestCase):
 
   def test_correct_reaching_limit(self):
     with self.assertRaisesRegex(
-        errors.CodeError, 'Cannot correct code after 1 attempts'
+        pg.coding.CodeError, 'Cannot correct code after 1 attempts'
     ):
       correction.correct(
           inspect.cleandoc("""

langfun/core/coding/python/execution.py

@@ -13,50 +13,23 @@
 # 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]:
@@ -84,186 +57,20 @@ def evaluate(
     run. The value for the last line can be accessed by key '__result__'. Or the
     stdout as a str.
   """
-  # 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 returns_stdout:
-    return stdout.getvalue()
-  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():
-      # We use `kill` instead of `terminate` to release process resources
-      # right away.
-      p.kill()
-      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:
-    The return value of `func`.
-
-  Raises:
-    TimeoutError: If the execution time exceeds the timeout.
-    Exception: Exception  that are raised from `func`.
-  """
-  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,
@@ -301,8 +108,12 @@ def run(
     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,
-      returns_stdout=returns_stdout, outputs_intermediate=outputs_intermediate,
-      sandbox=sandbox, timeout=timeout
+  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,
   )

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,14 +60,14 @@ class EvaluateTest(unittest.TestCase):
         ),
         3,
     )
-    with self.assertRaisesRegex(errors.CodeError, 'ValueError'):
+    with self.assertRaisesRegex(execution.CodeError, 'ValueError'):
       execution.evaluate(
           """
           def foo():
             raise ValueError("intentional error")
           foo()
           """,
-          permission=permissions.CodePermission.ALL
+          permission=execution.CodePermission.ALL
       )
 
   def test_class_def(self):
@@ -82,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,
     )
@@ -100,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(
@@ -125,7 +122,7 @@ class EvaluateTest(unittest.TestCase):
     )
     ret = execution.evaluate(
         code,
-        permission=permissions.CodePermission.ALL,
+        permission=execution.CodePermission.ALL,
         outputs_intermediate=True,
     )
     self.assertEqual(
@@ -134,7 +131,7 @@ class EvaluateTest(unittest.TestCase):
     self.assertEqual(ret['__result__'], 3)
     ret = execution.evaluate(
         code,
-        permission=permissions.CodePermission.ALL,
+        permission=execution.CodePermission.ALL,
         returns_stdout=True,
     )
     self.assertEqual(ret, 'z is 1\n')
@@ -153,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,
     )
@@ -167,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
       )
 
 
@@ -187,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

@@ -227,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