1  # encoding: utf-8
   2  """A class for managing IPython extensions."""
   3  
   4  # Copyright (c) IPython Development Team.
   5  # Distributed under the terms of the Modified BSD License.
   6  
   7  import os
   8  import os.path
   9  import sys
  10  from importlib import import_module, reload
  11  
  12  from traitlets.config.configurable import Configurable
  13  from IPython.utils.path import ensure_dir_exists
  14  from traitlets import Instance
  15  
  16  
  17  #-----------------------------------------------------------------------------
  18  # Main class
  19  #-----------------------------------------------------------------------------
  20  
  21  BUILTINS_EXTS = {"storemagic": False, "autoreload": False}
  22  
  23  
  24  class ExtensionManager(Configurable):
  25      """A class to manage IPython extensions.
  26  
  27      An IPython extension is an importable Python module that has
  28      a function with the signature::
  29  
  30          def load_ipython_extension(ipython):
  31              # Do things with ipython
  32  
  33      This function is called after your extension is imported and the
  34      currently active :class:`InteractiveShell` instance is passed as
  35      the only argument.  You can do anything you want with IPython at
  36      that point, including defining new magic and aliases, adding new
  37      components, etc.
  38  
  39      You can also optionally define an :func:`unload_ipython_extension(ipython)`
  40      function, which will be called if the user unloads or reloads the extension.
  41      The extension manager will only call :func:`load_ipython_extension` again
  42      if the extension is reloaded.
  43  
  44      You can put your extension modules anywhere you want, as long as
  45      they can be imported by Python's standard import mechanism.
  46      """
  47  
  48      shell = Instance('IPython.core.interactiveshell.InteractiveShellABC', allow_none=True)
  49  
  50      def __init__(self, shell=None, **kwargs):
  51          super(ExtensionManager, self).__init__(shell=shell, **kwargs)
  52          self.loaded = set()
  53  
  54      def load_extension(self, module_str: str):
  55          """Load an IPython extension by its module name.
  56  
  57          Returns the string "already loaded" if the extension is already loaded,
  58          "no load function" if the module doesn't have a load_ipython_extension
  59          function, or None if it succeeded.
  60          """
  61          try:
  62              return self._load_extension(module_str)
  63          except ModuleNotFoundError:
  64              if module_str in BUILTINS_EXTS:
  65                  BUILTINS_EXTS[module_str] = True
  66                  return self._load_extension("IPython.extensions." + module_str)
  67              raise
  68  
  69      def _load_extension(self, module_str: str):
  70          if module_str in self.loaded:
  71              return "already loaded"
  72  
  73          assert self.shell is not None
  74  
  75          with self.shell.builtin_trap:
  76              if module_str not in sys.modules:
  77                  mod = import_module(module_str)
  78              mod = sys.modules[module_str]
  79              if self._call_load_ipython_extension(mod):
  80                  self.loaded.add(module_str)
  81              else:
  82                  return "no load function"
  83  
  84      def unload_extension(self, module_str: str):
  85          """Unload an IPython extension by its module name.
  86  
  87          This function looks up the extension's name in ``sys.modules`` and
  88          simply calls ``mod.unload_ipython_extension(self)``.
  89  
  90          Returns the string "no unload function" if the extension doesn't define
  91          a function to unload itself, "not loaded" if the extension isn't loaded,
  92          otherwise None.
  93          """
  94          if BUILTINS_EXTS.get(module_str, False) is True:
  95              module_str = "IPython.extensions." + module_str
  96          if module_str not in self.loaded:
  97              return "not loaded"
  98  
  99          if module_str in sys.modules:
 100              mod = sys.modules[module_str]
 101              if self._call_unload_ipython_extension(mod):
 102                  self.loaded.discard(module_str)
 103              else:
 104                  return "no unload function"
 105  
 106      def reload_extension(self, module_str: str):
 107          """Reload an IPython extension by calling reload.
 108  
 109          If the module has not been loaded before,
 110          :meth:`InteractiveShell.load_extension` is called. Otherwise
 111          :func:`reload` is called and then the :func:`load_ipython_extension`
 112          function of the module, if it exists is called.
 113          """
 114  
 115          if BUILTINS_EXTS.get(module_str, False) is True:
 116              module_str = "IPython.extensions." + module_str
 117  
 118          if (module_str in self.loaded) and (module_str in sys.modules):
 119              self.unload_extension(module_str)
 120              mod = sys.modules[module_str]
 121              reload(mod)
 122              if self._call_load_ipython_extension(mod):
 123                  self.loaded.add(module_str)
 124          else:
 125              self.load_extension(module_str)
 126  
 127      def _call_load_ipython_extension(self, mod):
 128          if hasattr(mod, 'load_ipython_extension'):
 129              mod.load_ipython_extension(self.shell)
 130              return True
 131  
 132      def _call_unload_ipython_extension(self, mod):
 133          if hasattr(mod, 'unload_ipython_extension'):
 134              mod.unload_ipython_extension(self.shell)
 135              return True
 136