# -*- coding: utf-8 -*- ########################################################################### # Copyright (c), The AiiDA team. All rights reserved. # # This file is part of the AiiDA code. # # # # The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### """ An IPython extension that provides a magic command to load basic aiida commands. This makes it much easier to start. Produces output in: * Plaintext (IPython [qt]console) * HTML (IPython notebook, ``nbconvert --to html``, ``--to slides``) * JSON (IPython notebook ``.ipynb`` files) * LaTeX (e.g. ``ipython nbconvert example.ipynb --to LaTeX --post PDF``) Notes on how to load it at start: https://ipython.org/ipython-doc/3/config/intro.html Usage ====== .. sourcecode:: ipython In [1]: %load_ext aiida_magic In [2]: %aiida """ from IPython import version_info, get_ipython from IPython.core import magic [docs]def add_to_ns(local_ns, name, obj): """ Add a new variable with name ``name`` and value ``obj`` to the namespace ``local_ns``, optionally showing a warning if we are hiding an existing variable. .. todo:: implement the warning. Example:: # assuming that local_ns is a dictionary, e.g. from locals() import sys add_to_ns(local_ns, 'sys', sys) """ if name in local_ns: pass local_ns[name] = obj [docs]@magic.magics_class class AiiDALoaderMagics(magic.Magics): """AiiDA magic loader.""" [docs] @magic.needs_local_scope @magic.line_magic def aiida(self, line='', local_ns=None): """Load AiiDA in ipython (checking if it was already loaded), and inserts in the namespace the main AiiDA classes (the same that are loaded in ``verdi shell``. Usage:: %aiida [optional parameters] .. todo:: implement parameters, e.g. for the profile to load. """ # pylint: disable=unused-argument,attribute-defined-outside-init from aiida.manage.configuration import load_profile from aiida.cmdline.utils.shell import get_start_namespace self.is_warning = False lcontent = line.strip() if lcontent: profile = load_profile(lcontent) else: profile = load_profile() self.current_state = f'Loaded AiiDA DB environment - profile name: {profile.name}.' user_ns = get_start_namespace() for key, value in user_ns.items(): add_to_ns(local_ns, key, value) return self [docs] def _repr_json_(self): """ Output in JSON format. """ from aiida.common import json obj = {'current_state': self.current_state} if version_info[0] >= 3: return obj return json.dumps(obj) [docs] def _repr_html_(self): """ Output in HTML format. """ html = '<p>' if self.is_warning: html += '<strong>' html += self.current_state if self.is_warning: html += '</strong>' html += '</p>' return html [docs] def _repr_latex_(self): """ Output in LaTeX format. """ if self.is_warning: latex = '\\emph{%s}\n' % self.current_state else: latex = f'{self.current_state}\n' return latex [docs] def _repr_pretty_(self, pretty_print, cycle): # pylint: disable=unused-argument """ Output in text format. """ if self.is_warning: warning_str = '** ' else: warning_str = '' text = f'{warning_str}{self.current_state}\n' pretty_print.text(text) [docs]def load_ipython_extension(ipython): """ Registers the %aiida IPython extension. .. deprecated:: v3.0.0 Use :py:func:`~aiida.tools.ipython.ipython_magics.register_ipython_extension` instead. """ register_ipython_extension(ipython) [docs]def register_ipython_extension(ipython=None): """ Registers the %aiida IPython extension. The %aiida IPython extension provides the same environment as the `verdi shell`. :param ipython: InteractiveShell instance. If omitted, the global InteractiveShell is used. """ if ipython is None: ipython = get_ipython() ipython.register_magics(AiiDALoaderMagics)