Zig strack traces

pydust/src/functions.zig

@@ -169,13 +169,11 @@ pub fn wrap(comptime func: anytype, comptime sig: Signature, comptime flags: c_i
                     @field(args, field.name) = try py.as(field.type, pyargs[i]);
                 }
 
-                var callArgs = if (sig.selfParam) |_| .{ self, args } else .{args};
-                const result = @call(.always_inline, func, callArgs);
+                const result = if (sig.selfParam) |_| func(self, args) else func(args);
 
                 return py.createOwned(result);
             } else {
-                var callArgs = if (sig.selfParam) |_| .{self} else .{};
-                const result = @call(.always_inline, func, callArgs);
+                const result = if (sig.selfParam) |_| func(self) else func();
                 return py.createOwned(result);
             }
         }
@@ -256,8 +254,7 @@ pub fn wrap(comptime func: anytype, comptime sig: Signature, comptime flags: c_i
             }
 
             const self = if (sig.selfParam) |Self| try py.as(Self, pyself) else null;
-            var callArgs = if (sig.selfParam) |_| .{ self, args } else .{args};
-            const result = @call(.always_inline, func, callArgs);
+            const result = if (sig.selfParam) |_| func(self, args) else func(args);
             return py.createOwned(result);
         }
     };

pydust/src/types/error.zig

@@ -10,6 +10,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+const builtin = @import("builtin");
 const std = @import("std");
 const ffi = @import("../ffi.zig");
 const py = @import("../pydust.zig");
@@ -93,6 +94,7 @@ const PyExc = struct {
 
     pub fn raise(comptime self: Self, message: [:0]const u8) PyError {
         ffi.PyErr_SetString(self.asPyObject().py, message.ptr);
+        try augmentTraceback();
         return PyError.Raised;
     }
 
@@ -109,4 +111,97 @@ const PyExc = struct {
     inline fn asPyObject(comptime self: Self) py.PyObject {
         return .{ .py = @field(ffi, "PyExc_" ++ self.name) };
     }
+
+    /// In debug mode, augment the Python traceback to include Zig stack frames.
+    /// Warning: hackery ahead!
+    fn augmentTraceback() PyError!void {
+        if (builtin.mode == .Debug) {
+            // First of all, grab the current Python exception
+            var ptype: ?*ffi.PyObject = undefined;
+            var pvalue: ?*ffi.PyObject = undefined;
+            var ptraceback: ?*ffi.PyObject = undefined;
+            ffi.PyErr_Fetch(&ptype, &pvalue, &ptraceback);
+
+            // Capture at most 32 stack frames above us.
+            var addresses: [32]usize = undefined;
+            var st: std.builtin.StackTrace = .{
+                .index = 0,
+                .instruction_addresses = &addresses,
+            };
+            std.debug.captureStackTrace(@returnAddress(), &st);
+
+            const debugInfo = std.debug.getSelfDebugInfo() catch return;
+
+            // Skip the first frame (this function) and the last frame (the trampoline entrypoint)
+            for (0..st.index) |idx| {
+                // std.debug.writeStackTrace subtracts 1 from the address - not sure why, but it gives accurate frames.
+                const address = st.instruction_addresses[idx] - 1;
+
+                // If we can't find info for the stack frame, then we skip this frame..
+                const module = debugInfo.getModuleForAddress(address) catch continue;
+                const symbol_info: std.debug.SymbolInfo = module.getSymbolAtAddress(debugInfo.allocator, address) catch continue;
+                defer symbol_info.deinit(debugInfo.allocator);
+                const line_info = symbol_info.line_info orelse continue;
+
+                // We also want to skip any Pydust internal frames, e.g. the function trampoline and also this current function!
+                if (std.mem.indexOf(u8, line_info.file_name, "/pydust/src/")) |_| {
+                    continue;
+                }
+
+                // Allocate a string of newlines.
+                // Since we wrap the error in a function, we have an addition "def foo()" line.
+                // In addition to lineno being zero-based, we have to subtract 2.
+                // This means that exceptions on line 1 will be off... but that's quite rare.
+                const nnewlines = if (line_info.line < 2) 0 else line_info.line - 2;
+                const newlines = try py.allocator.alloc(u8, nnewlines);
+                @memset(newlines, '\n');
+
+                // Setup a function we know will fail (with DivideByZero error)
+                const code = try std.fmt.allocPrintZ(
+                    py.allocator,
+                    "{s}def {s}():\n    1/0\n",
+                    .{ newlines, symbol_info.symbol_name },
+                );
+
+                // Compilation should succeed, but execution will fail.
+                const filename = try py.allocator.dupeZ(u8, line_info.file_name);
+                defer py.allocator.free(filename);
+                const compiled = ffi.Py_CompileString(code.ptr, filename.ptr, ffi.Py_file_input) orelse continue;
+
+                // Import the compiled code as a module and invoke the failing function
+                const module_name = try py.allocator.dupeZ(u8, symbol_info.compile_unit_name);
+                defer py.allocator.free(module_name);
+                const fake_module: py.PyObject = .{
+                    .py = ffi.PyImport_ExecCodeModule(module_name.ptr, compiled) orelse continue,
+                };
+
+                const func_name = try py.allocator.dupeZ(u8, symbol_info.symbol_name);
+                defer py.allocator.free(func_name);
+                const fake_function = try fake_module.get(func_name);
+                _ = fake_function.call(.{}, .{}) catch null;
+
+                // Grab our forced exception info.
+                // We can ignore qtype and qvalue, we just want to get the traceback object.
+                var qtype: ?*ffi.PyObject = undefined;
+                var qvalue: ?*ffi.PyObject = undefined;
+                var qtraceback: ?*ffi.PyObject = undefined;
+                ffi.PyErr_Fetch(&qtype, &qvalue, &qtraceback);
+                if (qtype) |q| py.decref(q);
+                if (qvalue) |q| py.decref(q);
+                std.debug.assert(qtraceback != null);
+
+                // Extract the traceback frame by calling into Python (Pytraceback isn't part of the Stable API)
+                const pytb = py.PyObject{ .py = qtraceback.? };
+                const frame = (try pytb.get("tb_frame")).py;
+
+                // Restore the original exception, augment it with the new frame, then fetch the new exception.
+                ffi.PyErr_Restore(ptype, pvalue, ptraceback);
+                _ = ffi.PyTraceBack_Here(@alignCast(@ptrCast(frame)));
+                ffi.PyErr_Fetch(&ptype, &pvalue, &ptraceback);
+            }
+
+            // Restore the latest the exception info
+            ffi.PyErr_Restore(ptype, pvalue, ptraceback);
+        }
+    }
 };