Improving class loading efficiency and considering dynamic type list

forte/data/data_store.py

@@ -22,7 +22,7 @@
 from sortedcontainers import SortedList
 from typing_inspect import get_origin, get_args, is_generic_type
 
-from forte.utils import get_class
+from forte.utils import get_class, get_class_nc
 from forte.utils.utils import get_full_module_name
 from forte.data.ontology.code_generation_objects import EntryTree
 from forte.data.ontology.ontology_code_generator import OntologyCodeGenerator
@@ -895,7 +895,7 @@ def _is_subclass(
             if cls_qualified_name in type_name_parent_class:
                 return True
             else:
-                entry_class = get_class(type_name)
+                entry_class = get_class_nc(type_name)
                 if issubclass(entry_class, cls):
                     type_name_parent_class.add(cls_qualified_name)
                     return True

forte/utils/utils.py

@@ -16,7 +16,7 @@
 """
 import sys
 import difflib
-from functools import wraps
+from functools import wraps, lru_cache
 from inspect import getfullargspec
 from pydoc import locate
 from typing import Dict, List, Optional, get_type_hints, Tuple
@@ -28,6 +28,7 @@
     "get_full_module_name",
     "get_class_name",
     "get_class",
+    "get_class_nc",
     "get_qual_name",
     "create_class_with_kwargs",
     "check_type",
@@ -78,8 +79,10 @@ def get_class_name(o, lower: bool = False) -> str:
         return o.__name__
 
 
-def get_class(full_class_name: str, module_paths: Optional[List[str]] = None):
-    r"""Returns the class based on class name.
+def get_class_nc(
+    full_class_name: str, module_paths: Optional[List[str]] = None
+):
+    r"""Returns the class based on class name, not cached.
 
     Args:
         full_class_name (str): Name or full path to the class.
@@ -127,6 +130,57 @@ def get_class(full_class_name: str, module_paths: Optional[List[str]] = None):
     return class_
 
 
+@lru_cache()
+def cached_locate(name_to_locate_class):
+    return locate(name_to_locate_class)
+
+
+def get_class(full_class_name: str, module_paths: Optional[List[str]] = None):
+    r"""Returns the class based on class name, with cache to improve speed.
+
+    Args:
+        class_name (str): Name or full path to the class.
+        module_paths (list): Paths to candidate modules to search for the
+            class. This is used if the class cannot be located solely based on
+            ``class_name``. The first module in the list that contains the class
+            is used.
+
+    Returns:
+        The target class.
+
+    """
+    class_ = cached_locate(full_class_name)
+    if (class_ is None) and (module_paths is not None):
+        for module_path in module_paths:
+            class_ = cached_locate(".".join([module_path, full_class_name]))
+            if class_ is not None:
+                break
+
+    # Try to find classes that are dynamically loaded, class_ will still be None if failed.
+    if class_ is None:
+        try:
+            module_name, class_name = full_class_name.rsplit(".", 1)
+            try:
+                class_ = getattr(sys.modules[module_name], class_name)
+            except (AttributeError, KeyError):
+                # ignore when cannot find the module in sys.modules or cannot find the class
+                # in the module.
+                pass
+        except ValueError:
+            # ignoring when the full class name doesn't have multiple parts.
+            pass
+
+    if class_ is None:
+        if module_paths:
+            raise ValueError(
+                f"Class not found in {module_paths}: {full_class_name}"
+            )
+        else:
+            raise ValueError(f"Class not found in {full_class_name}")
+
+    return class_
+
+
 def get_qual_name(o: object, lower: bool = False) -> str:
     r"""Returns the qualified name of an object ``o``.