Add ignored file (?)

compiler/bindings/python/iree/build/main.py

@@ -0,0 +1,244 @@
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from typing import Any, IO
+
+import argparse
+import importlib
+import importlib.util
+from pathlib import Path
+import sys
+
+from iree.build.executor import BuildEntrypoint, Entrypoint, Executor
+
+__all__ = [
+    "iree_build_main",
+    "load_build_module",
+]
+
+
+def iree_build_main(
+    module="__main__",
+    args: list[str] | None = None,
+    stdout: IO | None = None,
+    stderr: IO | None = None,
+):
+    """Make a build module invoke iree.build on itself when run.
+
+    Typically, if you have a module that declares build entrypoints, you will
+    add a stanza at the end:
+
+    .. code-block:: python
+        from iree.build import *
+
+        if __name__ == "__main__":
+            iree_build_main()
+    """
+    main = CliMain(module=module, args=args, stdout=stdout, stderr=stderr)
+    main.run()
+
+
+def load_build_module(module_path: Path | str):
+    """Loads a build module by path, evaling and returning it."""
+    spec = importlib.util.spec_from_file_location("__iree_build__", module_path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+class CliMain:
+    """Composes command line programs."""
+
+    def __init__(
+        self,
+        *,
+        args: list[str] | None = None,
+        module=None,
+        stdout: IO | None = None,
+        stderr: IO | None = None,
+    ):
+        self.stdout = stdout if stdout is not None else sys.stdout
+        self.stderr = stderr if stderr is not None else sys.stderr
+        if args is None:
+            args = sys.argv[1:]
+        if module is not None and isinstance(module, str):
+            module = __import__(module)
+        module = module
+
+        p = self.arg_parser = argparse.ArgumentParser(
+            description="IREE program build driver"
+        )
+        if module is None:
+            args, self.top_module = self._resolve_module_arguments(args)
+        else:
+            self.top_module = ModuleWrapper(module)
+
+        p.add_argument(
+            "--output-dir",
+            type=Path,
+            default=Path.cwd(),
+            help="Output directory for the build tree (defaults to current directory)",
+        )
+
+        cmd_group_desc = p.add_argument_group(
+            title="Build command",
+            description="Selects a build sub-command to invoke (default '--build')",
+        )
+        cmd_group = cmd_group_desc.add_mutually_exclusive_group()
+        cmd_group.add_argument(
+            "--build",
+            dest="command",
+            action="store_const",
+            const=self.build_command,
+            help="Executes build actions",
+        )
+        cmd_group.add_argument(
+            "--list",
+            dest="command",
+            action="store_const",
+            const=self.list_command,
+            help="Lists top level build actions",
+        )
+
+        cmd_group.add_argument(
+            "--list-all",
+            dest="command",
+            action="store_const",
+            const=self.list_all_command,
+            help="Lists all build actions",
+        )
+
+        p.add_argument(
+            "action_path",
+            nargs="*",
+            help="Paths of actions to build (default to top-level actions)",
+        )
+
+        self._define_action_arguments(p)
+        self.args = self.arg_parser.parse_args(args)
+
+    def abort(self):
+        sys.exit(1)
+
+    def _define_action_arguments(self, p: argparse.ArgumentParser):
+        user_group = p.add_argument_group("Action defined options")
+        for ep in self.top_module.entrypoints.values():
+            for cl_arg in ep.cl_args():
+                cl_arg.define_arg(user_group)
+
+    def _resolve_module_arguments(
+        self, args: list[str]
+    ) -> tuple[list[str], "ModuleWrapper"]:
+        p = argparse.ArgumentParser(
+            add_help=False,
+            usage="python -m iree.build [-m] build_module [... additional module specific options ...]",
+            prog="python -m iree.build",
+        )
+        # Invoked as a standalone tool: need the user to specify the
+        # module.
+        p.add_argument(
+            "-m",
+            dest="parse_as_module",
+            action="store_true",
+            help="Interpret the build definitions argument as a module (vs a file)",
+        )
+        p.add_argument(
+            "build_module",
+            help="The Python file or module from which to load build definitions",
+        )
+
+        bootstrap_args, rem_args = p.parse_known_args(args)
+        # Resolve from arguments.
+        is_module = bootstrap_args.parse_as_module or _is_module_like_str(
+            bootstrap_args.build_module
+        )
+        if is_module:
+            try:
+                top_module = ModuleWrapper.load_module(bootstrap_args.build_module)
+            except ModuleNotFoundError as e:
+                print(
+                    f"ERROR: Module '{bootstrap_args.build_module}' not found: {e}",
+                    file=self.stderr,
+                )
+                self.abort()
+        else:
+            top_module = ModuleWrapper.load_py_file(bootstrap_args.build_module)
+        return rem_args, top_module
+
+    def _create_executor(self) -> Executor:
+        executor = Executor(self.args.output_dir, self.args, stderr=self.stderr)
+        executor.analyze(*self.top_module.entrypoints.values())
+        return executor
+
+    def run(self):
+        command = self.args.command
+        if command is None:
+            command = self.build_command
+        command()
+
+    def build_command(self):
+        executor = self._create_executor()
+
+        if not self.args.action_path:
+            # Default to all.
+            build_actions = list(executor.entrypoints)
+        else:
+            # Look up each requested and add it.
+            build_actions = []
+            for action_path in self.args.action_path:
+                try:
+                    build_actions.append(executor.all[action_path])
+                except KeyError:
+                    all_paths = "\n".join(executor.all.keys())
+                    print(
+                        f"ERROR: Action '{action_path}' not found. Available: \n{all_paths}",
+                        file=self.stderr,
+                    )
+                    self.abort()
+        executor.build(*build_actions)
+
+        for build_action in build_actions:
+            if isinstance(build_action, BuildEntrypoint):
+                for output in build_action.outputs:
+                    print(f"{output.get_fs_path()}", file=self.stdout)
+
+    def list_command(self):
+        executor = self._create_executor()
+        for ep in executor.entrypoints:
+            print(ep.path, file=self.stdout)
+
+    def list_all_command(self):
+        executor = self._create_executor()
+        for name in executor.all.keys():
+            if name:
+                print(name, file=self.stdout)
+
+
+class ModuleWrapper:
+    """Wraps a raw, loaded module with access to discovered details."""
+
+    def __init__(self, mod):
+        self.mod = mod
+        self.entrypoints = self._collect_entrypoints()
+
+    @staticmethod
+    def load_module(module_name: str) -> "ModuleWrapper":
+        return ModuleWrapper(importlib.import_module(module_name))
+
+    @staticmethod
+    def load_py_file(module_path: Path | str) -> "ModuleWrapper":
+        return ModuleWrapper(load_build_module(str(module_path)))
+
+    def _collect_entrypoints(self) -> dict[str, Entrypoint]:
+        results: dict[str, Entrypoint] = {}
+        for attr_name, attr_value in self.mod.__dict__.items():
+            if isinstance(attr_value, Entrypoint):
+                results[attr_name] = attr_value
+        return results
+
+
+def _is_module_like_str(s: str) -> bool:
+    return "/" not in s and "\\" not in s and not s.endswith(".py")