Ianhelle/lazy-import-init-2023-09-26 (#717)

* Switch init py modules to use lazy import function. * Fixed doc string in lazy importer * Fixing recursive import bug * Fixing pylint and mypy errors * Minor refactor of nbinit to clean up main init_notebook function * removing commented out code
microsoft · Oct 6, 2023 · 7ad1280 · 7ad1280
1 parent df22c8b
commit 7ad1280
Show file tree

Hide file tree

Showing 30 changed files with 372 additions and 288 deletions.
diff --git a/msticpy/__init__.py b/msticpy/__init__.py
@@ -133,6 +133,7 @@
 )
 from .common.utility import search_name as search
 from .init.logging import set_logging_level, setup_logging
+from .lazy_importer import lazy_import
 
 __version__ = VERSION
 __author__ = "Ian Hellen, Pete Bryan, Ashwin Patil"
@@ -144,71 +145,25 @@
 if not os.environ.get("KQLMAGIC_EXTRAS_REQUIRES"):
     os.environ["KQLMAGIC_EXTRAS_REQUIRES"] = "jupyter-basic"
 
-_STATIC_ATTRIBS = list(locals().keys())
-
-_DEFAULT_IMPORTS = {
-    "az_connect": "msticpy.auth.azure_auth",
-    "current_providers": "msticpy.init.nbinit",
-    "ContextLookup": "msticpy.context.contextlookup",
-    "GeoLiteLookup": "msticpy.context.geoip",
-    "init_notebook": "msticpy.init.nbinit",
-    "reset_ipython_exception_handler": "msticpy.init.nbinit",
-    "MicrosoftSentinel": "msticpy.context.azure",
-    "MpConfigEdit": "msticpy.config.mp_config_edit",
-    "MpConfigFile": "msticpy.config.mp_config_file",
-    "QueryProvider": "msticpy.data",
-    "TILookup": "msticpy.context.tilookup",
-    "TimeSpan": "msticpy.common.timespan",
-    "WorkspaceConfig": "msticpy.common.wsconfig",
-    "entities": "msticpy.datamodel",
-    "Pivot": "msticpy.init.pivot",
+_LAZY_IMPORTS = {
+    "msticpy.auth.azure_auth.az_connect",
+    "msticpy.common.timespan.TimeSpan",
+    "msticpy.common.wsconfig.WorkspaceConfig",
+    "msticpy.config.mp_config_edit.MpConfigEdit",
+    "msticpy.config.mp_config_file.MpConfigFile",
+    "msticpy.context.azure.MicrosoftSentinel",
+    "msticpy.context.contextlookup.ContextLookup",
+    "msticpy.context.geoip.GeoLiteLookup",
+    "msticpy.context.tilookup.TILookup",
+    "msticpy.data.QueryProvider",
+    "msticpy.datamodel.entities",
+    "msticpy.init.nbinit.current_providers",
+    "msticpy.init.nbinit.init_notebook",
+    "msticpy.init.nbinit.reset_ipython_exception_handler",
+    "msticpy.init.pivot.Pivot",
 }
 
-
-def __getattr__(attrib: str) -> Any:
-    """
-    Import and return an attribute of MSTICPy.
-
-    Parameters
-    ----------
-    attrib : str
-        The attribute name
-
-    Returns
-    -------
-    Any
-        The attribute value.
-
-    Raises
-    ------
-    AttributeError
-        No attribute found.
-
-    """
-    if attrib in _DEFAULT_IMPORTS:
-        try:
-            return getattr(importlib.import_module(_DEFAULT_IMPORTS[attrib]), attrib)
-        except (MsticpyImportExtraError, MsticpyImportExtraError):
-            raise
-        except (ImportError, MsticpyException) as err:
-            warnings.warn(f"Unable to import module for 'msticpy.{attrib}'")
-            print(
-                f"WARNING. The msticpy attribute '{attrib}' is not loadable.",
-                "You may need to install one or more additional dependencies.\n",
-                "Please check the exception details below for more information.",
-                "\n".join(
-                    traceback.format_exception(
-                        type(err), value=err, tb=err.__traceback__
-                    )
-                ),
-            )
-            raise AttributeError(f"msticpy failed to load '{attrib}'") from err
-    raise AttributeError(f"msticpy has no attribute '{attrib}'")
-
-
-def __dir__():
-    """Return attribute list."""
-    return sorted(set(_STATIC_ATTRIBS + list(_DEFAULT_IMPORTS)))
+module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS)
 
 
 def load_plugins(plugin_paths: Union[str, Iterable[str]]):

diff --git a/msticpy/common/utility/package.py b/msticpy/common/utility/package.py
@@ -323,7 +323,7 @@ def init_dir(static_attribs: List[str], dynamic_imports: Dict[str, str]):
     return sorted(set(static_attribs + list(dynamic_imports)))
 
 
