nighthawk-python 0.1.0__py3-none-any.whl

nighthawk/natural/blocks.py

@@ -0,0 +1,279 @@
+from __future__ import annotations
+
+import ast
+import re
+import textwrap
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import yaml
+
+from ..errors import NaturalParseError
+
+_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+_BINDING_PATTERN = re.compile(r"<(:?)([A-Za-z_][A-Za-z0-9_]*)>")
+
+
+@dataclass(frozen=True)
+class NaturalBlock:
+    kind: Literal["docstring", "inline"]
+    text: str
+    input_bindings: tuple[str, ...]
+    output_bindings: tuple[str, ...]
+    lineno: int
+
+
+def is_natural_sentinel(text: str) -> bool:
+    return text.startswith("natural\n")
+
+
+def extract_program(text: str) -> str:
+    if not is_natural_sentinel(text):
+        raise NaturalParseError("Missing natural sentinel")
+    program = text.removeprefix("natural\n")
+    return textwrap.dedent(program)
+
+
+def extract_bindings(program: str) -> tuple[tuple[str, ...], tuple[str, ...]]:
+    inputs: list[str] = []
+    outputs: list[str] = []
+    for match in _BINDING_PATTERN.finditer(program):
+        is_output = match.group(1) == ":"
+        name = match.group(2)
+        if not _IDENTIFIER_PATTERN.match(name):
+            raise NaturalParseError(f"Invalid binding name: {name!r}")
+        if is_output:
+            outputs.append(name)
+        else:
+            inputs.append(name)
+
+    def deduplicate(names: list[str]) -> tuple[str, ...]:
+        seen: set[str] = set()
+        ordered: list[str] = []
+        for name in names:
+            if name in seen:
+                continue
+            seen.add(name)
+            ordered.append(name)
+        return tuple(ordered)
+
+    return deduplicate(inputs), deduplicate(outputs)
+
+
+_JOINED_STRING_FORMATTED_VALUE_PLACEHOLDER = "\x00"
+
+
+def _joined_string_first_literal_or_none(joined_string: ast.JoinedStr) -> str | None:
+    if not joined_string.values:
+        return None
+    first = joined_string.values[0]
+    if not isinstance(first, ast.Constant) or not isinstance(first.value, str):
+        return None
+    return first.value
+
+
+def _joined_string_is_natural_sentinel(joined_string: ast.JoinedStr) -> bool:
+    first_literal = _joined_string_first_literal_or_none(joined_string)
+    if first_literal is None:
+        return False
+    return is_natural_sentinel(first_literal)
+
+
+def _joined_string_scan_text(joined_string: ast.JoinedStr, *, formatted_value_placeholder: str) -> str:
+    parts: list[str] = []
+    for part in joined_string.values:
+        if isinstance(part, ast.Constant) and isinstance(part.value, str):
+            parts.append(part.value)
+        else:
+            parts.append(formatted_value_placeholder)
+    return "".join(parts)
+
+
+def _validate_joined_string_bindings_do_not_span_formatted_values(joined_string: ast.JoinedStr) -> None:
+    """Validate that no binding marker spans a formatted value boundary."""
+    boundary_marked_text = _joined_string_scan_text(
+        joined_string,
+        formatted_value_placeholder=_JOINED_STRING_FORMATTED_VALUE_PLACEHOLDER,
+    )
+
+    if re.search(r"<[^>]*" + _JOINED_STRING_FORMATTED_VALUE_PLACEHOLDER + r"[^>]*>", boundary_marked_text):
+        raise NaturalParseError("Binding marker must not span formatted value boundary in inline f-string Natural block")
+
+
+def _extract_program_and_bindings_from_joined_string(joined_string: ast.JoinedStr) -> tuple[str, tuple[str, ...], tuple[str, ...]]:
+    _validate_joined_string_bindings_do_not_span_formatted_values(joined_string)
+
+    scan_text = _joined_string_scan_text(joined_string, formatted_value_placeholder="")
+    program = extract_program(scan_text)
+    input_bindings, output_bindings = extract_bindings(program)
+    return program, input_bindings, output_bindings
+
+
+def find_natural_blocks(func_source: str) -> tuple[NaturalBlock, ...]:
+    """Parse function source text and return Natural blocks (docstring + inline)."""
+
+    try:
+        module = ast.parse(func_source)
+    except SyntaxError as e:
+        raise NaturalParseError(str(e)) from e
+
+    blocks: list[NaturalBlock] = []
+
+    func_def: ast.FunctionDef | ast.AsyncFunctionDef | None = None
+    for node in module.body:
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            func_def = node
+            break
+    if func_def is None:
+        raise NaturalParseError("No function definition found")
+
+    docstring_text = ast.get_docstring(func_def, clean=False)
+    if docstring_text and is_natural_sentinel(docstring_text):
+        program = extract_program(docstring_text)
+        input_bindings, output_bindings = extract_bindings(program)
+        blocks.append(
+            NaturalBlock(
+                kind="docstring",
+                text=program,
+                input_bindings=input_bindings,
+                output_bindings=output_bindings,
+                lineno=getattr(func_def, "lineno", 1),
+            )
+        )
+
+    start_index = 0
+    if func_def.body:
+        first_statement = func_def.body[0]
+        if isinstance(first_statement, ast.Expr) and isinstance(first_statement.value, ast.Constant) and isinstance(first_statement.value.value, str):
+            start_index = 1
+
+    for statement in func_def.body[start_index:]:
+        if not isinstance(statement, ast.Expr):
+            continue
+
+        value = statement.value
+
+        if isinstance(value, ast.Constant) and isinstance(value.value, str):
+            text = value.value
+            if is_natural_sentinel(text):
+                program = extract_program(text)
+                input_bindings, output_bindings = extract_bindings(program)
+                blocks.append(
+                    NaturalBlock(
+                        kind="inline",
+                        text=program,
+                        input_bindings=input_bindings,
+                        output_bindings=output_bindings,
+                        lineno=getattr(statement, "lineno", 1),
+                    )
+                )
+
+        if isinstance(value, ast.JoinedStr) and _joined_string_is_natural_sentinel(value):
+            program, input_bindings, output_bindings = _extract_program_and_bindings_from_joined_string(value)
+            blocks.append(
+                NaturalBlock(
+                    kind="inline",
+                    text=program,
+                    input_bindings=input_bindings,
+                    output_bindings=output_bindings,
+                    lineno=getattr(statement, "lineno", 1),
+                )
+            )
+
+    return tuple(blocks)
+
+
+_FRONTMATTER_STEP_KINDS = ("pass", "return", "break", "continue", "raise")
+
+
+def validate_frontmatter_deny(frontmatter: dict[str, object]) -> tuple[str, ...]:
+    """Validate a parsed frontmatter mapping and return the denied step kinds.
+
+    Raises:
+        NaturalParseError: If the frontmatter contains unknown keys, unknown
+            step kind names, or has an invalid ``deny`` structure.
+    """
+    if not frontmatter:
+        return ()
+
+    allowed_keys = {"deny"}
+    unknown_keys = set(frontmatter.keys()) - allowed_keys
+    if unknown_keys:
+        unknown_key_list = ", ".join(sorted(str(k) for k in unknown_keys))
+        raise NaturalParseError(f"Unknown frontmatter keys: {unknown_key_list}")
+
+    if "deny" not in frontmatter:
+        raise NaturalParseError("Frontmatter must include 'deny'")
+
+    deny_value = frontmatter["deny"]
+    if not isinstance(deny_value, list) or not all(isinstance(item, str) for item in deny_value):
+        raise NaturalParseError("Frontmatter 'deny' must be a YAML sequence of strings")
+
+    if len(deny_value) == 0:
+        raise NaturalParseError("Frontmatter 'deny' must not be empty")
+
+    denied: list[str] = []
+    for item in deny_value:
+        if item not in _FRONTMATTER_STEP_KINDS:
+            raise NaturalParseError(f"Unknown denied step kind: {item}")
+        if item not in denied:
+            denied.append(item)
+
+    return tuple(denied)
+
+
+def parse_frontmatter(processed_natural_program: str) -> tuple[str, dict[str, Any]]:
+    """Parse and strip YAML frontmatter from a Natural program.
+
+    Frontmatter is recognized when the first non-blank line is ``---`` and a
+    matching closing ``---`` line follows.  The YAML content between the
+    delimiters must be a mapping.
+
+    Returns:
+        A tuple of (program_text_without_frontmatter, parsed_mapping).
+        When no frontmatter is present the mapping is empty.
+
+    Raises:
+        NaturalParseError: If the frontmatter is syntactically invalid.
+    """
+    lines = processed_natural_program.splitlines(keepends=True)
+    if not lines:
+        return processed_natural_program, {}
+
+    start_index: int | None = None
+    for i, line in enumerate(lines):
+        if line.strip(" \t\r\n") == "":
+            continue
+        start_index = i
+        break
+
+    if start_index is None:
+        return processed_natural_program, {}
+
+    first_line = lines[start_index]
+    if first_line not in ("---\n", "---"):
+        return processed_natural_program, {}
+
+    closing_index: int | None = None
+    for i, line in enumerate(lines[start_index + 1 :], start=start_index + 1):
+        if line in ("---\n", "---"):
+            closing_index = i
+            break
+
+    if closing_index is None:
+        return processed_natural_program, {}
+
+    yaml_text = "".join(lines[start_index + 1 : closing_index])
+    if yaml_text.strip() == "":
+        return processed_natural_program, {}
+
+    try:
+        loaded = yaml.safe_load(yaml_text)
+    except yaml.YAMLError as e:
+        raise NaturalParseError(f"Frontmatter YAML parsing failed: {e}") from e
+
+    if not isinstance(loaded, dict):
+        raise NaturalParseError("Frontmatter YAML must be a mapping")
+
+    instructions_without_frontmatter = "".join(lines[closing_index + 1 :])
+    return instructions_without_frontmatter, loaded

