From 7abc775ec20a1e2fc047a03e78b0e16f7ba21f79 Mon Sep 17 00:00:00 2001
From: Nicholas Gates <nick@nickgates.com>
Date: Fri, 13 Oct 2023 11:24:41 +0100
Subject: [PATCH] Allow state in root modules

---
 pydust/src/builtins.zig | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/pydust/src/builtins.zig b/pydust/src/builtins.zig
index d5223a31..2bb6b0f8 100644
--- a/pydust/src/builtins.zig
+++ b/pydust/src/builtins.zig
@@ -294,18 +294,21 @@ fn lift(comptime PydustStruct: type) !py.PyObject {
     // Grab the qualified name, importing the root module first.
     comptime var qualName = State.getIdentifier(PydustStruct).qualifiedName;
 
-    const root = try import(qualName[0]);
-    defer root.decref();
+    var mod = try import(qualName[0]);
 
     // Recursively resolve submodules / nested classes
-    var mod = root;
-    inline for (qualName[1 .. qualName.len - 1]) |part| {
-        mod = try mod.get(part);
-        defer mod.decref(); // Inline loop so this runs at the end of the function.
+    if (comptime qualName.len > 1) {
+        inline for (qualName[1 .. qualName.len - 1]) |part| {
+            defer mod.decref();
+            mod = try mod.get(part);
+        }
+
+        defer mod.decref();
+        mod = try mod.get(qualName[qualName.len - 1]);
     }
 
     // Grab the attribute using the final part of the qualified name.
-    return mod.get(qualName[qualName.len - 1]);
+    return mod;
 }
 
 const testing = std.testing;