-def lazy_import(module: str, attrib: str, call: bool = False):
+def delayed_import(module: str, attrib: str, call: bool = False):
     """Import attribute from module on demand."""
     attribute = None
 

diff --git a/msticpy/config/__init__.py b/msticpy/config/__init__.py
@@ -12,21 +12,12 @@
 It use the ipywidgets package.
 
 """
-from ..common.utility.package import init_dir, init_getattr
+from ..lazy_importer import lazy_import
 
-_STATIC_ATTRIBS = list(locals().keys())
-_DEFAULT_IMPORTS = {
-    "MpConfigControls": "msticpy.config.mp_config_control",
-    "MpConfigEdit": "msticpy.config.mp_config_edit",
-    "MpConfigFile": "msticpy.config.mp_config_file",
+_LAZY_IMPORTS = {
+    "msticpy.config.mp_config_control.MpConfigControls",
+    "msticpy.config.mp_config_edit.MpConfigEdit",
+    "msticpy.config.mp_config_file.MpConfigFile",
 }
 
-
-def __getattr__(attrib: str):
-    """Import and a dynamic attribute of module."""
-    return init_getattr(__name__, _DEFAULT_IMPORTS, attrib)
-
-
-def __dir__():
-    """Return attribute list."""
-    return init_dir(_STATIC_ATTRIBS, _DEFAULT_IMPORTS)
+module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS)
diff --git a/msticpy/config/ce_azure_sentinel.py b/msticpy/config/ce_azure_sentinel.py
@@ -9,7 +9,7 @@
 import ipywidgets as widgets
 
 from .._version import VERSION
-from ..common.utility.package import lazy_import
+from ..common.utility.package import delayed_import
 
 # from ..context.azure.sentinel_core import MicrosoftSentinel
 from .ce_common import (
@@ -24,7 +24,7 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
-ms_sentinel = lazy_import("msticpy.context.azure.sentinel_core", "MicrosoftSentinel")
+ms_sentinel = delayed_import("msticpy.context.azure.sentinel_core", "MicrosoftSentinel")
 
 
 # pylint: disable=too-many-ancestors

diff --git a/msticpy/config/ce_user_defaults.py b/msticpy/config/ce_user_defaults.py
@@ -67,7 +67,8 @@ def __init__(self, mp_controls: MpConfigControls):
 
         """
         super().__init__(mp_controls)
-        from ..data import DataEnvironment  # pylint: disable=import-outside-toplevel
+        # pylint: disable=import-outside-toplevel
+        from ..data.core.query_defns import DataEnvironment
 
         self._data_env_enum = DataEnvironment
 

diff --git a/msticpy/config/mp_config_file.py b/msticpy/config/mp_config_file.py
@@ -25,16 +25,13 @@
 except ImportError:
     _KEYVAULT = False
 
-try:
-    from ..context.azure.sentinel_core import MicrosoftSentinel
-
-    _SENTINEL = True
-except ImportError:
-    _SENTINEL = False
 from ..common.pkg_config import current_config_path, refresh_config, validate_config
+from ..common.utility.package import delayed_import
 from .comp_edit import CompEditDisplayMixin, CompEditStatusMixin
 from .file_browser import FileBrowser
 
