Skip to content

Commit

Permalink
Add ignored file (?)
Browse files Browse the repository at this point in the history
  • Loading branch information
stellaraccident committed Oct 31, 2024
1 parent 8502cdd commit 73d9120
Showing 1 changed file with 244 additions and 0 deletions.
244 changes: 244 additions & 0 deletions compiler/bindings/python/iree/build/main.py
Original file line number Diff line number Diff line change
@@ -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")

0 comments on commit 73d9120

Please sign in to comment.