feat(printer): Print Contextual Stmts in ToPython (#44)

Author: potatomashed

include/mlc/printer/ast.h

@@ -10,29 +10,32 @@ namespace printer {
 using mlc::core::ObjectPath;
 
 struct PrinterConfigObj : public Object {
+  bool def_free_var = true;
   int32_t indent_spaces = 2;
   int8_t print_line_numbers = 0;
   int32_t num_context_lines = -1;
   mlc::List<ObjectPath> path_to_underline;
 
   PrinterConfigObj() = default;
-  explicit PrinterConfigObj(int32_t indent_spaces, int8_t print_line_numbers, int32_t num_context_lines,
-                            mlc::List<ObjectPath> path_to_underline)
-      : indent_spaces(indent_spaces), print_line_numbers(print_line_numbers), num_context_lines(num_context_lines),
-        path_to_underline(path_to_underline) {}
+  explicit PrinterConfigObj(bool def_free_var, int32_t indent_spaces, int8_t print_line_numbers,
+                            int32_t num_context_lines, mlc::List<ObjectPath> path_to_underline)
+      : def_free_var(def_free_var), indent_spaces(indent_spaces), print_line_numbers(print_line_numbers),
+        num_context_lines(num_context_lines), path_to_underline(path_to_underline) {}
   MLC_DEF_DYN_TYPE(MLC_EXPORTS, PrinterConfigObj, Object, "mlc.printer.PrinterConfig");
 };
 
 struct PrinterConfig : public ObjectRef {
   MLC_DEF_OBJ_REF(MLC_EXPORTS, PrinterConfig, PrinterConfigObj, ObjectRef)
+      .Field("def_free_var", &PrinterConfigObj::def_free_var)
       .Field("indent_spaces", &PrinterConfigObj::indent_spaces)
       .Field("print_line_numbers", &PrinterConfigObj::print_line_numbers)
       .Field("num_context_lines", &PrinterConfigObj::num_context_lines)
       .Field("path_to_underline", &PrinterConfigObj::path_to_underline)
-      .StaticFn("__init__", InitOf<PrinterConfigObj, int32_t, int8_t, int32_t, mlc::List<ObjectPath>>);
-  explicit PrinterConfig(int32_t indent_spaces = 2, int8_t print_line_numbers = 0, int32_t num_context_lines = -1,
-                         mlc::List<ObjectPath> path_to_underline = {})
-      : PrinterConfig(PrinterConfig::New(indent_spaces, print_line_numbers, num_context_lines, path_to_underline)) {}
+      .StaticFn("__init__", InitOf<PrinterConfigObj, bool, int32_t, int8_t, int32_t, mlc::List<ObjectPath>>);
+  explicit PrinterConfig(bool def_free_var = true, int32_t indent_spaces = 2, int8_t print_line_numbers = 0,
+                         int32_t num_context_lines = -1, mlc::List<ObjectPath> path_to_underline = {})
+      : PrinterConfig(PrinterConfig::New(def_free_var, indent_spaces, print_line_numbers, num_context_lines,
+                                         path_to_underline)) {}
 };
 
 } // namespace printer
@@ -82,6 +85,17 @@ struct ExprObj : public ::mlc::Object {
 }; // struct ExprObj
 
 struct Expr : public ::mlc::printer::Node {
+  Expr Attr(mlc::Str name) const { return this->get()->Attr(name); }
+  Expr Index(mlc::List<::mlc::printer::Expr> idx) const { return this->get()->Index(idx); }
+  Expr Call(mlc::List<::mlc::printer::Expr> args) const { return this->get()->Call(args); }
+  Expr CallKw(mlc::List<::mlc::printer::Expr> args, mlc::List<::mlc::Str> kwargs_keys,
+              mlc::List<::mlc::printer::Expr> kwargs_values) const {
+    return this->get()->CallKw(args, kwargs_keys, kwargs_values);
+  }
+  Expr AddPath(mlc::core::ObjectPath p) {
+    this->get()->source_paths->push_back(p);
+    return *this;
+  }
   MLC_DEF_OBJ_REF(MLC_EXPORTS, Expr, ExprObj, ::mlc::printer::Node)
       .Field("source_paths", &ExprObj::source_paths)
       .StaticFn("__init__", ::mlc::InitOf<ExprObj, ::mlc::List<::mlc::core::ObjectPath>>)

include/mlc/printer/ir_printer.h

@@ -139,6 +139,17 @@ struct IRPrinterObj : public Object {
     return ret;
   }
 
+  template <typename T> mlc::List<T> ApplyToList(const UList &list, const ObjectPath &p) const {
+    // TODO: expose a Python interface
+    int64_t n = list->size();
+    mlc::List<T> args;
+    args.reserve(n);
+    for (int64_t i = 0; i < n; ++i) {
+      args->push_back(this->operator()(list[i], p->WithListIndex(i)));
+    }
+    return args;
+  }
+
   void FramePush(const ObjectRef &frame) {
     frames.push_back(frame);
     frame_vars[frame] = mlc::UList();
@@ -178,14 +189,28 @@ struct IRPrinter : public ObjectRef {
   explicit IRPrinter(PrinterConfig cfg, mlc::Dict<Any, VarInfo> obj2info, mlc::Dict<Str, int64_t> defined_names,
                      mlc::UList frames, mlc::UDict frame_vars)
       : IRPrinter(IRPrinter::New(cfg, obj2info, defined_names, frames, frame_vars)) {}
-
 }; // struct IRPrinter
 
 inline Str ToPython(const ObjectRef &obj, const PrinterConfig &cfg) {
   IRPrinter printer(cfg);
-  printer->FramePush(DefaultFrame());
+  DefaultFrame frame;
+  printer->FramePush(frame);
   Node ret = ::mlc::Lib::IRPrint(obj, printer, ObjectPath::Root());
   printer->FramePop();
+  if (frame->stmts->empty()) {
+    return ret->ToPython(cfg);
+  }
+  if (const auto *block = ret.as<StmtBlockObj>()) {
+    // TODO: support List::insert by iterator
+    frame->stmts->insert(frame->stmts.size(), block->stmts->begin(), block->stmts->end());
+  } else if (const auto *expr = ret.as<ExprObj>()) {
+    frame->stmts->push_back(ExprStmt(mlc::List<ObjectPath>{}, Optional<Str>{}, Expr(expr)));
+  } else if (const auto *stmt = ret.as<StmtObj>()) {
+    frame->stmts->push_back(Stmt(stmt));
+  } else {
+    MLC_THROW(ValueError) << "Unsupported type: " << ret;
+  }
+  ret = StmtBlock(mlc::List<ObjectPath>{}, Optional<Str>{}, frame->stmts);
   return ret->ToPython(cfg);
 }
 

python/mlc/printer/ast.py

@@ -9,6 +9,7 @@
 
 @mlcd.c_class("mlc.printer.PrinterConfig")
 class PrinterConfig(Object):
+    def_free_var: bool = True
     indent_spaces: int = 2
     print_line_numbers: int = 0
     num_context_lines: int = -1

tests/python/test_printer_ast.py

@@ -1,11 +1,13 @@
+from __future__ import annotations
+
 import itertools
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 import mlc.printer as mlcp
 import pytest
 
 if TYPE_CHECKING:
-    import _pytest
+    from _pytest.mark import ParameterSet
 
 
 @pytest.mark.parametrize(
@@ -645,7 +647,7 @@ def test_print_expr_stmt_doc() -> None:
     ],
     ids=itertools.count(),
 )
-def test_print_assert_doc(msg: Optional[mlcp.ast.Expr], expected: str) -> None:
+def test_print_assert_doc(msg: mlcp.ast.Expr | None, expected: str) -> None:
     test = mlcp.ast.Literal(True)
     doc = mlcp.ast.Assert(test, msg)
     assert doc.to_python().strip() == expected.strip()
@@ -774,7 +776,7 @@ def test_print_function_doc(
     args: list[mlcp.ast.Assign],
     decorators: list[mlcp.ast.Id],
     body: list[mlcp.ast.Stmt],
-    return_type: Optional[mlcp.ast.Expr],
+    return_type: mlcp.ast.Expr | None,
     expected: str,
 ) -> None:
     doc = mlcp.ast.Function(mlcp.ast.Id("func"), args, decorators, return_type, body)
@@ -1103,7 +1105,7 @@ def test_print_invalid_multiline_doc_comment(doc: mlcp.ast.Stmt) -> None:
     assert "cannot have newline" in str(e.value)
 
 
-def generate_expr_precedence_test_cases() -> list["_pytest.mark.ParameterSet"]:
+def generate_expr_precedence_test_cases() -> list[ParameterSet]:
     x = mlcp.ast.Id("x")
     y = mlcp.ast.Id("y")
     z = mlcp.ast.Id("z")