Use node.calcfunction decorator (#15)

Add node.calcfunction and node.workfunction decorators.
aiidateam · Dec 8, 2023 · b22467f · b22467f
1 parent dd04076
commit b22467f
Show file tree

Hide file tree

Showing 34 changed files with 4,299 additions and 1,174 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -36,7 +36,7 @@ jobs:
 
         strategy:
             matrix:
-                python-version: ['3.10']
+                python-version: ['3.9', '3.10', '3.11']
 
         services:
             postgres:

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -5,6 +5,11 @@ repos:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
+-   repo: https://github.com/pycqa/flake8
+    rev: '6.0.0'
+    hooks:
+    -   id: flake8
+        args: ['--max-line-length=121', '--ignore=F821, F722, E203, W503']
 -   repo: https://github.com/psf/black
     rev: 22.3.0
     hooks:

diff --git a/README.md b/README.md
@@ -33,7 +33,7 @@ Here is a detailed comparison between the ``WorkTree`` with two AiiDA built-in w
 ## Installation
 
 ```console
-    pip install --upgrade aiida-worktree
+    pip install git+https://github.com/superstar54/aiida-worktree
 ```
 
 
@@ -46,17 +46,14 @@ Create nodes from `calcfunction`.
 
 ```python
 from aiida_worktree import node
-from aiida.engine import calcfunction
 
 # define add calcfunction node
-@node()
-@calcfunction
+@node.calcfunction()
 def add(x, y):
     return x + y
 
 # define multiply calcfunction node
-@node()
-@calcfunction
+@node.calcfunction()
 def multiply(x, y):
     return x*y
 

diff --git a/aiida_worktree/__init__.py b/aiida_worktree/__init__.py
@@ -2,4 +2,7 @@
 from .node import Node
 from .decorator import node, build_node
 
+
 __version__ = "0.0.4"
+
+__all__ = ["WorkTree", "Node", "node", "build_node"]
diff --git a/aiida_worktree/cli/__init__.py b/aiida_worktree/cli/__init__.py
@@ -3,3 +3,6 @@
 The commands need to be imported here for them to be registered with the top-level command group.
 """
 from aiida_worktree.cli import cmd_tree
+
+
+__all__ = ["cmd_tree"]
diff --git a/aiida_worktree/decorator.py b/aiida_worktree/decorator.py
@@ -30,7 +30,6 @@ def add_input_recursive(inputs, port, prefix=None):
 
 def build_node(ndata):
     """Register a node from a AiiDA component."""
-    from aiida.engine import calcfunction
     from node_graph.decorator import create_node
     from aiida_worktree.node import Node
     import cloudpickle as pickle
@@ -250,20 +249,28 @@ class NodeDecoratorCollection:
 
     __call__: Any = node  # Alias '@node' to '@node.node'.
 
+    @staticmethod
+    def calcfunction(**kwargs):
+        def decorator(func):
+            # First, apply the calcfunction decorator
+            calcfunc_decorated = calcfunction(func)
+            # Then, apply node decorator
+            node_decorated = node(**kwargs)(calcfunc_decorated)
 
-node = NodeDecoratorCollection()
+            return node_decorated
+
+        return decorator
+
+    @staticmethod
+    def workfunction(**kwargs):
+        def decorator(func):
+            # First, apply the workfunction decorator
+            calcfunc_decorated = workfunction(func)
+            node_decorated = node(**kwargs)(calcfunc_decorated)
 
-if __name__ == "__main__":
-    from aiida.engine import calcfunction
-    from aiida_worktree.decorator import node
+            return node_decorated
 
-    @node(
-        identifier="MyAdd",
-        outputs=[["General", "result"]],
-        executor_type="calcfunction",
-    )
-    @calcfunction
-    def myadd(x, y):
-        return x + y
+        return decorator
 
-    print(myadd.node)
+
+node = NodeDecoratorCollection()
diff --git a/aiida_worktree/engine/worktree.py b/aiida_worktree/engine/worktree.py
@@ -19,7 +19,7 @@
 
 
 from aiida.engine.processes.exit_code import ExitCode
-from aiida.engine.processes.process import Process, ProcessState
+from aiida.engine.processes.process import Process
 
 from aiida.engine.processes.workchains.awaitable import (
     Awaitable,
@@ -30,9 +30,6 @@
 from aiida.engine.processes.workchains.workchain import Protect, WorkChainSpec
 from aiida.engine import run_get_node
 
-from os.path import splitext
-import yaml
-
 
 if t.TYPE_CHECKING:
     from aiida.engine.runners import Runner  # pylint: disable=unused-import
@@ -275,8 +272,10 @@ def _do_step(self) -> t.Any:
 
         # If the worktree is finished or the result is an ExitCode, we exit by returning
         if finished:
-            result = self.finalize()
-            return result
+            if isinstance(result, ExitCode):
+                return result
+            else:
+                return self.finalize()
 
         if self._awaitables:
             return Wait(self._do_step, "Waiting before next step")
@@ -684,8 +683,6 @@ def run_nodes(self, names):
                 self.to_context(**{name: process})
             elif node["metadata"]["node_type"] in ["worktree"]:
                 # process = run_get_node(executor, *args, **kwargs)
-                from aiida_worktree.utils import merge_properties
-                from aiida.orm.utils.serialize import serialize
 
                 print("node  type: worktree.")
                 wt = self.run_executor(executor, args, kwargs, var_args, var_kwargs)
@@ -859,9 +856,6 @@ def check_parent_state(self, name):
                         "SKIPPED",
                         "FAILED",
                     ]:
-                        # print(
-                        #     f"    {name}: Input node {link['from_node']}, {self.ctx.nodes[link['from_node']]['state']} ."
-                        # )
                         ready = False
                         return (
                             ready,

diff --git a/aiida_worktree/executors/qe.py b/aiida_worktree/executors/qe.py
@@ -1,6 +1,4 @@
-import time
-from aiida import engine, orm, load_profile
-from aiida.orm import load_code
+from aiida import load_profile
 
 load_profile()
 

diff --git a/aiida_worktree/nodes/builtin.py b/aiida_worktree/nodes/builtin.py
@@ -1,5 +1,4 @@
 from aiida_worktree.node import Node
-from aiida_worktree.executors.builtin import GatherWorkChain
 
 
 class AiiDAGather(Node):

diff --git a/aiida_worktree/nodes/qe.py b/aiida_worktree/nodes/qe.py
@@ -1,5 +1,4 @@
 from aiida_worktree.node import Node
-from aiida import orm
 
 
 class AiiDAKpoint(Node):

diff --git a/aiida_worktree/utils/analysis.py b/aiida_worktree/utils/analysis.py
@@ -26,7 +26,6 @@ def save(self):
         """
         self.build_node_link()
         self.build_connectivity()
-        print("exist_in_db: {}".format(self.exist_in_db()))
         if self.exist_in_db() or self.restart_process is not None:
             new_nodes, modified_nodes, update_metadata = self.check_diff(
                 self.restart_process

diff --git a/aiida_worktree/worktree.py b/aiida_worktree/worktree.py
@@ -43,7 +43,6 @@ def run(self):
         """
         from aiida_worktree.engine.worktree import WorkTree as WorkTreeEngine
         from aiida_worktree.utils import merge_properties
-        from aiida.orm.utils.serialize import serialize
         from aiida.manage import manager
 
         # One can not run again if the process is alreay created. otherwise, a new process node will
@@ -253,7 +252,6 @@ def play_nodes(self, nodes):
 
     def reset(self):
         """Reset the worktree."""
-        from aiida.engine.processes import control
 
         self.process = None
         for node in self.nodes:

diff --git a/docs/source/blog.ipynb b/docs/source/blog.ipynb
@@ -108,14 +108,12 @@
     "from aiida.engine import calcfunction\n",
     "\n",
     "# define add node\n",
-    "@node()\n",
-    "@calcfunction\n",
+    "@node.calcfunction()\n",
     "def add(x, y):\n",
     "    return x + y\n",
     "\n",
     "# define multiply node\n",
-    "@node()\n",
-    "@calcfunction\n",
+    "@node.calcfunction()\n",
     "def multiply(x, y):\n",
     "    return x*y\n",
     "\n"

diff --git a/docs/source/concept/node.rst b/docs/source/concept/node.rst
@@ -15,8 +15,7 @@ Decorate any Python function using the `node` decorator. To use the power of Aii
    from aiida.engine import calcfunction
 
    # define add calcfunction node
-   @node()
-   @calcfunction
+   @node.calcfunction()
    def add(x, y):
       return x + y
 

diff --git a/docs/source/howto/continue_finished_worktree.ipynb b/docs/source/howto/continue_finished_worktree.ipynb
diff --git a/docs/source/howto/ctx.ipynb b/docs/source/howto/ctx.ipynb
@@ -133,8 +133,7 @@
     "from aiida.orm import Int\n",
     "\n",
     "# define add node\n",
-    "@node()\n",
-    "@calcfunction\n",
+    "@node.calcfunction()\n",
     "def add(x, y):\n",
     "    return x + y\n",
     "\n",
@@ -188,7 +187,75 @@
    "outputs": [
     {
      "data": {
-      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.43.0 (0)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"224pt\" height=\"419pt\"\n viewBox=\"0.00 0.00 224.00 419.48\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 415.48)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-415.48 220,-415.48 220,4 -4,4\"/>\n<!-- N4231 -->\n<g id=\"node1\" class=\"node\">\n<title>N4231</title>\n<polygon fill=\"#e38851\" stroke=\"red\" stroke-width=\"6\" points=\"216,-291.74 0,-291.74 0,-238.74 216,-238.74 216,-291.74\"/>\n<text text-anchor=\"middle\" x=\"108\" y=\"-276.54\" font-family=\"Times,serif\" font-size=\"14.00\">WorkTree: test_worktree_ctx (4231)</text>\n<text text-anchor=\"middle\" x=\"108\" y=\"-261.54\" font-family=\"Times,serif\" font-size=\"14.00\">State: finished</text>\n<text text-anchor=\"middle\" x=\"108\" y=\"-246.54\" font-family=\"Times,serif\" font-size=\"14.00\">Exit Code: 0</text>\n</g>\n<!-- N4233 -->\n<g id=\"node3\" class=\"node\">\n<title>N4233</title>\n<polygon fill=\"#de707f\" fill-opacity=\"0.466667\" stroke=\"black\" stroke-width=\"0\" points=\"155.5,-172.74 60.5,-172.74 60.5,-119.74 155.5,-119.74 155.5,-172.74\"/>\n<text text-anchor=\"middle\" x=\"108\" y=\"-157.54\" font-family=\"Times,serif\" font-size=\"14.00\">add (4233)</text>\n<text text-anchor=\"middle\" x=\"108\" y=\"-142.54\" font-family=\"Times,serif\" font-size=\"14.00\">State: finished</text>\n<text text-anchor=\"middle\" x=\"108\" y=\"-127.54\" font-family=\"Times,serif\" font-size=\"14.00\">Exit Code: 0</text>\n</g>\n<!-- N4231&#45;&gt;N4233 -->\n<g id=\"edge3\" class=\"edge\">\n<title>N4231&#45;&gt;N4233</title>\n<path fill=\"none\" stroke=\"#000000\" stroke-dasharray=\"1,5\" d=\"M108,-238.56C108,-222.26 108,-200.92 108,-182.97\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"111.5,-182.86 108,-172.86 104.5,-182.86 111.5,-182.86\"/>\n<text text-anchor=\"middle\" x=\"147.5\" y=\"-209.54\" font-family=\"Times,serif\" font-size=\"14.00\">CALL_CALC</text>\n<text text-anchor=\"middle\" x=\"147.5\" y=\"-194.54\" font-family=\"Times,serif\" font-size=\"14.00\">add1</text>\n</g>\n<!-- N4230 -->\n<g id=\"node2\" class=\"node\">\n<title>N4230</title>\n<ellipse fill=\"#8cd499\" stroke=\"black\" stroke-width=\"0\" cx=\"108\" cy=\"-384.61\" rx=\"50.41\" ry=\"26.74\"/>\n<text text-anchor=\"middle\" x=\"108\" y=\"-388.41\" font-family=\"Times,serif\" font-size=\"14.00\">Int (4230)</text>\n<text text-anchor=\"middle\" x=\"108\" y=\"-373.41\" font-family=\"Times,serif\" font-size=\"14.00\">value: 2</text>\n</g>\n<!-- N4230&#45;&gt;N4231 -->\n<g id=\"edge1\" class=\"edge\">\n<title>N4230&#45;&gt;N4231</title>\n<path fill=\"none\" stroke=\"#000000\" stroke-dasharray=\"5,2\" d=\"M108,-357.56C108,-341.19 108,-319.85 108,-301.91\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"111.5,-301.81 108,-291.81 104.5,-301.81 111.5,-301.81\"/>\n<text text-anchor=\"middle\" x=\"152.5\" y=\"-328.54\" font-family=\"Times,serif\" font-size=\"14.00\">INPUT_WORK</text>\n<text text-anchor=\"middle\" x=\"152.5\" y=\"-313.54\" font-family=\"Times,serif\" font-size=\"14.00\">nt__ctx__x</text>\n</g>\n<!-- N4234 -->\n<g id=\"node4\" class=\"node\">\n<title>N4234</title>\n<ellipse fill=\"#8cd499\" stroke=\"black\" stroke-width=\"0\" cx=\"108\" cy=\"-26.87\" rx=\"50.41\" ry=\"26.74\"/>\n<text text-anchor=\"middle\" x=\"108\" y=\"-30.67\" font-family=\"Times,serif\" font-size=\"14.00\">Int (4234)</text>\n<text text-anchor=\"middle\" x=\"108\" y=\"-15.67\" font-family=\"Times,serif\" font-size=\"14.00\">value: 5</text>\n</g>\n<!-- N4233&#45;&gt;N4234 -->\n<g id=\"edge2\" class=\"edge\">\n<title>N4233&#45;&gt;N4234</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M108,-119.48C108,-103.24 108,-82.01 108,-64.09\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"111.5,-63.97 108,-53.97 104.5,-63.97 111.5,-63.97\"/>\n<text text-anchor=\"middle\" x=\"134\" y=\"-90.54\" font-family=\"Times,serif\" font-size=\"14.00\">CREATE</text>\n<text text-anchor=\"middle\" x=\"134\" y=\"-75.54\" font-family=\"Times,serif\" font-size=\"14.00\">result</text>\n</g>\n</g>\n</svg>\n",
+      "image/svg+xml": [
+       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
+       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
+       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
+       "<!-- Generated by graphviz version 2.43.0 (0)\n",
+       " -->\n",
+       "<!-- Title: %3 Pages: 1 -->\n",
+       "<svg width=\"224pt\" height=\"419pt\"\n",
+       " viewBox=\"0.00 0.00 224.00 419.48\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
+       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 415.48)\">\n",
+       "<title>%3</title>\n",
+       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-415.48 220,-415.48 220,4 -4,4\"/>\n",
+       "<!-- N4231 -->\n",
+       "<g id=\"node1\" class=\"node\">\n",
+       "<title>N4231</title>\n",
+       "<polygon fill=\"#e38851\" stroke=\"red\" stroke-width=\"6\" points=\"216,-291.74 0,-291.74 0,-238.74 216,-238.74 216,-291.74\"/>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-276.54\" font-family=\"Times,serif\" font-size=\"14.00\">WorkTree: test_worktree_ctx (4231)</text>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-261.54\" font-family=\"Times,serif\" font-size=\"14.00\">State: finished</text>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-246.54\" font-family=\"Times,serif\" font-size=\"14.00\">Exit Code: 0</text>\n",
+       "</g>\n",
+       "<!-- N4233 -->\n",
+       "<g id=\"node3\" class=\"node\">\n",
+       "<title>N4233</title>\n",
+       "<polygon fill=\"#de707f\" fill-opacity=\"0.466667\" stroke=\"black\" stroke-width=\"0\" points=\"155.5,-172.74 60.5,-172.74 60.5,-119.74 155.5,-119.74 155.5,-172.74\"/>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-157.54\" font-family=\"Times,serif\" font-size=\"14.00\">add (4233)</text>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-142.54\" font-family=\"Times,serif\" font-size=\"14.00\">State: finished</text>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-127.54\" font-family=\"Times,serif\" font-size=\"14.00\">Exit Code: 0</text>\n",
+       "</g>\n",
+       "<!-- N4231&#45;&gt;N4233 -->\n",
+       "<g id=\"edge3\" class=\"edge\">\n",
+       "<title>N4231&#45;&gt;N4233</title>\n",
+       "<path fill=\"none\" stroke=\"#000000\" stroke-dasharray=\"1,5\" d=\"M108,-238.56C108,-222.26 108,-200.92 108,-182.97\"/>\n",
+       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"111.5,-182.86 108,-172.86 104.5,-182.86 111.5,-182.86\"/>\n",
+       "<text text-anchor=\"middle\" x=\"147.5\" y=\"-209.54\" font-family=\"Times,serif\" font-size=\"14.00\">CALL_CALC</text>\n",
+       "<text text-anchor=\"middle\" x=\"147.5\" y=\"-194.54\" font-family=\"Times,serif\" font-size=\"14.00\">add1</text>\n",
+       "</g>\n",
+       "<!-- N4230 -->\n",
+       "<g id=\"node2\" class=\"node\">\n",
+       "<title>N4230</title>\n",
+       "<ellipse fill=\"#8cd499\" stroke=\"black\" stroke-width=\"0\" cx=\"108\" cy=\"-384.61\" rx=\"50.41\" ry=\"26.74\"/>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-388.41\" font-family=\"Times,serif\" font-size=\"14.00\">Int (4230)</text>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-373.41\" font-family=\"Times,serif\" font-size=\"14.00\">value: 2</text>\n",
+       "</g>\n",
+       "<!-- N4230&#45;&gt;N4231 -->\n",
+       "<g id=\"edge1\" class=\"edge\">\n",
+       "<title>N4230&#45;&gt;N4231</title>\n",
+       "<path fill=\"none\" stroke=\"#000000\" stroke-dasharray=\"5,2\" d=\"M108,-357.56C108,-341.19 108,-319.85 108,-301.91\"/>\n",
+       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"111.5,-301.81 108,-291.81 104.5,-301.81 111.5,-301.81\"/>\n",
+       "<text text-anchor=\"middle\" x=\"152.5\" y=\"-328.54\" font-family=\"Times,serif\" font-size=\"14.00\">INPUT_WORK</text>\n",
+       "<text text-anchor=\"middle\" x=\"152.5\" y=\"-313.54\" font-family=\"Times,serif\" font-size=\"14.00\">nt__ctx__x</text>\n",
+       "</g>\n",
+       "<!-- N4234 -->\n",
+       "<g id=\"node4\" class=\"node\">\n",
+       "<title>N4234</title>\n",
+       "<ellipse fill=\"#8cd499\" stroke=\"black\" stroke-width=\"0\" cx=\"108\" cy=\"-26.87\" rx=\"50.41\" ry=\"26.74\"/>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-30.67\" font-family=\"Times,serif\" font-size=\"14.00\">Int (4234)</text>\n",
+       "<text text-anchor=\"middle\" x=\"108\" y=\"-15.67\" font-family=\"Times,serif\" font-size=\"14.00\">value: 5</text>\n",
+       "</g>\n",
+       "<!-- N4233&#45;&gt;N4234 -->\n",
+       "<g id=\"edge2\" class=\"edge\">\n",
+       "<title>N4233&#45;&gt;N4234</title>\n",
+       "<path fill=\"none\" stroke=\"#000000\" d=\"M108,-119.48C108,-103.24 108,-82.01 108,-64.09\"/>\n",
+       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"111.5,-63.97 108,-53.97 104.5,-63.97 111.5,-63.97\"/>\n",
+       "<text text-anchor=\"middle\" x=\"134\" y=\"-90.54\" font-family=\"Times,serif\" font-size=\"14.00\">CREATE</text>\n",
+       "<text text-anchor=\"middle\" x=\"134\" y=\"-75.54\" font-family=\"Times,serif\" font-size=\"14.00\">result</text>\n",
+       "</g>\n",
+       "</g>\n",
+       "</svg>\n"
+      ],
       "text/plain": [
        "<graphviz.graphs.Digraph at 0x7f0f74ab60e0>"
       ]

diff --git a/docs/source/howto/for.ipynb b/docs/source/howto/for.ipynb