jobflow 0.1.13__tar.gz → 0.1.15__tar.gz

{jobflow-0.1.13/src/jobflow.egg-info → jobflow-0.1.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jobflow
-Version: 0.1.13
+Version: 0.1.15
 Summary: jobflow is a library for writing computational workflows
 Author-email: Alex Ganose <alexganose@gmail.com>
 License: modified BSD
@@ -9,34 +9,67 @@ Project-URL: repository, https://github.com/materialsproject/jobflow
 Project-URL: documentation, https://materialsproject.github.io/jobflow/
 Project-URL: changelog, https://github.com/materialsproject/jobflow/blob/main/CHANGELOG.md
 Keywords: high-throughput,workflow
-Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Information Technology
 Classifier: Intended Audience :: Science/Research
 Classifier: Intended Audience :: System Administrators
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Database :: Front-Ends
 Classifier: Topic :: Other/Nonlisted Topic
 Classifier: Topic :: Scientific/Engineering
-Requires-Python: >=3.8
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: PyYAML
+Requires-Dist: maggma>=0.57.0
+Requires-Dist: monty>=2023.9.25
+Requires-Dist: networkx
+Requires-Dist: pydantic-settings>=2.0.3
+Requires-Dist: pydantic>=2.0.1
+Requires-Dist: pydash
 Provides-Extra: docs
+Requires-Dist: autodoc_pydantic==2.0.1; extra == "docs"
+Requires-Dist: furo==2023.9.10; extra == "docs"
+Requires-Dist: ipython==8.18.1; extra == "docs"
+Requires-Dist: myst_parser==2.0.0; extra == "docs"
+Requires-Dist: nbsphinx==0.9.3; extra == "docs"
+Requires-Dist: sphinx-copybutton==0.5.2; extra == "docs"
+Requires-Dist: sphinx==7.2.6; extra == "docs"
 Provides-Extra: dev
+Requires-Dist: pre-commit>=2.12.1; extra == "dev"
 Provides-Extra: tests
+Requires-Dist: moto==4.2.11; extra == "tests"
+Requires-Dist: pytest-cov==4.1.0; extra == "tests"
+Requires-Dist: pytest==7.4.3; 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
-License-File: LICENSE
+Requires-Dist: FireWorks==2.0.3; extra == "strict"
+Requires-Dist: PyYAML==6.0.1; extra == "strict"
+Requires-Dist: maggma==0.58.0; extra == "strict"
+Requires-Dist: matplotlib==3.8.2; extra == "strict"
+Requires-Dist: monty==2023.11.3; extra == "strict"
+Requires-Dist: moto==4.2.11; extra == "strict"
+Requires-Dist: networkx==3.2.1; extra == "strict"
+Requires-Dist: pydantic-settings==2.1.0; extra == "strict"
+Requires-Dist: pydantic==2.5.2; extra == "strict"
+Requires-Dist: pydash==7.0.6; extra == "strict"
+Requires-Dist: pydot==1.4.2; extra == "strict"
+Requires-Dist: typing-extensions==4.8.0; extra == "strict"
 
 # jobflow
 
-<a href="https://github.com/materialsproject/jobflow/actions?query=workflow%3Atesting"><img alt="code coverage" src="https://img.shields.io/github/actions/workflow/status/materialsproject/jobflow/testing.yml?branch=main&label=tests"></a>
-<a href="https://codecov.io/gh/materialsproject/jobflow/"><img alt="code coverage" src="https://img.shields.io/codecov/c/gh/materialsproject/jobflow/main"></a>
-<a href="https://pypi.org/project/jobflow"><img alt="pypi version" src="https://img.shields.io/pypi/v/jobflow?color=blue"></a>
-<img alt="supported python versions" src="https://img.shields.io/pypi/pyversions/jobflow">
+[![tests](https://img.shields.io/github/actions/workflow/status/materialsproject/jobflow/testing.yml?branch=main&label=tests)](https://github.com/materialsproject/jobflow/actions?query=workflow%3Atesting)
+[![code coverage](https://img.shields.io/codecov/c/gh/materialsproject/jobflow/main)](https://codecov.io/gh/materialsproject/jobflow/)
+[![pypi version](https://img.shields.io/pypi/v/jobflow?color=blue)](https://pypi.org/project/jobflow/)
+![supported python versions](https://img.shields.io/pypi/pyversions/jobflow)
 
 [Documentation](https://materialsproject.github.io/jobflow/) | [PyPI](https://pypi.org/project/jobflow/) | [GitHub](https://github.com/materialsproject/jobflow)
 
@@ -109,7 +142,7 @@ the jobs is determined automatically and can be visualised using the flow graph.
 
 ## Installation
 
-The jobflow is a Python 3.8+ library and can be installed using pip.
+`jobflow` is a Python 3.9+ library and can be installed using `pip`.
 
 ```bash
 pip install jobflow

{jobflow-0.1.13 → jobflow-0.1.15}/README.md

@@ -1,9 +1,9 @@
 # jobflow
 
-<a href="https://github.com/materialsproject/jobflow/actions?query=workflow%3Atesting"><img alt="code coverage" src="https://img.shields.io/github/actions/workflow/status/materialsproject/jobflow/testing.yml?branch=main&label=tests"></a>
-<a href="https://codecov.io/gh/materialsproject/jobflow/"><img alt="code coverage" src="https://img.shields.io/codecov/c/gh/materialsproject/jobflow/main"></a>
-<a href="https://pypi.org/project/jobflow"><img alt="pypi version" src="https://img.shields.io/pypi/v/jobflow?color=blue"></a>
-<img alt="supported python versions" src="https://img.shields.io/pypi/pyversions/jobflow">
+[![tests](https://img.shields.io/github/actions/workflow/status/materialsproject/jobflow/testing.yml?branch=main&label=tests)](https://github.com/materialsproject/jobflow/actions?query=workflow%3Atesting)
+[![code coverage](https://img.shields.io/codecov/c/gh/materialsproject/jobflow/main)](https://codecov.io/gh/materialsproject/jobflow/)
+[![pypi version](https://img.shields.io/pypi/v/jobflow?color=blue)](https://pypi.org/project/jobflow/)
+![supported python versions](https://img.shields.io/pypi/pyversions/jobflow)
 
 [Documentation](https://materialsproject.github.io/jobflow/) | [PyPI](https://pypi.org/project/jobflow/) | [GitHub](https://github.com/materialsproject/jobflow)
 
@@ -76,7 +76,7 @@ the jobs is determined automatically and can be visualised using the flow graph.
 
 ## Installation
 
-The jobflow is a Python 3.8+ library and can be installed using pip.
+`jobflow` is a Python 3.9+ library and can be installed using `pip`.
 
 ```bash
 pip install jobflow

{jobflow-0.1.13 → jobflow-0.1.15}/pyproject.toml

@@ -11,55 +11,57 @@ license = { text = "modified BSD" }
 authors = [{ name = "Alex Ganose", email = "alexganose@gmail.com" }]
 dynamic = ["version"]
 classifiers = [
-    "Development Status :: 2 - Pre-Alpha",
+    "Development Status :: 5 - Production/Stable",
     "Intended Audience :: Information Technology",
     "Intended Audience :: Science/Research",
     "Intended Audience :: System Administrators",
     "Operating System :: OS Independent",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.9",
     "Topic :: Database :: Front-Ends",
     "Topic :: Other/Nonlisted Topic",
     "Topic :: Scientific/Engineering",
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 dependencies = [
     "PyYAML",
-    "maggma>=0.38.1",
-    "monty>=2021.5.9",
+    "maggma>=0.57.0",
+    "monty>=2023.9.25",
     "networkx",
-    "pydantic",
+    "pydantic-settings>=2.0.3",
+    "pydantic>=2.0.1",
     "pydash",
 ]
 
 [project.optional-dependencies]
 docs = [
-    "autodoc_pydantic==1.9.0",
-    "furo==2023.7.26",
-    "ipython==8.14.0",
+    "autodoc_pydantic==2.0.1",
+    "furo==2023.9.10",
+    "ipython==8.18.1",
     "myst_parser==2.0.0",
-    "nbsphinx==0.9.2",
+    "nbsphinx==0.9.3",
     "sphinx-copybutton==0.5.2",
-    "sphinx==7.1.2",
+    "sphinx==7.2.6",
 ]
 dev = ["pre-commit>=2.12.1"]
-tests = ["pytest-cov==4.1.0", "pytest==7.4.0"]
+tests = ["moto==4.2.11", "pytest-cov==4.1.0", "pytest==7.4.3"]
 vis = ["matplotlib", "pydot"]
 fireworks = ["FireWorks"]
 strict = [
     "FireWorks==2.0.3",
     "PyYAML==6.0.1",
-    "maggma==0.53.0",
-    "matplotlib==3.7.2",
-    "monty==2023.8.8",
-    "moto==4.1.14",
-    "networkx==3.1",
-    "pydantic==1.10.9",
+    "maggma==0.58.0",
+    "matplotlib==3.8.2",
+    "monty==2023.11.3",
+    "moto==4.2.11",
+    "networkx==3.2.1",
+    "pydantic-settings==2.1.0",
+    "pydantic==2.5.2",
     "pydash==7.0.6",
     "pydot==1.4.2",
-    "typing-extensions==4.7.1",
+    "typing-extensions==4.8.0",
 ]
 
 [project.urls]
@@ -80,7 +82,7 @@ max-line-length = 88
 max-doc-length = 88
 select = "C, E, F, W, B"
 extend-ignore = "E203, W503, E501, F401, RST21"
-min-python-version = "3.8.0"
+min-python-version = "3.9.0"
 docstring-convention = "numpy"
 rst-roles = "class, func, ref, obj"
 
@@ -116,29 +118,55 @@ exclude_lines = [
 ]
 
 [tool.ruff]
-target-version = "py38"
+target-version = "py39"
 ignore-init-module-imports = true
 select = [
-    "B",   # flake8-bugbear
-    "C4",  # flake8-comprehensions
-    "D",   # pydocstyle
-    "E",   # pycodestyle
-    "F",   # pyflakes
-    "I",   # isort
-    "PLE", # pylint error
-    "PLW", # pylint warning
-    "Q",   # flake8-quotes
-    "RUF", # Ruff-specific rules
-    "SIM", # flake8-simplify
-    "TID", # tidy imports
-    "UP",  # pyupgrade
-    "W",   # pycodestyle
-    "YTT", # flake8-2020
+    "B",      # flake8-bugbear
+    "C4",     # flake8-comprehensions
+    "D",      # pydocstyle
+    "E",      # pycodestyle error
+    "EXE",    # flake8-executable
+    "F",      # pyflakes
+    "FA",     # flake8-future-annotations
+    "FBT003", # boolean-positional-value-in-call
+    "FLY",    # flynt
+    "I",      # isort
+    "ICN",    # flake8-import-conventions
+    "ISC",    # flake8-implicit-str-concat
+    "PD",     # pandas-vet
+    "PERF",   # perflint
+    "PIE",    # flake8-pie
+    "PL",     # pylint
+    "PT",     # flake8-pytest-style
+    "PYI",    # flakes8-pyi
+    "Q",      # flake8-quotes
+    "RET",    # flake8-return
+    "RSE",    # flake8-raise
+    "RUF",    # Ruff-specific rules
+    "SIM",    # flake8-simplify
+    "SLOT",   # flake8-slots
+    "TCH",    # flake8-type-checking
+    "TID",    # flake8-tidy-imports
+    "UP",     # pyupgrade
+    "W",      # pycodestyle warning
+    "YTT",    # flake8-2020
+]
+ignore = [
+    "B028",
+    "PLR0911", # too-many-return-statements
+    "PLR0912", # too-many-branches
+    "PLR0913", # too-many-arguments
+    "PLR0915", # too-many-statements
+    "PLW0603",
+    "RUF013",
 ]
-ignore = ["B028", "PLW0603", "RUF013"]
 pydocstyle.convention = "numpy"
 isort.known-first-party = ["jobflow"]
 
 [tool.ruff.per-file-ignores]
+# F401: unused import
 "__init__.py" = ["F401"]
-"**/tests/*" = ["D"]
+# D: pydocstyle
+# PLR2004: magic-value-comparison
+# PT004: pytest-missing-fixture-name-underscore
+"**/tests/*" = ["D", "PLR2004", "PT004"]

jobflow-0.1.15/src/jobflow/_version.py

@@ -0,0 +1,7 @@
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("jobflow")
+except PackageNotFoundError:  # pragma: no cover
+    # package is not installed
+    __version__ = ""

{jobflow-0.1.13 → jobflow-0.1.15}/src/jobflow/core/flow.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import logging
 import warnings
 from copy import deepcopy
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from monty.json import MSONable
 
@@ -14,13 +14,13 @@ from jobflow.core.reference import find_and_get_references
 from jobflow.utils import ValueEnum, contains_flow_or_job, suuid
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Iterator
+    from collections.abc import Iterator, Sequence
+    from typing import Any, Callable
 
     from networkx import DiGraph
 
     from jobflow import Job
 
-__all__ = ["JobOrder", "Flow", "get_flow"]
 
 logger = logging.getLogger(__name__)
 
@@ -166,7 +166,7 @@ class Flow(MSONable):
             raise TypeError(
                 f"Flow can only contain Job or Flow objects, not {type(value).__name__}"
             )
-        jobs = list(self.jobs)
+        jobs = list(self)
         jobs[idx] = value  # type: ignore[index, assignment]
         self.jobs = tuple(jobs)
 
@@ -188,10 +188,10 @@ class Flow(MSONable):
 
     def __sub__(self, other: Flow | Job) -> Flow:
         """Remove a job or subflow from the flow."""
-        if other not in self.jobs:
+        if other not in self:
             raise ValueError(f"{other!r} not found in flow")
         new_flow = deepcopy(self)
-        new_flow.jobs = tuple([job for job in new_flow.jobs if job != other])
+        new_flow.jobs = tuple([job for job in new_flow if job != other])
         return new_flow
 
     def __repr__(self, level: int = 0, prefix: str = "") -> str:
@@ -201,8 +201,8 @@ class Flow(MSONable):
         _prefix = f"{prefix}." if prefix else ""
         job_reprs = "\n".join(
             f"{indent}{_prefix}{i}. "
-            f"{j.__repr__(level + 1, f'{_prefix}{i}') if isinstance(j, Flow) else j}"
-            for i, j in enumerate(self.jobs, 1)
+            f"{jb.__repr__(level + 1, f'{_prefix}{i}') if isinstance(jb, Flow) else jb}"
+            for i, jb in enumerate(self, 1)
         )
         return f"Flow({name=}, {uuid=})\n{job_reprs}"
 
@@ -298,7 +298,7 @@ class Flow(MSONable):
             The uuids of all Jobs in the Flow (including nested Flows).
         """
         uuids: list[str] = []
-        for job in self.jobs:
+        for job in self:
             if isinstance(job, Flow):
                 uuids.extend(job.job_uuids)
             else:
@@ -316,7 +316,7 @@ class Flow(MSONable):
             The uuids of all Jobs and Flows in the Flow (including nested Flows).
         """
         uuids: list[str] = []
-        for job in self.jobs:
+        for job in self:
             if isinstance(job, Flow):
                 uuids.extend(job.all_uuids)
             uuids.append(job.uuid)
@@ -336,7 +336,7 @@ class Flow(MSONable):
 
         import networkx as nx
 
-        graph = nx.compose_all([job.graph for job in self.jobs])
+        graph = nx.compose_all([job.graph for job in self])
 
         for node in graph:
             node_props = graph.nodes[node]
@@ -346,7 +346,7 @@ class Flow(MSONable):
         if self.order == JobOrder.LINEAR:
             # add fake edges between jobs to force linear order
             edges = []
-            for job_a, job_b in nx.utils.pairwise(self.jobs):
+            for job_a, job_b in nx.utils.pairwise(self):
                 if isinstance(job_a, Flow):
                     leaves = [v for v, d in job_a.graph.out_degree() if d == 0]
                 else:
@@ -474,7 +474,7 @@ class Flow(MSONable):
         >>> flow.update_kwargs({"number": 10}, name_filter="add")
         >>> flow.update_kwargs({"number": 10}, function_filter=add)
         """
-        for job in self.jobs:
+        for job in self:
             job.update_kwargs(
                 update,
                 name_filter=name_filter,
@@ -573,7 +573,7 @@ class Flow(MSONable):
         ...     {"number": 10}, class_filter=AddMaker, nested=False
         ... )
         """
-        for job in self.jobs:
+        for job in self:
             job.update_maker_kwargs(
                 update,
                 name_filter=name_filter,
@@ -598,7 +598,7 @@ class Flow(MSONable):
         else:
             self.name += append_str
 
-        for job in self.jobs:
+        for job in self:
             job.append_name(append_str, prepend=prepend)
 
     def update_metadata(
@@ -647,7 +647,7 @@ class Flow(MSONable):
 
         >>> flow.update_metadata({"tag": "addition_job"})
         """
-        for job in self.jobs:
+        for job in self:
             job.update_metadata(
                 update,
                 name_filter=name_filter,
@@ -717,7 +717,7 @@ class Flow(MSONable):
 
         >>> flow.update_config({"manager_config": {"_fworker": "myfworker"}})
         """
-        for job in self.jobs:
+        for job in self:
             job.update_config(
                 config,
                 name_filter=name_filter,
@@ -756,8 +756,8 @@ class Flow(MSONable):
                 self.hosts.extend(hosts_uuids)
         else:
             hosts_uuids = [self.uuid]
-        for j in self.jobs:
-            j.add_hosts_uuids(hosts_uuids, prepend=prepend)
+        for job in self:
+            job.add_hosts_uuids(hosts_uuids, prepend=prepend)
 
     def add_jobs(self, jobs: Job | Flow | Sequence[Flow | Job]) -> None:
         """
@@ -794,7 +794,8 @@ class Flow(MSONable):
                     f"current Flow ({self.uuid})"
                 )
             job_ids.add(job.uuid)
-            job.add_hosts_uuids(hosts)
+            if job.host != self.uuid:
+                job.add_hosts_uuids(hosts)
         self._jobs += tuple(jobs)
 
     def remove_jobs(self, indices: int | list[int]):
@@ -810,12 +811,12 @@ class Flow(MSONable):
         """
         if not isinstance(indices, (list, tuple)):
             indices = [indices]
-        if any(i < 0 or i >= len(self.jobs) for i in indices):
+        if any(idx < 0 or idx >= len(self) for idx in indices):
             raise ValueError(
                 "Only indices between 0 and the number of the jobs are accepted"
             )
 
-        new_jobs = tuple(j for i, j in enumerate(self.jobs) if i not in indices)
+        new_jobs = tuple(job for idx, job in enumerate(self) if idx not in indices)
         uuids: set = set()
         for job in new_jobs:
             if isinstance(job, Flow):

{jobflow-0.1.13 → jobflow-0.1.15}/src/jobflow/core/job.py

@@ -13,7 +13,8 @@ from jobflow.core.reference import OnMissing, OutputReference
 from jobflow.utils.uuid import suuid
 
 if typing.TYPE_CHECKING:
-    from typing import Any, Callable, Hashable, Sequence
+    from collections.abc import Hashable, Sequence
+    from typing import Any, Callable
 
     from networkx import DiGraph
     from pydantic import BaseModel
@@ -22,8 +23,6 @@ if typing.TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-__all__ = ["job", "Job", "Response", "JobConfig", "store_inputs"]
-
 
 @dataclass
 class JobConfig(MSONable):
@@ -560,6 +559,7 @@ class Job(MSONable):
 
         from jobflow import CURRENT_JOB
         from jobflow.core.flow import get_flow
+        from jobflow.core.schemas import JobStoreDocument
 
         index_str = f", {self.index}" if self.index != 1 else ""
         logger.info(f"Starting job - {self.name} ({self.uuid}{index_str})")
@@ -633,15 +633,15 @@ class Job(MSONable):
             ) from err
 
         save = {k: "output" if v is True else v for k, v in self._kwargs.items()}
-        data = {
-            "uuid": self.uuid,
-            "index": self.index,
-            "output": output,
-            "completed_at": datetime.now().isoformat(),
-            "metadata": self.metadata,
-            "hosts": self.hosts,
-            "name": self.name,
-        }
+        data: JobStoreDocument = JobStoreDocument(
+            uuid=self.uuid,
+            index=self.index,
+            output=output,
+            completed_at=datetime.now().isoformat(),
+            metadata=self.metadata,
+            hosts=self.hosts,
+            name=self.name,
+        )
         store.update(data, key=["uuid", "index"], save=save)
 
         CURRENT_JOB.reset()
@@ -1321,6 +1321,7 @@ def prepare_replace(
         store_output_job.index = current_job.index + 1
         store_output_job.metadata = current_job.metadata
         store_output_job.output_schema = current_job.output_schema
+        store_output_job._kwargs = current_job._kwargs
         replace.add_jobs(store_output_job)
 
     elif isinstance(replace, Job):

{jobflow-0.1.13 → jobflow-0.1.15}/src/jobflow/core/maker.py

@@ -12,8 +12,6 @@ if typing.TYPE_CHECKING:
 
     import jobflow
 
-__all__ = ["Maker"]
-
 
 @dataclass
 class Maker(MSONable):

{jobflow-0.1.13 → jobflow-0.1.15}/src/jobflow/core/reference.py

@@ -4,24 +4,18 @@ from __future__ import annotations
 
 import contextlib
 import typing
-from typing import Any, Sequence
+from typing import Any
 
 from monty.json import MontyDecoder, MontyEncoder, MSONable, jsanitize
 from pydantic import BaseModel
-from pydantic.utils import lenient_issubclass
+from pydantic.v1.utils import lenient_issubclass
 
 from jobflow.utils.enum import ValueEnum
 
 if typing.TYPE_CHECKING:
-    import jobflow
+    from collections.abc import Sequence
 
-__all__ = [
-    "OnMissing",
-    "OutputReference",
-    "resolve_references",
-    "find_and_resolve_references",
-    "find_and_get_references",
-]
+    import jobflow
 
 
 class OnMissing(ValueEnum):
@@ -95,14 +89,14 @@ class OutputReference(MSONable):
         uuid: str,
         attributes: tuple[tuple[str, Any], ...] = (),
         output_schema: type[BaseModel] = None,
-    ):
+    ) -> None:
         super().__init__()
         self.uuid = uuid
         self.attributes = attributes
         self.output_schema = output_schema
 
         for attr_type, attr in attributes:
-            if attr_type not in ("a", "i"):
+            if attr_type not in {"a", "i"}:
                 raise ValueError(
                     f"Unrecognised attribute type '{attr_type}' for attribute '{attr}'"
                 )
@@ -165,11 +159,12 @@ class OutputReference(MSONable):
         if on_missing == OnMissing.ERROR and index not in cache[self.uuid]:
             istr = f" ({index})" if index is not None else ""
             raise ValueError(
-                f"Could not resolve reference - {self.uuid}{istr} not in store or cache"
+                f"Could not resolve reference - {self.uuid}{istr} not in store or "
+                f"{index=}, {cache=}"
             )
-        elif on_missing == OnMissing.NONE and index not in cache[self.uuid]:
+        if on_missing == OnMissing.NONE and index not in cache[self.uuid]:
             return None
-        elif on_missing == OnMissing.PASS and index not in cache[self.uuid]:
+        if on_missing == OnMissing.PASS and index not in cache[self.uuid]:
             return self
 
         data = cache[self.uuid][index]
@@ -182,7 +177,11 @@ class OutputReference(MSONable):
 
         for attr_type, attr in self.attributes:
             # i means index else use attribute access
-            data = data[attr] if attr_type == "i" else getattr(data, attr)
+            data = (
+                data[attr]
+                if attr_type == "i" or isinstance(data, dict)
+                else getattr(data, attr)
+            )
 
         return data
 
@@ -206,12 +205,11 @@ class OutputReference(MSONable):
         if inplace:
             self.uuid = uuid
             return self
-        else:
-            from copy import deepcopy
+        from copy import deepcopy
 
-            new_reference = deepcopy(self)
-            new_reference.uuid = uuid
-            return new_reference
+        new_reference = deepcopy(self)
+        new_reference.uuid = uuid
+        return new_reference
 
     def __getitem__(self, item) -> OutputReference:
         """Index the reference."""
@@ -269,7 +267,7 @@ class OutputReference(MSONable):
         """Return a hash of the reference."""
         return hash(str(self))
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         """Test for equality against another reference."""
         if isinstance(other, OutputReference):
             return (
@@ -291,7 +289,7 @@ class OutputReference(MSONable):
         """Serialize the reference as a dict."""
         schema = self.output_schema
         schema_dict = MontyEncoder().default(schema) if schema is not None else None
-        data = {
+        return {
             "@module": self.__class__.__module__,
             "@class": type(self).__name__,
             "@version": None,
@@ -299,7 +297,6 @@ class OutputReference(MSONable):
             "attributes": self.attributes,
             "output_schema": schema_dict,
         }
-        return data
 
 
 def resolve_references(
@@ -382,7 +379,7 @@ def find_and_get_references(arg: Any) -> tuple[OutputReference, ...]:
         # if the argument is a reference then stop there
         return (arg,)
 
-    elif isinstance(arg, (float, int, str, bool)):
+    if isinstance(arg, (float, int, str, bool)):
         # argument is a primitive, we won't find a reference here
         return ()
 
@@ -438,7 +435,7 @@ def find_and_resolve_references(
         # if the argument is a reference then stop there
         return arg.resolve(store, cache=cache, on_missing=on_missing)
 
-    elif isinstance(arg, (float, int, str, bool)):
+    if isinstance(arg, (float, int, str, bool)):
         # argument is a primitive, we won't find a reference here
         return arg
 
@@ -503,7 +500,7 @@ def validate_schema_access(
         raise AttributeError(f"{schema.__name__} does not have attribute '{item}'.")
 
     subschema = None
-    item_type = schema.__fields__[item].outer_type_
+    item_type = schema.model_fields[item].annotation
     if lenient_issubclass(item_type, BaseModel):
         subschema = item_type