From 42fd90734b9e528c4852024f77fafe2fe79b4885 Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Tue, 19 May 2026 11:13:05 +0100
Subject: [PATCH 01/10] First changes to metatransformation docs

---
 src/psyclone/docstring_parser.py            | 44 +++++++++++++++++-
 src/psyclone/tests/docstring_parser_test.py | 51 ++++++++++++++++-----
 src/psyclone/tests/utils_test.py            | 13 ++----
 src/psyclone/utils.py                       | 51 ++++++++++++++++++---
 4 files changed, 132 insertions(+), 27 deletions(-)

diff --git a/src/psyclone/docstring_parser.py b/src/psyclone/docstring_parser.py
index 7e886b3ba4..272c65403d 100644
--- a/src/psyclone/docstring_parser.py
+++ b/src/psyclone/docstring_parser.py
@@ -171,6 +171,7 @@ class DocstringData():
     arguments: OrderedDict
     raises: list
     returns: ReturnsData
+    sub_arguments: OrderedDict[str, OrderedDict]
 
     def add_data(self, docstring_element:
                  Union[ArgumentData, RaisesData, ReturnsData, str]) -> None:
@@ -200,6 +201,31 @@ def add_data(self, docstring_element:
                 f"'{docstring_element}'."
             )
 
+    def add_subarguments(self, cls_name: str, other_data: DocstringData,
+                         replace_args: bool = False):
+        '''Add the other DocstringData's arguments as sub arguments to be
+        referenced via cls_name.
+
+        :param cls_name: The class name to use for the sub arguments.
+        :param other_data: The DocstringData object whose arguments to add
+            as sub arguments for this object.
+        :param replace_args: whether to replace duplicate sub arguments if
+            found.
+        '''
+        # If the class isn't already in the sub arguments list, we just
+        # add a copy of the other data's argument dicts to the sub_arguments.
+        if cls_name not in self.sub_arguments:
+            self.sub_arguments[cls_name] = other_data.arguments.copy()
+            return
+        # Otherwise we need to add them manually.
+        for arg in other_data.arguments:
+            # If the arg is already present and we're not overwriting then
+            # skip.
+            if (arg in self.sub_arguments[cls_name].keys() and
+                    not replace_args):
+                continue
+            self.sub_arguments[cls_name][arg] = other_data.arguments[arg]
+
     def merge(self, other_data, replace_desc: bool = False,
               replace_args: bool = False, replace_returns: bool = False):
         '''
@@ -239,6 +265,22 @@ def merge(self, other_data, replace_desc: bool = False,
                 and other_data.returns is not None):
             self.returns = other_data.returns
 
+        # Merge the sub_arguments.
+        for subarg in other_data.sub_arguments:
+            # If the sub argument isn't in our own sub arguments, make a copy
+            # of the sub argument.
+            if subarg not in self.sub_arguments.keys():
+                self.sub_arguments[subarg] = other_data.sub_arguments[subarg].copy()
+            else:
+                # Otherwise we merge the sub arguments in the same fashion we
+                # merge the arguments.
+                for arg in other_data.sub_arguments[subarg]:
+                    if (arg in self.sub_arguments[subarg].keys()
+                            and not replace_args):
+                        continue
+                    self.sub_arguments[subarg][arg] = \
+                        other_data.sub_arguments[subarg][arg]
+
     def gen_docstring(
             self, indentation: str = "    ",
             function: Union[None, Callable[..., Any]] = None
@@ -375,7 +417,7 @@ def create_from_object(cls, obj: Any):
         rtype = None
         docstring_data = DocstringData(
             desc="", arguments=OrderedDict(), raises=[],
-            returns=None
+            returns=None, sub_arguments=OrderedDict()
         )
 
         docstring_data.add_data(desc_chunk)
diff --git a/src/psyclone/tests/docstring_parser_test.py b/src/psyclone/tests/docstring_parser_test.py
index 316b0219ce..646524830a 100644
--- a/src/psyclone/tests/docstring_parser_test.py
+++ b/src/psyclone/tests/docstring_parser_test.py
@@ -76,19 +76,21 @@ def test_docstringdata_base():
     arguments = OrderedDict()
     returns = ReturnsData(desc="desc", datatype="datatype")
     raises = []
+    subargs = OrderedDict()
 
     docdata = DocstringData(desc="desc", arguments=arguments, raises=raises,
-                            returns=returns)
+                            returns=returns, sub_arguments=subargs)
     assert docdata.desc == "desc"
     assert docdata.arguments is arguments
     assert docdata.raises is raises
     assert docdata.returns is returns
+    assert docdata.sub_arguments is subargs
 
 
 def test_docstringdata_add_data():
     'Test the add_data function of the DocstringData dataclass.'''
     docdata = DocstringData(desc=None, arguments=OrderedDict(), raises=[],
-                            returns=None)
+                            returns=None, sub_arguments=OrderedDict())
 
     docdata.add_data("desc")
     assert docdata.desc == "desc"
@@ -124,10 +126,10 @@ def test_docstringdata_add_data():
 def test_docstringdata_merge():
     '''Test the merge function of the DocstringData dataclass.'''
     docdata = DocstringData(desc=None, arguments=OrderedDict(), raises=[],
-                            returns=None)
+                            returns=None, sub_arguments=OrderedDict())
 
     docdata2 = DocstringData(desc=None, arguments=OrderedDict(), raises=[],
-                             returns=None)
+                             returns=None, sub_arguments=OrderedDict())
     adata = ArgumentData(name="name", datatype="datatype",
                          desc="desc", inline_type=True)
     docdata2.add_data(adata)
@@ -142,7 +144,7 @@ def test_docstringdata_merge():
 
     # Check we don't overwrite arguments without replace set.
     docdata3 = DocstringData(desc=None, arguments=OrderedDict(), raises=[],
-                             returns=None)
+                             returns=None, sub_arguments=OrderedDict())
     adata3 = ArgumentData(name="name", datatype="datatype",
                           desc="desc", inline_type=True)
     docdata3.add_data(adata3)
@@ -153,13 +155,13 @@ def test_docstringdata_merge():
 
     # Merge a description.
     docdata4 = DocstringData(desc="desc", arguments=OrderedDict(), raises=[],
-                             returns=None)
+                             returns=None, sub_arguments=OrderedDict())
     docdata.merge(docdata4)
     assert docdata.desc == "desc"
 
     # Don't overwrite without param
     docdata5 = DocstringData(desc="desc2", arguments=OrderedDict(), raises=[],
-                             returns=None)
+                             returns=None, sub_arguments=OrderedDict())
     docdata.merge(docdata5)
     assert docdata.desc == "desc"
     docdata.merge(docdata5, replace_desc=True)
