refactor: defer heavy imports — blackjax/jax, sqlalchemy, IPython lazy (#1505)

autofit/__init__.py

@@ -12,7 +12,6 @@
 from . import mock as m
 from .non_linear.grid.grid_search import GridSearch as SearchGridSearch
 from .aggregator.base import AggBase
-from .database.aggregator.aggregator import GridSearchAggregator
 from .graphical.expectation_propagation.history import EPHistory
 from .graphical.declarative.factor.analysis import AnalysisFactor
 from .graphical.declarative.factor.analysis import EPAnalysisFactor
@@ -34,12 +33,10 @@
 from .non_linear.samples import Sample
 from .non_linear.samples import load_from_table
 from .non_linear.samples import SamplesStored
-from .database.aggregator import Aggregator
 from .aggregator.summary.aggregate_csv import AggregateCSV
 from .aggregator.summary.aggregate_csv import ValueType
 from .aggregator.summary.aggregate_images import AggregateImages
 from .aggregator.summary.aggregate_fits import AggregateFITS
-from .database.aggregator import Query
 from autofit.aggregator.fit_interface import Fit
 from .aggregator.search_output import SearchOutput
 from .mapper import prior
@@ -95,7 +92,6 @@
 from .non_linear.search.mcmc.emcee.search import Emcee
 from .non_linear.search.mcmc.zeus.search import Zeus
 from .non_linear.search.nest.nautilus.search import Nautilus
-from .non_linear.search.nest.nss.search import NSS
 from .non_linear.search.nest.dynesty.search.dynamic import DynestyDynamic
 from .non_linear.search.nest.dynesty.search.static import DynestyStatic
 from .non_linear.search.mle.drawer.search import Drawer
@@ -136,7 +132,6 @@
 from autofit.mapper.prior.arithmetic.compound import Log10
 
 from . import example as ex
-from . import database as db
 
 
 for type_ in (
@@ -201,3 +196,34 @@ def save_abc(pickler, obj):
     is_test_mode,
     test_mode_level,
 )
+
+# Lazy attributes (PEP 562): NSS pulls blackjax -> jax, and the database
+# aggregator pulls sqlalchemy + the declarative models — together over a
+# second of import time that most sessions never use.
+_LAZY_ATTRS = {
+    "NSS": ("autofit.non_linear.search.nest.nss.search", "NSS"),
+    "Aggregator": ("autofit.database.aggregator", "Aggregator"),
+    "Query": ("autofit.database.aggregator", "Query"),
+    "GridSearchAggregator": (
+        "autofit.database.aggregator.aggregator",
+        "GridSearchAggregator",
+    ),
+    "db": ("autofit.database", None),
+}
+
+
+def __getattr__(name):
+    try:
+        module_name, attr = _LAZY_ATTRS[name]
+    except KeyError:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    import importlib
+
+    module = importlib.import_module(module_name)
+    value = module if attr is None else getattr(module, attr)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(globals()) | set(_LAZY_ATTRS))

autofit/database/sqlalchemy_.py

@@ -4,44 +4,100 @@
 
 Sufficient interface is implemented to permit import of SQLAlchemy based classes without
 any error. If any attempt is made to use those classes a meaningful warning is returned.
