mlir/utils/jupyter/mlir_opt_kernel/install.py

   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
   5 import os
   6 import argparse
   7
   8 from jupyter_client.kernelspec import KernelSpecManager
   9
  10
  11 def 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
  20 def _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
  29 def 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
  50 if __name__ == "__main__":
  51     main()