@@ -167,7 +169,7 @@ def test_docstringdata_merge():
 
     # Merge raises
     docdata6 = DocstringData(desc="desc", arguments=OrderedDict(), raises=[],
-                             returns=None)
+                             returns=None, sub_arguments=OrderedDict())
     rdata = RaisesData(desc="desc2", exception="Error")
     docdata6.add_data(rdata)
     docdata.merge(docdata6)
@@ -175,7 +177,7 @@ def test_docstringdata_merge():
 
     # Merge returns
     docdata7 = DocstringData(desc="desc", arguments=OrderedDict(), raises=[],
-                             returns=None)
+                             returns=None, sub_arguments=OrderedDict())
     rdata = ReturnsData(desc="desc", datatype="datatype")
     docdata7.add_data(rdata)
     docdata.merge(docdata7)
@@ -183,7 +185,7 @@ def test_docstringdata_merge():
 
     # Don't overwrite without param
     docdata8 = DocstringData(desc="desc", arguments=OrderedDict(), raises=[],
-                             returns=None)
+                             returns=None, sub_arguments=OrderedDict())
     rdata2 = ReturnsData(desc="desc2", datatype="datatype")
     docdata8.add_data(rdata2)
     docdata.merge(docdata8)
@@ -191,6 +193,33 @@ def test_docstringdata_merge():
     docdata.merge(docdata8, replace_returns=True)
     assert docdata.returns is rdata2
 
+    # Test merging of sub arguments.
+    subargs1 = {"Subclass1": {"opt1": "Test opt1"}}
+    subargs2 = {"Subclass2": {"opt2": "Test opt2"}}
+    docdata = DocstringData(desc="desc", arguments=OrderedDict(), raises=[],
+                            returns=None, sub_arguments=subargs1)
+    docdata2 = DocstringData(desc="desc", arguments=OrderedDict(), raises=[],
+                             returns=None, sub_arguments=subargs2)
+
+    docdata.merge(docdata2)
+    assert docdata.sub_arguments["Subclass2"] == {"opt2": "Test opt2"}
+    # Modify subargs2[Subclass2] to show its not modifying our docdata.
+    subargs2["Subclass2"]["opt3"] = "Test opt3"
+    assert docdata.sub_arguments["Subclass2"] == {"opt2": "Test opt2"}
+
+    subargs3 = {"Subclass1": {"opt1": "Not test opt1",
+                              "opt3": "Test opt 3"}}
+    docdata3 = DocstringData(desc="desc", arguments=OrderedDict(), raises=[],
+                             returns=None, sub_arguments=subargs3)
+    # Check the merging without replace_args doesn't change things.
+    docdata.merge(docdata3)
+    assert docdata.sub_arguments["Subclass1"]["opt1"] == "Test opt1"
+    assert docdata.sub_arguments["Subclass1"]["opt3"] == "Test opt 3"
+
+    # Check that merging with replace_args does change things.
+    docdata.merge(docdata3, replace_args=True)
+    assert docdata.sub_arguments["Subclass1"]["opt1"] == "Not test opt1"
+
 
 def dummy_function(typed_arg: int, untyped_arg):
     '''Dummy function for testing functionality.'''
@@ -316,7 +345,7 @@ def test_DocstringData_gen_docstring_():
 
     # Check we get nothing for an empty DocstringData
     docdata = DocstringData(desc=None, arguments=OrderedDict(), raises=[],
-                            returns=None)
+                            returns=None, sub_arguments=OrderedDict())
     output = docdata.gen_docstring()
     assert output == ""
 
diff --git a/src/psyclone/tests/utils_test.py b/src/psyclone/tests/utils_test.py
index 99ff13f289..d08df02c26 100644
--- a/src/psyclone/tests/utils_test.py
+++ b/src/psyclone/tests/utils_test.py
@@ -99,7 +99,8 @@ def test_transformation_doc_wrapper_non_transformation():
 def test_transformation_doc_wrapper_single_inheritance():
     '''Test the transformation_doc_wrapper.'''
 
-    # Create a base transformation class
+    # Createa base transformation class
+    @transformation_documentation_wrapper(inherit=False)
     class BaseTrans(Transformation):
 
         def validate(self, node, opt1, opt2, **kwargs):
@@ -116,6 +117,7 @@ def apply(self, node, opt1: bool = False, opt2=None, **kwargs):
             :type opt2: opt2 type.
             '''
 
+    @transformation_documentation_wrapper(inherit=False)
     class InheritingTrans(BaseTrans):
 
         def validate(self, node, opt3, **kwargs):
@@ -130,18 +132,11 @@ def apply(self, node, opt3: int = 1, **kwargs):
             :param opt3: opt3 docstring.
             '''
 
-    assert "opt2" not in BaseTrans.validate.__doc__
-
-    transformation_documentation_wrapper(BaseTrans, inherit=False)
-
     assert ":param bool opt1: opt1 docstring." in BaseTrans.validate.__doc__
     assert ":param opt2: opt2 docstring." in BaseTrans.validate.__doc__
     assert ":type opt2: opt2 type." in BaseTrans.validate.__doc__
 
-    assert "opt2" not in InheritingTrans.apply.__doc__
-    assert "opt3" not in InheritingTrans.validate.__doc__
-    transformation_documentation_wrapper(InheritingTrans, inherit=False)
-
+    # Test that the option worked correctly.
     assert (":param int opt3: opt3 docstring." in
             InheritingTrans.validate.__doc__)
 
diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index 9a7393cb02..0ca83ccbe1 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -88,12 +88,31 @@ def stringify_annotation(annotation) -> str:
     return str(annotation)
 
 
-def transformation_documentation_wrapper(cls, *args, inherit=True, **kwargs):
+def transformation_documentation_wrapper(*args, inherit=True,
+                                         add_subtransformations: bool = True,
+                                         **kwargs):
     '''
     Updates the apply and validate methods' docstrings for the supplied cls,
-    according to the value of inherit.
+    according to the value of inherit. args is either a length 1 set argument
+    containing the class to be wrapper, or a length 0 argument set if
+    options are specified on the transformation_docstring_wrapper that is
+    handled by python. This works due to:
+
+    >>> @transformation_documentation_wrapper(inherit=True)
+    >>> class myclass():
+    >>>     pass
+
+    essentially being:
+
+    >>> transformation_documentation_wrapper(inherit=True)(myclass)
+
+    whilst the decorator without any argument is simply:
+
+    >>> transformation_documentation_wrapper(myclass)
+
+    We use this to vary the behaviour of the wrapper slightly depending on
+    whether any arguments are present.
 
-    :param Class cls: The class whose docstrings are to be updated.
     :param inherit: whether to inherit argument docstrings from cls' parent's
                     apply method. If the provided argument is a list, instead
                     the docstrings are updated from each class included in
@@ -102,6 +121,8 @@ def transformation_documentation_wrapper(cls, *args, inherit=True, **kwargs):
                     Transformation's validate docstring from its own
                     apply docstring.
     :type inherit: Union[list[Class], bool]
+    :param add_subtransformations: Whether to add parameter docstrings from
+        sub transformations used by this Transformation.
     '''
     # List of argument doctrings to never inherit.
     _uninheritable_args = ["options"]
