From 4e7db009930e93596d6ceea327e417cb60998c58 Mon Sep 17 00:00:00 2001
From: Valentin Berlier <berlier.v@gmail.com>
Date: Wed, 6 Dec 2023 02:29:45 +0100
Subject: [PATCH] feat: add codegen cli command

---
 bolt/commands.py | 18 +++++++++++++++++-
 bolt/module.py   |  4 ++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/bolt/commands.py b/bolt/commands.py
index ca9a4fb..b8a2bac 100644
--- a/bolt/commands.py
+++ b/bolt/commands.py
@@ -4,15 +4,31 @@
 from typing import Iterable, Optional, Tuple
 
 import click
-from beet import Project
+from beet import ErrorMessage, Project
 from beet.core.utils import format_directory
 from beet.toolchain.cli import beet, message_fence
+from mecha import AstCacheBackend
 
 from bolt import AstMemo, MemoRegistry, MemoStorage
 
 pass_project = click.make_pass_decorator(Project)  # type: ignore
 
 
+@beet.command()
+@pass_project
+@click.argument("filename", type=click.Path(exists=True, dir_okay=False))
+def codegen(project: Project, filename: str):
+    """Inspect cached bolt codegen."""
+    source_path = Path(filename).resolve()
+    ast_path = project.cache["mecha"].get_path(f"{source_path}-ast")
+
+    try:
+        with ast_path.open("rb") as f:
+            click.echo(AstCacheBackend().load_data(f)["codegen"], nl=False)
+    except Exception as exc:
+        raise ErrorMessage(f'No cached codegen for "{filename}".') from exc
+
+
 @beet.command()
 @pass_project
 @click.option(
diff --git a/bolt/module.py b/bolt/module.py
index 90c263a..3aabad8 100644
--- a/bolt/module.py
+++ b/bolt/module.py
@@ -102,6 +102,7 @@ class CompiledModule:
 
     ast: AstRoot
     code: Optional[CodeType]
+    python: Optional[str]
     refs: List[Any]
     dependencies: Set[str]
     prelude_imports: List[AstPrelude]
@@ -268,6 +269,7 @@ def __getitem__(self, current: Union[TextFileBase[Any], str]) -> CompiledModule:
         module = CompiledModule(
             ast=compilation_unit.ast,
             code=code,
+            python=result.source,
             refs=result.refs,
             dependencies=result.dependencies,
             prelude_imports=result.prelude_imports,
@@ -470,6 +472,7 @@ def load(self, f: BufferedReader) -> AstRoot:
         self.modules.registry[self.modules.database.current] = CompiledModule(
             ast=data["ast"],
             code=marshal.load(f),
+            python=data["codegen"],
             refs=data["refs"],
             dependencies=data["dependencies"],
             prelude_imports=data["prelude_imports"],
@@ -491,6 +494,7 @@ def dump(self, node: AstRoot, f: BufferedWriter):
         self.dump_data(
             {
                 "ast": module.ast,
+                "codegen": module.python,
                 "refs": module.refs,
                 "dependencies": module.dependencies,
                 "prelude_imports": module.prelude_imports,