pw_build: Support more expressions

- Add <TARGET_FILE_IF_EXISTS(target)>, which evalutes to a target's
  output file only if the file exists.
- Add <TARGET_OBJECTS(target)>, which evaluates to zero or more separate
  arguments with the object files from the target.
- Expand python_runner_test.py to cover expression expansion.

Change-Id: Ibfd07274b8b325066e5bfd0096dd4a3d2b366e93
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/15229
Commit-Queue: Wyatt Hepler <hepler@google.com>
Reviewed-by: Alexei Frolov <frolv@google.com>
diff --git a/pw_build/docs.rst b/pw_build/docs.rst
index c16eab1..16e9d9f 100644
--- a/pw_build/docs.rst
+++ b/pw_build/docs.rst
@@ -116,7 +116,7 @@
 BUILD.gn files. This allows build code to use GN labels without having to worry
 about converting them to files.
 
-Currently, only one expression is supported.
+The following expressions are supported:
 
 .. describe:: <TARGET_FILE(gn_target)>
 
@@ -145,6 +145,50 @@
   converts the provided GN path or list of paths to be relative to the build
   directory, from which all build commands and scripts are executed.
 
+.. describe:: <TARGET_FILE_IF_EXISTS(gn_target)>
+
+  ``TARGET_FILE_IF_EXISTS`` evaluates to the output file of the provided GN
+  target, if the output file exists. If the output file does not exist, the
+  entire argument that includes this expression is omitted, even if there is
+  other text or another expression.
+
+  For example, consider this expression:
+
+  .. code::
+
+    "--database=<TARGET_FILE_IF_EXISTS(//alpha/bravo)>"
+
+  If the ``//alpha/bravo`` target file exists, this might expand to the
+  following:
+
+  .. code::
+
+    "--database=/home/User/project/out/obj/alpha/bravo/bravo.elf"
+
+  If the ``//alpha/bravo`` target file does not exist, the entire
+  ``--database=`` argument is omitted from the script arguments.
+
+.. describe:: <TARGET_OBJECTS(gn_target)>
+
+  Evaluates to the object files of the provided GN target. Expands to a separate
+  argument for each object file. If the target has no object files, the argument
+  is omitted entirely. Because it does not expand to a single expression, the
+  ``<TARGET_OBJECTS(...)>`` expression may not have leading or trailing text.
+
+  For example, the expression
+
+  .. code::
+
+    "<TARGET_OBJECTS(//foo/bar:a_source_set)>"
+
+  might expand to multiple separate arguments:
+
+  .. code::
+
+    "/home/User/project_root/out/obj/foo/bar/a_source_set.file_a.cc.o"
+    "/home/User/project_root/out/obj/foo/bar/a_source_set.file_b.cc.o"
+    "/home/User/project_root/out/obj/foo/bar/a_source_set.file_c.cc.o"
+
 **Example**
 
 .. code::
diff --git a/pw_build/py/pw_build/python_runner.py b/pw_build/py/pw_build/python_runner.py
index 3679041..81bd4ec 100755
--- a/pw_build/py/pw_build/python_runner.py
+++ b/pw_build/py/pw_build/python_runner.py
@@ -14,22 +14,20 @@
 """Script that preprocesses a Python command then runs it.
 
 This script evaluates expressions in the Python command's arguments then invokes
-the command. Only one expression is supported currently:
-
-  <TARGET_FILE(gn_target)> -- gets the target output file (e.g. .elf, .a,, .so)
-      for a GN target; raises an error for targets with no output file, such as
-      a source_set or group
+the command.
 """
 
 import argparse
 from dataclasses import dataclass
+import enum
 import logging
 from pathlib import Path
 import re
 import shlex
 import subprocess
 import sys
-from typing import Callable, Dict, Iterator, List, NamedTuple, Optional, Tuple
+from typing import Callable, Dict, Iterable, Iterator, List, NamedTuple
+from typing import Optional, Tuple
 
 _LOG = logging.getLogger(__name__)
 
@@ -269,43 +267,110 @@
     """An error occurred while parsing an expression."""
 
 
