TheRock/compiler/amd-llvm/mlir/utils/jupyter/mlir_opt_kernel/install.py

# Part of the LLVM Project, 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

import os
import argparse

from jupyter_client.kernelspec import KernelSpecManager


def install_my_kernel_spec(user=True, prefix=None):
    """Install the kernel spec for user in given prefix."""
    print("Installing mlir-opt IPython kernel spec")
    pkgroot = os.path.dirname(__file__)
    KernelSpecManager().install_kernel_spec(
        os.path.join(pkgroot, "assets"), "mlir", user=user, prefix=prefix
    )


def _is_root():
    """Returns whether the current user is root."""
    try:
        return os.geteuid() == 0
    except AttributeError:
        # Return false wherever unknown.
        return False


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Install KernelSpec for MlirOpt Kernel"
    )
    prefix_locations = parser.add_mutually_exclusive_group()

    prefix_locations.add_argument(
        "--user", help="Install in user home directory", action="store_true"
    )
    prefix_locations.add_argument(
        "--prefix", help="Install directory prefix", default=None
    )

    args = parser.parse_args(argv)

    user = args.user or not _is_root()
    prefix = args.prefix

    install_my_kernel_spec(user=user, prefix=prefix)


if __name__ == "__main__":
    main()