Semi-automatically populate __all__ in __init__.py files

bayesflow/adapters/__init__.py

@@ -1,2 +1,6 @@
 from . import transforms
 from .adapter import Adapter
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=["transforms"])

bayesflow/adapters/transforms/__init__.py

@@ -16,3 +16,7 @@
 from .standardize import Standardize
 from .to_array import ToArray
 from .transform import Transform
+
+from ...utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=["transforms"])

bayesflow/approximators/__init__.py

@@ -1,3 +1,7 @@
 from .approximator import Approximator
 from .continuous_approximator import ContinuousApproximator
 from .model_comparison_approximator import ModelComparisonApproximator
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/benchmarks/__init__.py

@@ -1,2 +1,6 @@
 from .simulators import TwoMoons
 from .simulators import LotkaVolterra
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/datasets/__init__.py

@@ -2,3 +2,7 @@
 from .online_dataset import OnlineDataset
 from .disk_dataset import DiskDataset
 from .rounds_dataset import RoundsDataset
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/diagnostics/__init__.py

@@ -12,3 +12,7 @@
     recovery,
     z_score_contraction,
 )
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/distributions/__init__.py

@@ -1,3 +1,7 @@
 from .distribution import Distribution
 from .diagonal_normal import DiagonalNormal
 from .diagonal_student_t import DiagonalStudentT
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/metrics/__init__.py

@@ -1,2 +1,7 @@
+from . import functional
 from .maximum_mean_discrepancy import MaximumMeanDiscrepancy
 from .root_mean_squard_error import RootMeanSquaredError
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=["functional"])

bayesflow/metrics/functional/__init__.py

@@ -1,2 +1,6 @@
 from .maximum_mean_discrepancy import maximum_mean_discrepancy
 from .root_mean_squared_error import root_mean_squared_error
+
+from ...utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/networks/__init__.py

@@ -9,3 +9,7 @@
 from .lstnet import LSTNet
 from .summary_network import SummaryNetwork
 from .transformers import SetTransformer, TimeSeriesTransformer, FusionTransformer
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/simulators/__init__.py

@@ -5,3 +5,7 @@
 from .model_comparison_simulator import ModelComparisonSimulator
 from .simulator import Simulator
 from .two_moons import TwoMoons
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/types/__init__.py

@@ -1,2 +1,6 @@
 from .shape import Shape
 from .tensor import Tensor
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/utils/__init__.py

@@ -68,3 +68,7 @@
 )
 from .validators import check_lengths_same
 from .workflow_utils import find_inference_network, find_summary_network
+
+from ._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=["keras_utils", "logging", "numpy_utils"])

bayesflow/utils/_docs/__init__.py

@@ -0,0 +1 @@
+from ._populate_all import _add_imports_to_all

bayesflow/utils/_docs/_populate_all.py

@@ -0,0 +1,21 @@
+import inspect
+
+
+def _add_imports_to_all(include_modules: bool | list[str] = False, exclude: list[str] | None = None):
+    """Add all global variables to __all__"""
+    assert type(include_modules) in [bool, list]
+    exclude = exclude or []
+    calling_module = inspect.stack()[1]
+    local_stack = calling_module[0]
+    global_vars = local_stack.f_globals
+    all_vars = global_vars["__all__"] if "__all__" in global_vars else []
+    included_vars = []
+    for var_name in set(global_vars.keys()):
+        if inspect.ismodule(global_vars[var_name]):
+            if include_modules is True and var_name not in exclude and not var_name.startswith("_"):
+                included_vars.append(var_name)
+            elif isinstance(include_modules, list) and var_name in include_modules:
+                included_vars.append(var_name)
+        elif var_name not in exclude and not var_name.startswith("_"):
+            included_vars.append(var_name)
+    global_vars["__all__"] = list(set(all_vars).union(included_vars))

bayesflow/workflows/__init__.py

@@ -1 +1,5 @@
 from .basic_workflow import BasicWorkflow
+
+from ..utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

docsrc/source/api/bayesflow.rst

@@ -12,6 +12,7 @@
     bayesflow.datasets
     bayesflow.diagnostics
     bayesflow.distributions
+    bayesflow.experimental
     bayesflow.metrics
     bayesflow.networks
     bayesflow.simulators

docsrc/source/conf.py

@@ -90,7 +90,7 @@
 # do not ignore __all__, use it to determine public members
 autosummary_ignore_module_all = False
 # include imported members in autosummary
-autosummary_imported_members = True
+autosummary_imported_members = False
 # selects content to insert into the main body of an autoclass directive.
 autoclass_content = "both"