jobflow 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

jobflow/__init__.py

@@ -1,7 +1,7 @@
 """Jobflow is a package for writing dynamic and connected workflows."""
 
 from jobflow._version import __version__
-from jobflow.core.flow import Flow, JobOrder
+from jobflow.core.flow import Flow, JobOrder, flow
 from jobflow.core.job import Job, JobConfig, Response, job
 from jobflow.core.maker import Maker
 from jobflow.core.reference import OnMissing, OutputReference

jobflow/core/flow.py

@@ -4,6 +4,8 @@ from __future__ import annotations
 
 import logging
 import warnings
+from contextlib import contextmanager
+from contextvars import ContextVar
 from copy import deepcopy
 from typing import TYPE_CHECKING
 
@@ -155,6 +157,12 @@ class Flow(MSONable):
         self.add_jobs(jobs)
         self.output = output
 
+        # If we're running inside a `DecoratedFlow`, add *this* Flow to the
+        # context.
+        current_flow_children_list = _current_flow_context.get()
+        if current_flow_children_list is not None:
+            current_flow_children_list.append(self)
+
     def __len__(self) -> int:
         """Get the number of jobs or subflows in the flow."""
         return len(self.jobs)
@@ -828,7 +836,7 @@ class Flow(MSONable):
             if job.host is not None and job.host != self.uuid:
                 raise ValueError(
                     f"{type(job).__name__} {job.name} ({job.uuid}) already belongs "
-                    f"to another flow."
+                    f"to another flow: {job.host}."
                 )
             if job.uuid in job_ids:
                 raise ValueError(
@@ -921,3 +929,104 @@ def get_flow(
             )
 
     return flow
+
+
+class DecoratedFlow(Flow):
+    """A DecoratedFlow is a Flow that is returned on using the @flow decorator."""
+
+    def __init__(self, fn, *args, **kwargs):
+        from jobflow import Maker
+
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+
+        # Collect the jobs and flows that are used in the function
+        children_list = []
+        with flow_build_context(children_list):
+            output = self.fn(*self.args, **self.kwargs)
+
+        # From the collected items, remove those that have already been assigned
+        # to another Flow during the call of the function.
+        # This handles the case where a Flow object is instantiated inside
+        # the decorated function
+        children_list = [c for c in children_list if c.host is None]
+
+        name = getattr(self.fn, "__qualname__", self.fn.__name__)
+
+        # if decorates a make() in a Maker use that as a name
+        if (
+            len(self.args) > 0
+            and name.split(".")[-1] == "make"
+            and getattr(args[0], self.fn.__name__, None)
+            and isinstance(args[0], Maker)
+        ):
+            name = args[0].name
+
+        if isinstance(output, (jobflow.Job, jobflow.Flow)):
+            warnings.warn(
+                f"@flow decorated function '{name}' contains a Flow or"
+                f"Job as an output. Usually the output should be the output of"
+                f"a Job or another Flow (e.g. job.output). Replacing the"
+                f"output of the @flow with the output of the Flow/Job."
+                f"If this message is unexpected then double check the outputs"
+                f"of your @flow decorated function.",
+                stacklevel=2,
+            )
+            output = output.output
+
+        super().__init__(name=name, jobs=children_list, output=output)
+
+
+def flow(fn):
+    """
+    Turn a function into a DecoratedFlow object.
+
+    Parameters
+    ----------
+        fn (Callable): The function to be wrapped in a DecoratedFlow object.
+
+    Returns
+    -------
+        Callable: A wrapper function that, when called, creates and returns
+        an instance of DecoratedFlow initialized with the provided function
+        and its arguments.
+    """
+    from functools import wraps
+
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        return DecoratedFlow(fn, *args, **kwargs)
+
+    return wrapper
+
+
+@contextmanager
+def flow_build_context(children_list):
+    """Provide a context manager for flows.
+
+    Provides a context manager for setting and resetting the `Job` and `Flow`
+    objects in the current flow context.
+
+    Parameters
+    ----------
+        children_list: The `Job` or `Flow` objects that are part of the current
+        flow context.
+
+    Yields
+    ------
+        None: Temporarily sets the provided `Job` or `Flow` objects as
+        belonging to the current flow context within the managed block.
+
+    Raises
+    ------
+        None
+    """
+    token = _current_flow_context.set(children_list)
+    try:
+        yield
+    finally:
+        _current_flow_context.reset(token)
+
+
+_current_flow_context = ContextVar("current_flow_context", default=None)

jobflow/core/job.py

@@ -11,6 +11,7 @@ from typing import cast, overload
 from monty.json import MSONable, jsanitize
 from typing_extensions import Self
 
+from jobflow.core.flow import _current_flow_context
 from jobflow.core.reference import OnMissing, OutputReference
 from jobflow.utils.uid import suid
 
@@ -384,6 +385,30 @@ class Job(MSONable):
                 stacklevel=2,
             )
 
