better type annotation fir id lookup

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()
 

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()