aiida-pythonjob 0.3.1__tar.gz → 0.3.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (63) hide show
  1. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/PKG-INFO +26 -20
  2. aiida_pythonjob-0.3.2/pyproject.toml +114 -0
  3. aiida_pythonjob-0.3.2/setup.cfg +4 -0
  4. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/__init__.py +4 -3
  5. aiida_pythonjob-0.3.2/src/aiida_pythonjob/calculations/__init__.py +4 -0
  6. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/calculations/pyfunction.py +17 -12
  7. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/calculations/pythonjob.py +2 -2
  8. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/serializer.py +18 -9
  9. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/decorator.py +4 -7
  10. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/launch.py +16 -31
  11. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/parsers/pythonjob.py +10 -9
  12. aiida_pythonjob-0.3.2/src/aiida_pythonjob/parsers/utils.py +125 -0
  13. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/utils.py +94 -168
  14. aiida_pythonjob-0.3.2/src/aiida_pythonjob.egg-info/PKG-INFO +99 -0
  15. aiida_pythonjob-0.3.2/src/aiida_pythonjob.egg-info/SOURCES.txt +38 -0
  16. aiida_pythonjob-0.3.2/src/aiida_pythonjob.egg-info/dependency_links.txt +1 -0
  17. aiida_pythonjob-0.3.2/src/aiida_pythonjob.egg-info/entry_points.txt +21 -0
  18. aiida_pythonjob-0.3.2/src/aiida_pythonjob.egg-info/requires.txt +24 -0
  19. aiida_pythonjob-0.3.2/src/aiida_pythonjob.egg-info/top_level.txt +1 -0
  20. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_parser.py +13 -44
  21. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_pyfunction.py +1 -1
  22. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_pythonjob.py +1 -1
  23. aiida_pythonjob-0.3.1/.github/workflows/ci.yml +0 -91
  24. aiida_pythonjob-0.3.1/.github/workflows/python-publish.yml +0 -40
  25. aiida_pythonjob-0.3.1/.gitignore +0 -164
  26. aiida_pythonjob-0.3.1/.pre-commit-config.yaml +0 -13
  27. aiida_pythonjob-0.3.1/.readthedocs.yml +0 -25
  28. aiida_pythonjob-0.3.1/docs/Makefile +0 -20
  29. aiida_pythonjob-0.3.1/docs/environment.yml +0 -7
  30. aiida_pythonjob-0.3.1/docs/gallery/autogen/GALLERY_HEADER.rst +0 -3
  31. aiida_pythonjob-0.3.1/docs/gallery/autogen/pyfunction.py +0 -191
  32. aiida_pythonjob-0.3.1/docs/gallery/autogen/pythonjob.py +0 -491
  33. aiida_pythonjob-0.3.1/docs/make.bat +0 -35
  34. aiida_pythonjob-0.3.1/docs/requirements.txt +0 -10
  35. aiida_pythonjob-0.3.1/docs/source/conf.py +0 -161
  36. aiida_pythonjob-0.3.1/docs/source/index.rst +0 -52
  37. aiida_pythonjob-0.3.1/docs/source/installation.rst +0 -44
  38. aiida_pythonjob-0.3.1/docs/source/tutorial/dft.ipynb +0 -1291
  39. aiida_pythonjob-0.3.1/docs/source/tutorial/html/atomization_energy.html +0 -290
  40. aiida_pythonjob-0.3.1/docs/source/tutorial/html/pythonjob_eos_emt.html +0 -290
  41. aiida_pythonjob-0.3.1/docs/source/tutorial/index.rst +0 -11
  42. aiida_pythonjob-0.3.1/pyproject.toml +0 -169
  43. aiida_pythonjob-0.3.1/src/aiida_pythonjob/calculations/__init__.py +0 -3
  44. aiida_pythonjob-0.3.1/tests/conftest.py +0 -18
  45. aiida_pythonjob-0.3.1/tests/inputs_folder/another_input.txt +0 -1
  46. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/LICENSE +0 -0
  47. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/README.md +0 -0
  48. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/calculations/utils.py +0 -0
  49. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/config.py +0 -0
  50. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/__init__.py +0 -0
  51. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/atoms.py +0 -0
  52. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/data_wrapper.py +0 -0
  53. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/deserializer.py +0 -0
  54. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/jsonable_data.py +0 -0
  55. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/pickled_data.py +0 -0
  56. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/data/utils.py +0 -0
  57. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/parsers/__init__.py +0 -0
  58. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/src/aiida_pythonjob/ports_adapter.py +0 -0
  59. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_create_env.py +0 -0
  60. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_data.py +0 -0
  61. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_entry_points.py +0 -0
  62. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_serializer.py +0 -0
  63. {aiida_pythonjob-0.3.1 → aiida_pythonjob-0.3.2}/tests/test_utils.py +0 -0