+        # If we're running inside a `DecoratedFlow`, add *this* Job to the
+        # context.
+        current_flow_children_list = _current_flow_context.get()
+        if current_flow_children_list is not None:
+            current_flow_children_list.append(self)
+
+    def __getitem__(self, key: Any) -> OutputReference:
+        """
+        Get the corresponding `OutputReference` for the `Job`.
+
+        This is for when it is indexed like a dictionary or list.
+
+        Parameters
+        ----------
+        key
+            The index/key.
+
+        Returns
+        -------
+        OutputReference
+            The equivalent of `Job.output[k]`
+        """
+        return self.output[key]
+
     def __repr__(self):
         """Get a string representation of the job."""
         name, uuid = self.name, self.uuid
@@ -1279,6 +1304,21 @@ class Response(typing.Generic[T]):
         Response
             The job response controlling the data to store and flow execution options.
         """
+        # If the Job returns another Job, or something that can be interpreted
+        # as an iterable of jobs, interpret it as a replace.
+        from jobflow import Flow
+
+        def is_job_or_flow(x):
+            return isinstance(x, Job | Flow)
+
+        should_replace = is_job_or_flow(job_returns)
+        if isinstance(job_returns, (list, tuple)):
+            should_replace = all(is_job_or_flow(resp) for resp in job_returns)
+
+        if should_replace:
+            job_returns = Response(replace=job_returns)
+
+        # only apply output schema if there is no replace.
         if isinstance(job_returns, Response):
             if job_returns.replace is None:
                 # only apply output schema if there is no replace.

jobflow/core/reference.py

@@ -514,12 +514,47 @@ def validate_schema_access(
         The BaseModel class associated with the item, if any.
     """
     schema_dict = schema.model_json_schema()
-    if item not in schema_dict["properties"]:
+    item_in_schema = item in schema_dict["properties"]
+    property_like = isinstance(item, str) and has_property_like(schema, item)
+    if not item_in_schema and not property_like:
         raise AttributeError(f"{schema.__name__} does not have attribute '{item}'.")
 
     subschema = None
-    item_type = schema.model_fields[item].annotation
-    if lenient_issubclass(item_type, BaseModel):
-        subschema = item_type
+    if item_in_schema:
+        item_type = schema.model_fields[item].annotation
+        if lenient_issubclass(item_type, BaseModel):
+            subschema = item_type
 
     return True, subschema
+
+
+def has_property_like(obj_type: type, name: str) -> bool:
+    """
+    Check if a class has an attribute and if it is property-like.
+
+    Parameters
+    ----------
+    obj_type
+        The class that needs to be checked
+    name
+        The name of the attribute to be verified
+
+    Returns
+    -------
+    bool
+        True if the property corresponding to the name is property-like.
+    """
+    if not hasattr(obj_type, name):
+        return False
+
+    attr = getattr(obj_type, name)
+
+    if isinstance(attr, property):
+        return True
+
+    if callable(attr):
+        return False
+
+    # Check for custom property-like descriptors with __get__ but not callable
+    # If not, is not property-like.
+    return hasattr(attr, "__get__")

jobflow/core/store.py

@@ -21,9 +21,17 @@ if typing.TYPE_CHECKING:
 
     from jobflow.core.schemas import JobStoreDocument
 
-    obj_type = Union[str, Enum, type[MSONable], list[Union[Enum, str, type[MSONable]]]]
-    save_type = Optional[dict[str, obj_type]]
-    load_type = Union[bool, dict[str, Union[bool, obj_type]]]
+    obj_save_type = Union[
+        str,
+        Enum,
+        type[MSONable],
+        list[Union[Enum, str, type[MSONable], list[Union[str, type[MSONable]]]]],
+    ]
+    save_type = Optional[dict[str, obj_save_type]]
+    obj_load_type = Union[
+        str, Enum, type[MSONable], list[Union[Enum, str, type[MSONable]]]
+    ]
+    load_type = Union[bool, dict[str, Union[bool, obj_load_type]]]
 
 
 T = typing.TypeVar("T", bound="JobStore")
@@ -43,6 +51,9 @@ class JobStore(Store):
         Which items to save in additional stores when uploading documents. Given as a
         mapping of ``{store name: store_type}`` where ``store_type`` can be a dictionary
         key (string or enum), an :obj:`.MSONable` class, or a list of keys/classes.
