From 478501f76929d80df44a61caa6ea9e67087239e0 Mon Sep 17 00:00:00 2001 From: "Andrew S. Rosen" Date: Tue, 26 Sep 2023 23:11:33 -1000 Subject: [PATCH] Include validation of Maggma store in Pydantic settings (#1002) --- src/quacc/settings.py | 8 +++++++- src/quacc/wflow/db.py | 10 +++------- tests/settings/test_settings.py | 3 +-- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index 7c363c3674..a673a79ae3 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -7,6 +7,8 @@ from shutil import which from typing import List, Optional, Union +from maggma.core import Store +from monty.json import MontyDecoder from pydantic import BaseSettings, Field, root_validator, validator from quacc.calculators.presets import vasp as vasp_defaults @@ -94,7 +96,7 @@ class QuaccSettings(BaseSettings): # --------------------------- # Data Store Settings # --------------------------- - PRIMARY_STORE: str = Field( + PRIMARY_STORE: Union[str, Store] = Field( None, description=( "String-based JSON representation of the primary Maggma data store " @@ -268,6 +270,10 @@ def resolve_and_make_paths(cls, v): def expand_paths(cls, v): return v.expanduser() + @validator("PRIMARY_STORE") + def generate_store(cls, v): + return MontyDecoder().decode(v) if v else None + class Config: """Pydantic config settings.""" diff --git a/src/quacc/wflow/db.py b/src/quacc/wflow/db.py index 11ed3e82bb..be96aebe1d 100644 --- a/src/quacc/wflow/db.py +++ b/src/quacc/wflow/db.py @@ -6,8 +6,6 @@ import warnings from typing import TYPE_CHECKING -from monty.json import MontyDecoder - if TYPE_CHECKING: from maggma.core import Store @@ -66,7 +64,7 @@ def covalent_to_db( store.update(docs, key="dispatch_id") -def results_to_db(store: Store | str, results: dict | list[dict]) -> None: +def results_to_db(store: Store, results: dict | list[dict]) -> None: """ Store the results of a quacc recipe in a user-specified Maggma Store. A UUID will be generated for each entry. @@ -74,8 +72,7 @@ def results_to_db(store: Store | str, results: dict | list[dict]) -> None: Parameters ---------- store - The Maggma Store object to store the results in or a str representation - of a Maggma Store (taken from `.to_json()`) + The Maggma Store object to store the results in results The output summary dictionary or list of dictionaries from a quacc recipe @@ -84,11 +81,10 @@ def results_to_db(store: Store | str, results: dict | list[dict]) -> None: ------- None """ - if isinstance(store, str): - store = MontyDecoder().decode(store) if isinstance(results, dict): results = [results] + for result in results: result["uuid"] = str(uuid.uuid4()) diff --git a/tests/settings/test_settings.py b/tests/settings/test_settings.py index 7d90b3477c..ae6ee7253a 100644 --- a/tests/settings/test_settings.py +++ b/tests/settings/test_settings.py @@ -42,8 +42,7 @@ def test_file(monkeypatch, tmpdir): def test_store(tmpdir): tmpdir.chdir() - store = MemoryStore() - SETTINGS.PRIMARY_STORE = store.to_json() + SETTINGS.PRIMARY_STORE = MemoryStore() atoms = bulk("Cu") static_job(atoms)