nighthawk/natural/decorator.py

@@ -0,0 +1,302 @@
+from __future__ import annotations
+
+import ast
+import inspect
+import logging
+import sys
+import textwrap
+from collections.abc import Awaitable, Callable
+from functools import wraps
+from typing import Any, cast
+
+from ..runtime.runner import Runner, StepEnvelope
+from ..runtime.scoping import get_step_executor
+from ..runtime.step_context import python_cell_scope, python_name_scope
+from ..tools.registry import call_scope
+from .blocks import find_natural_blocks
+from .transform import transform_module_ast
+
+type NaturalFunctionCallable = Callable[..., Any]
+
+
+class _RunnerProxy:
+    @staticmethod
+    def run_step(
+        natural_program: str,
+        input_binding_names: list[str],
+        output_binding_names: list[str],
+        binding_name_to_type: dict[str, object],
+        return_annotation: object,
+        is_in_loop: bool,
+    ) -> StepEnvelope:
+        caller_frame = sys._getframe(1)
+        current_step_executor = get_step_executor()
+        runner = Runner(current_step_executor)
+        return runner.run_step(
+            natural_program,
+            input_binding_names,
+            output_binding_names,
+            binding_name_to_type,
+            return_annotation,
+            is_in_loop,
+            caller_frame=caller_frame,
+        )
+
+    @staticmethod
+    async def run_step_async(
+        natural_program: str,
+        input_binding_names: list[str],
+        output_binding_names: list[str],
+        binding_name_to_type: dict[str, object],
+        return_annotation: object,
+        is_in_loop: bool,
+    ) -> StepEnvelope:
+        caller_frame = sys._getframe(1)
+        current_step_executor = get_step_executor()
+        runner = Runner(current_step_executor)
+        return await runner.run_step_async(
+            natural_program,
+            input_binding_names,
+            output_binding_names,
+            binding_name_to_type,
+            return_annotation,
+            is_in_loop,
+            caller_frame=caller_frame,
+        )
+
+
+def _extract_inline_fstring_name_set(function_source: str, *, function_name: str) -> set[str]:
+    """Extract names referenced in f-string expressions of inline Natural blocks."""
+    try:
+        module = ast.parse(function_source)
+    except SyntaxError:
+        return set()
+
+    function_def: ast.FunctionDef | ast.AsyncFunctionDef | None = None
+    for node in module.body:
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == function_name:
+            function_def = node
+            break
+    if function_def is None:
+        return set()
+
+    names: set[str] = set()
+
+    class Visitor(ast.NodeVisitor):
+        def visit_Name(self, node: ast.Name) -> None:  # noqa: N802
+            names.add(node.id)
+
+    visitor = Visitor()
+
+    for statement in function_def.body:
+        if not isinstance(statement, ast.Expr):
+            continue
+        value = statement.value
+        if not isinstance(value, ast.JoinedStr):
+            continue
+
+        first_part: ast.expr | None = value.values[0] if value.values else None
+        if not isinstance(first_part, ast.Constant) or not isinstance(first_part.value, str):
+            continue
+        if not first_part.value.startswith("natural\n"):
+            continue
+
+        for part in value.values:
+            if isinstance(part, ast.FormattedValue):
+                visitor.visit(part.value)
+
+    return names
+
+
+def _build_capture_name_set(source: str, function_name: str) -> set[str]:
+    """Build the set of names that need to be captured from the enclosing scope."""
+    capture_name_set: set[str] = set()
+    try:
+        for block in find_natural_blocks(source):
+            capture_name_set.update(block.input_bindings)
+            capture_name_set.update(block.output_bindings)
+        capture_name_set.update(_extract_inline_fstring_name_set(source, function_name=function_name))
+    except Exception as exception:
+        logging.getLogger("nighthawk").warning("Failed to extract capture names for %s: %s", function_name, exception)
+        capture_name_set = set()
+    return capture_name_set
+
+
+def _build_transformed_factory_module(
+    *,
+    transformed_module: ast.Module,
+    function_name: str,
+    name_to_value: dict[str, object],
+) -> ast.Module:
+    """Build a factory-function module that captures enclosing-scope values via closure."""
+    transformed_function_def: ast.FunctionDef | ast.AsyncFunctionDef | None = None
+    for node in transformed_module.body:
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == function_name:
+            transformed_function_def = node
+            break
+
+    if transformed_function_def is None:
+        raise RuntimeError("Transformed function not found in transformed module")
+
+    captured_value_name = "__nh_captured_values__"
+    factory_name = "__nh_factory__"
+
+    factory_body: list[ast.stmt] = []
+    for name in sorted(name_to_value.keys()):
+        factory_body.append(
+            ast.Assign(
+                targets=[ast.Name(id=name, ctx=ast.Store())],
+                value=ast.Subscript(
+                    value=ast.Name(id=captured_value_name, ctx=ast.Load()),
+                    slice=ast.Constant(name),
+                    ctx=ast.Load(),
+                ),
+            )
+        )
+
+    factory_body.append(transformed_function_def)
+    factory_body.append(ast.Return(value=ast.Name(id=function_name, ctx=ast.Load())))
+
+    factory_function_def = ast.FunctionDef(
+        name=factory_name,
+        args=ast.arguments(
+            posonlyargs=[],
+            args=[ast.arg(arg=captured_value_name)],
+            kwonlyargs=[],
+            kw_defaults=[],
+            defaults=[],
+        ),
+        body=factory_body,
+        decorator_list=[],
+        returns=None,
+        type_comment=None,
+    )
+
+    factory_module = ast.Module(body=[factory_function_def], type_ignores=[])
+    ast.fix_missing_locations(factory_module)
+    return factory_module
+
+
+def natural_function(func: NaturalFunctionCallable | None = None) -> NaturalFunctionCallable:
+    """Transform a function containing Natural blocks into an executable Natural function.
+
+    Parses the function source to find Natural blocks, rewrites the AST to
+    delegate block execution to the active step executor at runtime.
+
+    Args:
+        func: The function to transform. Can be omitted for use as a bare
+            decorator.
+
+    Example:
+        ```python
+        @nighthawk.natural_function
+        def summarize(text: str) -> str:
+            '''natural
+            Summarize <text> in one sentence.
+            -> <:result>
+            '''
+            return result
+        ```
+    """
+    if func is None:
+        return lambda f: natural_function(f)  # type: ignore[return-value]
+
+    if isinstance(func, staticmethod):
+        decorated_static_function = natural_function(func.__func__)
+        return cast(NaturalFunctionCallable, staticmethod(decorated_static_function))
+
+    if isinstance(func, classmethod):
+        decorated_class_function = natural_function(func.__func__)
+        return cast(NaturalFunctionCallable, classmethod(decorated_class_function))
+
+    lines, starting_line_number = inspect.getsourcelines(func)
+    source = textwrap.dedent("".join(lines))
+
+    try:
+        original_module = ast.parse(source)
+        for node in original_module.body:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func.__name__:
+                node.decorator_list = []
+                break
+        ast.increment_lineno(original_module, starting_line_number - 1)
+    except Exception as exception:
+        logging.getLogger("nighthawk").warning("Failed to parse original module AST for %s: %s", func.__name__, exception)
+        original_module = ast.Module(body=[], type_ignores=[])
+
+    capture_name_set = _build_capture_name_set(source, func.__name__)
+
+    definition_frame = inspect.currentframe()
+    name_to_value: dict[str, object] = {}
+    if definition_frame is not None and definition_frame.f_back is not None:
+        caller_frame = definition_frame.f_back
+        if caller_frame.f_code.co_name != "<module>":
+            for name in capture_name_set:
+                if name in caller_frame.f_locals:
+                    name_to_value[name] = caller_frame.f_locals[name]
+
+    captured_name_tuple = tuple(sorted(capture_name_set))
+
+    transformed_module = transform_module_ast(original_module, captured_name_tuple=captured_name_tuple)
+
+    filename = inspect.getsourcefile(func) or "<nighthawk>"
+
+    factory_module = _build_transformed_factory_module(
+        transformed_module=transformed_module,
+        function_name=func.__name__,
+        name_to_value=name_to_value,
+    )
+    code = compile(factory_module, filename, "exec")
+
+    globals_namespace: dict[str, object] = dict(func.__globals__)
+    globals_namespace["__nighthawk_runner__"] = _RunnerProxy()
+    from .blocks import extract_program as _nh_extract_program
+
+    globals_namespace["__nh_extract_program__"] = _nh_extract_program
+    globals_namespace["__nh_python_cell_scope__"] = python_cell_scope
+
+    module_namespace: dict[str, object] = {}
+    exec(code, globals_namespace, module_namespace)
+
+    factory = module_namespace.get("__nh_factory__")
+    if not callable(factory):
+        raise RuntimeError("Transformed factory not found after compilation")
+
+    transformed = factory(name_to_value)
+    if not callable(transformed):
+        raise RuntimeError("Transformed function not found after factory execution")
+
+    transformed_freevar_name_set = set(transformed.__code__.co_freevars)
+    captured_name_set = set(name_to_value.keys())
+
+    unexpected_freevar_name_set = transformed_freevar_name_set - captured_name_set
+    allowed_unexpected_freevar_name_set = {func.__name__}
+    if not unexpected_freevar_name_set.issubset(allowed_unexpected_freevar_name_set):
+        raise RuntimeError(
+            f"Transformed function freevars do not match captured names. freevars={transformed.__code__.co_freevars!r} captured={tuple(sorted(name_to_value.keys()))!r}"
+        )
+
+    if transformed.__closure__ is None and name_to_value:
+        raise RuntimeError("Transformed function closure is missing for captured names")
+
+    if inspect.iscoroutinefunction(func):
+        transformed_async = cast(Callable[..., Awaitable[Any]], transformed)
+
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with call_scope():
+                if name_to_value:
+                    with python_name_scope(name_to_value):
+                        return await transformed_async(*args, **kwargs)
+                return await transformed_async(*args, **kwargs)
+
+        return cast(NaturalFunctionCallable, async_wrapper)  # type: ignore[return-value]
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        with call_scope():
+            if name_to_value:
+                with python_name_scope(name_to_value):
+                    return transformed(*args, **kwargs)
+            return transformed(*args, **kwargs)
+
+    return cast(NaturalFunctionCallable, wrapper)  # type: ignore[return-value]