Coverage for cosmolayer/__init__.py: 85%

1""" 

2Differentiable COSMO-Type Activity Coefficient Layer 

3""" 

4 

5import importlib 

6from typing import TYPE_CHECKING 

7 

8from . import store 

9from ._version import __version__ 

10 

11if TYPE_CHECKING: 

12 from . import cosmosac 

13 from .cosmodata import ( 

14 MixtureDatapoint, 

15 MixtureInferenceDataset, 

16 MixtureTrainingDataset, 

17 ) 

18 from .cosmolayer import CosmoLayer 

19 from .cosmolightning import LogGammaLightningModule 

20 from .cosmosolver import CosmoSolver 

21 

22__all__ = [ 

23 "__version__", 

24 "cosmosac", 

25 "store", 

26 "CosmoLayer", 

27 "LogGammaLightningModule", 

28 "CosmoSolver", 

29 "MixtureDatapoint", 

30 "MixtureInferenceDataset", 

31 "MixtureTrainingDataset", 

32] 

33 

34# Submodules and attributes that pull in heavy optional dependencies (torch, 

35# lightning, rdkit-adjacent chains) are imported lazily, on first access, 

36# rather than eagerly at package-import time. This keeps `import cosmolayer` 

37# (and cheap subpackages like `cosmolayer.store`) fast. 

38_LAZY_SUBMODULES = frozenset({"cosmosac"}) 

39_LAZY_ATTRS = { 

40 "CosmoLayer": ".cosmolayer", 

41 "LogGammaLightningModule": ".cosmolightning", 

42 "CosmoSolver": ".cosmosolver", 

43 "MixtureDatapoint": ".cosmodata", 

44 "MixtureInferenceDataset": ".cosmodata", 

45 "MixtureTrainingDataset": ".cosmodata", 

46} 

47 

48 

49def __getattr__(name: str) -> object: 

50 if name in _LAZY_SUBMODULES: 

51 module = importlib.import_module(f".{name}", __name__) 

52 globals()[name] = module 

53 return module 

54 if name in _LAZY_ATTRS: 

55 module = importlib.import_module(_LAZY_ATTRS[name], __name__) 

56 value = getattr(module, name) 

57 globals()[name] = value 

58 return value 

59 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

60 

61 

62def __dir__() -> list[str]: 

63 return sorted(set(globals()) | set(_LAZY_ATTRS) | _LAZY_SUBMODULES)