@@ -125,7 +146,7 @@ def update_func_docstring(func, added_parameters: DocstringData) -> None:
         func_data.merge(added_parameters, replace_args=False)
         func.__doc__ = func_data.gen_docstring(function=func)
 
-    def wrapper():
+    def wrapper(cls):
         '''
         The wrapping function of the decorator.
 
@@ -141,7 +162,8 @@ def wrapper():
         if isinstance(inherit, list):
             added_parameters = DocstringData(
                 desc=None, arguments=OrderedDict(),
-                raises=[], returns=None)
+                raises=[], returns=None,
+                sub_arguments=OrderedDict())
             for superclass in inherit:
                 inherited_params = \
                     DocstringData.create_from_object(superclass.apply)
@@ -153,16 +175,33 @@ def wrapper():
                  )
         else:
             added_parameters = None
+        if add_subtransformations and len(cls._SUB_TRANSFORMATIONS > 0):
+            if added_parameters = None:
+                added_parameters = DocstringData(
+                    desc=None, arguments=OrderedDict(),
+                    raises=[], returns=None,
+                    sub_arguments=OrderedDict())
+            for trans in cls.SUB_TRANSFORMATION:
+                inherited_params = \
+                    DocstringData.create_from_object(trans.apply)
+                added_parameters.add_subarguments(trans.__name__,
+                                                  inherited_params)
+
+
         if added_parameters is not None:
             # Remove any arguments we don't want to inherit.
             for arg in list(added_parameters.arguments.keys()):
                 if arg in _uninheritable_args:
                     del added_parameters.arguments[arg]
             update_func_docstring(cls.apply, added_parameters)
+
         # Update the validate docstring
         added_parameters = DocstringData.create_from_object(cls.apply)
         if added_parameters is not None:
             update_func_docstring(cls.validate, added_parameters)
 
         return cls
-    return wrapper(*args, **kwargs)
+    if len(args) > 0:
+        return wrapper(*args)
+    else:
+        return wrapper

From c8f464cc816884e413e3b0aec991b3a6c5515ece Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Wed, 20 May 2026 09:56:05 +0100
Subject: [PATCH 02/10] Updates

---
 src/psyclone/docstring_parser.py | 49 ++++++++++++++++++++++++++------
 src/psyclone/utils.py            | 10 ++++---
 2 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/src/psyclone/docstring_parser.py b/src/psyclone/docstring_parser.py
index 272c65403d..041c9e01d0 100644
--- a/src/psyclone/docstring_parser.py
+++ b/src/psyclone/docstring_parser.py
@@ -59,6 +59,8 @@
 Marcin Kurczewski - https://github.com/rr-/docstring_parser.
 '''
 
+from __future__ import annotations
+
 from collections import OrderedDict
 from dataclasses import dataclass
 import inspect
@@ -85,16 +87,22 @@ class ArgumentData():
     desc: str
     inline_type: bool
 
-    def gen_docstring(self, function: Union[None, Callable[..., Any]] = None)\
-            -> str:
+    def gen_docstring(self, function: Union[None, Callable[..., Any]] = None,
+                      cls_name: str = "") -> str:
         '''
         :param function: The function who the generated docstring will be for.
                     Default option is None. If no function is supplied, there
                     can be no type annotation and so the type information is
                     included inline in the :param or in a separate :type
                     entry.
+        :param cls_name: The cls_name to use for saying which subclass (or
+            otherwise) this argument is from. Used for documenting
+            sub transformation options.
+
         :returns: The docstring represented by this ArgumentData.
         '''
+        if cls_name:
+            cls_name = f"(Option used for {cls_name}) "
         rstr = ":param "
         if function:
             # If the argument is in function's parameter list and has a type
@@ -103,13 +111,13 @@ def gen_docstring(self, function: Union[None, Callable[..., Any]] = None)\
             val = signature.parameters.get(self.name)
             if (val is not None and val.annotation is not
                     inspect.Parameter.empty):
-                rstr += f"{self.name}: {self.desc}"
+                rstr += f"{self.name}: {cls_name}{self.desc}"
                 return rstr
 
         if self.inline_type:
-            rstr += f"{self.datatype} {self.name}: {self.desc}"
+            rstr += f"{self.datatype} {self.name}: {cls_name}{self.desc}"
         else:
-            rstr += f"{self.name}: {self.desc}{os.linesep}"
+            rstr += f"{self.name}: {cls_name}{self.desc}{os.linesep}"
             rstr += f":type {self.name}: {self.datatype}"
 
         return rstr
@@ -201,7 +209,7 @@ def add_data(self, docstring_element:
                 f"'{docstring_element}'."
             )
 
