xref: /llvm-project/mlir/utils/jupyter/mlir_opt_kernel/install.py (revision f9008e6366c2496b1ca1785b891d5578174ad63e)
1# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2# See https://llvm.org/LICENSE.txt for license information.
3# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5import os
6import argparse
7
8from jupyter_client.kernelspec import KernelSpecManager
9
10
11def install_my_kernel_spec(user=True, prefix=None):
12    """Install the kernel spec for user in given prefix."""
13    print("Installing mlir-opt IPython kernel spec")
14    pkgroot = os.path.dirname(__file__)
15    KernelSpecManager().install_kernel_spec(
16        os.path.join(pkgroot, "assets"), "mlir", user=user, prefix=prefix
17    )
18
19
20def _is_root():
21    """Returns whether the current user is root."""
22    try:
23        return os.geteuid() == 0
24    except AttributeError:
25        # Return false wherever unknown.
26        return False
27
28
29def main(argv=None):
30    parser = argparse.ArgumentParser(
31        description="Install KernelSpec for MlirOpt Kernel"
32    )
33    prefix_locations = parser.add_mutually_exclusive_group()
34
35    prefix_locations.add_argument(
36        "--user", help="Install in user home directory", action="store_true"
37    )
38    prefix_locations.add_argument(
39        "--prefix", help="Install directory prefix", default=None
40    )
41
42    args = parser.parse_args(argv)
43
44    user = args.user or not _is_root()
45    prefix = args.prefix
46
47    install_my_kernel_spec(user=user, prefix=prefix)
48
49
50if __name__ == "__main__":
51    main()
52