+        If the list of keys/classes itself contains a list/tuple, this will be treated
+        as the path to the save key. Can be used if a key is duplicate in the output
+        and only a single occurrence shall be put in the additional store.
     load
         Which items to load from additional stores when querying documents. Given as a
         mapping of ``{store name: store_type}`` where ``store_type`` can be `True``, in
@@ -304,7 +315,10 @@ class JobStore(Store):
                 locations = []
                 for store_name, store_save in save_keys.items():
                     for save_key in store_save:
-                        locations.extend(find_key(doc, save_key, include_end=True))
+                        if isinstance(save_key, (list, tuple)):
+                            locations.append(list(save_key))
+                        else:
+                            locations.extend(find_key(doc, save_key, include_end=True))
 
                     locations = get_root_locations(locations)
                     objects = [get(doc, list(loc)) for loc in locations]
@@ -726,7 +740,7 @@ def _prepare_load(
 
 def _prepare_save(
     save: bool | save_type,
-) -> dict[str, list[str | type[MSONable]]]:
+) -> dict[str, list[str | type[MSONable] | list[str | type[MSONable]]]]:
     """Standardize save type."""
     from enum import Enum
 

{jobflow-0.2.1.dist-info → jobflow-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jobflow
-Version: 0.2.1
+Version: 0.3.0
 Summary: jobflow is a library for writing computational workflows
 Author-email: Alex Ganose <a.ganose@imperial.ac.uk>
 License: modified BSD
@@ -36,35 +36,35 @@ Provides-Extra: ulid
 Requires-Dist: python-ulid; extra == "ulid"
 Provides-Extra: docs
 Requires-Dist: autodoc_pydantic==2.2.0; extra == "docs"
-Requires-Dist: furo==2025.9.25; extra == "docs"
-Requires-Dist: ipython==9.3.0; extra == "docs"
+Requires-Dist: furo==2025.12.19; extra == "docs"
+Requires-Dist: ipython==9.9.0; extra == "docs"
 Requires-Dist: myst_parser==4.0.1; extra == "docs"
-Requires-Dist: nbsphinx==0.9.7; extra == "docs"
+Requires-Dist: nbsphinx==0.9.8; extra == "docs"
 Requires-Dist: sphinx-copybutton==0.5.2; extra == "docs"
 Requires-Dist: sphinx==8.1.3; extra == "docs"
 Provides-Extra: dev
 Requires-Dist: pre-commit>=2.12.1; extra == "dev"
 Requires-Dist: typing_extensions; python_version < "3.11" and extra == "dev"
 Provides-Extra: tests
-Requires-Dist: moto==5.1.16; extra == "tests"
+Requires-Dist: moto==5.1.20; extra == "tests"
 Requires-Dist: pytest-cov==7.0.0; extra == "tests"
-Requires-Dist: pytest==8.4.2; extra == "tests"
+Requires-Dist: pytest==9.0.2; extra == "tests"
 Provides-Extra: vis
 Requires-Dist: matplotlib; extra == "vis"
 Requires-Dist: pydot; extra == "vis"
 Provides-Extra: fireworks
 Requires-Dist: FireWorks; extra == "fireworks"
 Provides-Extra: strict
-Requires-Dist: FireWorks==2.0.6; extra == "strict"
+Requires-Dist: FireWorks==2.0.8; extra == "strict"
 Requires-Dist: PyYAML==6.0.3; extra == "strict"
 Requires-Dist: maggma==0.72.0; extra == "strict"
-Requires-Dist: matplotlib==3.10.7; extra == "strict"
+Requires-Dist: matplotlib==3.10.8; extra == "strict"
 Requires-Dist: monty==2025.3.3; extra == "strict"
-Requires-Dist: moto==5.1.16; extra == "strict"
+Requires-Dist: moto==5.1.20; extra == "strict"
 Requires-Dist: networkx==3.4.2; extra == "strict"
-Requires-Dist: pydantic-settings==2.11.0; extra == "strict"
-Requires-Dist: pydantic==2.12.3; extra == "strict"
-Requires-Dist: pydash==8.0.5; extra == "strict"
+Requires-Dist: pydantic-settings==2.12.0; extra == "strict"
+Requires-Dist: pydantic==2.12.5; extra == "strict"
+Requires-Dist: pydash==8.0.6; extra == "strict"
 Requires-Dist: pydot==4.0.1; extra == "strict"
 Requires-Dist: python-ulid==3.1.0; extra == "strict"
 Requires-Dist: typing-extensions==4.15.0; extra == "strict"

{jobflow-0.2.1.dist-info → jobflow-0.3.0.dist-info}/RECORD

@@ -1,15 +1,15 @@
-jobflow/__init__.py,sha256=l7o10BaqEQWw5aZziWRg40PsIAgQ4lrlluXs9hIv2mg,570
+jobflow/__init__.py,sha256=gKxdw3Yi4gjq6218fMBTHSBnVq-kZ4uZ9fLoKbC1xr0,576
 jobflow/_version.py,sha256=Ym07PBD7sAmpqVpX8tuzWma3P_Hv6KXbDKXWkw8OwaI,205
 jobflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jobflow/settings.py,sha256=yinHpix-DPwzcBhQCO8zDXFuv048rmBgpYMPa9wcs8c,6094
 jobflow/core/__init__.py,sha256=3sx5t1Gysejc4c_fPrhvCjPUg0p_384Zko8ms2c_NnY,98
-jobflow/core/flow.py,sha256=AWseNK4P42CiiPykl4kud0e-SJYxDJoqM6I_n77w4_M,30854
-jobflow/core/job.py,sha256=ogKour02EF6uvSYZz4F4zaqiM6bUiaCwwtboF0WprPk,49002
+jobflow/core/flow.py,sha256=ernJoGkts5hi3ruReJHDqioB5N2wpl6zjiwCQqhqGYM,34337
+jobflow/core/job.py,sha256=9Tqv4u9Fc_3Ou1qBizHLvrgLB_WjVKcyjodlvMAnAMg,50319
 jobflow/core/maker.py,sha256=WhsYw2wDNVIyAEeRUoikOQMzXzHuXfFVwXrJpwGCD1E,11162
-jobflow/core/reference.py,sha256=2xqNwyo57wFPLMLPprXa5h8f_cOWcyCHmDsfVEl5pRk,17101
+jobflow/core/reference.py,sha256=ZATBhaB9HDhPujsfscOcHw7nJf7QfKvozvII47_q1UU,18012
 jobflow/core/schemas.py,sha256=Oi5-PnZpI8S9jSY7Q4f8H7xUybbRZDXlgugeVewVsrA,968
 jobflow/core/state.py,sha256=IGJTtmpotDKEcgDEnsT5x20ZeyvQT68Mr3teTjkgYnM,709
-jobflow/core/store.py,sha256=PAtAhZTLU1b17OD66rqlo8_JP1p-IzH8c2WJwWd1Q2o,26997
+jobflow/core/store.py,sha256=njeKh_F-sHGt-hk0lsETmgK9H1I7hAqCeJY4LzStAxo,27627
 jobflow/managers/__init__.py,sha256=KkA5cVDe2os2_2aTa8eiB9SnkGLZNybcci-Lo4tbaWM,55
 jobflow/managers/fireworks.py,sha256=xKPcL1BX59jnF6LpyEwcsxvCCShLqep-9UHRyv9FmYE,7146
 jobflow/managers/local.py,sha256=iI1QI6Dg7P7C3vVd4Mt1xXAmcQ-o37Ox_gQRnafsBHE,6142
@@ -21,8 +21,8 @@ jobflow/utils/graph.py,sha256=kweAowEzv8ZGjJ9KZvJ-G5ueAqGPWbU0z7Xd4e-q8no,6596
 jobflow/utils/log.py,sha256=4-_1OUSQ8I4Q6CgQ4pxNBeveBdpXla7nibZNF7Vk3pw,1110
 jobflow/utils/uid.py,sha256=hNkpJ5AYhKd_sPWE_iGPcn6YD_AyizKX_swWksFr-_M,2537
 jobflow/utils/uuid.py,sha256=m0fInOs1yklpavf7jId85luDOHjZqfIIUzEtd89Bk_s,440
-jobflow-0.2.1.dist-info/licenses/LICENSE,sha256=jUEiENfZNQZh9RE9ixtUWgVkLRD85ScZ6iv1WREf19w,2418
-jobflow-0.2.1.dist-info/METADATA,sha256=4B39RQEMtBLt7_K2rghR9VfZ4GfxgFHIig9iixQRdaw,10039
-jobflow-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jobflow-0.2.1.dist-info/top_level.txt,sha256=IanNooU88OupQPDrWnT0rbL3E27P2wEy7Jsfx9_j8zc,8
-jobflow-0.2.1.dist-info/RECORD,,
+jobflow-0.3.0.dist-info/licenses/LICENSE,sha256=jUEiENfZNQZh9RE9ixtUWgVkLRD85ScZ6iv1WREf19w,2418
+jobflow-0.3.0.dist-info/METADATA,sha256=5lPrcdt01_rFBFjVSCxQpbHvH7N9JQwSOaODeG-Z6Po,10040
+jobflow-0.3.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+jobflow-0.3.0.dist-info/top_level.txt,sha256=IanNooU88OupQPDrWnT0rbL3E27P2wEy7Jsfx9_j8zc,8
+jobflow-0.3.0.dist-info/RECORD,,

{jobflow-0.2.1.dist-info → jobflow-0.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any