-    def add_subarguments(self, cls_name: str, other_data: DocstringData,
+    def add_subarguments(self, cls_name: str, other_data: "DocstringData",
                          replace_args: bool = False):
         '''Add the other DocstringData's arguments as sub arguments to be
         referenced via cls_name.
@@ -216,9 +224,14 @@ def add_subarguments(self, cls_name: str, other_data: DocstringData,
         # add a copy of the other data's argument dicts to the sub_arguments.
         if cls_name not in self.sub_arguments:
             self.sub_arguments[cls_name] = other_data.arguments.copy()
+            # Remove the first argument, which is always the node or nodes
+            # access.
+            first_arg = list(self.sub_arguments[cls_name].keys())[0]
+            self.sub_arguments[cls_name].pop(first_arg, None)
+            print(self.sub_arguments[cls_name])
             return
         # Otherwise we need to add them manually.
-        for arg in other_data.arguments:
+        for arg in other_data.arguments[1:]:
             # If the arg is already present and we're not overwriting then
             # skip.
             if (arg in self.sub_arguments[cls_name].keys() and
@@ -226,7 +239,7 @@ def add_subarguments(self, cls_name: str, other_data: DocstringData,
                 continue
             self.sub_arguments[cls_name][arg] = other_data.arguments[arg]
 
-    def merge(self, other_data, replace_desc: bool = False,
+    def merge(self, other_data: "DocstringData", replace_desc: bool = False,
               replace_args: bool = False, replace_returns: bool = False):
         '''
         Merges the other_data DocstringData object into this one.
@@ -237,7 +250,6 @@ def merge(self, other_data, replace_desc: bool = False,
         always.
 
         :param other_data: the DocstringData object to merge into this object.
-        :type other_data: :py:class:`psyclone.docstring_parser.DocstringData`
         :param replace_desc: whether to replace the desc with that of
                              other_data.
         :param replace_args: whether to replace duplicate arguments with that
@@ -327,6 +339,24 @@ def gen_docstring(
             else:
                 argstring = indentation + argstring
             argstrings.append(argstring)
+        for cls in self.sub_arguments:
+            for arg in self.sub_arguments[cls]:
+                argstring = self.sub_arguments[cls][arg].gen_docstring(
+                    cls_name=cls
+                )
+                if os.linesep in argstring:
+                    lines = argstring.split(os.linesep)
+                    argstring = indentation + lines[0] + os.linesep
+                    for line in lines[1:]:
+                        if ":type" not in line:
+                            argstring += indentation*2 + line + os.linesep
+                        else:
+                            argstring += indentation + line + os.linesep
+                    # Remove the last newline character
+                    argstring = argstring[:-1]
+                else:
+                    argstring = indentation + argstring
+                argstrings.append(argstring)
 
         raisestrings = []
         for element in self.raises:
@@ -354,6 +384,7 @@ def gen_docstring(
         else:
             returnstring = ""
 
+
         docstring = description
         docstring += os.linesep.join(argstrings)
         if len(argstrings) > 0:
diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index 0ca83ccbe1..3f64a5e00a 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -37,6 +37,7 @@
 
 '''This module provides generic utility functions.'''
 
+
 from collections import OrderedDict
 import sys
 from psyclone.errors import InternalError
@@ -127,7 +128,7 @@ def transformation_documentation_wrapper(*args, inherit=True,
     # List of argument doctrings to never inherit.
     _uninheritable_args = ["options"]
 
-    def update_func_docstring(func, added_parameters: DocstringData) -> None:
+    def update_func_docstring(func, added_parameters: "DocstringData") -> None:
         '''
         Adds the docstrings specified in added_parameters to the
         docstring of func.
@@ -175,17 +176,18 @@ def wrapper(cls):
                  )
         else:
             added_parameters = None
-        if add_subtransformations and len(cls._SUB_TRANSFORMATIONS > 0):
-            if added_parameters = None:
+        if add_subtransformations and len(cls._SUB_TRANSFORMATIONS) > 0:
+            if added_parameters is None:
                 added_parameters = DocstringData(
                     desc=None, arguments=OrderedDict(),
                     raises=[], returns=None,
                     sub_arguments=OrderedDict())
-            for trans in cls.SUB_TRANSFORMATION:
+            for trans in cls._SUB_TRANSFORMATIONS:
                 inherited_params = \
                     DocstringData.create_from_object(trans.apply)
                 added_parameters.add_subarguments(trans.__name__,
                                                   inherited_params)
+                print("?")
 
 
         if added_parameters is not None:

From 5c2b4e1fdb4b1be406900d138e4f4ff823191afd Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Wed, 20 May 2026 13:26:28 +0100
Subject: [PATCH 03/10] Initial draft for PR

---
 src/psyclone/docstring_parser.py | 1 -
 src/psyclone/utils.py            | 7 +++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/psyclone/docstring_parser.py b/src/psyclone/docstring_parser.py
index 041c9e01d0..c725a0d446 100644
--- a/src/psyclone/docstring_parser.py
+++ b/src/psyclone/docstring_parser.py
@@ -228,7 +228,6 @@ def add_subarguments(self, cls_name: str, other_data: "DocstringData",
             # access.
             first_arg = list(self.sub_arguments[cls_name].keys())[0]
             self.sub_arguments[cls_name].pop(first_arg, None)
-            print(self.sub_arguments[cls_name])
             return
         # Otherwise we need to add them manually.
         for arg in other_data.arguments[1:]:
diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index 3f64a5e00a..9153f0d06f 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -126,7 +126,7 @@ def transformation_documentation_wrapper(*args, inherit=True,
         sub transformations used by this Transformation.
     '''
     # List of argument doctrings to never inherit.
-    _uninheritable_args = ["options"]
+    _uninheritable_args = ["options", "nodes", "node_list", "node"]
 
     def update_func_docstring(func, added_parameters: "DocstringData") -> None:
         '''
@@ -187,7 +187,6 @@ def wrapper(cls):
                     DocstringData.create_from_object(trans.apply)
                 added_parameters.add_subarguments(trans.__name__,
                                                   inherited_params)
-                print("?")
 
 
         if added_parameters is not None:
@@ -195,6 +194,10 @@ def wrapper(cls):
             for arg in list(added_parameters.arguments.keys()):
                 if arg in _uninheritable_args:
                     del added_parameters.arguments[arg]
+            for trans in cls._SUB_TRANSFORMATIONS:
+                for arg in added_parameters.sub_arguments[trans.__name__]:
+                    if arg in _uninheritable_args:
+                        del added_parameters.sub_arguments[trans.__name__][arg]
             update_func_docstring(cls.apply, added_parameters)
 
         # Update the validate docstring

From 5b174a39ad01d029fdf05432b85ceb0a4b84a717 Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Wed, 20 May 2026 13:27:09 +0100
Subject: [PATCH 04/10] Linting

---
 src/psyclone/docstring_parser.py | 4 ++--
 src/psyclone/utils.py            | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/psyclone/docstring_parser.py b/src/psyclone/docstring_parser.py
index c725a0d446..b32ac7c625 100644
--- a/src/psyclone/docstring_parser.py
+++ b/src/psyclone/docstring_parser.py
@@ -281,7 +281,8 @@ def merge(self, other_data: "DocstringData", replace_desc: bool = False,
             # If the sub argument isn't in our own sub arguments, make a copy
             # of the sub argument.
             if subarg not in self.sub_arguments.keys():
-                self.sub_arguments[subarg] = other_data.sub_arguments[subarg].copy()
+                self.sub_arguments[subarg] = \
+                    other_data.sub_arguments[subarg].copy()
             else:
                 # Otherwise we merge the sub arguments in the same fashion we
                 # merge the arguments.
@@ -383,7 +384,6 @@ def gen_docstring(
         else:
             returnstring = ""
 
-
         docstring = description
         docstring += os.linesep.join(argstrings)
         if len(argstrings) > 0:
diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index 9153f0d06f..e100efa2dd 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -188,7 +188,6 @@ def wrapper(cls):
                 added_parameters.add_subarguments(trans.__name__,
                                                   inherited_params)
 
-
         if added_parameters is not None:
             # Remove any arguments we don't want to inherit.
             for arg in list(added_parameters.arguments.keys()):

From 992bb15f0a9c85304f853ac65bc36fbb6d116cdb Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Thu, 21 May 2026 15:58:34 +0100
Subject: [PATCH 05/10] Test new features

---
 src/psyclone/docstring_parser.py            |  6 +-
 src/psyclone/tests/docstring_parser_test.py | 65 +++++++++++++++
 src/psyclone/tests/utils_test.py            | 92 +++++++++++++++++++++
 src/psyclone/utils.py                       | 14 +++-
 4 files changed, 168 insertions(+), 9 deletions(-)

diff --git a/src/psyclone/docstring_parser.py b/src/psyclone/docstring_parser.py
index b32ac7c625..786e0d7762 100644
--- a/src/psyclone/docstring_parser.py
+++ b/src/psyclone/docstring_parser.py
@@ -224,13 +224,9 @@ def add_subarguments(self, cls_name: str, other_data: "DocstringData",
         # add a copy of the other data's argument dicts to the sub_arguments.
         if cls_name not in self.sub_arguments:
             self.sub_arguments[cls_name] = other_data.arguments.copy()
-            # Remove the first argument, which is always the node or nodes
-            # access.
-            first_arg = list(self.sub_arguments[cls_name].keys())[0]
-            self.sub_arguments[cls_name].pop(first_arg, None)
             return
         # Otherwise we need to add them manually.
-        for arg in other_data.arguments[1:]:
+        for arg in other_data.arguments:
             # If the arg is already present and we're not overwriting then
             # skip.
             if (arg in self.sub_arguments[cls_name].keys() and
diff --git a/src/psyclone/tests/docstring_parser_test.py b/src/psyclone/tests/docstring_parser_test.py
index 646524830a..8ecd520d6f 100644
--- a/src/psyclone/tests/docstring_parser_test.py
+++ b/src/psyclone/tests/docstring_parser_test.py
@@ -623,3 +623,68 @@ def test_function(param: DocstringData):
         assert isinstance(data, ArgumentData)
         assert (data.datatype ==
                 "<class 'psyclone.docstring_parser.DocstringData'>")
+
+
+def test_subarguments():
+    '''
+    Test that we can add sub arguments to a docstring data as expected.
+    '''
+
+    def docstringobj(arg1: int):
+        '''
+        description.
+
+        :param arg1: my param
+        '''
+        pass
+
+    def subobject(arg1: int, arg2: int):
+        '''
+        subobject description
+
+        :param arg1: sub param
+        :param arg2: arg2
+        '''
+        pass
+
+    def subobject2(arg3: int, arg2: int):
+        '''
+        subobject2 description
+
+        :param arg2: subobj2 arg2
+        :param arg3: arg3
+        '''
+
+    doc1 = DocstringData.create_from_object(docstringobj)
+
+    # Should have no sub arguments
+    assert doc1.sub_arguments == {}
+
+    doc2 = DocstringData.create_from_object(subobject)
+    assert "arg1" in doc2.arguments
+    assert "arg2" in doc2.arguments
+    assert len(doc2.arguments) == 2
+
+    # Add doc2 as subarguments to doc1
+    doc1.add_subarguments("subobject", doc2, replace_args=False)
+
+    # Should only have arg2 in the sub arguments
+    assert len(doc1.sub_arguments["subobject"]) == 2
+    assert doc1.sub_arguments["subobject"]["arg1"].desc == "sub param"
+    assert doc1.sub_arguments["subobject"]["arg2"].desc == "arg2"
+
+    doc3 = DocstringData.create_from_object(subobject2)
+    # Add doc3 also as subobject to doc1 without replace_args
+    doc1.add_subarguments("subobject", doc3, replace_args=False)
+    assert len(doc1.sub_arguments["subobject"]) == 3
+    assert doc1.sub_arguments["subobject"]["arg1"].desc == "sub param"
+    assert doc1.sub_arguments["subobject"]["arg2"].desc == "arg2"
+    assert doc1.sub_arguments["subobject"]["arg3"].desc == "arg3"
+
+    # Add doc 3 as a subojbject with replace_args
+    # Add doc3 also as subobject to doc1 without replace_args
+    doc1.add_subarguments("subobject", doc3, replace_args=True)
+    assert len(doc1.sub_arguments["subobject"]) == 3
+    assert doc1.sub_arguments["subobject"]["arg1"].desc == "sub param"
+    assert doc1.sub_arguments["subobject"]["arg2"].desc == "subobj2 arg2"
+    assert doc1.sub_arguments["subobject"]["arg3"].desc == "arg3"
diff --git a/src/psyclone/tests/utils_test.py b/src/psyclone/tests/utils_test.py
index d08df02c26..33526b532a 100644
--- a/src/psyclone/tests/utils_test.py
+++ b/src/psyclone/tests/utils_test.py
@@ -354,3 +354,95 @@ def func(temp: bool, temp2: Union[bool, int]):
             anno = stringify_annotation(v.annotation)
             # Python >= 3.14 uses the second format
             assert "typing.Union[bool, int]" == anno or "bool | int" == anno
+
+
+def test_transformation_doc_wrapper_subtrans():
+    '''Test the transformation doc wrapper works correctly for
+    subtransformations.'''
+
+    class SubTrans1(Transformation):
+
+        def validate(self, node, opt3, **kwargs):
+            '''
+            Sub validate docstring
+            '''
+
+        def apply(self, node, opt3: int = 1, **kwargs):
+            '''
+            Sub apply docstring
+
+            :param opt3: opt3 docstring.
+            '''
+
+    class SubTrans2(Transformation):
+
+        def validate(self, node, opt3, **kwargs):
+            '''
+            Sub validate docstring
+            '''
+
+        def apply(self, node, opt3: int = 1, **kwargs):
+            '''
+            Sub apply docstring
+
+            :param opt3: opt3 docstring.
+            '''
+
+    # Create a base transformation class
+    @transformation_documentation_wrapper(add_subtransformations=False)
+    class BaseTrans(Transformation):
+        _SUB_TRANSFORMATIONS = [SubTrans1, SubTrans2]
+
+        def validate(self, node, **kwargs):
+            '''
+            Super validate docstring
+            '''
+
+        def apply(self, node, opt1: bool = False, opt2=None,
+                  **kwargs):
+            '''
+            Super apply docstring
+
+            :param opt1: opt1 docstring.
+            :param opt2: opt2 docstring.
+            :type opt2: opt2 type.
+            '''
+            pass
+
+    # With add_subtransformations=False we shouldn't get any of the SubTrans
+    # arguments.
+    assert "opt3" not in BaseTrans.apply.__doc__
+
+    # Create a base transformation class
+    @transformation_documentation_wrapper()
+    class BaseTrans(Transformation):
+        _SUB_TRANSFORMATIONS = [SubTrans1, SubTrans2]
+
+        def validate(self, node, **kwargs):
+            '''
+            Super validate docstring
+            '''
+
+        def apply(self, node, opt1: bool = False, opt2=None,
+                  **kwargs):
+            '''
+            Super apply docstring
+
+            :param opt1: opt1 docstring.
+            :param opt2: opt2 docstring.
+            :type opt2: opt2 type.
+            '''
+            pass
+
+    # Disable some flake8 for this string, as empty lines in output
+    # contain whitespace.
+    correct = """Super apply docstring
+    
+    
+    :param opt1: opt1 docstring.
+    :param opt2: opt2 docstring.
+    :type opt2: opt2 type.
+    :param int opt3: (Option used for SubTrans1) opt3 docstring.
+    :param int opt3: (Option used for SubTrans2) opt3 docstring.\
+"""  # noqa: W293
+    assert correct in BaseTrans.apply.__doc__
diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index e100efa2dd..3f94897362 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -193,10 +193,16 @@ def wrapper(cls):
             for arg in list(added_parameters.arguments.keys()):
                 if arg in _uninheritable_args:
                     del added_parameters.arguments[arg]
-            for trans in cls._SUB_TRANSFORMATIONS:
-                for arg in added_parameters.sub_arguments[trans.__name__]:
-                    if arg in _uninheritable_args:
-                        del added_parameters.sub_arguments[trans.__name__][arg]
+            if add_subtransformations:
+                for trans in cls._SUB_TRANSFORMATIONS:
+                    for arg in list(
+                            added_parameters.sub_arguments[
+                                trans.__name__
+                            ].keys()
+                    ):
+                        if arg in _uninheritable_args:
+                            del added_parameters.sub_arguments[
+                                    trans.__name__][arg]
             update_func_docstring(cls.apply, added_parameters)
 
         # Update the validate docstring

From 56b281bae2749c18a3bea728338f4c4ad569fb9e Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Thu, 21 May 2026 16:17:15 +0100
Subject: [PATCH 06/10] Fixes to finalise coverage

---
 src/psyclone/tests/docstring_parser_test.py |  5 +--
 src/psyclone/tests/utils_test.py            | 42 +++++++++++++++++++--
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/src/psyclone/tests/docstring_parser_test.py b/src/psyclone/tests/docstring_parser_test.py
index 8ecd520d6f..5960c12e88 100644
--- a/src/psyclone/tests/docstring_parser_test.py
+++ b/src/psyclone/tests/docstring_parser_test.py
@@ -636,7 +636,6 @@ def docstringobj(arg1: int):
 
         :param arg1: my param
         '''
-        pass
 
     def subobject(arg1: int, arg2: int):
         '''
@@ -645,14 +644,14 @@ def subobject(arg1: int, arg2: int):
         :param arg1: sub param
         :param arg2: arg2
         '''
-        pass
 
-    def subobject2(arg3: int, arg2: int):
+    def subobject2(arg3, arg2: int):
         '''
         subobject2 description
 
         :param arg2: subobj2 arg2
         :param arg3: arg3
+        :type arg3: int
         '''
 
     doc1 = DocstringData.create_from_object(docstringobj)
diff --git a/src/psyclone/tests/utils_test.py b/src/psyclone/tests/utils_test.py
index 33526b532a..c6c339864d 100644
--- a/src/psyclone/tests/utils_test.py
+++ b/src/psyclone/tests/utils_test.py
@@ -367,11 +367,12 @@ def validate(self, node, opt3, **kwargs):
             Sub validate docstring
             '''
 
-        def apply(self, node, opt3: int = 1, **kwargs):
+        def apply(self, node, opt3=1, **kwargs):
             '''
             Sub apply docstring
 
             :param opt3: opt3 docstring.
+            :type opt3: int
             '''
 
     class SubTrans2(Transformation):
@@ -407,7 +408,6 @@ def apply(self, node, opt1: bool = False, opt2=None,
             :param opt2: opt2 docstring.
             :type opt2: opt2 type.
             '''
-            pass
 
     # With add_subtransformations=False we shouldn't get any of the SubTrans
     # arguments.
@@ -432,7 +432,6 @@ def apply(self, node, opt1: bool = False, opt2=None,
             :param opt2: opt2 docstring.
             :type opt2: opt2 type.
             '''
-            pass
 
     # Disable some flake8 for this string, as empty lines in output
     # contain whitespace.
@@ -442,7 +441,42 @@ def apply(self, node, opt1: bool = False, opt2=None,
     :param opt1: opt1 docstring.
     :param opt2: opt2 docstring.
     :type opt2: opt2 type.
-    :param int opt3: (Option used for SubTrans1) opt3 docstring.
+    :param opt3: (Option used for SubTrans1) opt3 docstring.
+    :type opt3: int
+    :param int opt3: (Option used for SubTrans2) opt3 docstring.\
+"""  # noqa: W293
+    assert correct in BaseTrans.apply.__doc__
+
+    # Test behaviour still is consistant with inherit=False
+    @transformation_documentation_wrapper(inherit=False)
+    class BaseTrans(Transformation):
+        _SUB_TRANSFORMATIONS = [SubTrans1, SubTrans2]
+
+        def validate(self, node, **kwargs):
+            '''
+            Super validate docstring
+            '''
+
+        def apply(self, node, opt1: bool = False, opt2=None,
+                  **kwargs):
+            '''
+            Super apply docstring
+
+            :param opt1: opt1 docstring.
+            :param opt2: opt2 docstring.
+            :type opt2: opt2 type.
+            '''
+
+    # Disable some flake8 for this string, as empty lines in output
+    # contain whitespace.
+    correct = """Super apply docstring
+    
+    
+    :param opt1: opt1 docstring.
+    :param opt2: opt2 docstring.
+    :type opt2: opt2 type.
+    :param opt3: (Option used for SubTrans1) opt3 docstring.
+    :type opt3: int
     :param int opt3: (Option used for SubTrans2) opt3 docstring.\
 """  # noqa: W293
     assert correct in BaseTrans.apply.__doc__

From 697423b8984209918f6fc23d6f27766fac27046f Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Mon, 1 Jun 2026 15:11:13 +0100
Subject: [PATCH 07/10] Updates for review

---
 src/psyclone/docstring_parser.py |  2 +-
 src/psyclone/tests/utils_test.py | 11 +++---
 src/psyclone/utils.py            | 57 +++++++++++++++++++++-----------
 3 files changed, 43 insertions(+), 27 deletions(-)

diff --git a/src/psyclone/docstring_parser.py b/src/psyclone/docstring_parser.py
index 786e0d7762..1c356c55fc 100644
--- a/src/psyclone/docstring_parser.py
+++ b/src/psyclone/docstring_parser.py
@@ -102,7 +102,7 @@ def gen_docstring(self, function: Union[None, Callable[..., Any]] = None,
         :returns: The docstring represented by this ArgumentData.
         '''
         if cls_name:
-            cls_name = f"(Option used for {cls_name}) "
+            cls_name = f"(Option provided for {cls_name}) "
         rstr = ":param "
         if function:
             # If the argument is in function's parameter list and has a type
diff --git a/src/psyclone/tests/utils_test.py b/src/psyclone/tests/utils_test.py
index c6c339864d..2e87577ae5 100644
--- a/src/psyclone/tests/utils_test.py
+++ b/src/psyclone/tests/utils_test.py
@@ -99,7 +99,7 @@ def test_transformation_doc_wrapper_non_transformation():
 def test_transformation_doc_wrapper_single_inheritance():
     '''Test the transformation_doc_wrapper.'''
 
-    # Createa base transformation class
+    # Create base transformation class
     @transformation_documentation_wrapper(inherit=False)
     class BaseTrans(Transformation):
 
@@ -204,7 +204,6 @@ def apply(self, node, opt3: int = 1, **kwargs):
         InheritingTrans,
         inherit=[BaseTrans1, BaseTrans2]
     )
-    print(InheritingTrans.apply.__doc__)
     assert "param bool opt1: opt1 docstring." in InheritingTrans.apply.__doc__
     assert "param bool opt2: opt2 docstring." in InheritingTrans.apply.__doc__
     assert ("param bool opt1: opt1 docstring."
@@ -441,9 +440,9 @@ def apply(self, node, opt1: bool = False, opt2=None,
     :param opt1: opt1 docstring.
     :param opt2: opt2 docstring.
     :type opt2: opt2 type.
-    :param opt3: (Option used for SubTrans1) opt3 docstring.
+    :param opt3: (Option provided for SubTrans1) opt3 docstring.
     :type opt3: int
-    :param int opt3: (Option used for SubTrans2) opt3 docstring.\
+    :param int opt3: (Option provided for SubTrans2) opt3 docstring.\
 """  # noqa: W293
     assert correct in BaseTrans.apply.__doc__
 
@@ -475,8 +474,8 @@ def apply(self, node, opt1: bool = False, opt2=None,
     :param opt1: opt1 docstring.
     :param opt2: opt2 docstring.
     :type opt2: opt2 type.
-    :param opt3: (Option used for SubTrans1) opt3 docstring.
+    :param opt3: (Option provided for SubTrans1) opt3 docstring.
     :type opt3: int
-    :param int opt3: (Option used for SubTrans2) opt3 docstring.\
+    :param int opt3: (Option provided for SubTrans2) opt3 docstring.\
 """  # noqa: W293
     assert correct in BaseTrans.apply.__doc__
diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index 3f94897362..bc895c4bb6 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -37,13 +37,14 @@
 
 '''This module provides generic utility functions.'''
 
-
+from typing import Type
 from collections import OrderedDict
 import sys
 from psyclone.errors import InternalError
 from psyclone.docstring_parser import (
-    DocstringData,
+    DocstringData, ReturnsData
 )
+from psyclone.psyGen import Transformation
 
 
 def within_virtual_env():
@@ -89,29 +90,46 @@ def stringify_annotation(annotation) -> str:
     return str(annotation)
 
 
-def transformation_documentation_wrapper(*args, inherit=True,
+def transformation_documentation_wrapper(*args,
+                                         inherit: list[Type[Transformation]]
+                                         | bool = True,
                                          add_subtransformations: bool = True,
                                          **kwargs):
     '''
-    Updates the apply and validate methods' docstrings for the supplied cls,
-    according to the value of inherit. args is either a length 1 set argument
-    containing the class to be wrapper, or a length 0 argument set if
-    options are specified on the transformation_docstring_wrapper that is
-    handled by python. This works due to:
+    Updates the apply and validate methods' docstrings for the decorated
+    class, according to the value of inherit. Should be used as a
+    decorator on a Transformation subclass.
+
+    *args is either a length 1 list of
+    arguments containing the class to be wrapped, or a length 0 argument set
+    when options are specified on the transformation_docstring_wrapper that is
+    handled by python. The length 0 set would happen with a case like this:
 
     >>> @transformation_documentation_wrapper(inherit=True)
-    >>> class myclass():
-    >>>     pass
+    ... class mytrans(Transformation):
+    ...     pass
+
+    as this code is equivalent to:
+
+    >>> mytrans = transformation_documentation_wrapper(inherit=True)(mytrans)
+
+    For this case *args is empty, as only the inherit argument is provided
+    to the transformation_documentation_wrapper call.
+
+    Without arguments to the decorator:
 
-    essentially being:
+    >>> @transformation_documentation_wrapper
+    ... class mytrans(Transformation):
+    ...    pass
 
-    >>> transformation_documentation_wrapper(inherit=True)(myclass)
+    the resultant code is the same as writing:
 
-    whilst the decorator without any argument is simply:
+    >>> mytrans = transformation_documentation_wrapper(mytrans)
 
-    >>> transformation_documentation_wrapper(myclass)
+    In this case *args contains the wrapped class in *args, which needs to be
+    passed manually to the sub-function inside the wrapper.
 
-    We use this to vary the behaviour of the wrapper slightly depending on
+    This allows us to vary the behaviour of the wrapper slightly depending on
     whether any arguments are present.
 
     :param inherit: whether to inherit argument docstrings from cls' parent's
@@ -121,7 +139,6 @@ def transformation_documentation_wrapper(*args, inherit=True,
                     inherit is False, the wrapper will just update the
                     Transformation's validate docstring from its own
                     apply docstring.
-    :type inherit: Union[list[Class], bool]
     :param add_subtransformations: Whether to add parameter docstrings from
         sub transformations used by this Transformation.
     '''
@@ -162,8 +179,8 @@ def wrapper(cls):
             )
         if isinstance(inherit, list):
             added_parameters = DocstringData(
-                desc=None, arguments=OrderedDict(),
-                raises=[], returns=None,
+                desc="", arguments=OrderedDict(),
+                raises=[], returns=ReturnsData("", None),
                 sub_arguments=OrderedDict())
             for superclass in inherit:
                 inherited_params = \
@@ -179,8 +196,8 @@ def wrapper(cls):
         if add_subtransformations and len(cls._SUB_TRANSFORMATIONS) > 0:
             if added_parameters is None:
                 added_parameters = DocstringData(
-                    desc=None, arguments=OrderedDict(),
-                    raises=[], returns=None,
+                    desc="", arguments=OrderedDict(),
+                    raises=[], returns=ReturnsData("", None),
                     sub_arguments=OrderedDict())
             for trans in cls._SUB_TRANSFORMATIONS:
                 inherited_params = \

From c7590f8fb2f1d9660d53e4e76b899bca35ef4eb1 Mon Sep 17 00:00:00 2001
From: LonelyCat124 <3043914+LonelyCat124@users.noreply.github.com.>
Date: Mon, 1 Jun 2026 15:20:59 +0100
Subject: [PATCH 08/10] Fix 3.9 typing issue

---
 src/psyclone/utils.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index bc895c4bb6..432ba2f54c 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -37,7 +37,7 @@
 
 '''This module provides generic utility functions.'''
 
-from typing import Type
+from typing import Type, Union
 from collections import OrderedDict
 import sys
 from psyclone.errors import InternalError
@@ -91,8 +91,9 @@ def stringify_annotation(annotation) -> str:
 
 
 def transformation_documentation_wrapper(*args,
-                                         inherit: list[Type[Transformation]]
-                                         | bool = True,
+                                         inherit:
+                                         Union[list[Type[Transformation]],
+                                               bool] = True,
                                          add_subtransformations: bool = True,
                                          **kwargs):
     '''

From 62d086161e399dde6852771fb26077ba94ea4624 Mon Sep 17 00:00:00 2001
From: Sergi Siso <sergiesg@gmail.com>
Date: Tue, 2 Jun 2026 15:12:15 +0100
Subject: [PATCH 09/10] Update formating and doctrings of kwargs documentation
 code

---
 src/psyclone/docstring_parser.py |  5 ++---
 src/psyclone/tests/utils_test.py | 15 ++++-----------
 src/psyclone/utils.py            | 16 ++++++++--------
 3 files changed, 14 insertions(+), 22 deletions(-)

diff --git a/src/psyclone/docstring_parser.py b/src/psyclone/docstring_parser.py
index 1c356c55fc..c9f7f1267c 100644
--- a/src/psyclone/docstring_parser.py
+++ b/src/psyclone/docstring_parser.py
@@ -95,9 +95,8 @@ def gen_docstring(self, function: Union[None, Callable[..., Any]] = None,
                     can be no type annotation and so the type information is
                     included inline in the :param or in a separate :type
                     entry.
-        :param cls_name: The cls_name to use for saying which subclass (or
-            otherwise) this argument is from. Used for documenting
-            sub transformation options.
+        :param cls_name: The cls_name to use for saying which argument is this
+            from. Used for documenting sub transformation options.
 
         :returns: The docstring represented by this ArgumentData.
         '''
diff --git a/src/psyclone/tests/utils_test.py b/src/psyclone/tests/utils_test.py
index 2e87577ae5..06d2dbb122 100644
--- a/src/psyclone/tests/utils_test.py
+++ b/src/psyclone/tests/utils_test.py
@@ -434,16 +434,14 @@ def apply(self, node, opt1: bool = False, opt2=None,
 
     # Disable some flake8 for this string, as empty lines in output
     # contain whitespace.
-    correct = """Super apply docstring
-    
-    
+    correct = """    Super apply docstring\n    \n    \n\
     :param opt1: opt1 docstring.
     :param opt2: opt2 docstring.
     :type opt2: opt2 type.
     :param opt3: (Option provided for SubTrans1) opt3 docstring.
     :type opt3: int
     :param int opt3: (Option provided for SubTrans2) opt3 docstring.\
-"""  # noqa: W293
+"""
     assert correct in BaseTrans.apply.__doc__
 
     # Test behaviour still is consistant with inherit=False
@@ -466,16 +464,11 @@ def apply(self, node, opt1: bool = False, opt2=None,
             :type opt2: opt2 type.
             '''
 
-    # Disable some flake8 for this string, as empty lines in output
-    # contain whitespace.
-    correct = """Super apply docstring
-    
-    
+    correct = """Super apply docstring\n    \n    \n\
     :param opt1: opt1 docstring.
     :param opt2: opt2 docstring.
     :type opt2: opt2 type.
     :param opt3: (Option provided for SubTrans1) opt3 docstring.
     :type opt3: int
-    :param int opt3: (Option provided for SubTrans2) opt3 docstring.\
-"""  # noqa: W293
+    :param int opt3: (Option provided for SubTrans2) opt3 docstring."""
     assert correct in BaseTrans.apply.__doc__
diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py
index 432ba2f54c..5989a18221 100644
--- a/src/psyclone/utils.py
+++ b/src/psyclone/utils.py
@@ -97,14 +97,14 @@ def transformation_documentation_wrapper(*args,
                                          add_subtransformations: bool = True,
                                          **kwargs):
     '''
-    Updates the apply and validate methods' docstrings for the decorated
-    class, according to the value of inherit. Should be used as a
-    decorator on a Transformation subclass.
+    This function should be used as a decorator on a Transformation subclass.
+    It updates the apply and validate methods' docstrings for the decorated
+    class, according to the value of inherit.
 
-    *args is either a length 1 list of
-    arguments containing the class to be wrapped, or a length 0 argument set
-    when options are specified on the transformation_docstring_wrapper that is
-    handled by python. The length 0 set would happen with a case like this:
+    *args is either a length 1 list of arguments containing the class to be
+    wrapped, or a length 0 argument set when options are specified on the
+    transformation_docstring_wrapper that is handled by python. The length 0
+    set would happen with a case like this:
 
     >>> @transformation_documentation_wrapper(inherit=True)
     ... class mytrans(Transformation):
@@ -115,7 +115,7 @@ def transformation_documentation_wrapper(*args,
     >>> mytrans = transformation_documentation_wrapper(inherit=True)(mytrans)
 
     For this case *args is empty, as only the inherit argument is provided
-    to the transformation_documentation_wrapper call.
+    to the transformation_documentation_wrapper call and is through kwargs.
 
     Without arguments to the decorator:
 

From 275003d3954736c7d6ec61ef9272d1e16e96f7aa Mon Sep 17 00:00:00 2001
From: Sergi Siso <sergiesg@gmail.com>
Date: Tue, 2 Jun 2026 15:15:02 +0100
Subject: [PATCH 10/10] Update changelog

---
 changelog | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/changelog b/changelog
index d1013cf66a..62754b4f5b 100644
--- a/changelog
+++ b/changelog
@@ -1,3 +1,6 @@
+   27) PR #3438 for #3435. Add support for sub_transformations in the kwargs
+   documentation functionality.
+
    26) PR #3443 for #3243. Unifies NEMO parallelisation scripts to reduce the
    amount of code duplication between all the parallelisation scripts, and
    adds a README for the NEMO scripts directory.