diff --git a/corrlib/find.py b/corrlib/find.py index 4c51e05..e4ee735 100644 --- a/corrlib/find.py +++ b/corrlib/find.py @@ -8,6 +8,7 @@ from .tools import k2m, get_db_file from .tracker import get from typing import Any, Optional from pathlib import Path +import datetime as dt def _project_lookup_by_alias(db: Path, alias: str) -> str: @@ -62,8 +63,37 @@ def _project_lookup_by_id(db: Path, uuid: str) -> list[tuple[str, str]]: return results -def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None, parameters: Optional[str]=None, - created_before: Optional[str]=None, created_after: Optional[Any]=None, updated_before: Optional[Any]=None, updated_after: Optional[Any]=None) -> pd.DataFrame: +def _time_filter(results: pd.DataFrame, created_before: Optional[str]=None, created_after: Optional[Any]=None, updated_before: Optional[Any]=None, updated_after: Optional[Any]=None) -> pd.DataFrame: + drops = [] + for ind in len(results): + result = results.iloc[ind] + created_at = dt.datetime.fromisoformat(result['created_at']) + updated_at = dt.datetime.fromisoformat(result['updated_at']) + + if created_before is not None: + created_before = dt.datetime.fromisoformat(created_before) + if created_before < created_at: + drops.append(ind) + continue + if created_after is not None: + created_after = dt.datetime.fromisoformat(created_after) + if created_before > created_at: + drops.append(ind) + continue + if updated_before is not None: + updated_before = dt.datetime.fromisoformat(updated_before) + if updated_before < updated_at: + drops.append(ind) + continue + if updated_after is not None: + updated_after = dt.datetime.fromisoformat(updated_after) + if updated_after > updated_at: + drops.append(ind) + continue + return results.drop(drops) + + +def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None, parameters: Optional[str]=None) -> pd.DataFrame: """ Look up a correlator record in the database by the data given to the method. @@ -105,14 +135,6 @@ def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project search_expr += f" AND code = '{code}'" if parameters: search_expr += f" AND parameters = '{parameters}'" - if created_before: - search_expr += f" AND created_at < '{created_before}'" - if created_after: - search_expr += f" AND created_at > '{created_after}'" - if updated_before: - search_expr += f" AND updated_at < '{updated_before}'" - if updated_after: - search_expr += f" AND updated_at > '{updated_after}'" conn = sqlite3.connect(db) results = pd.read_sql(search_expr, conn) conn.close() @@ -236,7 +258,9 @@ def find_record(path: Path, ensemble: str, correlator_name: str, code: str, proj 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, created_before=created_before, created_after=created_after, updated_before=updated_before, updated_after=updated_after) + results = _db_lookup(db, ensemble, correlator_name,code, project, parameters=parameters) + if Any([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":