6 files changed, 112 insertions, 67 deletions
diff --git a/sources/shiboken6/ApiExtractor/apiextractor.cpp b/sources/shiboken6/ApiExtractor/apiextractor.cpp
index d18758fe8..6958c8c8f 100644
--- a/sources/shiboken6/ApiExtractor/apiextractor.cpp
+++ b/sources/shiboken6/ApiExtractor/apiextractor.cpp
@@ -607,7 +607,10 @@ static void getCode(QStringList &code, const CodeSnipList &codeSnips)
 
 static void getCode(QStringList &code, const TypeEntry *type)
 {
-    getCode(code, type->codeSnips());
+    if (type->isComplex())
+        getCode(code, static_cast<const ComplexTypeEntry *>(type)->codeSnips());
+    else if (type->isTypeSystem())
+        getCode(code, static_cast<const TypeSystemTypeEntry *>(type)->codeSnips());
 
     CustomConversion *customConversion = type->customConversion();
     if (!customConversion)
diff --git a/sources/shiboken6/ApiExtractor/complextypeentry.h b/sources/shiboken6/ApiExtractor/complextypeentry.h
index fd17fe572..b1402c8e4 100644
--- a/sources/shiboken6/ApiExtractor/complextypeentry.h
+++ b/sources/shiboken6/ApiExtractor/complextypeentry.h
@@ -69,6 +69,18 @@ public:
     void addFunctionModification(const FunctionModification &functionModification);
     FunctionModificationList functionModifications(const QString &signature) const;
 
+    const CodeSnipList &codeSnips() const;
+    CodeSnipList &codeSnips();
+    void setCodeSnips(const CodeSnipList &codeSnips);
+    void addCodeSnip(const CodeSnip &codeSnip);
+
+    void setDocModification(const DocModificationList& docMods);
+    DocModificationList docModifications() const;
+
+    /// Extra includes for function arguments determined by the meta builder.
+    const IncludeList &argumentIncludes() const;
+    void addArgumentInclude(const Include &newInclude);
+
     AddedFunctionList addedFunctions() const;
     void setAddedFunctions(const AddedFunctionList &addedFunctions);
     void addNewFunction(const AddedFunctionPtr &addedFunction);
diff --git a/sources/shiboken6/ApiExtractor/typesystem.cpp b/sources/shiboken6/ApiExtractor/typesystem.cpp
index 9888de31a..a818b576c 100644
--- a/sources/shiboken6/ApiExtractor/typesystem.cpp
+++ b/sources/shiboken6/ApiExtractor/typesystem.cpp
@@ -64,10 +64,7 @@ public:
     QString m_targetLangPackage;
     mutable QString m_cachedTargetLangName; // "Foo.Bar"
     mutable QString m_cachedTargetLangEntryName; // "Bar"
-    CodeSnipList m_codeSnips;
-    DocModificationList m_docModifications;
     IncludeList m_extraIncludes;
-    IncludeList m_argumentIncludes;
     Include m_include;
     QString m_targetConversionRule;
     QVersionNumber m_version;
@@ -111,36 +108,6 @@ TypeEntry::TypeEntry(TypeEntryPrivate *d) : m_d(d)
 
 TypeEntry::~TypeEntry() = default;
 
-const CodeSnipList &TypeEntry::codeSnips() const
-{
-    return m_d->m_codeSnips;
-}
-
-CodeSnipList &TypeEntry::codeSnips()
-{
-    return m_d->m_codeSnips;
-}
-
-void TypeEntry::setCodeSnips(const CodeSnipList &codeSnips)
-{
-    m_d->m_codeSnips = codeSnips;
-}
-
-void TypeEntry::addCodeSnip(const CodeSnip &codeSnip)
-{
-    m_d->m_codeSnips << codeSnip;
-}
-
-void TypeEntry::setDocModification(const DocModificationList &docMods)
-{
-    m_d->m_docModifications << docMods;
-}
-
-DocModificationList TypeEntry::docModifications() const
-{
-    return m_d->m_docModifications;
-}
-
 const IncludeList &TypeEntry::extraIncludes() const
 {
     return m_d->m_extraIncludes;
@@ -157,17 +124,6 @@ void TypeEntry::addExtraInclude(const Include &newInclude)
         m_d->m_extraIncludes.append(newInclude);
 }
 
-const IncludeList &TypeEntry::argumentIncludes() const
-{
-    return m_d->m_argumentIncludes;
-}
-
-void TypeEntry::addArgumentInclude(const Include &newInclude)
-{
-    if (!m_d->m_argumentIncludes.contains(newInclude))
-        m_d->m_argumentIncludes.append(newInclude);
-}
-
 Include TypeEntry::include() const
 {
     return m_d->m_include;
@@ -746,6 +702,7 @@ class TypeSystemTypeEntryPrivate : public TypeEntryPrivate
 public:
     using TypeEntryPrivate::TypeEntryPrivate;
 
+    CodeSnipList m_codeSnips;
     TypeSystem::SnakeCase m_snakeCase = TypeSystem::SnakeCase::Disabled;
 };
 
@@ -766,6 +723,24 @@ TypeEntry *TypeSystemTypeEntry::clone() const
     return new TypeSystemTypeEntry(new TypeSystemTypeEntryPrivate(*d));
 }
 
