From 3409bcbdce47bb4fe0df568a5b296fe2e5e726a0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henning=20Jan=C3=9Fen?= <he.janssen@fz-juelich.de>
Date: Mon, 25 Apr 2022 10:39:15 +0200
Subject: [PATCH] Add ipython extension subpackage

Added small initial example, which contains a cell magic exporting all
nodes from a ipython/jupyter runtime into a group and yaml file

Example::
  %load_ext aiida_jutools
  %export_nodes <group_name>

 This will create a group `notebook-exports/<group_name>` and a yaml file
 `<group_name>_export.yml`, the variable names (names are preferred over
 the ipython hacky references, e.g. _12) are shown in the yaml file and
 in the extras of the node
---
 aiida_jutools/__init__.py         | 11 ++++++
 aiida_jutools/ipython/__init__.py |  2 +
 aiida_jutools/ipython/export.py   | 62 +++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+)
 create mode 100644 aiida_jutools/ipython/__init__.py
 create mode 100644 aiida_jutools/ipython/export.py

diff --git a/aiida_jutools/__init__.py b/aiida_jutools/__init__.py
index 8ba8ed5..0073769 100644
--- a/aiida_jutools/__init__.py
+++ b/aiida_jutools/__init__.py
@@ -45,3 +45,14 @@
 #
 # The potential problem and the solution stated above, if it becomes one, applies to other JuDFTTeam plugins as well,
 # should they start using aiida-jutools as common codebase (aiida-fleur, aiida-spirit, aiida-spex, ...).
+
+def load_ipython_extension(ipython):
+    """
+    Register ipython extensions
+    """
+    from aiida_jutools.ipython import NotebookAiidaExport
+    import masci_tools
+
+    #Register extensions from masci-tools
+    masci_tools.load_ipython_extension(ipython)
+    ipython.register_magics(NotebookAiidaExport)
\ No newline at end of file
diff --git a/aiida_jutools/ipython/__init__.py b/aiida_jutools/ipython/__init__.py
new file mode 100644
index 0000000..92d58ae
--- /dev/null
+++ b/aiida_jutools/ipython/__init__.py
@@ -0,0 +1,2 @@
+
+from .export import NotebookAiidaExport
\ No newline at end of file
diff --git a/aiida_jutools/ipython/export.py b/aiida_jutools/ipython/export.py
new file mode 100644
index 0000000..bfd7ddd
--- /dev/null
+++ b/aiida_jutools/ipython/export.py
@@ -0,0 +1,62 @@
+"""
+Magics class for exporting aiida nodes from a jupyer notebook environment
+"""
+from IPython.core.magic import magics_class, line_magic, needs_local_scope
+from aiida import orm
+from collections import defaultdict
+import re
+import yaml
+
+IPYTHON_VARIABLE = re.compile('_+[0-9]*')
+
+@magics_class
+class NotebookAiidaExport:
+
+    @line_magic
+    @needs_local_scope
+    def export_nodes(self, line, local_ns):
+
+        args = line.split(' ', maxsplit=1)
+        group_name = args[0]
+
+        if len(args) == 2:
+            if args[1].lower() not in ('pk', 'uuid'):
+                raise ValueError('Invalid value for export format (Either pk or uuid)')
+            export_format = args[1].lower()
+        else:
+            export_format = 'uuid'
+        
+        #Get all AiiDA nodes in the local namespace
+        aiida_nodes = defaultdict(set)
+        for name, node in local_ns.items():
+            if isinstance(node, orm.Node):
+                aiida_nodes[node].add(name)
+        
+        aiida_nodes_one_name = {}
+        #Now choose the names and prefer proper variable names over _12, etc. (ipython utilities)
+        export_dict = defaultdict(list)
+        for node, names in aiida_nodes.items():
+            class_name = getattr(node, 'process_class', node.__class__).__name__
+            identifier = node.pk if export_format == 'pk' else node.uuid
+            if len(names) == 1:
+                name = names.pop()
+                export_dict[class_name].append([name, identifier])
+                aiida_nodes_one_name[name] = node
+            elif len({name for name in names if not re.fullmatch(IPYTHON_VARIABLE, name)}) != 0:
+                name = {name for name in names if not re.fullmatch(IPYTHON_VARIABLE, name)}.pop()
+                export_dict[class_name].append([name, identifier])
+                aiida_nodes_one_name[name] = node
+            else:
+                name = names.pop()
+                export_dict[class_name].append([name, identifier])
+                aiida_nodes_one_name[name] = node
+                
+        
+        group, _ = orm.Group.objects.get_or_create(f'notebook-exports/{group_name}')
+        for name, node in aiida_nodes_one_name.items():
+            node.extras['notebook-variable'] = name
+            node.extras['notebook-origin'] = group_name
+            group.add_nodes(node)
+        
+        with open(f'{group_name}-nodes.yml', 'w') as file:
+            yaml.dump(dict(export_dict), file)