Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           0 : from functools import wraps
       5           0 : from typing import Any, Callable
       6             : 
       7           0 : from pairinteraction import (
       8             :     _wrapped as pi,
       9             : )
      10           0 : from pairinteraction_gui.app import Application
      11             : 
      12             : 
      13           0 : class DatabaseMissingError(Exception):
      14           0 :     def __init__(self, err: RuntimeError) -> None:
      15           0 :         super().__init__(str(err))
      16           0 :         if not self.is_database_missing_error(err):
      17           0 :             raise ValueError("The message must contain 'Table' and 'not found' to be a DatabaseMissingError.")
      18           0 :         table = next(w for w in str(err).split(" ") if "states" in w)
      19           0 :         self.species = table.replace("_states", "")
      20             : 
      21           0 :     @classmethod
      22           0 :     def is_database_missing_error(cls, err: RuntimeError) -> bool:
      23           0 :         return "Table" in str(err) and "not found" in str(err)
      24             : 
      25             : 
      26           0 : class NoStateFoundError(Exception):
      27           0 :     pass
      28             : 
      29             : 
      30           0 : def catch_download_missing(func: Callable[..., Any]) -> Callable[..., Any]:
      31           0 :     @wraps(func)
      32           0 :     def wrapper_func(*args: Any, **kwargs: Any) -> Any:
      33           0 :         try:
      34           0 :             return func(*args, **kwargs)
      35           0 :         except RuntimeError as err:
      36           0 :             if DatabaseMissingError.is_database_missing_error(err):
      37           0 :                 Application.signals.ask_download_database.emit(DatabaseMissingError(err).species)
      38           0 :                 return func(*args, **kwargs)
      39           0 :             raise err
      40             : 
      41           0 :     return wrapper_func
      42             : 
      43             : 
      44           0 : @catch_download_missing
      45           0 : def get_ket_atom(species: str, **qns: float) -> pi.KetAtom:
      46           0 :     try:
      47           0 :         return pi.KetAtom(species, **qns)  # type: ignore
      48           0 :     except ValueError as err:
      49           0 :         if "No state found" in str(err) or "quantum number m must be" in str(err):
      50           0 :             raise NoStateFoundError(str(err)) from err
      51           0 :         raise err