Add __repr_pretty__ support

cmake/deps.cmake

@@ -134,10 +134,7 @@ if(CODON_JUPYTER)
         NAME cppzmq
         URL https://github.com/zeromq/cppzmq/archive/refs/tags/v4.8.1.tar.gz
         VERSION 4.8.1
-        OPTION "CPPZMQ_BUILD_TESTS OFF")
-    if(cppzmq_ADDED)
-        set_target_properties(unit_tests PROPERTIES EXCLUDE_FROM_ALL ON)
-    endif()
+        OPTIONS "CPPZMQ_BUILD_TESTS OFF")
     CPMAddPackage(
         NAME xtl
         GITHUB_REPOSITORY "xtensor-stack/xtl"

codon/compiler/jit.cpp

@@ -84,8 +84,8 @@ public:
 };
 } // namespace
 
-JIT::JIT(const std::string &argv0)
-    : compiler(std::make_unique<Compiler>(argv0, /*debug=*/true)) {
+JIT::JIT(const std::string &argv0, const std::string &mode)
+    : compiler(std::make_unique<Compiler>(argv0, /*debug=*/true)), mode(mode) {
   if (auto e = Engine::create()) {
     engine = std::move(e.get());
   } else {
@@ -176,6 +176,17 @@ llvm::Expected<std::string> JIT::exec(const std::string &code) {
   auto sctx = cache->imports[MAIN_IMPORT].ctx;
   auto preamble = std::make_shared<ast::SimplifyVisitor::Preamble>();
   try {
+    auto *e = node->getSuite()
+                  ? const_cast<ast::SuiteStmt *>(node->getSuite())->lastInBlock()
+                  : &node;
+    if (e)
+      if (auto ex = (*e)->getExpr()) {
+        *e = std::make_shared<ast::PrintStmt>(
+            std::vector<ast::ExprPtr>{std::make_shared<ast::CallExpr>(
+                std::make_shared<ast::IdExpr>("_jit_display"), ex->expr,
+                std::make_shared<ast::StringExpr>(mode))},
+            false);
+      }
     auto s = ast::SimplifyVisitor(sctx, preamble).transform(node);
     auto simplified = std::make_shared<ast::SuiteStmt>();
     for (auto &s : preamble->globals)

codon/compiler/jit.h

@@ -19,9 +19,10 @@ class JIT {
 private:
   std::unique_ptr<Compiler> compiler;
   std::unique_ptr<Engine> engine;
+  std::string mode;
 
 public:
-  explicit JIT(const std::string &argv0);
+  explicit JIT(const std::string &argv0, const std::string &mode = "");
 
   Compiler *getCompiler() const { return compiler.get(); }
   Engine *getEngine() const { return engine.get(); }

codon/parser/ast/stmt.cpp

@@ -52,6 +52,16 @@ void SuiteStmt::flatten(StmtPtr s, std::vector<StmtPtr> &stmts) {
       stmts.push_back(ss);
   }
 }
+StmtPtr *SuiteStmt::lastInBlock() {
+  if (stmts.empty())
+    return nullptr;
+  if (auto s = const_cast<SuiteStmt *>(stmts.back()->getSuite())) {
+    auto l = s->lastInBlock();
+    if (l)
+      return l;
+  }
+  return &(stmts.back());
+}
 
 std::string BreakStmt::toString(int) const { return "(break)"; }
 ACCEPT_IMPL(BreakStmt, ASTVisitor);

codon/parser/ast/stmt.h

@@ -92,6 +92,7 @@ struct SuiteStmt : public Stmt {
   const Stmt *firstInBlock() const override {
     return stmts.empty() ? nullptr : stmts[0]->firstInBlock();
   }
+  StmtPtr *lastInBlock();
 
   /// Flatten all nested SuiteStmt objects that do not own a block in the statement
   /// vector. This is shallow flattening.

extra/jupyter/src/codon.cpp

@@ -61,7 +61,7 @@ nl::json CodonJupyter::execute_request_impl(int execution_counter, const string
 }
 
 void CodonJupyter::configure_impl() {
-  jit = std::make_unique<codon::jit::JIT>(argv0);
+  jit = std::make_unique<codon::jit::JIT>(argv0, "jupyter");
   llvm::cantFail(jit->init());
 }
 

stdlib/internal/__init__.codon

@@ -25,6 +25,7 @@ from internal.types.collections.dict import *
 
 import internal.c_stubs as _C
 from internal.builtin import *
+from internal.builtin import _jit_display
 from internal.box import Box
 from internal.str import *
 

stdlib/internal/builtin.codon

@@ -329,3 +329,14 @@ class int:
             raise ValueError("invalid literal for int() with base " + str(base) + ": " + s)
 
         return result
+
+
+def _jit_display(x, s: Static[str]):
+    if hasattr(x, "__repr_pretty__") and s == "jupyter":
+        return x.__repr_pretty__()
+    elif hasattr(x, "__repr__"):
+        return x.__repr__()
+    elif hasattr(x, "__str__"):
+        return x.__str__()
+    else:
+        return ''