@@ -1,8 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: aiida-pythonjob
3
- Version: 0.3.1
3
+ Version: 0.3.2
4
4
  Summary: Run Python functions on a remote computer.
5
- Project-URL: Source, https://github.com/aiidateam/aiida-pythonjob
6
5
  Author-email: Xing Wang <xingwang1991@gmail.com>
7
6
  License: MIT License
8
7
 
@@ -25,33 +24,40 @@ License: MIT License
25
24
  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
26
25
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
27
26
  SOFTWARE.
28
- License-File: LICENSE
27
+
28
+ Project-URL: Source, https://github.com/aiidateam/aiida-pythonjob
29
29
  Keywords: aiida,plugin
30
- Classifier: Development Status :: 3 - Alpha
31
- Classifier: Framework :: AiiDA
30
+ Classifier: Programming Language :: Python
32
31
  Classifier: Intended Audience :: Science/Research
33
32
  Classifier: License :: OSI Approved :: MIT License
34
33
  Classifier: Natural Language :: English
35
- Classifier: Programming Language :: Python
34
+ Classifier: Development Status :: 3 - Alpha
35
+ Classifier: Framework :: AiiDA
36
36
  Requires-Python: >=3.9
37
+ Description-Content-Type: text/markdown
38
+ License-File: LICENSE
37
39
  Requires-Dist: aiida-core<3,>=2.3
38
40
  Requires-Dist: ase
39
41
  Requires-Dist: cloudpickle
40
- Requires-Dist: node-graph~=0.2.23
41
- Provides-Extra: dev
42
- Requires-Dist: hatch; extra == 'dev'
43
- Provides-Extra: docs
44
- Requires-Dist: furo; extra == 'docs'
45
- Requires-Dist: markupsafe<2.1; extra == 'docs'
46
- Requires-Dist: nbsphinx; extra == 'docs'
47
- Requires-Dist: sphinx; extra == 'docs'
48
- Requires-Dist: sphinx-gallery; extra == 'docs'
49
- Requires-Dist: sphinx-rtd-theme; extra == 'docs'
50
- Requires-Dist: sphinxcontrib-contentui; extra == 'docs'
51
- Requires-Dist: sphinxcontrib-details-directive; extra == 'docs'
42
+ Requires-Dist: node-graph>=0.3.0
43
+ Provides-Extra: test
44
+ Requires-Dist: pgtest>=1.3.1,~=1.3; extra == "test"
45
+ Requires-Dist: coverage~=7.0; extra == "test"
46
+ Requires-Dist: pytest~=7.0; extra == "test"
47
+ Requires-Dist: pytest-cov~=4.1; extra == "test"
48
+ Requires-Dist: ipdb; extra == "test"
52
49
  Provides-Extra: pre-commit
53
- Requires-Dist: pre-commit~=3.5; extra == 'pre-commit'
54
- Description-Content-Type: text/markdown
50
+ Requires-Dist: pre-commit~=3.5; extra == "pre-commit"
51
+ Provides-Extra: docs
52
+ Requires-Dist: sphinx_rtd_theme; extra == "docs"
53
+ Requires-Dist: sphinx; extra == "docs"
54
+ Requires-Dist: sphinxcontrib-contentui; extra == "docs"
55
+ Requires-Dist: sphinxcontrib-details-directive; extra == "docs"
56
+ Requires-Dist: sphinx-gallery; extra == "docs"
57
+ Requires-Dist: furo; extra == "docs"
58
+ Requires-Dist: markupsafe<2.1; extra == "docs"
59
+ Requires-Dist: nbsphinx; extra == "docs"
60
+ Dynamic: license-file
55
61
 
