S0L009/Luna-GNN-Scorer-InferenceAPI

1"""2Firebase database utilities.3"""4 5import re6import copy7from datetime import datetime, timedelta, timezone8 9import firebase_admin10from firebase_admin import credentials11from firebase_admin import firestore12 13from typing import Union, Any, Optional14 15def norm_field(16    field: str17):18    """19    Normalize a field name for use in Firestore.20 21    Parameters22    ----------23    field : str24        The field name to normalize.25 26    Returns27    -------28    str29        The normalized field name. Only lowercase alphanumeric30        characters are allowed.31    """32    field = field.lower()33    return re.sub(r'[^a-z0-9]', '', field)34 35def get_firestore_obj(path):36    """37    Return a Firestore client.38 39    Parameters40    ----------41    path : str42        Path to the service account JSON file.43 44    Returns45    -------46    db : Firestore Client47        The Firestore client.48    """49    cred = credentials.Certificate(path)50 51    # Initialize the Firebase app52    firebase_admin.initialize_app(cred)53 54    # Reference Firestore55    db = firestore.client()56    return db57 58def generic_doc_ref(59        db,60        path: list[tuple["collection", "document"]],61        value: bool = False62    ):63    """64    Return a Firestore document or collection reference, optionally with data.65 66    Parameters67    ----------68    db : Firestore Client69        The Firestore client.70    path : list[tuple["collection", "document"]]71        A list of tuples, where each tuple contains a collection name and a document ID.72        The document ID may be None, in which case the function will return a CollectionReference.73    value : bool, optional74        If True: Return the document data as a dictionary by calling get() and to_dict(); or if75        a document ID is None, returns the list of documents in the collection as an iterable by76        calling stream().77        78        If False: Return the reference directly. Default is False.79 80    Returns81    -------82    doc_ref : DocumentReference|CollectionReference or tuple[Any, DocumentReference|CollectionReference]83        The Firestore reference, or a tuple of the returned data and reference.84    """85    x = db86    for collection_name, doc_id in path:87        if doc_id is None:88            x = x.collection(collection_name)89            if value:90                return x.stream(), x91            return x92        x = x.collection(collection_name).document(doc_id)93    if value:94        return x.get().to_dict(), x95    return x96 97def safe_doc_get(98        db,99        path: Union[list[tuple["collection", "document"]], "docref"],100        fallback_generator: Any101    ):102    """103    Safely retrieve a Firestore document as a dictionary, using a fallback generator if104    the document does not exist.105 106    Parameters107    ----------108    db : Firestore Client109        The Firestore client.110    path : list[tuple["collection", "document"]]|"docref"111        A list of tuples representing the path to the document in the database, 112        or a direct document reference.113    fallback_generator : Any114        A callable that generates a fallback value if the document does not exist.115        If fallback_generator is a Dict, then we will iteratively check through each116        key to ensure it exists.117 118    Returns119    -------120    dict121        The document data as a dictionary if the document exists, 122        otherwise the result of the fallback_generator.123    """124    if isinstance(path, list):125        ref = generic_doc_ref(db=db, path=path, value=False)126    else:127        ref = path128    129    doc = ref.get()130    fallback = fallback_generator()131    132    if doc.exists:133        info = doc.to_dict()134        if isinstance(fallback, dict):135            for key in fallback:136                if key not in info:137                    info[key] = fallback[key]138        return info139    return fallback140 141 142def generator_wrap(obj, use_copy=False):143    """144    Wrap an object into a generator function that yields the object.145 146    This is useful for converting an object into a generator that can be147    used with functions that expect a generator.148 149    Parameters150    ----------151    obj : Any152        The object to wrap.153    use_copy : bool, optional154        If True, use a copy of the object instead of the original, by default False.155 156    Returns157    -------158    Callable[[], Any]159        A generator that yields the object.160    """161    def _gen_obj():162        if use_copy:163            return copy.deepcopy(obj)164        else:165            return obj166    return _gen_obj167 168def safe_doc_update(169        db,170        path: Union[list[tuple["collection", "document"]], "docref"],171        data: Any172    ):173    """174    Safely update a Firestore document. If the document does not exist, it will be created.175    176    Parameters177    ----------178    db : Firestore Client179        The Firestore client.180    path : list[tuple["collection", "document"]]|"docref"181        A list of tuples representing the path to the document in the database, 182        or a direct document reference.183    data : Any184        The document data to update.185    """186    if isinstance(path, list):187        ref = generic_doc_ref(db=db, path=path, value=False)188    else:189        ref = path190    if ref.get().exists:191        ref.update(data)192    else:193        ref.set(data)194 195_TIMESTAMP_FORMAT = r"%d/%m/%Y-%H:%M:%S"196 197def timestamp(stamp: str = None):198    """199    Get or parse a datetime object in the format "%d/%m/%Y-%H:%M:%S".200 201    Parameters202    ----------203    stamp : str, optional204        A string representing a datetime object in the format "%d/%m/%Y-%H:%M:%S".205        If None, return the current datetime object, by default None.206        Otherwise, parse the string into a datetime object.207 208    Returns209    -------210    Union[datetime, str]211        The string or parsed datetime object.212    """213    if stamp is None:214        now = datetime.now(timezone.utc)215        return now.strftime(_TIMESTAMP_FORMAT)216    dt = datetime.strptime(stamp, _TIMESTAMP_FORMAT)217    return datetime(dt.year, dt.month, dt.day, dt.hour, dt.minute, dt.second, tzinfo=timezone.utc)218 219def days_since(last: datetime):220    """221    Calculate the number of days since a given datetime object.222 223    Parameters224    ----------225    last : datetime226        The datetime object to calculate the number of days since.227 228    Returns229    -------230    int231        The number of days since the given datetime object.232    """233    return (datetime.now(timezone.utc) - last).seconds / 86400234 235def key_bundle(236    *args,237    separator: str = '_',238):239    """240    Bundle multiple key values into a string.241 242    Parameters243    ----------244    *args : Any245        The values to bundle. Each must have a __str__() method246        that does not output separator.247    separator : str248        The separator to use between the values. Must be length 1.249        Defaults to '_'.250 251    Returns252    -------253    str254        The bundled string.255    """256    if len(separator) != 1:257        raise ValueError(f"separator must be length 1; got {len(separator)}.")258    259    def check(key):260        if separator in key:261            raise ValueError(f"separator {separator} found in key {key}.")262        return key263    264    return separator.join([check(str(key)) for key in args])265