diff --git a/capellambse_context_diagrams/collectors/custom.py b/capellambse_context_diagrams/collectors/custom.py
index fde8399..6169ea1 100644
--- a/capellambse_context_diagrams/collectors/custom.py
+++ b/capellambse_context_diagrams/collectors/custom.py
@@ -4,12 +4,9 @@
 """This module defines the collector for the CustomDiagram."""
 from __future__ import annotations
 
-import builtins
-import collections.abc as cabc
 import copy
 import typing as t
 
-import capellambse
 import capellambse.model as m
 
 from .. import _elkjs, context
@@ -33,7 +30,7 @@ class CustomCollector:
 
     def __init__(
         self,
-        diagram: context.CustomDiagram,
+        diagram: context.ContextDiagram,
         params: dict[str, t.Any],
     ) -> None:
         self.diagram = diagram
diff --git a/docs/custom_diagram.md b/docs/custom_diagram.md
index 3f4f7de..98d1c1c 100644
--- a/docs/custom_diagram.md
+++ b/docs/custom_diagram.md
@@ -5,7 +5,8 @@
 
 # Custom Diagram
 
-`Custom diagram`s let's you create custom diagrams based on the data in the model. You define the data collection using an iterable and `Custom diagram` takes care of the rest.
+`Custom diagram`s let you create custom diagrams based on the data in the model. You define the data collection using an iterable, and `Custom diagram` takes care of the rest.
+
 You can access `.custom_diagram` on any supported model element.
 
 ??? example "Custom Diagram of `PP 1 `"
@@ -13,19 +14,21 @@ You can access `.custom_diagram` on any supported model element.
     ``` py
     import capellambse
 
-    visited = set()
-
     def _collector(
         target: m.ModelElement,
     ) -> cabc.Iterator[m.ModelElement]:
-        if target.uuid in visited:
-            return
-        visited.add(target.uuid)
-        for link in target.links:
-            yield link
-            yield from _collector(link.source)
-            yield from _collector(link.target)
-
+        visited = set()
+        def collector(
+            target: m.ModelElement,
+        ) -> cabc.Iterator[m.ModelElement]:
+            if target.uuid in visited:
+                return
+            visited.add(target.uuid)
+            for link in target.links:
+                yield link
+                yield from collector(link.source)
+                yield from collector(link.target)
+        yield from collector(target)
 
     model = capellambse.MelodyModel("tests/data/ContextDiagram.aird")
     obj = model.by_uuid("c403d4f4-9633-42a2-a5d6-9e1df2655146")