utility for loading class

Author: aditya0by0

chebai/ensemble/_scripts/_ensemble_run_script.py

@@ -1,28 +1,11 @@
-import importlib
 from typing import Any, Dict, Type
 
 import yaml
 from jsonargparse import ArgumentParser
 
-from ._base import EnsembleBase
+from chebai.ensemble._utils import _load_class
 
-
-def load_class(class_path: str) -> Type[EnsembleBase]:
-    """
-    Dynamically imports and returns a class from a full dotted path.
-
-    Args:
-        class_path (str): Full module path to the class (e.g., 'my_package.module.MyClass').
-
-    Returns:
-        Type[EnsembleBase]: The imported class object.
-
-    Raises:
-        ModuleNotFoundError, AttributeError: If module or class cannot be loaded.
-    """
-    module_path, class_name = class_path.rsplit(".", 1)
-    module = importlib.import_module(module_path)
-    return getattr(module, class_name)
+from .._base import EnsembleBase
 
 
 def load_config_and_instantiate(config_path: str) -> EnsembleBase:
@@ -44,7 +27,7 @@ def load_config_and_instantiate(config_path: str) -> EnsembleBase:
     class_path: str = config["class_path"]
     init_args: Dict[str, Any] = config.get("init_args", {})
 
-    cls = load_class(class_path)
+    cls = _load_class(class_path)
 
     if not issubclass(cls, EnsembleBase):
         raise TypeError(f"{cls} must be subclass of EnsembleBase")

chebai/ensemble/_utils.py

@@ -0,0 +1,7 @@
+import importlib
+
+
+def _load_class(class_path):
+    module_path, class_name = class_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)

chebai/ensemble/_wrappers/_base.py

@@ -136,10 +136,3 @@ def _predict_from_list_of_smiles(self, smiles_list: list) -> dict:
     @abstractmethod
     def _predict_from_data_file(self, data_file_path: str) -> dict:
         pass
-
-    @staticmethod
-    def _load_class(class_path):
-        class_name = class_path.split(".")[-1]
-        module_path = ".".join(class_path.split(".")[:-1])
-        module = importlib.import_module(module_path)
-        return getattr(module, class_name)

chebai/ensemble/_wrappers/_neural_network.py

@@ -8,15 +8,13 @@
 from chebai.preprocessing.reader import DataReader
 
 from .._constants import MODEL_CKPT_PATH, READER_CLS_PATH, READER_KWARGS
+from .._utils import _load_class
 from ._base import BaseWrapper
 
 
 class NNWrapper(BaseWrapper):
 
-    def __init__(
-        self,
-        **kwargs,
-    ):
+    def __init__(self, **kwargs):
         self._validate_model_configs(**kwargs)
         super().__init__(**kwargs)
 
@@ -28,7 +26,7 @@ def __init__(
             else dict()
         )
 
-        reader_cls: Type[DataReader] = self._load_class(self._reader_class_path)
+        reader_cls: Type[DataReader] = _load_class(self._reader_class_path)
         assert issubclass(reader_cls, DataReader), ""
         self._reader = reader_cls(**self._reader_kwargs)
         self._collator = reader_cls.COLLATOR()