From 96505249d6cdfa12c9a6ac693f8beb7d487c8233 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Sun, 2 Nov 2025 14:45:29 +0800
Subject: [PATCH] fix #401

---
 docs/C-API/introduction.md  |  6 +++++-
 include/pocketpy/pocketpy.h |  7 +++++--
 src/interpreter/ceval.c     | 16 +++++++++++-----
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/docs/C-API/introduction.md b/docs/C-API/introduction.md
index be744d81..88a7e66d 100644
--- a/docs/C-API/introduction.md
+++ b/docs/C-API/introduction.md
@@ -55,4 +55,8 @@ Mark a function that can raise an exception on failure.
 
 ### `PY_RETURN` macro
 
-Mark a function that can store a value in `py_retval()` on success.
\ No newline at end of file
+Mark a function that can store a value in `py_retval()` on success.
+
+### `PY_MAYBENULL` macro
+
+Mark a variable or callback function that may be `NULL`.
diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h
index d3969699..62767e5a 100644
--- a/include/pocketpy/pocketpy.h
+++ b/include/pocketpy/pocketpy.h
@@ -49,6 +49,7 @@ typedef struct c11_sv {
 
 #define PY_RAISE
 #define PY_RETURN
+#define PY_MAYBENULL
 
 /// A generic reference to a python object.
 typedef py_TValue* py_Ref;
@@ -79,7 +80,7 @@ typedef struct py_Callbacks {
     /// Used by `__import__` to load a source module.
     char* (*importfile)(const char*);
     /// Called before `importfile` to lazy-import a C module.
-    py_GlobalRef (*lazyimport)(const char*);
+    PY_MAYBENULL py_GlobalRef (*lazyimport)(const char*);
     /// Used by `print` to output a string.
     void (*print)(const char*);
     /// Flush the output buffer of `print`.
@@ -87,7 +88,9 @@ typedef struct py_Callbacks {
     /// Used by `input` to get a character.
     int (*getchr)();
     /// Used by `gc.collect()` to mark extra objects for garbage collection.
-    void (*gc_mark)(void (*f)(py_Ref val, void* ctx), void* ctx);
+    PY_MAYBENULL void (*gc_mark)(void (*f)(py_Ref val, void* ctx), void* ctx);
+    /// Used by `PRINT_EXPR` bytecode.
+    PY_MAYBENULL bool (*displayhook)(py_Ref val) PY_RAISE;
 } py_Callbacks;
 
 /// A struct contains the application-level callbacks.
diff --git a/src/interpreter/ceval.c b/src/interpreter/ceval.c
index 246557e6..0133ff0f 100644
--- a/src/interpreter/ceval.c
+++ b/src/interpreter/ceval.c
@@ -152,15 +152,21 @@ __NEXT_STEP:
             *THIRD() = tmp;
             DISPATCH();
         }
-        case OP_PRINT_EXPR:
-            if(TOP()->type != tp_NoneType) {
-                bool ok = py_repr(TOP());
+        case OP_PRINT_EXPR: {
+            if(self->callbacks.displayhook) {
+                bool ok = self->callbacks.displayhook(TOP());
                 if(!ok) goto __ERROR;
-                self->callbacks.print(py_tostr(&self->last_retval));
-                self->callbacks.print("\n");
+            } else {
+                if(TOP()->type != tp_NoneType) {
+                    bool ok = py_repr(TOP());
+                    if(!ok) goto __ERROR;
+                    self->callbacks.print(py_tostr(&self->last_retval));
+                    self->callbacks.print("\n");
+                }
             }
             POP();
             DISPATCH();
+        }
         /*****************************************/
         case OP_LOAD_CONST: {
             PUSH(c11__at(py_TValue, &frame->co->consts, byte.arg));