-def _target_output_file(paths: GnPaths, target_name: str) -> str:
-    target = TargetInfo(paths, target_name)
+class _ArgAction(enum.Enum):
+    APPEND = 0
+    OMIT = 1
+    EMIT_NEW = 2
 
-    if not target.artifact:
+
+class _Expression:
+    def __init__(self, match: re.Match, ending: int):
+        self._match = match
+        self._ending = ending
+
+    @property
+    def string(self):
+        return self._match.string
+
+    @property
+    def end(self) -> int:
+        return self._ending + len(_ENDING)
+
+    def contents(self) -> str:
+        return self.string[self._match.end():self._ending]
+
+    def expression(self) -> str:
+        return self.string[self._match.start():self.end]
+
+
+_Actions = Iterator[Tuple[_ArgAction, str]]
+
+
+def _target_file(paths: GnPaths, expr: _Expression) -> _Actions:
+    target = TargetInfo(paths, expr.contents())
+
+    if target.artifact is None:
         raise ExpressionError(f'Target {target} has no output file!')
 
-    return str(target.artifact)
+    yield _ArgAction.APPEND, str(target.artifact)
 
 
-_FUNCTIONS: Dict['str', Callable[[GnPaths, str], str]] = {
-    'TARGET_FILE': _target_output_file,
+def _target_file_if_exists(paths: GnPaths, expr: _Expression) -> _Actions:
+    (_, file), = _target_file(paths, expr)
+    if Path(file).exists():
+        yield _ArgAction.APPEND, file
+    else:
+        yield _ArgAction.OMIT, ''
+
+
+def _target_objects(paths: GnPaths, expr: _Expression) -> _Actions:
+    if expr.expression() != expr.string:
+        raise ExpressionError(
+            f'The expression "{expr.expression()}" in "{expr.string}" may '
+            'expand to multiple arguments, so it cannot be used alongside '
+            'other text or expressions')
+
+    for obj in TargetInfo(paths, expr.contents()).object_files:
+        yield _ArgAction.EMIT_NEW, str(obj)
+
+
+_FUNCTIONS: Dict['str', Callable[[GnPaths, _Expression], _Actions]] = {
+    'TARGET_FILE': _target_file,
+    'TARGET_FILE_IF_EXISTS': _target_file_if_exists,
+    'TARGET_OBJECTS': _target_objects,
 }
 
 _START_EXPRESSION = re.compile(fr'<({"|".join(_FUNCTIONS)})\(')
+_ENDING = ')>'
 
 
-def _expand_expressions(paths: GnPaths, string: str) -> Iterator[str]:
-    pos = None
+def _expand_arguments(paths: GnPaths, string: str) -> _Actions:
+    pos = 0
 
     for match in _START_EXPRESSION.finditer(string):
-        yield string[pos:match.start()]
+        if pos != match.start():
+            yield _ArgAction.APPEND, string[pos:match.start()]
 
-        pos = string.find(')>', match.end())
-        if pos == -1:
-            raise ExpressionError('Parse error: no terminating ")>" '
+        ending = string.find(_ENDING, match.end())
+        if ending == -1:
+            raise ExpressionError(f'Parse error: no terminating "{_ENDING}" '
                                   f'was found for "{string[match.start():]}"')
 
-        yield _FUNCTIONS[match.group(1)](paths, string[match.end():pos])
+        expression = _Expression(match, ending)
+        yield from _FUNCTIONS[match.group(1)](paths, expression)
 
-        pos += 2  # skip the terminating ')>'
+        pos = expression.end
 
-    yield string[pos:]
+    if pos < len(string):
+        yield _ArgAction.APPEND, string[pos:]
 
 
-def expand_expressions(paths: GnPaths, arg: str) -> str:
-    """Expands <FUNCTION(...)> expressions."""
-    return ''.join(_expand_expressions(paths, arg))
+def expand_expressions(paths: GnPaths, arg: str) -> Iterable[str]:
+    """Expands <FUNCTION(...)> expressions; yields zero or more arguments."""
+    if arg == '':
+        return ['']
+
+    expanded_args: List[List[str]] = [[]]
+
+    for action, piece in _expand_arguments(paths, arg):
+        if action is _ArgAction.OMIT:
+            return []
+
+        expanded_args[-1].append(piece)
+        if action is _ArgAction.EMIT_NEW:
+            expanded_args.append([])
+
+    return (''.join(arg) for arg in expanded_args if arg)
 
 
 def main(
@@ -334,7 +399,8 @@
 
     command = [sys.executable]
     try:
-        command += (expand_expressions(paths, arg) for arg in original_cmd[1:])
+        for arg in original_cmd[1:]:
+            command += expand_expressions(paths, arg)
     except ExpressionError as err:
         _LOG.error('%s: %s', sys.argv[0], err)
         return 1
diff --git a/pw_build/py/python_runner_test.py b/pw_build/py/python_runner_test.py
index bf5f297..ecf78d3 100755
--- a/pw_build/py/python_runner_test.py
+++ b/pw_build/py/python_runner_test.py
@@ -19,7 +19,8 @@
 import tempfile
 import unittest
 
-from pw_build.python_runner import GnPaths, Label, TargetInfo
+from pw_build.python_runner import ExpressionError, GnPaths, Label, TargetInfo
+from pw_build.python_runner import expand_expressions
 
 TEST_PATHS = GnPaths(Path('/gn_root'), Path('/gn_root/out'),
                      Path('/gn_root/some/cwd'), '//toolchains/cool:ToolChain')
@@ -152,23 +153,29 @@
 '''
 
 
+def _create_ninja_files():
+    tempdir = tempfile.TemporaryDirectory(prefix='pw_build_test_')
+
+    module = Path(tempdir.name, 'out', 'fake_toolchain', 'obj', 'fake_module')
+    os.makedirs(module)
+    module.joinpath('fake_test.ninja').write_text(NINJA_EXECUTABLE)
+    module.joinpath('fake_source_set.ninja').write_text(NINJA_SOURCE_SET)
+    module.joinpath('fake_no_objects.ninja').write_text('\n')
+
+    outdir = Path(tempdir.name, 'out', 'fake_toolchain', 'obj', 'fake_module')
+
+    paths = GnPaths(root=Path(tempdir.name),
+                    build=Path(tempdir.name, 'out'),
+                    cwd=Path(tempdir.name, 'some', 'module'),
+                    toolchain='//tools:fake_toolchain')
+
+    return tempdir, outdir, paths
+
+
 class TargetTest(unittest.TestCase):
     """Tests querying GN target information."""
     def setUp(self):
-        self._tempdir = tempfile.TemporaryDirectory(prefix='pw_build_test')
-        self._paths = GnPaths(root=Path(self._tempdir.name),
-                              build=Path(self._tempdir.name, 'out'),
-                              cwd=Path(self._tempdir.name, 'some', 'module'),
-                              toolchain='//tools:fake_toolchain')
-
-        module = Path(self._tempdir.name, 'out', 'fake_toolchain', 'obj',
-                      'fake_module')
-        os.makedirs(module)
-        module.joinpath('fake_test.ninja').write_text(NINJA_EXECUTABLE)
-        module.joinpath('fake_source_set.ninja').write_text(NINJA_SOURCE_SET)
-
-        self._outdir = Path(self._tempdir.name, 'out', 'fake_toolchain', 'obj',
-                            'fake_module')
+        self._tempdir, self._outdir, self._paths = _create_ninja_files()
 
     def tearDown(self):
         self._tempdir.cleanup()
@@ -199,5 +206,118 @@
                          self._outdir / 'test' / 'fake_test.elf')
 
 
+class ExpandExpressionsTest(unittest.TestCase):
+    """Tests expansion of expressions like <TARGET_FILE(//foo)>."""
+    def setUp(self):
+        self._tempdir, self._outdir, self._paths = _create_ninja_files()
+
+    def tearDown(self):
+        self._tempdir.cleanup()
+
+    def _path(self, *segments: str, create: bool = False) -> str:
+        path = Path(self._outdir, *segments)
+        if create:
+            os.makedirs(path.parent)
+            path.touch()
+        else:
+            assert not path.exists()
+        return str(path)
+
+    def test_empty(self):
+        self.assertEqual(list(expand_expressions(self._paths, '')), [''])
+
+    def test_no_expressions(self):
+        self.assertEqual(list(expand_expressions(self._paths, 'foobar')),
+                         ['foobar'])
+        self.assertEqual(
+            list(expand_expressions(self._paths, '<NOT_AN_EXPRESSION()>')),
+            ['<NOT_AN_EXPRESSION()>'])
+
+    def test_incomplete_expression(self):
+        for incomplete_expression in [
+                '<TARGET_FILE(',
+                '<TARGET_FILE(//foo)',
+                '<TARGET_FILE(//foo>',
+                '<TARGET_FILE(//foo) >',
+                '--arg=<TARGET_FILE_IF_EXISTS(//foo) Hello>',
+        ]:
+            with self.assertRaises(ExpressionError):
+                expand_expressions(self._paths, incomplete_expression)
+
+    def test_target_file(self):
+        path = self._path('test', 'fake_test.elf')
+
+        for expr, expected in [
+            ('<TARGET_FILE(//fake_module:fake_test)>', path),
+            ('--arg=<TARGET_FILE(//fake_module:fake_test)>', f'--arg={path}'),
+            ('--argument=<TARGET_FILE(//fake_module:fake_test)>;'
+             '<TARGET_FILE(//fake_module:fake_test)>',
+             f'--argument={path};{path}'),
+        ]:
+            self.assertEqual(list(expand_expressions(self._paths, expr)),
+                             [expected])
+
+    def test_target_file_if_exists(self):
+        path = self._path('test', 'fake_test.elf', create=True)
+
+        for expr, expected in [
+            ('<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>', path),
+            ('--arg=<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>',
+             f'--arg={path}'),
+            ('--argument=<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>;'
+             '<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>',
+             f'--argument={path};{path}'),
+        ]:
+            self.assertEqual(list(expand_expressions(self._paths, expr)),
+                             [expected])
+
+    def test_target_file_if_exists_arg_omitted(self):
+        for expr in [
+                '<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>',
+                '--arg=<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>',
+                '--argument=<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>;'
+                '<TARGET_FILE_IF_EXISTS(//fake_module:fake_test)>',
+        ]:
+            self.assertEqual(list(expand_expressions(self._paths, expr)), [])
+
+    def test_target_objects(self):
+        self.assertEqual(
+            set(
+                expand_expressions(
+                    self._paths,
+                    '<TARGET_OBJECTS(//fake_module:fake_source_set)>')), {
+                        self._path('fake_source_set.file_a.cc.o'),
+                        self._path('fake_source_set.file_b.c.o')
+                    })
+        self.assertEqual(
+            set(
+                expand_expressions(
+                    self._paths, '<TARGET_OBJECTS(//fake_module:fake_test)>')),
+            {
+                self._path('fake_test.fake_test.cc.o'),
+                self._path('fake_test.fake_test_c.c.o')
+            })
+
+    def test_target_objects_no_objects(self):
+        self.assertEqual(
+            list(
+                expand_expressions(
+                    self._paths,
+                    '<TARGET_OBJECTS(//fake_module:fake_no_objects)>')), [])
+
+    def test_target_objects_other_content_in_arg(self):
+        for arg in [
+                '--foo=<TARGET_OBJECTS(//fake_module:fake_no_objects)>',
+                '<TARGET_OBJECTS(//fake_module:fake_no_objects)>bar',
+                '--foo<TARGET_OBJECTS(//fake_module:fake_no_objects)>bar',
+                '<TARGET_OBJECTS(//fake_module:fake_no_objects)>'
+                '<TARGET_OBJECTS(//fake_module:fake_no_objects)>',
+                '<TARGET_OBJECTS(//fake_module:fake_source_set)>'
+                '<TARGET_OBJECTS(//fake_module:fake_source_set)>',
+        ]:
+            with self.assertRaises(ExpressionError):
+                expand_expressions(self._paths, arg)
+
+
 if __name__ == '__main__':
     unittest.main()