From d435ccb69d2b57d640c920d0e696023a9795ab56 Mon Sep 17 00:00:00 2001
From: FeepingCreature <feepingcreature@gmail.com>
Date: Mon, 4 Nov 2024 21:39:44 +0100
Subject: [PATCH] Work around PyQt memory leak on Krita::documents() and
 Node::childNodes().

---
 ai_diffusion/document.py | 12 ++++++++----
 ai_diffusion/layer.py    | 12 ++++++++----
 ai_diffusion/util.py     | 16 +++++++++++++++-
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/ai_diffusion/document.py b/ai_diffusion/document.py
index 27fcdfbc6..8fbd5bddc 100644
--- a/ai_diffusion/document.py
+++ b/ai_diffusion/document.py
@@ -11,6 +11,7 @@
 from .layer import Layer, LayerManager, LayerType
 from .pose import Pose
 from .localization import translate as _
+from .util import acquire_elements
 
 
 class Document(QObject):
@@ -120,7 +121,10 @@ def __init__(self, krita_document: krita.Document):
     @classmethod
     def active(cls):
         if doc := Krita.instance().activeDocument():
-            if doc not in Krita.instance().documents() or doc.activeNode() is None:
+            if (
+                doc not in acquire_elements(Krita.instance().documents())
+                or doc.activeNode() is None
+            ):
                 return None
             id = doc.rootNode().uniqueId().toString()
             return cls._instances.get(id) or KritaDocument(doc)
@@ -237,7 +241,7 @@ def current_time(self):
 
     @property
     def is_valid(self):
-        return self._doc in Krita.instance().documents()
+        return self._doc in acquire_elements(Krita.instance().documents())
 
     @property
     def is_active(self):
@@ -303,14 +307,14 @@ def update(self):
 
         layer = cast(krita.VectorLayer, layer.node)
         pose = self._layers.setdefault(layer.uniqueId(), Pose(doc.extent))
-        self._update(layer, layer.shapes(), pose, doc.resolution)
+        self._update(layer, acquire_elements(layer.shapes()), pose, doc.resolution)
 
     def add_character(self, layer: krita.VectorLayer):
         doc = KritaDocument.active()
         assert doc is not None
         pose = self._layers.setdefault(layer.uniqueId(), Pose(doc.extent))
         svg = Pose.create_default(doc.extent, pose.people_count).to_svg()
-        shapes = layer.addShapesFromSvg(svg)
+        shapes = acquire_elements(layer.addShapesFromSvg(svg))
         self._update(layer, shapes, pose, doc.resolution)
 
     def _update(
diff --git a/ai_diffusion/layer.py b/ai_diffusion/layer.py
index a01771825..e0d5ca63d 100644
--- a/ai_diffusion/layer.py
+++ b/ai_diffusion/layer.py
@@ -6,7 +6,7 @@
 from PyQt5.QtGui import QImage
 
 from .image import Extent, Bounds, Image
-from .util import ensure, maybe, client_logger as log
+from .util import acquire_elements, ensure, maybe, client_logger as log
 from . import eventloop
 
 
@@ -143,7 +143,11 @@ def parent_layer(self):
 
     @property
     def child_layers(self):
-        return [self._manager.wrap(child) for child in self._node.childNodes() if _is_real(child)]
+        return [
+            self._manager.wrap(child)
+            for child in acquire_elements(self._node.childNodes())
+            if _is_real(child)
+        ]
 
     @property
     def is_root(self):
@@ -199,7 +203,7 @@ def get_mask(self, bounds: Bounds | None = None, time: int | None = None):
 
     def move_to_top(self):
         parent = self._node.parentNode()
-        if parent.childNodes()[-1] == self._node:
+        if acquire_elements(parent.childNodes())[-1] == self._node:
             return  # already top-most layer
         with RestoreActiveLayer(self._manager):
             parent.removeChildNode(self.node)
@@ -556,7 +560,7 @@ def __bool__(self):
 
 
 def traverse_layers(node: krita.Node, type_filter: list[str] | None = None):
-    for child in node.childNodes():
+    for child in acquire_elements(node.childNodes()):
         type = child.type()
         if _is_real(type) and (not type_filter or type in type_filter):
             yield child
diff --git a/ai_diffusion/util.py b/ai_diffusion/util.py
index 0f4cd47b2..c2cb1c2fb 100644
--- a/ai_diffusion/util.py
+++ b/ai_diffusion/util.py
@@ -14,10 +14,12 @@
 import statistics
 import zipfile
 from typing import Any, Callable, Iterable, Optional, Sequence, TypeVar
-from PyQt5.QtCore import QStandardPaths
+from PyQt5 import sip
+from PyQt5.QtCore import QObject, QStandardPaths
 
 T = TypeVar("T")
 R = TypeVar("R")
+QOBJECT = TypeVar("QOBJECT", bound=QObject)
 
 is_windows = sys.platform.startswith("win")
 is_macos = sys.platform == "darwin"
@@ -243,3 +245,15 @@ def _extract_member(self, member, targetpath, pwd):
 
 
 ZipFile = LongPathZipFile if is_windows else zipfile.ZipFile
+
+
+def acquire_elements(l: list[QOBJECT]) -> list[QOBJECT]:
+    # Many Pykrita functions return a `QList<QObject*>` where the objects are
+    # allocated for the caller. SIP does not handle this case and just leaks
+    # the objects outright. Fix this by taking explicit ownership of the objects.
+    # Note: ONLY call this if you are confident that the Pykrita function
+    # allocates the list members!
+    for obj in l:
+        if obj is not None:
+            sip.transferback(obj)
+    return l