56
62
  # AiiDA-PythonJob
57
63
  [![PyPI version](https://badge.fury.io/py/aiida-pythonjob.svg)](https://badge.fury.io/py/aiida-pythonjob)
@@ -0,0 +1,114 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "aiida-pythonjob"
7
+ dynamic = ["version"] # read from src/aiida_pythonjob/__init__.py via setuptools
8
+ description = "Run Python functions on a remote computer."
9
+ authors = [{name = "Xing Wang", email = "xingwang1991@gmail.com"}]
10
+ readme = "README.md"
11
+ license = {file = "LICENSE"}
12
+ classifiers = [
13
+ "Programming Language :: Python",
14
+ "Intended Audience :: Science/Research",
15
+ "License :: OSI Approved :: MIT License",
16
+ "Natural Language :: English",
17
+ "Development Status :: 3 - Alpha",
18
+ "Framework :: AiiDA"
19
+ ]
20
+ keywords = ["aiida", "plugin"]
21
+ requires-python = ">=3.9"
22
+ dependencies = [
23
+ "aiida-core>=2.3,<3",
24
+ "ase",
25
+ "cloudpickle",
26
+ "node-graph>=0.3.0",
27
+ ]
28
+
29
+ [project.optional-dependencies]
30
+ test = [
31
+ "pgtest~=1.3,>=1.3.1",
32
+ "coverage~=7.0",
33
+ "pytest~=7.0",
34
+ "pytest-cov~=4.1",
35
+ "ipdb"
36
+ ]
37
+ pre-commit = [
38
+ "pre-commit~=3.5",
39
+ ]
40
+ docs = [
41
+ "sphinx_rtd_theme",
42
+ "sphinx",
43
+ "sphinxcontrib-contentui",
44
+ "sphinxcontrib-details-directive",
45
+ "sphinx-gallery",
46
+ "furo",
47
+ "markupsafe<2.1",
48
+ "nbsphinx"
49
+ ]
50
+
51
+ [project.urls]
52
+ Source = "https://github.com/aiidateam/aiida-pythonjob"
53
+
54
+ [project.entry-points."aiida.data"]
55
+ "pythonjob.jsonable_data" = "aiida_pythonjob.data.jsonable_data:JsonableData"
56
+ "pythonjob.pickled_data" = "aiida_pythonjob.data.pickled_data:PickledData"
57
+ "pythonjob.ase.atoms.Atoms" = "aiida_pythonjob.data.atoms:AtomsData"
58
+ "pythonjob.builtins.int" = "aiida.orm.nodes.data.int:Int"
59
+ "pythonjob.builtins.float" = "aiida.orm.nodes.data.float:Float"
60
+ "pythonjob.builtins.str" = "aiida.orm.nodes.data.str:Str"
61
+ "pythonjob.builtins.bool" = "aiida.orm.nodes.data.bool:Bool"
62
+ "pythonjob.builtins.list" = "aiida_pythonjob.data.data_wrapper:List"
63
+ "pythonjob.builtins.dict" = "aiida_pythonjob.data.data_wrapper:Dict"
64
+ "pythonjob.numpy.float32" = "aiida.orm.nodes.data.float:Float"
65
+ "pythonjob.numpy.float64" = "aiida.orm.nodes.data.float:Float"
66
+ "pythonjob.numpy.int64" = "aiida.orm.nodes.data.int:Int"
67
+ "pythonjob.numpy.bool_" = "aiida.orm.nodes.data.bool:Bool"
68
+ "pythonjob.numpy.ndarray" = "aiida_pythonjob.data.data_wrapper:ArrayData"
69
+
70
+ [project.entry-points."aiida.calculations"]
71
+ "pythonjob.pythonjob" = "aiida_pythonjob.calculations.pythonjob:PythonJob"
72
+
73
+ [project.entry-points."aiida.parsers"]
74
+ "pythonjob.pythonjob" = "aiida_pythonjob.parsers.pythonjob:PythonJobParser"
75
+
76
+ [tool.setuptools]
77
+ package-dir = {"" = "src"}
78
+
79
+ [tool.setuptools.packages.find]
80
+ where = ["src"]
81
+ include = ["aiida_pythonjob*"]
82
+
83
+ [tool.setuptools.dynamic]
84
+ version = {attr = "aiida_pythonjob.__version__"}
85
+
86
+ [tool.pytest.ini_options]
87
+ python_files = "test_*.py example_*.py"
88
+ addopts = "--pdbcls=IPython.terminal.debugger:TerminalPdb"
89
+ filterwarnings = [
90
+ "ignore::DeprecationWarning:aiida:",
91
+ "ignore:Creating AiiDA configuration folder:",
92
+ "ignore::DeprecationWarning:plumpy:",
93
+ "ignore::DeprecationWarning:yaml:",
94
+ ]
95
+
96
+ [tool.coverage.run]
97
+ source = ["src/aiida_pythonjob"]
98
+
99
+ [tool.ruff]
100
+ line-length = 120
101
+
102
+ [tool.ruff.lint]
103
+ ignore = [
104
+ "F403",
105
+ "F405",
106
+ "PLR0911",
107
+ "PLR0912",
108
+ "PLR0913",
109
+ "PLR0915",
110
+ "PLR2004",
111
+ "RUF005",
112
+ "RUF012"
113
+ ]
114
+ select = ["E", "W", "F", "I", "N", "PLC", "PLE", "PLR", "PLW", "RUF"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -1,16 +1,17 @@
1
1
  """AiiDA plugin that run Python function on remote computers."""
2
2
 
3
- __version__ = "0.3.1"
3
+ __version__ = "0.3.2"
4
4
 
5
- from node_graph import spec
5
+ from node_graph import socket_spec as spec
6
6
 
7
- from .calculations import PythonJob
7
+ from .calculations import PyFunction, PythonJob
8
8
  from .decorator import pyfunction
9
9
  from .launch import prepare_pythonjob_inputs
10
10
  from .parsers import PythonJobParser
11
11
 
12
12
  __all__ = (
13
13
  "PythonJob",
14
+ "PyFunction",
14
15
  "pyfunction",
15
16
  "PickledData",
16
17
  "prepare_pythonjob_inputs",
@@ -0,0 +1,4 @@
1
+ from .pyfunction import PyFunction
2
+ from .pythonjob import PythonJob
3
+
4
+ __all__ = ("PythonJob", "PyFunction")
@@ -17,12 +17,13 @@ from aiida.orm import (
17
17
  Str,
18
18
  to_aiida_type,
19
19
  )
20
+ from node_graph.socket_spec import SocketSpec
20
21
 
21
22
  __all__ = ("PyFunction",)
22
23
 
23
24
 
24
25
  class PyFunction(Process):
25
- """"""
26
+ """Run a Python function in-process, using SocketSpec for I/O."""
26
27
 
27
28
  _node_class = CalcFunctionNode
28
29
 
@@ -53,8 +54,8 @@ class PyFunction(Process):
53
54
  """Define the process specification, including its inputs, outputs and known exit codes."""
54
55
  super().define(spec)
55
56
  spec.input_namespace("function_data", dynamic=True, required=True)
56
- spec.input("function_data.output_ports", valid_type=Dict, serializer=to_aiida_type, required=False)
57
- spec.input("function_data.input_ports", valid_type=Dict, serializer=to_aiida_type, required=False)
57
+ spec.input("function_data.outputs_spec", valid_type=Dict, serializer=to_aiida_type, required=False)
58
+ spec.input("function_data.inputs_spec", valid_type=Dict, serializer=to_aiida_type, required=False)
58
59
  spec.input("process_label", valid_type=Str, serializer=to_aiida_type, required=False)
59
60
  spec.input_namespace("function_inputs", valid_type=Data, required=False)
60
61
  spec.input(
@@ -160,11 +161,13 @@ class PyFunction(Process):
160
161
  else:
161
162
  self.deserializers = None
162
163
 
164
+ # Build input namespace (raw Python) from AiiDA inputs using the declared SocketSpec
163
165
  inputs = dict(self.inputs.function_inputs or {})
164
166
  try:
167
+ inputs_spec = SocketSpec.from_dict(self.node.inputs.function_data.inputs_spec.get_dict())
165
168
  inputs = deserialize_ports(
166
169
  serialized_data=inputs,
167
- port_schema=self.node.inputs.function_data.input_ports.get_dict(),
170
+ port_schema=inputs_spec,
168
171
  deserializers=self.deserializers,
169
172
  )
170
173
  except Exception as exception:
@@ -174,6 +177,7 @@ class PyFunction(Process):
174
177
  exception=exception_message, traceback=traceback_str
175
178
  )
176
179
 
180
+ # Execute user function
177
181
  try:
178
182
  results = self.func(**inputs)
179
183
  except Exception as exception:
@@ -182,16 +186,18 @@ class PyFunction(Process):
182
186
  return self.exit_codes.ERROR_FUNCTION_EXECUTION_FAILED.format(
183
187
  exception=exception_message, traceback=traceback_str
184
188
  )
189
+
190
+ # Parse & output
185
191
  return self.parse(results)
186
192
 
187
193
  def parse(self, results):
188
- """Parse the results of the function."""
189
- from aiida_pythonjob.utils import parse_outputs
194
+ """Parse the results of the function and attach outputs."""
195
+ from aiida_pythonjob.parsers.utils import parse_outputs
190
196
 
191
- self.output_ports = self.node.inputs.function_data.output_ports.get_dict()
192
- exit_code = parse_outputs(
197
+ outputs_spec = SocketSpec.from_dict(self.node.inputs.function_data.outputs_spec.get_dict())
198
+ outputs, exit_code = parse_outputs(
193
199
  results,
194
- output_ports=self.output_ports,
200
+ output_spec=outputs_spec,
195
201
  exit_codes=self.exit_codes,
196
202
  logger=self.logger,
197
203
  serializers=self.serializers,
@@ -199,8 +205,7 @@ class PyFunction(Process):
199
205
  if exit_code:
200
206
  return exit_code
201
207
  # Store the outputs
202
- for name, port in self.output_ports["ports"].items():
203
- if "value" in port:
204
- self.out(name, port["value"])
208
+ for name, value in (outputs or {}).items():
209
+ self.out(name, value)
205
210
 
206
211
  return ExitCode()
@@ -45,8 +45,8 @@ class PythonJob(CalcJob):
45
45
  """Define the process specification, including its inputs, outputs and known exit codes."""
46
46
  super().define(spec)
47
47
  spec.input_namespace("function_data", dynamic=True, required=True)
48
- spec.input("function_data.input_ports", valid_type=Dict, serializer=to_aiida_type, required=False)
49
- spec.input("function_data.output_ports", valid_type=Dict, serializer=to_aiida_type, required=False)
48
+ spec.input("function_data.inputs_spec", valid_type=Dict, serializer=to_aiida_type, required=False)
49
+ spec.input("function_data.outputs_spec", valid_type=Dict, serializer=to_aiida_type, required=False)
50
50
  spec.input("process_label", valid_type=Str, serializer=to_aiida_type, required=False)
51
51
  spec.input_namespace("function_inputs", valid_type=Data, required=False)
52
52
  spec.input(
@@ -18,23 +18,32 @@ def atoms_to_structure_data(structure):
18
18
 
19
19
 
20
20
  def get_serializers_from_entry_points() -> dict:
21
- # Retrieve the entry points for 'aiida.data' and store them in a dictionary
22
- eps = entry_points()
21
+ """Retrieve the entry points for 'aiida.data' and store them in a dictionary."""
22
+ eps_all = entry_points()
23
23
  if sys.version_info >= (3, 10):
24
- group = eps.select(group="aiida.data")
24
+ group = eps_all.select(group="aiida.data")
25
25
  else:
26
- group = eps.get("aiida.data", [])
27
- eps = {}
28
- for ep in group:
26
+ group = eps_all.get("aiida.data", [])
27
+
28
+ # By converting the group to a set, we remove accidental duplicates
29
+ # where the same EntryPoint object is discovered twice. Legitimate
30
+ # competing entry points from different packages will remain.
31
+ unique_group = set(group)
32
+
33
+ serializers = {}
34
+ for ep in unique_group:
29
35
  # split the entry point name by first ".", and check the last part
30
36
  key = ep.name.split(".", 1)[-1]
37
+
31
38
  # skip key without "." because it is not a module name for a data type
32
39
  if "." not in key:
33
40
  continue
34
- eps.setdefault(key, [])
41
+
42
+ serializers.setdefault(key, [])
35
43
  # get the path of the entry point value and replace ":" with "."
36
- eps[key].append(ep.value.replace(":", "."))
37
- return eps
44
+ serializers[key].append(ep.value.replace(":", "."))
45
+
46
+ return serializers
38
47
 
39
48
 
40
49
  def get_serializers() -> dict:
@@ -5,11 +5,12 @@ import logging
5
5
  import signal
6
6
  import sys
7
7
  import typing as t
8
- from typing import Any, Mapping
8
+ from typing import List
9
9
 
10
10
  from aiida.engine.processes.functions import FunctionType, get_stack_size
11
11
  from aiida.manage import get_manager
12
12
  from aiida.orm import ProcessNode
13
+ from node_graph.socket_spec import SocketSpec
13
14
 
14
15
  from aiida_pythonjob.calculations.pyfunction import PyFunction
15
16
  from aiida_pythonjob.launch import create_inputs, prepare_pyfunction_inputs
@@ -19,8 +20,8 @@ LOGGER = logging.getLogger(__name__)
19
20
 
20
21
  # The following code is modified from the aiida-core.engine.processes.functions module
21
22
  def pyfunction(
22
- inputs: t.Optional[Mapping[str, Any]] = None,
23
- outputs: t.Optional[t.List[Mapping[str, Any]]] = None,
23
+ inputs: t.Optional[SocketSpec | List[str]] = None,
24
+ outputs: t.Optional[t.List[SocketSpec | List[str]]] = None,
24
25
  ) -> t.Callable[[FunctionType], FunctionType]:
25
26
  """The base function decorator to create a FunctionProcess out of a normal python function.
26
27
 
@@ -62,8 +63,6 @@ def pyfunction(
62
63
  # # Remove all the known inputs from the kwargs
63
64
  outputs_spec = kwargs.pop("outputs_spec", None) or outputs
64
65
  inputs_spec = kwargs.pop("inputs_spec", None) or inputs
65
- input_ports = kwargs.pop("input_ports", None)
66
- output_ports = kwargs.pop("output_ports", None)
67
66
  metadata = kwargs.pop("metadata", None)
68
67
  function_data = kwargs.pop("function_data", None)
69
68
  deserializers = kwargs.pop("deserializers", None)
@@ -77,8 +76,6 @@ def pyfunction(
77
76
  function_inputs=function_inputs,
78
77
  inputs_spec=inputs_spec,
79
78
  outputs_spec=outputs_spec,
80
- input_ports=input_ports,
81
- output_ports=output_ports,
82
79
  metadata=metadata,
83
80
  process_label=process_label,
84
81
  function_data=function_data,
@@ -5,9 +5,9 @@ import os
5
5
  from typing import Any, Callable, Dict, Optional, Union
6
6
 
7
7
  from aiida import orm
8
- from node_graph.nodes.utils import generate_input_sockets, generate_output_sockets
8
+ from node_graph.node_spec import BaseHandle
9
+ from node_graph.socket_spec import infer_specs_from_callable
9
10
 
10
- from .ports_adapter import inputs_sockets_to_ports, outputs_sockets_to_ports
11
11
  from .utils import build_function_data, get_or_create_code, serialize_ports
12
12
 
13
13
 
@@ -16,8 +16,6 @@ def prepare_pythonjob_inputs(
16
16
  function_inputs: Optional[Dict[str, Any]] = None,
17
17
  inputs_spec: Optional[type] = None,
18
18
  outputs_spec: Optional[type] = None,
19
- output_ports: Optional[Dict[str, Any]] = None,
20
- input_ports: Optional[Dict[str, Any]] = None,
21
19
  code: Optional[orm.AbstractCode] = None,
22
20
  command_info: Optional[Dict[str, str]] = None,
23
21
  computer: Union[str, orm.Computer] = "localhost",
@@ -36,6 +34,8 @@ def prepare_pythonjob_inputs(
36
34
  raise ValueError("Either function or function_data must be provided")
37
35
  if function is not None and function_data is not None:
38
36
  raise ValueError("Only one of function or function_data should be provided")
37
+ if isinstance(function, BaseHandle):
38
+ function = function._func
39
39
  # if function is a function, inspect it and get the source code
40
40
  if function is not None and inspect.isfunction(function):
41
41
  function_data = build_function_data(function, register_pickle_by_value=register_pickle_by_value)
@@ -60,20 +60,12 @@ def prepare_pythonjob_inputs(
60
60
  if code is None:
61
61
  command_info = command_info or {}
62
62
  code = get_or_create_code(computer=computer, **command_info)
63
- # outputs
64
- if not output_ports:
65
- node_outputs = generate_output_sockets(function or (lambda **_: None), outputs=outputs_spec)
66
- output_ports = outputs_sockets_to_ports(node_outputs)
67
- # inputs
68
- if not input_ports:
69
- node_inputs = generate_input_sockets(function or (lambda **_: None), inputs=inputs_spec)
70
- input_ports = inputs_sockets_to_ports(node_inputs)
71
-
72
- function_data["output_ports"] = output_ports
73
- function_data["input_ports"] = input_ports
63
+ in_spec, out_spec = infer_specs_from_callable(function, inputs=inputs_spec, outputs=outputs_spec)
64
+ function_data["inputs_spec"] = in_spec.to_dict()
65
+ function_data["outputs_spec"] = out_spec.to_dict()
74
66
  # serialize kwargs against the (nested) input schema
75
67
  function_inputs = function_inputs or {}
76
- function_inputs = serialize_ports(python_data=function_inputs, port_schema=input_ports, serializers=serializers)
68
+ function_inputs = serialize_ports(python_data=function_inputs, port_schema=in_spec, serializers=serializers)
77
69
  # replace "." with "__dot__" in the keys of a dictionary
78
70
  if deserializers:
79
71
  deserializers = orm.Dict({k.replace(".", "__dot__"): v for k, v in deserializers.items()})
@@ -116,8 +108,6 @@ def prepare_pyfunction_inputs(
116
108
  function_inputs: Optional[Dict[str, Any]] = None,
117
109
  inputs_spec: Optional[type] = None,
118
110
  outputs_spec: Optional[type] = None,
119
- output_ports: Optional[Dict[str, Any]] = None,
120
- input_ports: Optional[Dict[str, Any]] = None,
121
111
  metadata: Optional[Dict[str, Any]] = None,
122
112
  process_label: Optional[str] = None,
123
113
  function_data: dict | None = None,
@@ -126,13 +116,15 @@ def prepare_pyfunction_inputs(
126
116
  register_pickle_by_value: bool = False,
127
117
  **kwargs: Any,
128
118
  ) -> Dict[str, Any]:
129
- """Prepare the inputs for PythonJob"""
119
+ """Prepare the inputs for PyFunction."""
130
120
  import types
131
121
 
132
122
  if function is None and function_data is None:
133
123
  raise ValueError("Either function or function_data must be provided")
134
124
  if function is not None and function_data is not None:
135
125
  raise ValueError("Only one of function or function_data should be provided")
126
+ if isinstance(function, BaseHandle):
127
+ function = function._func
136
128
  # if function is a function, inspect it and get the source code
137
129
  if function is not None:
138
130
  if inspect.isfunction(function):
@@ -141,20 +133,13 @@ def prepare_pyfunction_inputs(
141
133
  raise NotImplementedError("Built-in functions are not supported yet")
142
134
  else:
143
135
  raise ValueError("Invalid function type")
144
- # outputs
145
- if not output_ports:
146
- node_outputs = generate_output_sockets(function or (lambda **_: None), outputs=outputs_spec)
147
- output_ports = outputs_sockets_to_ports(node_outputs)
148
- # inputs
149
- if not input_ports:
150
- node_inputs = generate_input_sockets(function or (lambda **_: None), inputs=inputs_spec)
151
- input_ports = inputs_sockets_to_ports(node_inputs)
152
-
153
- function_data["output_ports"] = output_ports
154
- function_data["input_ports"] = input_ports
136
+ # spec
137
+ in_spec, out_spec = infer_specs_from_callable(function, inputs=inputs_spec, outputs=outputs_spec)
138
+ function_data["inputs_spec"] = in_spec.to_dict()
139
+ function_data["outputs_spec"] = out_spec.to_dict()
155
140
  # serialize the kwargs into AiiDA Data
156
141
  function_inputs = function_inputs or {}
157
- function_inputs = serialize_ports(python_data=function_inputs, port_schema=input_ports, serializers=serializers)
142
+ function_inputs = serialize_ports(python_data=function_inputs, port_schema=in_spec, serializers=serializers)
158
143
  # replace "." with "__dot__" in the keys of a dictionary
159
144
  if deserializers:
160
145
  deserializers = orm.Dict({k.replace(".", "__dot__"): v for k, v in deserializers.items()})
@@ -3,8 +3,9 @@
3
3
  import json
4
4
 
5
5
  from aiida.parsers.parser import Parser
6
+ from node_graph.socket_spec import SocketSpec
6
7
 
7
- from aiida_pythonjob.utils import parse_outputs
8
+ from .utils import parse_outputs
8
9
 
9
10
  # Map error_type from script.py to exit code label
10
11
  ERROR_TYPE_TO_EXIT_CODE = {
@@ -22,8 +23,9 @@ class PythonJobParser(Parser):
22
23
  def parse(self, **kwargs):
23
24
  import pickle
24
25
 
25
- # Read output_ports specification
26
- self.output_ports = self.node.inputs.function_data.output_ports.get_dict()
26
+ # Read outputs SocketSpec
27
+ spec_dict = self.node.inputs.function_data.outputs_spec.get_dict()
28
+ self.outputs_spec = SocketSpec.from_dict(spec_dict)
27
29
 
28
30
  # load custom serializers
29
31
  if "serializers" in self.node.inputs and self.node.inputs.serializers:
@@ -53,7 +55,7 @@ class PythonJobParser(Parser):
53
55
  pass
54
56
  except json.JSONDecodeError as exc:
55
57
  self.logger.error(f"Error reading _error.json: {exc}")
56
- return self.exit_codes.ERROR_INVALID_OUTPUT # or a different exit code
58
+ return self.exit_codes.ERROR_INVALID_OUTPUT
57
59
 
58
60
  # 2) If we reach here, _error.json exists but is empty or doesn't exist at all -> no error recorded
59
61
  # Proceed with parsing results.pickle
@@ -61,9 +63,9 @@ class PythonJobParser(Parser):
61
63
  with self.retrieved.base.repository.open("results.pickle", "rb") as handle:
62
64
  results = pickle.load(handle)
63
65
 
64
- exit_code = parse_outputs(
66
+ outputs, exit_code = parse_outputs(
65
67
  results,
66
- output_ports=self.output_ports,
68
+ output_spec=self.outputs_spec,
67
69
  exit_codes=self.exit_codes,
68
70
  logger=self.logger,
69
71
  serializers=self.serializers,
@@ -72,9 +74,8 @@ class PythonJobParser(Parser):
72
74
  return exit_code
73
75
 
74
76
  # Store the outputs
75
- for name, port in self.output_ports["ports"].items():
76
- if "value" in port:
77
- self.out(name, port["value"])
77
+ for name, value in (outputs or {}).items():
78
+ self.out(name, value)
78
79
 
79
80
  except OSError:
80
81
  return self.exit_codes.ERROR_READING_OUTPUT_FILE