rewriter: create {add,rm}_extra_files commands

Add ability to mutate a target's `extra_files` list through the
rewriter.

The logic is copied from sources add/rm, but changes the `extra_files`
kwarg instead of the sources positional argument.
Has additional logic to handle creating the `extra_files` list if it
doesn't exist.

1 files changed, 96 insertions, 3 deletions

diff --git a/mesonbuild/rewriter.py b/mesonbuild/rewriter.py
index 67cd871..d816cce 100644
--- a/mesonbuild/rewriter.py
+++ b/mesonbuild/rewriter.py
@@ -46,7 +46,7 @@ def add_arguments(parser, formatter=None):
     tgt_parser.add_argument('--type', dest='tgt_type', choices=rewriter_keys['target']['target_type'][2], default='executable',
                             help='Type of the target to add (only for the "add_target" action)')
     tgt_parser.add_argument('target', help='Name or ID of the target')
-    tgt_parser.add_argument('operation', choices=['add', 'rm', 'add_target', 'rm_target', 'info'],
+    tgt_parser.add_argument('operation', choices=['add', 'rm', 'add_target', 'rm_target', 'add_extra_files', 'rm_extra_files', 'info'],
                             help='Action to execute')
     tgt_parser.add_argument('sources', nargs='*', help='Sources to add/remove')
 
@@ -308,7 +308,7 @@ rewriter_keys = {
     },
     'target': {
         'target': (str, None, None),
-        'operation': (str, None, ['src_add', 'src_rm', 'target_rm', 'target_add', 'info']),
+        'operation': (str, None, ['src_add', 'src_rm', 'target_rm', 'target_add', 'extra_files_add', 'extra_files_rm', 'info']),
         'sources': (list, [], None),
         'subdir': (str, '', None),
         'target_type': (str, 'executable', ['both_libraries', 'executable', 'jar', 'library', 'shared_library', 'shared_module', 'static_library']),
@@ -709,6 +709,91 @@ class Rewriter:
                 if root not in self.modified_nodes:
                     self.modified_nodes += [root]
 
+        elif cmd['operation'] == 'extra_files_add':
+            tgt_function: FunctionNode = target['node']
+            mark_array = True
+            try:
+                node = target['extra_files'][0]
+            except IndexError:
+                # Specifying `extra_files` with a list that flattens to empty gives an empty
+                # target['extra_files'] list, account for that.
+                try:
+                    extra_files_key = next(k for k in tgt_function.args.kwargs.keys() if isinstance(k, IdNode) and k.value == 'extra_files')
+                    node = tgt_function.args.kwargs[extra_files_key]
+                except StopIteration:
+                    # Target has no extra_files kwarg, create one
+                    node = ArrayNode(ArgumentNode(Token('', tgt_function.filename, 0, 0, 0, None, '[]')), tgt_function.end_lineno, tgt_function.end_colno, tgt_function.end_lineno, tgt_function.end_colno)
+                    tgt_function.args.kwargs[IdNode(Token('string', tgt_function.filename, 0, 0, 0, None, 'extra_files'))] = node
+                    mark_array = False
+                    if tgt_function not in self.modified_nodes:
+                        self.modified_nodes += [tgt_function]
+                target['extra_files'] = [node]
+            if isinstance(node, IdNode):
+                node = self.interpreter.assignments[node.value]
+                target['extra_files'] = [node]
+            if not isinstance(node, ArrayNode):
+                mlog.error('Target', mlog.bold(cmd['target']), 'extra_files argument must be a list', *self.on_error())
+                return self.handle_error()
+
+            # Generate the current extra files list
+            extra_files_list = []
+            for i in target['extra_files']:
+                for j in arg_list_from_node(i):
+                    if isinstance(j, StringNode):
+                        extra_files_list += [j.value]
+
+            # Generate the new String nodes
+            to_append = []
+            for i in sorted(set(cmd['sources'])):
+                if i in extra_files_list:
+                    mlog.log('  -- Extra file', mlog.green(i), 'is already defined for the target --> skipping')
+                    continue
+                mlog.log('  -- Adding extra file', mlog.green(i), 'at',
+                         mlog.yellow(f'{node.filename}:{node.lineno}'))
+                token = Token('string', node.filename, 0, 0, 0, None, i)
+                to_append += [StringNode(token)]
+
+            # Append to the AST at the right place
+            arg_node = node.args
+            arg_node.arguments += to_append
+
+            # Mark the node as modified
+            if arg_node not in to_sort_nodes:
+                to_sort_nodes += [arg_node]
+            # If the extra_files array is newly created, don't mark it as its parent function node already is,
+            # otherwise this would cause double modification.
+            if mark_array and node not in self.modified_nodes:
+                self.modified_nodes += [node]
+
+        elif cmd['operation'] == 'extra_files_rm':
+            # Helper to find the exact string node and its parent
+            def find_node(src):
+                for i in target['extra_files']:
+                    for j in arg_list_from_node(i):
+                        if isinstance(j, StringNode):
+                            if j.value == src:
+                                return i, j
+                return None, None
+
+            for i in cmd['sources']:
+                # Try to find the node with the source string
+                root, string_node = find_node(i)
+                if root is None:
+                    mlog.warning('  -- Unable to find extra file', mlog.green(i), 'in the target')
+                    continue
+
+                # Remove the found string node from the argument list
+                arg_node = root.args
+                mlog.log('  -- Removing extra file', mlog.green(i), 'from',
+                         mlog.yellow(f'{string_node.filename}:{string_node.lineno}'))
+                arg_node.arguments.remove(string_node)
+
+                # Mark the node as modified
+                if arg_node not in to_sort_nodes and not isinstance(root, FunctionNode):
+                    to_sort_nodes += [arg_node]
+                if root not in self.modified_nodes:
+                    self.modified_nodes += [root]
+
         elif cmd['operation'] == 'target_add':
             if target is not None:
                 mlog.error('Can not add target', mlog.bold(cmd['target']), 'because it already exists', *self.on_error())
@@ -756,9 +841,15 @@ class Rewriter:
                 for j in arg_list_from_node(i):
                     if isinstance(j, StringNode):
                         src_list += [j.value]
+            extra_files_list = []
+            for i in target['extra_files']:
+                for j in arg_list_from_node(i):
+                    if isinstance(j, StringNode):
+                        extra_files_list += [j.value]
             test_data = {
                 'name': target['name'],
-                'sources': src_list
+                'sources': src_list,
+                'extra_files': extra_files_list
             }
             self.add_info('target', target['id'], test_data)
 
@@ -879,6 +970,8 @@ target_operation_map = {
     'rm': 'src_rm',
     'add_target': 'target_add',
     'rm_target': 'target_rm',
+    'add_extra_files': 'extra_files_add',
+    'rm_extra_files': 'extra_files_rm',
     'info': 'info',
 }