asyncio

codegen/README.md

@@ -0,0 +1 @@
+codegen using `ast` to generate `sync_api` from `async_api`

codegen/__main__.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+import re
+import argparse
+import contextlib
+import os
+from copy import deepcopy
+from typing import Any, TypedDict, cast, Optional, TYPE_CHECKING, Sequence
+import ast
+from pathlib import Path
+import json
+import subprocess
+import libcst
+
+if TYPE_CHECKING:
+    from _typeshed import StrPath
+
+
+PRE_CODEGEN_NAME = "IS_PRE_CODEGEN"
+STATICALLY_ASYNC_NAME = "IS_ASYNC"
+DYNAMICALLY_ASYNC_NAME = "IS_ASYNC"
+
+COMMENT_IDENTIFIER = "COMMENT"
+PREVIOUS_LINE_COMMENT_IDENTIFIER = "PREV_LINE_COMMENT"
+
+
+class CodegenConfig(TypedDict):
+    sync_target_directory: str
+    async_target_directory: str
+    exclude: list[str]
+    include_directories: list[str]
+
+
+class AsyncCodegenNodeTransformer(ast.NodeTransformer):
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
+        self.generic_visit(node)
+
+        if node.name.endswith("_prim_sync"):
+            return None
+
+        return node
+
+    def _is_statically_async_literal(self, node: ast.AST) -> bool:
+        return isinstance(node, ast.Constant) and node.value == STATICALLY_ASYNC_NAME
+
+    def _is_pre_codegen_literal(self, node: ast.AST) -> bool:
+        return isinstance(node, ast.Constant) and node.value == PRE_CODEGEN_NAME
+
+    def _match_static_condition(self, test: ast.AST) -> bool | None:
+        if self._is_statically_async_literal(test):
+            return True
+
+        if self._is_pre_codegen_literal(test):
+            return False
+
+        if (
+            isinstance(test, ast.UnaryOp)
+            and isinstance(test.op, ast.Not)
+            and self._is_statically_async_literal(test.operand)
+        ):
+            return False
+
+        if (
+            isinstance(test, ast.UnaryOp)
+            and isinstance(test.op, ast.Not)
+            and self._is_pre_codegen_literal(test.operand)
+        ):
+            return True
+
+        return None
+
+    def visit_If(self, node: ast.If) -> Any:
+        self.generic_visit(node)
+
+        if (condition_value := self._match_static_condition(node.test)) is not None:
+            return node.body if condition_value else node.orelse
+
+        return node
+
+
+class SyncCodegenNodeTransformer(ast.NodeTransformer):
+    def visit_Assign(self, node: ast.Assign) -> Any:
+        self.generic_visit(node)
+
+        if node.targets:
+            first_target = node.targets[0]
+            if isinstance(first_target, ast.Name) and first_target.id == DYNAMICALLY_ASYNC_NAME:
+                node.value = ast.Constant(value=False, kind=None)
+
+        return node
+
+    def visit_Await(self, node: ast.Await) -> Any:
+        self.generic_visit(node)
+        return node.value
+
+    def visit_AsyncFor(self, node: ast.AsyncFor) -> Any:
+        self.generic_visit(node)
+        new_node = ast.For(**{field: getattr(node, field) for field in node._fields})
+        return ast.copy_location(new_node, node)
+
+    def visit_AsyncWith(self, node: ast.AsyncWith) -> Any:
+        self.generic_visit(node)
+        new_node = ast.With(**{field: getattr(node, field) for field in node._fields})
+        return ast.copy_location(new_node, node)
+
+    def visit_comprehension(self, node: ast.comprehension) -> Any:
+        self.generic_visit(node)
+        node.is_async = 0
+        return node
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
+        self.generic_visit(node)
+
+        # primitive functions are implemented as sync and as async so the async variant can be dropped
+        if node.name.endswith("_prim"):
+            return None
+
+        new_node = ast.FunctionDef(**{field: getattr(node, field) for field in node._fields})
+
+        return ast.copy_location(new_node, node)
+
+    def _is_statically_async_literal(self, node: ast.AST) -> bool:
+        return isinstance(node, ast.Constant) and node.value == STATICALLY_ASYNC_NAME
+
+    def _is_pre_codegen_literal(self, node: ast.AST) -> bool:
+        return isinstance(node, ast.Constant) and node.value == PRE_CODEGEN_NAME
+
+    def _match_static_condition(self, test: ast.AST) -> bool | None:
+        if self._is_statically_async_literal(test):
+            return False
+
+        if self._is_pre_codegen_literal(test):
+            return False
+
+        if (
+            isinstance(test, ast.UnaryOp)
+            and isinstance(test.op, ast.Not)
+            and self._is_statically_async_literal(test.operand)
+        ):
+            return True
+
+        if (
+            isinstance(test, ast.UnaryOp)
+            and isinstance(test.op, ast.Not)
+            and self._is_pre_codegen_literal(test.operand)
+        ):
+            return True
+
+        return None
+
+    def visit_If(self, node: ast.If) -> Any:
+        self.generic_visit(node)
+
+        if (condition_value := self._match_static_condition(node.test)) is not None:
+            return node.body if condition_value else node.orelse
+
+        return node
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
+        self.generic_visit(node)
+
+        if node.name.endswith("_prim_sync"):
+            node.name = node.name.removesuffix("_sync")
+
+        return node
+
+
+class CommentTransformer(libcst.CSTTransformer):
+    @staticmethod
+    def _get_comment(node: libcst.CSTNode, target_name: str) -> str | None:
+        if isinstance(node, libcst.SimpleStatementLine) and len(node.body) == 1:
+            expr = node.body[0]
+            if (
+                isinstance(expr, libcst.Expr)
+                and isinstance((call := expr.value), libcst.Call)
+                and isinstance((func := call.func), libcst.Name)
+                and func.value == target_name
+                and len(call.args) == 1
+                and isinstance((comment_value := call.args[0].value), libcst.SimpleString)
+            ):
+                comment_value_string = comment_value.evaluated_value
+                if isinstance(comment_value_string, bytes):
+                    return comment_value_string.decode()
+                return comment_value_string
+        return None
+
+    @classmethod
+    def _process_statement_sequence(
+        cls, sequence: Sequence[libcst.BaseStatement]
+    ) -> tuple[list[libcst.BaseStatement], list[libcst.EmptyLine], libcst.Comment | None]:
+        new_body: list[libcst.BaseStatement] = []
+        pending_comments: list[libcst.EmptyLine] = []
+        header_comment: libcst.Comment | None = None
+        for statement in sequence:
+            prev_line_comment_text = cls._get_comment(statement, PREVIOUS_LINE_COMMENT_IDENTIFIER)
+            if prev_line_comment_text is not None:
+                comment = libcst.Comment(f"# {prev_line_comment_text}")
+                if new_body:
+                    prev_statement = new_body[-1]
+                    if isinstance(prev_statement, libcst.SimpleStatementLine):
+                        whitespace = libcst.TrailingWhitespace(
+                            whitespace=libcst.SimpleWhitespace(" "), comment=comment
+                        )
+                        new_body[-1] = prev_statement.with_changes(trailing_whitespace=whitespace)
+                    else:
+                        pending_comments.append(libcst.EmptyLine(indent=True, comment=comment))
+                elif header_comment is None:
+                    header_comment = comment
+                else:
+                    pending_comments.append(libcst.EmptyLine(indent=True, comment=comment))
+                continue
+            comment_text = cls._get_comment(statement, COMMENT_IDENTIFIER)
+            if comment_text is not None:
+                pending_comments.append(
+                    libcst.EmptyLine(indent=True, comment=libcst.Comment(f"# {comment_text}"))
+                )
+                continue
+            if (
+                pending_comments
+                and (leading_lines := getattr(statement, "leading_lines", None)) is not None
+            ):
+                statement = statement.with_changes(
+                    leading_lines=(*leading_lines, *pending_comments)
+                )
+                pending_comments = []
+            new_body.append(statement)
+        return (new_body, pending_comments, header_comment)
+
+    def leave_Module(
+        self, original_node: libcst.Module, updated_node: libcst.Module
+    ) -> libcst.Module:
+        new_body, pending_comments, header_comment = self._process_statement_sequence(
+            updated_node.body
+        )
+        if header_comment is not None:
+            pending_comments.insert(0, libcst.EmptyLine(indent=True, comment=header_comment))
+        new_footer = (*updated_node.footer, *pending_comments)
+        return updated_node.with_changes(body=new_body, footer=new_footer)
+
+    def leave_IndentedBlock(
+        self, original_node: libcst.IndentedBlock, updated_node: libcst.IndentedBlock
+    ) -> libcst.IndentedBlock:
+        new_body, pending_comments, header_comment = self._process_statement_sequence(
+            updated_node.body
+        )
+        new_footer = (*updated_node.footer, *pending_comments)
+        if not new_body:
+            new_body = [libcst.SimpleStatementLine([libcst.Pass()])]
+
+        changes: dict[str, Any] = {"body": new_body, "footer": new_footer}
+
+        if header_comment:
+            changes["header"] = libcst.TrailingWhitespace(
+                whitespace=libcst.SimpleWhitespace(" "), comment=header_comment
+            )
+
+        return updated_node.with_changes(**changes)
+
+
+def codegen_for_ast(ast: ast.AST) -> tuple[ast.AST, ast.AST]:
+    ast_2 = deepcopy(ast)
+    return (
+        SyncCodegenNodeTransformer().generic_visit(ast),
+        AsyncCodegenNodeTransformer().generic_visit(ast_2),
+    )
+
+
+def add_comments(code: str) -> str:
+    cst = libcst.parse_module(code)
+    transformer = CommentTransformer()
+    return cst.visit(transformer).code
+
+
+def codegen_for_file(file: Path) -> tuple[ast.AST, ast.AST]:
+    code = file.read_text()
+    return codegen_for_ast(ast.parse(code))
+
+
+def codegen_for_whole_directory(directory: "StrPath"):
+    directory = Path(directory).resolve()
+    items = {path.name: path for path in directory.iterdir()}
+    codegen_config: CodegenConfig
+    try:
+        codegen_config = cast(
+            "CodegenConfig", json.loads(items.pop("codegen_config.json").read_text())
+        )
+    except KeyError:
+        codegen_config = CodegenConfig(
+            sync_target_directory=str(
+                directory.with_stem(f"{directory.stem}_sync"),
+            ),
+            async_target_directory=str(
+                directory.with_stem(f"{directory.stem}_async"),
+            ),
+            exclude=[],
+            include_directories=[],
+        )
+    sync_target_directory = directory / codegen_config["sync_target_directory"]
+    async_target_directory = directory / codegen_config["async_target_directory"]
+    sync_target_directory.mkdir(parents=True, exist_ok=True)
+    async_target_directory.mkdir(parents=True, exist_ok=True)
+    exclusions = {(directory / exclusion).resolve() for exclusion in codegen_config["exclude"]}
+    for path in items.values():
+        path = path.resolve()
+        if path.suffix.lower() != ".py":
+            continue
+        if not path.is_file():
+            continue
+        if path in exclusions:
+            continue
+        (sync_ast, async_ast) = codegen_for_file(path)
+        (sync_code, async_code) = (
+            add_comments(ast.unparse(sync_ast)),
+            add_comments(ast.unparse(async_ast)),
+        )
+        (sync_target_directory / path.name).write_text(sync_code)
+        (async_target_directory / path.name).write_text(async_code)
+    subprocess.run(
+        [
+            "python",
+            "-m",
+            "ruff",
+            "format",
+            str(sync_target_directory.resolve()),
+            str(async_target_directory.resolve()),
+        ],
+        capture_output=True,
+        text=True,
+    )
+    for included_dir in codegen_config["include_directories"]:
+        codegen_for_whole_directory(directory / included_dir)
+
+
+class CodegenArgumentNamespace(argparse.Namespace):
+    targets: list[Path]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("targets", nargs="*", type=Path)
+    parsed = parser.parse_args(namespace=CodegenArgumentNamespace())
+    for target in parsed.targets:
+        codegen_for_whole_directory(target)
+
+
+if __name__ == "__main__":
+    main()