feat: enhance notebook loading utility

opvious/specifications/notebook.py

@@ -2,7 +2,7 @@
 import os
 import threading
 import types
-from typing import Optional
+from typing import Optional, Sequence
 import warnings
 
 from ..modeling import Model
@@ -12,15 +12,21 @@
 
 
 def load_notebook_models(
-    path: str, root: Optional[str] = None, allow_empty=False
+    path: str,
+    root: Optional[str] = None,
+    allow_empty=False,
+    include_classes=False,
+    include_symbols: Sequence[str] = (),
 ) -> types.SimpleNamespace:
     """Loads all models from a notebook
 
     Args:
         path: Path to the notebook, relative to `root` if present otherwise CWD
         root: Root path. If set to a file, its parent directory will be used
             (convenient for use with `__file__`).
-        allow_empty: Do not throw an error if no models were found
+        allow_empty: Do not throw an error if no exports were found
+        include_classes: Also include :class:`Model` subclasses
+        include_symbols: Also add values with these names
     """
     if root:
         root = os.path.realpath(root)
@@ -31,7 +37,10 @@ def load_notebook_models(
 
     # We run the import logic in a separate, fresh, thread since `importnb`
     # expects an inactive event loop if the notebook includes async statements.
-    t = _ImportThread(target=_populate_notebook_namespace, args=(path, ns))
+    t = _ImportThread(
+        target=_populate_notebook_namespace,
+        args=(path, ns, include_classes, set(include_symbols)),
+    )
     t.start()
     t.join()
 
@@ -57,7 +66,12 @@ def join(self):
             raise Exception("Notebook import failed") from self._exception
 
 
-def _populate_notebook_namespace(path: str, ns: types.SimpleNamespace) -> None:
+def _populate_notebook_namespace(
+    path: str,
+    ns: types.SimpleNamespace,
+    include_classes: bool,
+    include_symbols: set[str],
+) -> None:
     import asyncio
 
     loop = asyncio.new_event_loop()
@@ -92,7 +106,15 @@ def code(self, raw):
     count = 0
     for attr in dir(nb):
         value = getattr(nb, attr)
-        if isinstance(value, Model):
+        if (
+            isinstance(value, Model)
+            or (
+                include_classes
+                and isinstance(value, type)
+                and issubclass(value, Model)
+            )
+            or attr in include_symbols
+        ):
             count += 1
             setattr(ns, attr, value)
-    _logger.debug("Loaded %s model(s) from %s.", count, path)
+    _logger.debug("Loaded %s symbols(s) from %s.", count, path)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "opvious"
-version = "0.19.0rc4"
+version = "0.19.1rc1"
 description = "Opvious Python SDK"
 authors = ["Opvious Engineering <[email protected]>"]
 readme = "README.md"

tests/test_specifications.py

@@ -2,6 +2,7 @@
 import pytest
 
 
+from opvious.modeling import Model
 from opvious.specifications import (
     LocalSpecification,
     load_notebook_models,
@@ -29,7 +30,27 @@ async def test_inline_sources(self):
         assert "greaterThanBound" in spec.sources[0].text
 
     def test_load_notebook_models(self):
-        ns = load_notebook_models("notebooks/set-cover.ipynb", root=__file__)
+        ns = load_notebook_models(
+            "notebooks/set-cover.ipynb",
+            root=__file__,
+        )
         spec = ns.model.specification()
         text = spec.sources[0].text
         assert r"\S^d_\mathrm{sets}&: S" in text
+        assert not hasattr(ns, "SetCover")
+
+    def test_load_notebook_model_classes(self):
+        ns = load_notebook_models(
+            "notebooks/set-cover.ipynb",
+            root=__file__,
+            include_classes=True,
+        )
+        assert issubclass(ns.SetCover, Model)
+
+    def test_load_notebook_model_symbols(self):
+        ns = load_notebook_models(
+            "notebooks/set-cover.ipynb",
+            root=__file__,
+            include_symbols=["SetCover"],
+        )
+        assert issubclass(ns.SetCover, Model)