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

langfun/core/coding/python/parsing.py

@@ -13,205 +13,120 @@
 # limitations under the License.
 """Python code parsing."""
 
-import ast
 import inspect
 import io
 import re
 
-import langfun.core as lf
-from langfun.core.coding.python import errors
-from langfun.core.coding.python import permissions
-
-
-class PythonCodeParser(lf.Component):
-  """Python code parser with permission control."""
-
-  _ID_REGEX = re.compile('^[a-zA-Z_\\-]*$')
-
-  class _CodeValidator(ast.NodeVisitor):
-    """Python AST node visitor for ensuring code are permitted."""
-
-    def __init__(self, code: str, permission: permissions.CodePermission):
-      super().__init__()
-      self.code = code
-      self.permission = permission
-
-    def verify(
-        self,
-        node,
-        flag: permissions.CodePermission,
-        node_type,
-        error_message: str,
-    ) -> None:
-      if isinstance(node, node_type) and not (self.permission & flag):
-        raise SyntaxError(
-            error_message, (
-                '<generated-code>',
-                node.lineno,
-                node.col_offset,
-                self._code_line(node.lineno),
-                node.end_lineno,
-                node.end_col_offset,
-            ))
-
-    def _code_line(self, lineno):
-      return self.code.split('\n')[lineno - 1]
-
-    def generic_visit(self, node):
-      self.verify(
-          node,
-          permissions.CodePermission.CONDITION,
-          (ast.If, ast.Match),
-          'Condition is not allowed.',
-      )
-
-      self.verify(
-          node,
-          permissions.CodePermission.LOOP,
-          (ast.For, ast.While, ast.AsyncFor, ast.AsyncWith),
-          'Loop is not allowed.',
-      )
-
-      self.verify(
-          node,
-          permissions.CodePermission.EXCEPTION,
-          (ast.Try, ast.Raise, ast.Assert),
-          'Exception is not allowed.',
-      )
-
-      self.verify(
-          node,
-          permissions.CodePermission.CLASS_DEFINITION,
-          ast.ClassDef,
-          'Class definition is not allowed.',
-      )
-
-      self.verify(
-          node,
-          permissions.CodePermission.FUNCTION_DEFINITION,
-          (
-              ast.FunctionDef,
-              ast.AsyncFunctionDef,
-              ast.Return,
-              ast.Yield,
-              ast.YieldFrom,
-          ),
-          'Function definition is not allowed.',
-      )
-
-      self.verify(
-          node,
-          permissions.CodePermission.IMPORT,
-          (ast.Import, ast.ImportFrom),
-          '`import` is not allowed.',
-      )
-
-      super().generic_visit(node)
-
-  def parse(
-      self, code: str, permission: permissions.CodePermission
-  ) -> tuple[str, ast.AST]:
-    code = self.clean(code)
-    try:
-      parsed_code = ast.parse(code, mode='exec')
-      PythonCodeParser._CodeValidator(code, permission).visit(parsed_code)
-    except SyntaxError as e:
-      raise errors.CodeError(code, e) from e
-    return code, parsed_code
-
-  def clean(self, code_text: str) -> str:
-    # TODO(daiyip): Deal with markdown in docstrings.
-    code = io.StringIO()
-    quote_char = None
-    in_code = False
-    i = 0
-    in_comment = False
-    while i < len(code_text):
-      c = code_text[i]
-      # Detect code block separator (```).
-      if (not in_comment
-          and quote_char is None
-          and c == '`'
-          and code_text[i:i + 3] == '```'):
-        in_code = not in_code
-        if in_code:
-          i += 3
-          continue
-        else:
-          break
-
-      # Detect string literal boundary.
-      if (in_code
-          and not in_comment
-          and c in ('\'', '"')
-          and i > 0
-          and code_text[i - 1] != '\\'):
-        # Handle ''' and """.
-        if code_text[i: i + 3] == c * 3:
-          c = c * 3
-          i += 2
-
-        if quote_char is None:
-          quote_char = c
-        elif quote_char == c:
-          # NOTE(daiyip): at times, LM forgets to escape quotes inside a string.
-          # Thus we do some smart checking here to automatically correct such
-          # case. This logic here is pretty involved in handling special cases.
-          # We might want to revisit them later.
-
-          # Peek forward to see if it could be a valid string.
-          nt, nnt_start = _next_token(code_text, i + 1)
-          if nt in (',', '[', ']', '}', ')', '+', '*', '%', '\n', ':'):
-            end_quote = True
-          elif nt == ' ':
-            # Detect if . could be a method invocation.
-            # NOTE(daiyip): 'in' and 'not in' might have false positives. But
-            # given the chance is low, we do not complicate the reasoning logic
-            # for now.
-            nnt, _ = _next_token(code_text, nnt_start, skip_whitespace=True)
-            end_quote = nnt in ('+', '*', '%', '#', '[', 'in', 'not', ':')
-          elif nt == '.':
-            # Detect if . could be method invocation on string.
-            nnt, nnnt_start = _next_token(code_text, nnt_start)
-            nnnt, _ = _next_token(code_text, nnnt_start)
-            end_quote = nnt.isidentifier() and nnnt == '('
-          else:
-            end_quote = False
-
-          if end_quote:
-            quote_char = None
-          else:
-            c = f'\\{c}'
-      # Detect comment.
-      elif c == '#' and quote_char is None:
-        in_comment = True
-      # Detect end-of-comment.
-      elif c == '\n':
-        # NOTE(daiyip): deal with cases that LM forgot to escape linebreaks
-        # within strings.
-        if quote_char is not None:
-          # Only add \\ for ' and " (other than ''' and """).
-          if len(quote_char) == 1:
-            c = '\\n'
-        else:
-          in_comment = False
 