+
+``sa`` is a lazy proxy: the real ``sqlalchemy`` module is imported on first
+attribute access, not when this module is imported, so sessions that never use
+the database do not pay sqlalchemy's import cost. Modules that reference
+``sa.<attr>`` in function signatures must use ``from __future__ import
+annotations`` so the annotation does not trigger the import at definition time.
 """
 
-try:
-    import sqlalchemy as sa
-    from sqlalchemy.ext import declarative
-except ImportError:
-    class MockSQlAlchemy:
-        def __getattr__(self, item):
-            return self
 
-        def __call__(self, *args, **kwargs):
-            return self
+def fail():
+    raise ImportError(
+        "Please install SQLAlchemy to use the database"
+    )
 
-        def __mro_entries__(self, *args, **kwargs):
-            return tuple()
 
-        def declarative_base(self):
-            return MockBase
+class MockBase:
+    def __init__(self, *args, **kwargs):
+        fail()
+
+
+class MockSQlAlchemy:
+    def __getattr__(self, item):
+        return self
+
+    def __call__(self, *args, **kwargs):
+        return self
+
+    def __mro_entries__(self, *args, **kwargs):
+        return tuple()
+
+    def declarative_base(self):
+        return MockBase
+
+    def __getitem__(self, item):
+        fail()
+
+    def __setitem__(self, key, value):
+        fail()
+
 
-        def __getitem__(self, item):
-            fail()
+MockBase.metadata = MockSQlAlchemy()
+
+
+class _LazySQLAlchemy:
+    """
+    Defers ``import sqlalchemy`` until an attribute is first accessed, then
+    delegates every attribute to the real module (or to ``MockSQlAlchemy``
+    when sqlalchemy is not installed).
+    """
+
+    _target = None
+
+    def _load(self):
+        if _LazySQLAlchemy._target is None:
+            try:
+                import sqlalchemy
+
+                _LazySQLAlchemy._target = sqlalchemy
+            except ImportError:
+                _LazySQLAlchemy._target = MockSQlAlchemy()
+        return _LazySQLAlchemy._target
+
+    def __getattr__(self, item):
+        return getattr(self._load(), item)
+
+    def __call__(self, *args, **kwargs):
+        return self._load()(*args, **kwargs)
+
+    def __mro_entries__(self, *args, **kwargs):
+        target = self._load()
+        if isinstance(target, MockSQlAlchemy):
+            return tuple()
+        raise TypeError(f"{target!r} cannot be used as a base class")
 
-        def __setitem__(self, key, value):
-            fail()
+    def __getitem__(self, item):
+        return self._load()[item]
 
+    def __setitem__(self, key, value):
+        self._load()[key] = value
 
-    sa = MockSQlAlchemy()
-    declarative = sa
 
+class _LazyDeclarative(_LazySQLAlchemy):
+    def _load(self):
+        if _LazyDeclarative._target is None:
+            try:
+                from sqlalchemy.ext import declarative
 
-    def fail():
-        raise ImportError(
-            "Please install SQLAlchemy to use the database"
-        )
+                _LazyDeclarative._target = declarative
+            except ImportError:
+                _LazyDeclarative._target = MockSQlAlchemy()
+        return _LazyDeclarative._target
 
 
-    class MockBase:
-        def __init__(self, *args, **kwargs):
-            fail()
+_LazyDeclarative._target = None
 
-        metadata = sa
+sa = _LazySQLAlchemy()
+declarative = _LazyDeclarative()

autofit/non_linear/fitness.py

@@ -1,6 +1,5 @@
 import logging
 import numpy as np
-from IPython.display import clear_output
 import os
 import time
 
@@ -440,6 +439,8 @@ def manage_quick_update(self, parameters, log_likelihood):
 
         if self.quick_update_count >= self.iterations_per_quick_update:
 
+            from IPython.display import clear_output
+
             clear_output(wait=True)
 
             start_time = time.time()

autofit/non_linear/grid/grid_search/result_builder.py

@@ -1,9 +1,13 @@
-from typing import List, Union
+from __future__ import annotations
+
+from typing import List, Union, TYPE_CHECKING
 
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.non_linear.samples import Samples
-from autofit.database import Prior
 from autofit.non_linear.result import Result, Placeholder
+
+if TYPE_CHECKING:
+    from autofit.database import Prior
 from .job import JobResult
 from .result import GridSearchResult
 

autofit/non_linear/paths/database.py

@@ -1,16 +1,18 @@
+from __future__ import annotations
+
 import shutil
-from typing import Optional, Union
+from typing import Optional, Union, TYPE_CHECKING
 
 from autonerves.output import conditional_output, should_output
-from autofit.database.sqlalchemy_ import sa
 from .abstract import AbstractPaths
 import numpy as np
 
-from autofit.database.model import Fit
 from autonerves.dictable import to_dict, from_dict
-from autofit.database.aggregator.info import Info
 from autofit.non_linear.samples.summary import SamplesSummary
 
+if TYPE_CHECKING:
+    from autofit.database.model import Fit
+
 
 class DatabasePaths(AbstractPaths):
     def __init__(
@@ -108,6 +110,8 @@ def zip_remove(self):
         """
         Remove files from both the symlinked folder and the output directory
         """
+        from autofit.database.aggregator.info import Info
+
         self.session.commit()
         Info(self.session).write()
 
@@ -234,6 +238,9 @@ def remove_search_internal(self):
 
     @property
     def fit(self) -> Fit:
+        from autofit.database.model import Fit
+        from autofit.database.sqlalchemy_ import sa
+
         if self._fit is None:
             try:
                 self._fit = (
@@ -341,5 +348,7 @@ def save_all(self, info, *_, **kwargs):
         self.save_json("search", to_dict(self.search))
         self.save_json("model", to_dict(self.model))
 
+        from autofit.database.aggregator.info import Info
+
         self.session.commit()
         Info(self.session).write()

autofit/non_linear/search/abstract_search.py

@@ -16,14 +16,14 @@
 import psutil
 
 if TYPE_CHECKING:
+    from autofit.database.sqlalchemy_ import sa
     from autofit.non_linear.result import Result
 
 from autonerves import conf
 
 from autonerves.output import should_output
 
 from autofit import exc
-from autofit.database.sqlalchemy_ import sa
 from autofit.graphical import (
     MeanField,
     AnalysisFactor,

autofit/non_linear/search/mcmc/abstract_mcmc.py

@@ -1,12 +1,16 @@
-from typing import Optional
+from __future__ import annotations
+
+from typing import Optional, TYPE_CHECKING
 
 from autonerves import conf
-from autofit.database.sqlalchemy_ import sa
 from autofit.non_linear.search.abstract_search import NonLinearSearch
 from autofit.non_linear.initializer import Initializer, InitializerBall
 from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelationsSettings
 from autofit.non_linear.plot import corner_cornerpy
 
+if TYPE_CHECKING:
+    from autofit.database.sqlalchemy_ import sa
+
 class AbstractMCMC(NonLinearSearch):
 
     def __init__(

autofit/non_linear/search/mcmc/blackjax/nuts/search.py

@@ -1,14 +1,15 @@
+from __future__ import annotations
+
 import logging
 import os
 import pickle
 from pathlib import Path
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 import numpy as np
 
 from autonerves import conf
 
-from autofit.database.sqlalchemy_ import sa
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
 from autofit.non_linear.fitness import Fitness
 from autofit.non_linear.initializer import Initializer
@@ -19,6 +20,9 @@
 from autofit.non_linear.samples.mcmc import SamplesMCMC
 from autofit.non_linear.samples.sample import Sample
 
+if TYPE_CHECKING:
+    from autofit.database.sqlalchemy_ import sa
+
 logger = logging.getLogger(__name__)
 
 

autofit/non_linear/search/mcmc/emcee/search.py

@@ -1,13 +1,14 @@
+from __future__ import annotations
+
 import logging
 import os
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, Optional, TYPE_CHECKING
 
 import numpy as np
 
 from autonerves import conf
 
-from autofit.database.sqlalchemy_ import sa
 from autofit.mapper.model_mapper import ModelMapper
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
 from autofit.non_linear.fitness import Fitness
@@ -19,6 +20,9 @@
 from autofit.non_linear.samples.sample import Sample
 from autofit.non_linear.samples.mcmc import SamplesMCMC
 
+if TYPE_CHECKING:
+    from autofit.database.sqlalchemy_ import sa
+
 logger = logging.getLogger(__name__)
 
 

autofit/non_linear/search/mcmc/zeus/search.py

@@ -1,10 +1,11 @@
+from __future__ import annotations
+
 import logging
-from typing import Dict, Optional
+from typing import Dict, Optional, TYPE_CHECKING
 
 import numpy as np
 import os
 
-from autofit.database.sqlalchemy_ import sa
 from autofit.mapper.model_mapper import ModelMapper
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
 from autofit.non_linear.fitness import Fitness
@@ -16,6 +17,9 @@
 from autofit.non_linear.test_mode import is_test_mode
 from autofit.non_linear.samples.mcmc import SamplesMCMC
 
+if TYPE_CHECKING:
+    from autofit.database.sqlalchemy_ import sa
+
 logger = logging.getLogger(__name__)