+const CodeSnipList &TypeSystemTypeEntry::codeSnips() const
+{
+    S_D(const TypeSystemTypeEntry);
+    return d->m_codeSnips;
+}
+
+CodeSnipList &TypeSystemTypeEntry::codeSnips()
+{
+    S_D(TypeSystemTypeEntry);
+    return d->m_codeSnips;
+}
+
+void TypeSystemTypeEntry::addCodeSnip(const CodeSnip &codeSnip)
+{
+    S_D(TypeSystemTypeEntry);
+    d->m_codeSnips.append(codeSnip);
+}
+
 TypeSystem::SnakeCase TypeSystemTypeEntry::snakeCase() const
 {
     S_D(const TypeSystemTypeEntry);
@@ -1249,6 +1224,9 @@ public:
 
     AddedFunctionList m_addedFunctions;
     FunctionModificationList m_functionMods;
+    CodeSnipList m_codeSnips;
+    DocModificationList m_docModifications;
+    IncludeList m_argumentIncludes;
     QSet<QString> m_generateFunctions;
     FieldModificationList m_fieldMods;
     QList<TypeSystemProperty> m_properties;
@@ -1358,6 +1336,55 @@ FunctionModificationList ComplexTypeEntry::functionModifications(const QString &
     return lst;
 }
 
+const CodeSnipList &ComplexTypeEntry::codeSnips() const
+{
+    S_D(const ComplexTypeEntry);
+    return d->m_codeSnips;
+}
+
+CodeSnipList &ComplexTypeEntry::codeSnips()
+{
+    S_D(ComplexTypeEntry);
+    return d->m_codeSnips;
+}
+
+void ComplexTypeEntry::setCodeSnips(const CodeSnipList &codeSnips)
+{
+    S_D(ComplexTypeEntry);
+    d->m_codeSnips = codeSnips;
+}
+
+void ComplexTypeEntry::addCodeSnip(const CodeSnip &codeSnip)
+{
+    S_D(ComplexTypeEntry);
+    d->m_codeSnips << codeSnip;
+}
+
+void ComplexTypeEntry::setDocModification(const DocModificationList &docMods)
+{
+    S_D(ComplexTypeEntry);
+    d->m_docModifications << docMods;
+}
+
+DocModificationList ComplexTypeEntry::docModifications() const
+{
+    S_D(const ComplexTypeEntry);
+    return d->m_docModifications;
+}
+
+const IncludeList &ComplexTypeEntry::argumentIncludes() const
+{
+    S_D(const ComplexTypeEntry);
+    return d->m_argumentIncludes;
+}
+
+void ComplexTypeEntry::addArgumentInclude(const Include &newInclude)
+{
+    S_D(ComplexTypeEntry);
+    if (!d->m_argumentIncludes.contains(newInclude))
+        d->m_argumentIncludes.append(newInclude);
+}
+
 AddedFunctionList ComplexTypeEntry::addedFunctions() const
 {
     S_D(const ComplexTypeEntry);
@@ -2387,7 +2414,6 @@ void TypeEntry::formatDebug(QDebug &debug) const
     FORMAT_NONEMPTY_STRING("package", m_d->m_targetLangPackage)
     FORMAT_BOOL("stream", m_d->m_stream)
     FORMAT_BOOL("built-in", m_d->m_builtin)
-    FORMAT_LIST_SIZE("codeSnips", m_d->m_codeSnips)
     FORMAT_NONEMPTY_STRING("targetConversionRule", m_d->m_targetConversionRule)
     if (m_d->m_viewOn)
        debug << ", views=" << m_d->m_viewOn->name();
@@ -2433,6 +2459,7 @@ void ComplexTypeEntry::formatDebug(QDebug &debug) const
     FORMAT_NONEMPTY_STRING("hash", d->m_hashFunction)
     FORMAT_LIST_SIZE("addedFunctions", d->m_addedFunctions)
     formatList(debug, "functionMods", d->m_functionMods, ", ");
+    FORMAT_LIST_SIZE("codeSnips", d->m_codeSnips)
     FORMAT_LIST_SIZE("fieldMods", d->m_fieldMods)
 }
 
diff --git a/sources/shiboken6/ApiExtractor/typesystem.h b/sources/shiboken6/ApiExtractor/typesystem.h
index 186d82675..5e22af5f8 100644
--- a/sources/shiboken6/ApiExtractor/typesystem.h
+++ b/sources/shiboken6/ApiExtractor/typesystem.h
@@ -153,22 +153,10 @@ public:
     virtual bool isValue() const;
     virtual bool isComplex() const;
 
-    const CodeSnipList &codeSnips() const;
-    CodeSnipList &codeSnips();
-    void setCodeSnips(const CodeSnipList &codeSnips);
-    void addCodeSnip(const CodeSnip &codeSnip);
-
-    void setDocModification(const DocModificationList& docMods);
-    DocModificationList docModifications() const;
-
     const IncludeList &extraIncludes() const;
     void setExtraIncludes(const IncludeList &includes);
     void addExtraInclude(const Include &newInclude);
 
-    /// Extra includes for function arguments determined by the meta builder.
-    const IncludeList &argumentIncludes() const;
-    void addArgumentInclude(const Include &newInclude);
-
     Include include() const;
     void setInclude(const Include &inc);
 
diff --git a/sources/shiboken6/ApiExtractor/typesystemparser.cpp b/sources/shiboken6/ApiExtractor/typesystemparser.cpp
index da6cc8255..235c14324 100644
--- a/sources/shiboken6/ApiExtractor/typesystemparser.cpp
+++ b/sources/shiboken6/ApiExtractor/typesystemparser.cpp
@@ -881,7 +881,7 @@ bool TypeSystemParser::endElement(StackElement element)
                     toNative->setSourceType(m_context->db->findType(toNative->sourceTypeName()));
             }
         }
