Merge pull request #337 from vpratz/feat-doc-all

vpratz · web-flow · commit 5632da3bc534 · 2025-02-28T19:15:23.000+01:00
Semi-automatically populate `__all__` in `__init__.py` files
diff --git a/bayesflow/adapters/__init__.py b/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"])
diff --git a/bayesflow/adapters/transforms/__init__.py b/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"])
diff --git a/bayesflow/approximators/__init__.py b/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=[])
diff --git a/bayesflow/benchmarks/__init__.py b/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=[])
diff --git a/bayesflow/datasets/__init__.py b/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=[])
diff --git a/bayesflow/diagnostics/__init__.py b/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=[])
diff --git a/bayesflow/distributions/__init__.py b/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=[])
diff --git a/bayesflow/metrics/__init__.py b/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"])
diff --git a/bayesflow/metrics/functional/__init__.py b/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=[])
diff --git a/bayesflow/networks/__init__.py b/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=[])
diff --git a/bayesflow/simulators/__init__.py b/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=[])
diff --git a/bayesflow/types/__init__.py b/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=[])
diff --git a/bayesflow/utils/__init__.py b/bayesflow/utils/__init__.py
@@ -67,3 +67,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"])
diff --git a/bayesflow/utils/_docs/__init__.py b/bayesflow/utils/_docs/__init__.py
@@ -0,0 +1 @@
+from ._populate_all import _add_imports_to_all
diff --git a/bayesflow/utils/_docs/_populate_all.py b/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))
diff --git a/bayesflow/workflows/__init__.py b/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=[])
diff --git a/docsrc/source/api/bayesflow.rst b/docsrc/source/api/bayesflow.rst
@@ -12,6 +12,7 @@
     bayesflow.datasets
     bayesflow.diagnostics
     bayesflow.distributions
+    bayesflow.experimental
     bayesflow.metrics
     bayesflow.networks
     bayesflow.simulators
diff --git a/docsrc/source/conf.py b/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"
 

Original file line number	Diff line number	Diff line change
`@@ -12,3 +12,7 @@`
`12`	`12`	`recovery,`
`13`	`13`	`z_score_contraction,`
`14`	`14`	`)`
	`15`	`+`
	`16`	`+from ..utils._docs import _add_imports_to_all`
	`17`	`+`
	`18`	`+_add_imports_to_all(include_modules=[])`