+ms_sentinel = delayed_import("msticpy.context.azure.sentinel_core", "MicrosoftSentinel")
+
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
@@ -306,7 +303,7 @@ def get_workspace_from_url(url: str) -> Dict[str, Dict[str, str]]:
             workspace.
 
         """
-        return MicrosoftSentinel.get_workspace_details_from_url(url)
+        return ms_sentinel().get_workspace_details_from_url(url)
 
     def _show_sentinel_workspace(self, show: bool = True):
         """Fetch settings from Sentinel Portal URL."""

diff --git a/msticpy/context/__init__.py b/msticpy/context/__init__.py
@@ -6,10 +6,8 @@
 """Context Providers Subpackage."""
 from typing import Any
 
-# flake8: noqa: F403
 from ..common.utility import ImportPlaceholder
-from .geoip import GeoLiteLookup, IPStackLookup
-from .tilookup import TILookup
+from ..lazy_importer import lazy_import
 from .vtlookupv3 import VT3_AVAILABLE
 
 vtlookupv3: Any
@@ -20,3 +18,12 @@
     vtlookupv3 = ImportPlaceholder(  # type: ignore
         "vtlookupv3", ["vt-py", "vt-graph-api", "nest_asyncio"]
     )
+
+
+_LAZY_IMPORTS = {
+    "msticpy.context.geoip.GeoLiteLookup",
+    "msticpy.context.geoip.IPStackLookup",
+    "msticpy.context.tilookup.TILookup",
+}
+
+module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS)
diff --git a/msticpy/context/azure/__init__.py b/msticpy/context/azure/__init__.py
@@ -5,6 +5,11 @@
 # --------------------------------------------------------------------------
 """Data provider sub-package."""
 
-# flake8: noqa: F401
-from .azure_data import AzureData
-from .sentinel_core import MicrosoftSentinel
+from ...lazy_importer import lazy_import
+
+_LAZY_IMPORTS = {
+    "msticpy.context.azure.azure_data.AzureData",
+    "msticpy.context.azure.sentinel_core.MicrosoftSentinel",
+}
+
+module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS)
diff --git a/msticpy/context/azure/sentinel_dynamic_summary.py b/msticpy/context/azure/sentinel_dynamic_summary.py
@@ -15,7 +15,7 @@
 from ..._version import VERSION
 from ...common.exceptions import MsticpyAzureConnectionError, MsticpyParameterError
 from ...common.pkg_config import get_config, get_http_timeout
-from ...data import QueryProvider
+from ...data.core.data_providers import QueryProvider
 from .azure_data import get_api_headers
 
 # pylint: disable=unused-import

diff --git a/msticpy/context/azure/sentinel_workspaces.py b/msticpy/context/azure/sentinel_workspaces.py
@@ -18,7 +18,7 @@
 from ...common.data_utils import df_has_data
 from ...common.pkg_config import get_http_timeout
 from ...common.utility import mp_ua_header
-from ...data import QueryProvider
+from ...data.core.data_providers import QueryProvider
 
 __version__ = VERSION
 __author__ = "Ian Hellen"

diff --git a/msticpy/context/tiproviders/kql_base.py b/msticpy/context/tiproviders/kql_base.py
@@ -26,7 +26,7 @@
 from ...common.exceptions import MsticpyConfigError
 from ...common.utility import export
 from ...common.wsconfig import WorkspaceConfig
-from ...data import QueryProvider
+from ...data.core.data_providers import QueryProvider
 from ..lookup_result import LookupStatus
 from ..provider_base import generate_items
 from .ti_provider_base import ResultSeverity, TIProvider

diff --git a/msticpy/context/tiproviders/open_page_rank.py b/msticpy/context/tiproviders/open_page_rank.py
@@ -31,7 +31,11 @@
 
 @export
 class OPR(HttpTIProvider):
-    """Open PageRank Lookup."""
+    """
+    Open PageRank Lookup.
+
+    See https://www.domcop.com/openpagerank/what-is-openpagerank
+    """
 
     _BASE_URL = "https://openpagerank.com"
 
@@ -50,10 +54,6 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
         self._provider_name = self.__class__.__name__
-        print(
-            "Using Open PageRank.",
-            "See https://www.domcop.com/openpagerank/what-is-openpagerank",
-        )
 
     async def lookup_iocs_async(
         self,

diff --git a/msticpy/data/__init__.py b/msticpy/data/__init__.py
@@ -21,10 +21,16 @@
 
 """
 from .._version import VERSION
-from ..common.exceptions import MsticpyImportExtraError
 
-# flake8: noqa: F403
-from .core.data_providers import QueryProvider
-from .core.query_defns import DataEnvironment, DataFamily
+# from ..common.exceptions import MsticpyImportExtraError
+from ..lazy_importer import lazy_import
 
 __version__ = VERSION
+
+_LAZY_IMPORTS = {
+    "msticpy.data.core.data_providers.QueryProvider",
+    "msticpy.data.core.query_defns.DataEnvironment",
+    "msticpy.data.core.query_defns.DataFamily",
+}
+
+module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS)
diff --git a/msticpy/data/azure_data.py b/msticpy/data/azure_data.py
@@ -12,7 +12,7 @@
 
 # flake8: noqa: F403, F401
 # pylint: disable=unused-import
-from ..context.azure import AzureData
+from ..context.azure.azure_data import AzureData
 
 WARN_MSSG = (
     "This module has moved to msticpy.context.azure.azure_data\n"

diff --git a/msticpy/data/core/data_providers.py b/msticpy/data/core/data_providers.py
@@ -14,7 +14,7 @@
 from ..._version import VERSION
 from ...common.pkg_config import get_config
 from ...common.utility import export, valid_pyname
-from ...nbwidgets import QueryTime
+from ...nbwidgets.query_time import QueryTime
 from .. import drivers
 from ..drivers.driver_base import DriverBase, DriverProps
 from .param_extractor import extract_query_params

diff --git a/msticpy/data/core/query_provider_utils_mixin.py b/msticpy/data/core/query_provider_utils_mixin.py
@@ -9,7 +9,7 @@
 from typing import Dict, Iterable, List, NamedTuple, Optional, Pattern, Protocol, Union
 
 from ..._version import VERSION
-from ...common.utility.package import lazy_import
+from ...common.utility.package import delayed_import
 from ..drivers.driver_base import DriverBase
 from .query_defns import DataEnvironment
 from .query_source import QuerySource
@@ -18,7 +18,7 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
-query_browser = lazy_import("msticpy.vis.query_browser", "browse_queries")
+query_browser = delayed_import("msticpy.vis.query_browser", "browse_queries")
 
 
 # pylint: disable=too-few-public-methods