-        purgeEmptyCodeSnips(&top->entry->codeSnips());
+        purgeEmptyCodeSnips(&static_cast<TypeSystemTypeEntry *>(top->entry)->codeSnips());
         break;
     case StackElement::FunctionTypeEntry:
         TypeDatabase::instance()->addGlobalUserFunctionModifications(top->functionMods);
@@ -907,8 +907,11 @@ bool TypeSystemParser::endElement(StackElement element)
         centry->setAddedFunctions(centry->addedFunctions() + top->addedFunctions);
         centry->setFunctionModifications(centry->functionModifications() + top->functionMods);
         centry->setFieldModifications(centry->fieldModifications() + top->fieldMods);
-        centry->setCodeSnips(centry->codeSnips() + top->entry->codeSnips());
         centry->setDocModification(centry->docModifications() + top->docModifications);
+        if (top->entry->isComplex()) {
+            auto *cte = static_cast<const ComplexTypeEntry *>(top->entry);
+            centry->setCodeSnips(centry->codeSnips() + cte->codeSnips());
+        }
     }
     break;
 
@@ -964,8 +967,6 @@ bool TypeSystemParser::endElement(StackElement element)
         break;
 
     case StackElement::EnumTypeEntry:
-        top->entry->setDocModification(top->docModifications);
-        top->docModifications = DocModificationList();
         m_currentEnum = nullptr;
         break;
     case StackElement::Template:
@@ -1066,8 +1067,11 @@ CodeSnipAbstract *TypeSystemParser::injectCodeTarget(qsizetype offset) const
         return &funcMod.snips().last();
     }
     case ParserState::TypeEntryCodeInjection:
+        Q_ASSERT(top->entry->isComplex());
+        return &static_cast<ComplexTypeEntry *>(top->entry)->codeSnips().last();
     case ParserState::TypeSystemCodeInjection:
-        return &top->entry->codeSnips().last();
+        Q_ASSERT(top->entry->isTypeSystem());
+        return &static_cast<TypeSystemTypeEntry *>(top->entry)->codeSnips().last();
     case ParserState::Template:
         return m_templateEntry;
     default:
@@ -2917,14 +2921,21 @@ bool TypeSystemParser::parseInjectCode(const ConditionalStreamReader &,
     snip.position = position;
     snip.language = lang;
 
-    if (topElement == StackElement::ModifyFunction
-        || topElement == StackElement::AddFunction) {
+    switch (topElement) {
+    case StackElement::ModifyFunction:
+    case StackElement::AddFunction: {
         FunctionModification &mod = m_contextStack.top()->functionMods.last();
         mod.appendSnip(snip);
         if (!snip.code().isEmpty())
             mod.setModifierFlag(FunctionModification::CodeInjection);
-    } else {
-        m_contextStack.top()->entry->addCodeSnip(snip);
+    }
+    break;
+    case StackElement::Root:
+        static_cast<TypeSystemTypeEntry *>(m_contextStack.top()->entry)->addCodeSnip(snip);
+        break;
+    default:
+        static_cast<ComplexTypeEntry *>(m_contextStack.top()->entry)->addCodeSnip(snip);
+        break;
     }
     return true;
 }
diff --git a/sources/shiboken6/ApiExtractor/typesystemtypeentry.h b/sources/shiboken6/ApiExtractor/typesystemtypeentry.h
index 87d1101d3..25efec9cb 100644
--- a/sources/shiboken6/ApiExtractor/typesystemtypeentry.h
+++ b/sources/shiboken6/ApiExtractor/typesystemtypeentry.h
@@ -17,6 +17,10 @@ public:
     TypeSystem::SnakeCase snakeCase() const;
     void setSnakeCase(TypeSystem::SnakeCase sc);
 
+    const CodeSnipList &codeSnips() const;
+    CodeSnipList &codeSnips();
+    void addCodeSnip(const CodeSnip &codeSnip);
+
 protected:
     explicit TypeSystemTypeEntry(TypeEntryPrivate *d);
 };