+_ID_REGEX = re.compile('^[a-zA-Z_\\-]*$')
+
+
+def clean(code_text: str) -> str:
+  """Cleans up Python code.
+
+  LLM may generate code with markdown annotations, as well as minor syntax
+  errors. This function removes such annotations and fixes minor syntax errors
+  without extra LLM calls.
+
+  Args:
+    code_text: The code text to clean up.
+
+  Returns:
+    The cleaned up code text.
+  """
+  # TODO(daiyip): Deal with markdown in docstrings.
+  code = io.StringIO()
+  quote_char = None
+  in_code = False
+  i = 0
+  in_comment = False
+  while i < len(code_text):
+    c = code_text[i]
+    # Detect code block separator (```).
+    if (not in_comment
+        and quote_char is None
+        and c == '`'
+        and code_text[i:i + 3] == '```'):
+      in_code = not in_code
       if in_code:
-        code.write(c)
-
-      i += 1
-
-    code = code.getvalue()
-    if code:
-      pos = code.find('\n')
-      # Strip markdown code type. E.g. ```python
-      if pos > 0 and self._ID_REGEX.match(code[:pos]):
-        code = code[pos:]
-    else:
-      # Maybe-code that resides not within a code markdown block.
-      # Adding '\n' makes inspect.cleandoc to make right adjustment.
-      code = '\n' + code_text
-    return inspect.cleandoc(code).strip()
+        i += 3
+        continue
+      else:
+        break
+
+    # Detect string literal boundary.
+    if (in_code
+        and not in_comment
+        and c in ('\'', '"')
+        and i > 0
+        and code_text[i - 1] != '\\'):
+      # Handle ''' and """.
+      if code_text[i: i + 3] == c * 3:
+        c = c * 3
+        i += 2
+
+      if quote_char is None:
+        quote_char = c
+      elif quote_char == c:
+        # NOTE(daiyip): at times, LM forgets to escape quotes inside a string.
+        # Thus we do some smart checking here to automatically correct such
+        # case. This logic here is pretty involved in handling special cases.
+        # We might want to revisit them later.
+
+        # Peek forward to see if it could be a valid string.
+        nt, nnt_start = _next_token(code_text, i + 1)
+        if (len(c) == 3
+            or nt in (',', '[', ']', '}', ')', '+', '*', '%', '\n', ':')):
+          end_quote = True
+        elif nt == ' ':
+          # Detect if . could be a method invocation.
+          # NOTE(daiyip): 'in' and 'not in' might have false positives. But
+          # given the chance is low, we do not complicate the reasoning logic
+          # for now.
+          nnt, _ = _next_token(code_text, nnt_start, skip_whitespace=True)
+          end_quote = nnt in ('+', '*', '%', '#', '[', 'in', 'not', ':')
+        elif nt == '.':
+          # Detect if . could be method invocation on string.
+          nnt, nnnt_start = _next_token(code_text, nnt_start)
+          nnnt, _ = _next_token(code_text, nnnt_start)
+          end_quote = nnt.isidentifier() and nnnt == '('
+        else:
+          end_quote = False
+
+        if end_quote:
+          quote_char = None
+        else:
+          c = f'\\{c}'
+    # Detect comment.
+    elif c == '#' and quote_char is None:
+      in_comment = True
+    # Detect end-of-comment.
+    elif c == '\n':
+      # NOTE(daiyip): deal with cases that LM forgot to escape linebreaks
+      # within strings.
+      if quote_char is not None:
+        # Only add \\ for ' and " (other than ''' and """).
+        if len(quote_char) == 1:
+          c = '\\n'
+      else:
+        in_comment = False
+
+    if in_code:
+      code.write(c)
+
+    i += 1
+
+  code = code.getvalue()
+  if code:
+    pos = code.find('\n')
+    # Strip markdown code type. E.g. ```python
+    if pos > 0 and _ID_REGEX.match(code[:pos]):
+      code = code[pos:]
+  else:
+    # Maybe-code that resides not within a code markdown block.
+    # Adding '\n' makes inspect.cleandoc to make right adjustment.
+    code = '\n' + code_text
+  return inspect.cleandoc(code).strip()
 
 
 def _next_token(

langfun/core/coding/python/parsing_test.py

@@ -15,18 +15,16 @@
 
 import inspect
 import unittest
-from langfun.core.coding.python import errors
 from langfun.core.coding.python import parsing
-from langfun.core.coding.python import permissions
 
 
-class PythonCodeParserTest(unittest.TestCase):
+class CleanTest(unittest.TestCase):
 
   def assert_clean(self, code: str, cleaned_code: str, clean: bool = True):
     if clean:
       cleaned_code = inspect.cleandoc(cleaned_code)
     self.assertEqual(
-        parsing.PythonCodeParser().clean(code), cleaned_code
+        parsing.clean(code), cleaned_code
     )
 
   def test_clean(self):
@@ -272,107 +270,6 @@ class PythonCodeParserTest(unittest.TestCase):
         """
     )
 
-  def assert_allowed(self, code: str, permission: permissions.CodePermission):
-    _, ast = parsing.PythonCodeParser().parse(code, permission)
-    self.assertIsNotNone(ast)
-
-  def assert_not_allowed(
-      self, code: str, permission: permissions.CodePermission
-  ):
-    with self.assertRaisesRegex(errors.CodeError, '.* is not allowed'):
-      parsing.PythonCodeParser().parse(code, permission)
-
-  def test_parse_with_allowed_code(self):
-    self.assert_allowed(
-        """
-        x = y + 1
-        z = x + y
-        """,
-        permissions.CodePermission.BASIC,
-    )
-    self.assert_allowed(
-        """
-        if x > 0:
-          print(x)
-        """,
-        permissions.CodePermission.CONDITION,
-    )
-    self.assert_allowed(
-        """
-        for i in range(5):
-          print(i)
-        """,
-        permissions.CodePermission.LOOP,
-    )
-    self.assert_allowed(
-        """
-        assert x > 1
-        """,
-        permissions.CodePermission.EXCEPTION,
-    )
-    self.assert_allowed(
-        """
-        class A:
-          pass
-        """,
-        permissions.CodePermission.CLASS_DEFINITION,
-    )
-    self.assert_allowed(
-        """
-        def foo(x, y):
-          return x + y
-        """,
-        permissions.CodePermission.FUNCTION_DEFINITION,
-    )
-    self.assert_allowed(
-        """
-        import re
-        """,
-        permissions.CodePermission.IMPORT,
-    )
-
-  def test_parse_with_not_allowed_code(self):
-    self.assert_not_allowed(
-        """
-        if x > 0:
-          print(x)
-        """,
-        permissions.CodePermission.BASIC,
-    )
-    self.assert_not_allowed(
-        """
-        for i in range(5):
-          print(i)
-        """,
-        permissions.CodePermission.BASIC,
-    )
-    self.assert_not_allowed(
-        """
-        assert x > 1
-        """,
-        permissions.CodePermission.BASIC,
-    )
-    self.assert_not_allowed(
-        """
-        class A:
-          pass
-        """,
-        permissions.CodePermission.BASIC,
-    )
-    self.assert_not_allowed(
-        """
-        def foo(x, y):
-          return x + y
-        """,
-        permissions.CodePermission.BASIC,
-    )
-    self.assert_not_allowed(
-        """
-        import re
-        """,
-        permissions.CodePermission.BASIC,
-    )
-
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/component.py

@@ -73,7 +73,7 @@ class Component(pg.Object):
         field.value.set_default(attr_value)
         additional_fields.append(field)
     if additional_fields:
-      pg.symbolic.update_schema(cls, additional_fields)
+      cls.update_schema(additional_fields)
 
   def _on_bound(self):
     super()._on_bound()
@@ -210,6 +210,22 @@ def get_contextual_override(var_name: str) -> ContextualOverride | None:
   return _get_scoped_value(_global_tls, _CONTEXT_OVERRIDES, var_name)
 
 
+def context_value(var_name: str, default: Any = RAISE_IF_HAS_ERROR) -> Any:
+  """Returns the value of a variable defined in `lf.context`."""
+  override = get_contextual_override(var_name)
+  if override is None:
+    if default == RAISE_IF_HAS_ERROR:
+      raise KeyError(f'{var_name!r} does not exist in current context.')
+    return default
+  return override.value
+
+
+def all_contextual_values() -> dict[str, Any]:
+  """Returns all contextual values provided from `lf.context` in scope."""
+  overrides = getattr(_global_tls, _CONTEXT_OVERRIDES, {})
+  return {k: v.value for k, v in overrides.items()}
+
+
 @contextlib.contextmanager
 def _contextual_scope(
     tls: threading.local, tls_key, **variables
@@ -237,7 +253,9 @@ def _get_scoped_value(
   return scoped_values.get(var_name, default)
 
 
-class ContextualAttribute(pg.symbolic.ValueFromParentChain):
+class ContextualAttribute(
+    pg.symbolic.ValueFromParentChain, pg.views.HtmlTreeView.Extension
+):
   """Attributes whose values are inferred from the context of the component.
 
   Please see go/langfun-component#attribute-value-retrieval for details.
@@ -270,6 +288,55 @@ class ContextualAttribute(pg.symbolic.ValueFromParentChain):
     else:
       return pg.MISSING_VALUE
 
+  def _html_tree_view_content(
+      self,
+      *,
+      view: pg.views.HtmlTreeView,
+      parent: Any = None,
+      root_path: pg.KeyPath | None = None,
+      **kwargs,
+  ) -> pg.Html:
+    inferred_value = pg.MISSING_VALUE
+    if isinstance(parent, pg.Symbolic) and root_path:
+      inferred_value = parent.sym_inferred(root_path.key, pg.MISSING_VALUE)
+
+    if inferred_value is not pg.MISSING_VALUE:
+      kwargs.pop('name', None)
+      return view.render(
+          inferred_value, parent=self,
+          root_path=pg.KeyPath('<inferred>', root_path),
+          **view.get_passthrough_kwargs(**kwargs)
+      )
+    return pg.Html.element(
+        'div',
+        [
+            '(not available)',
+        ],
+        css_classes=['unavailable-contextual'],
+    )
+
+  def _html_tree_view_config(self) -> dict[str, Any]:
+    return pg.views.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            collapse_level=1,
+        )
+    )
+
+  @classmethod
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        .contextual-attribute {
+          color: purple;
+        }
+        .unavailable-contextual {
+          color: gray;
+          font-style: italic;
+        }
+        """
+    ]
+
 
 # NOTE(daiyip): Returning Any instead of `lf.ContextualAttribute` to avoid
 # pytype check error as `contextual()` can be assigned to any type.

langfun/core/component_test.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """Contextual component and app test."""
 
+import inspect
+from typing import Any
 import unittest
 import weakref
 
@@ -84,6 +86,12 @@ class ComponentContextTest(unittest.TestCase):
           lf.get_contextual_override('y'),
           lf.ContextualOverride(3, cascade=False, override_attrs=False),
       )
+      self.assertEqual(lf.context_value('x'), 3)
+      self.assertIsNone(lf.context_value('f', None))
+      with self.assertRaisesRegex(KeyError, '.* does not exist'):
+        lf.context_value('f')
+
+      self.assertEqual(lf.all_contextual_values(), dict(x=3, y=3, z=3))
 
       # Member attributes take precedence over `lf.context`.
       self.assertEqual(a1.x, 1)
@@ -291,6 +299,52 @@ class ContextualAttributeTest(unittest.TestCase):
     self.assertEqual(c.z, 3)
     self.assertEqual(b.z, 3)
 
+  def test_to_html(self):
+    class A(lf.Component):
+      x: int = 1
+      y: int = lf.contextual()
+
+    def assert_content(html, expected):
+      expected = inspect.cleandoc(expected).strip()
+      actual = html.content.strip()
+      if actual != expected:
+        print(actual)
+      self.assertEqual(actual.strip(), expected)
+
+    self.assertIn(
+        inspect.cleandoc(
+            """
+            .contextual-attribute {
+              color: purple;
+            }
+            .unavailable-contextual {
+              color: gray;
+              font-style: italic;
+            }
+            """
+        ),
+        A().to_html().style_section,
+    )
+
+    assert_content(
+        A().to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove a"><summary><div class="summary-title">A(...)</div></summary><div class="complex-value a"><details open class="pyglove int"><summary><div class="summary-name">x<span class="tooltip">x</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove contextual-attribute"><summary><div class="summary-name">y<span class="tooltip">y</span></div><div class="summary-title">ContextualAttribute(...)</div></summary><div class="unavailable-contextual">(not available)</div></details></div></details>
+        """
+    )
+
+    class B(lf.Component):
+      z: Any
+      y: int = 2
+
+    b = B(A())
+    assert_content(
+        b.z.to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove a"><summary><div class="summary-title">A(...)</div></summary><div class="complex-value a"><details open class="pyglove int"><summary><div class="summary-name">x<span class="tooltip">x</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove contextual-attribute"><summary><div class="summary-name">y<span class="tooltip">y</span></div><div class="summary-title">ContextualAttribute(...)</div></summary><span class="simple-value int">2</span></details></div></details>
+        """
+    )
+
 
 if __name__ == '__main__':
   unittest.main()