diff --git a/corrlib/find.py b/corrlib/find.py index 14b1772..dd3a9a6 100644 --- a/corrlib/find.py +++ b/corrlib/find.py @@ -6,9 +6,10 @@ import numpy as np from .input.implementations import codes from .tools import k2m, get_db_file from .tracker import get -from typing import Any, Optional +from typing import Any, Optional, Union from pathlib import Path import datetime as dt +from collections.abc import Callable def _project_lookup_by_alias(db: Path, alias: str) -> str: @@ -39,7 +40,7 @@ def _project_lookup_by_alias(db: Path, alias: str) -> str: return str(results[0][0]) -def _project_lookup_by_id(db: Path, uuid: str) -> list[tuple[str, str]]: +def _project_lookup_by_id(db: Path, uuid: str) -> list[tuple[str, ...]]: """ Return the project information available in the database by UUID. @@ -251,22 +252,31 @@ def sfcf_filter(results: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: return results.drop(drops) +def openQCD_filter(results:pd.DataFrame, **kwargs: Any) -> pd.DataFrame: + return results + + def find_record(path: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None, parameters: Optional[str]=None, - created_before: Optional[str]=None, created_after: Optional[str]=None, updated_before: Optional[str]=None, updated_after: Optional[str]=None, revision: Optional[str]=None, **kwargs: Any) -> pd.DataFrame: + created_before: Optional[str]=None, created_after: Optional[str]=None, updated_before: Optional[str]=None, updated_after: Optional[str]=None, + revision: Optional[str]=None, + customFilter: Optional[Callable[[pd.DataFrame], pd.DataFrame]] = None, + **kwargs: Any) -> pd.DataFrame: db_file = get_db_file(path) db = path / db_file if code not in codes: raise ValueError("Code " + code + "unknown, take one of the following:" + ", ".join(codes)) get(path, db_file) results = _db_lookup(db, ensemble, correlator_name,code, project, parameters=parameters) - if any(arg is not None for arg in [created_before, created_after, updated_before, updated_after]): + if any([arg is not None for arg in [created_before, created_after, updated_before, updated_after]]): results = _time_filter(results, created_before, created_after, updated_before, updated_after) if code == "sfcf": results = sfcf_filter(results, **kwargs) elif code == "openQCD": - pass + results = openQCD_filter(results, **kwargs) else: - raise Exception + raise ValueError(f"Code {code} is not known.") + if customFilter is not None: + results = customFilter(results) print("Found " + str(len(results)) + " result" + ("s" if len(results)>1 else "")) return results.reset_index() diff --git a/tests/find_test.py b/tests/find_test.py index da1bfc1..e895b85 100644 --- a/tests/find_test.py +++ b/tests/find_test.py @@ -70,7 +70,6 @@ def test_db_lookup(tmp_path: Path) -> None: c.execute("INSERT INTO backlogs (name, ensemble, code, path, project, parameters, parameter_file, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'), datetime('now'))", (corr, ensemble, code, meas_path, uuid, pars, parameter_file)) conn.commit() - conn.close() results = find._db_lookup(db, ensemble, corr, code) assert len(results) == 1 @@ -88,8 +87,38 @@ def test_db_lookup(tmp_path: Path) -> None: assert len(results) == 1 results = find._db_lookup(db, ensemble, corr, code, parameters = "{par_A: 3.0, par_B: 4.0}") assert results.empty - #results = find._db_lookup(db, ensemble, corr, code, project, parameters, created_before, created_after, updated_before, updated_after) - #results = find._db_lookup(db, ensemble, corr, code, project, parameters, created_before, created_after, updated_before, updated_after) + corr = "g_A" + ensemble = "SF_A" + code = "openQCD" + meas_path = "archive/SF_A/f_A/Project_A.json.gz::asdfasdfasdf" + uuid = "Project_A" + pars = "{par_A: 3.0, par_B: 4.0}" + parameter_file = "projects/Project_A/myinput.in" + c.execute("INSERT INTO backlogs (name, ensemble, code, path, project, parameters, parameter_file, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'), datetime('now'))", + (corr, ensemble, code, meas_path, uuid, pars, parameter_file)) + conn.commit() + corr = "f_A" + results = find._db_lookup(db, ensemble, corr, code) + assert len(results) == 1 + results = find._db_lookup(db, "SF_B", corr, code) + assert results.empty + results = find._db_lookup(db, ensemble, "g_A", code) + assert len(results) == 1 + results = find._db_lookup(db, ensemble, corr, "sfcf") + assert results.empty + results = find._db_lookup(db, ensemble, corr, code, project = "Project_A") + assert len(results) == 1 + results = find._db_lookup(db, ensemble, "g_A", code, project = "Project_A") + assert len(results) == 1 + results = find._db_lookup(db, ensemble, corr, code, project = "Project_B") + assert results.empty + results = find._db_lookup(db, ensemble, "g_A", code, project = "Project_B") + assert results.empty + results = find._db_lookup(db, ensemble, corr, code, parameters = pars) + assert results.empty + results = find._db_lookup(db, ensemble, "g_A", code, parameters = "{par_A: 3.0, par_B: 4.0}") + assert len(results) == 1 + conn.close()