feat: support exporting notebook models

src/opvious/__main__.py

@@ -6,9 +6,8 @@
 import sys
 from typing import Any, Mapping, Optional
 
-from .common import __version__
-from .client import Client
-from .specifications import load_notebook_models, LocalSpecification
+from . import __version__, Client, LocalSpecification, load_notebook_models
+from .modeling import Model
 
 _COMMAND = "python -m opvious"
 
@@ -23,6 +22,7 @@
     {_COMMAND} register-notebook PATH [MODEL]
         [-dn NAME] [-t TAGS] [--allow-empty]
     {_COMMAND} register-sources GLOB [-dn NAME] [-t TAGS]
+    {_COMMAND} export-notebook-model PATH [MODEL] [-o PATH]
     {_COMMAND} (-h | --help)
     {_COMMAND} --version
 
@@ -33,6 +33,7 @@
                         server
     -n, --name NAME     Formulation name. By default this name is inferred
                         from the file's name, omitting the extension
+    -o, --out PATH      Path where to store the exported model.
     -t, --tags TAGS     Comma-separated list of tags. By default only the
                         `latest` tag is added
     --version           Show SDK version
@@ -80,18 +81,11 @@ async def handle_notebook(
         name: Optional[str],
         allow_empty: bool,
     ) -> None:
-        sn = load_notebook_models(path, allow_empty=allow_empty)
-        if model_name is None:
-            model_names = list(sn.__dict__.keys())
-            if not self._dry_run and len(model_names) != 1:
-                raise Exception(f"Notebook has 0 or 2+ models ({model_names})")
-        else:
-            model_names = [model_name]
-        if name is None:
-            name = _default_name(path)
-        for model_name in model_names:
-            model = getattr(sn, model_name)
-            await self._handle(model.specification(), name)
+        models = _load_notebook_models(path, model_name)
+        if self._dry_run:
+            return
+        _name, model = _singleton_model(models)
+        await self._handle(model.specification(), name or _default_name(path))
 
     async def handle_sources(self, glob: str, name: Optional[str]) -> None:
         if name is None:
@@ -104,10 +98,53 @@ def _default_name(path: str) -> str:
     return os.path.splitext(os.path.basename(path))[0]
 
 
+def _load_notebook_models(
+    path: str,
+    model_name: Optional[str],
+) -> dict[str, Model]:
+    sn = load_notebook_models(path, allow_empty=True)
+    if model_name is None:
+        return {k: v for k, v in sn.__dict__.items() if isinstance(v, Model)}
+    return {model_name: getattr(sn, model_name)}
+
+
+def _singleton_model(models: dict[str, Model]) -> tuple[str, Model]:
+    if len(models) != 1:
+        raise Exception(
+            "Notebook has 0 or 2+ models, please specify a model "
+            "name to select one"
+        )
+    return next(iter(models.items()))
+
+
+async def _export_notebook_model(
+    client: Client,
+    notebook_path: str,
+    model_name: Optional[str] = None,
+    export_path: Optional[str] = None,
+) -> None:
+    models = _load_notebook_models(notebook_path, model_name)
+    name, model = _singleton_model(models)
+    if not export_path:
+        export_path = f"{name}.proto"
+    with open(export_path, "bw+") as writer:
+        await client.export_specification(model.specification(), writer)
+
+
 async def _run(args: Mapping[str, Any]) -> None:
     client = Client.from_environment()
     if not client:
         raise Exception("Missing OPVIOUS_ENDPOINT environment variable")
+
+    if args["export-notebook-model"]:
+        await _export_notebook_model(
+            client,
+            notebook_path=args["PATH"],
+            model_name=args["MODEL"],
+            export_path=args["--out"],
+        )
+        return
+
     handler = _SpecificationHandler(client, args["--tags"], args["--dry-run"])
     if args["register-notebook"]:
         await handler.handle_notebook(

src/opvious/client/handlers.py

@@ -5,6 +5,7 @@
 import logging
 from typing import (
     AsyncIterator,
+    BinaryIO,
     Iterable,
     Mapping,
     Optional,
@@ -48,6 +49,7 @@
     solve_strategy_to_json,
 )
 from ..executors import (
+    BinaryExecutorResult,
     Executor,
     JsonExecutorResult,
     JsonSeqExecutorResult,
@@ -69,6 +71,7 @@
     Problem,
     ProblemOutlineCache,
     ProblemOutlineGenerator,
+    ProblemTransformation,
     SolveInputsBuilder,
     feasible_outcome_details,
     log_progress,
@@ -193,6 +196,42 @@ async def annotate_specification(
             ]
         return specification.annotated(issues)
 
+    async def export_specification(
+        self,
+        specification: LocalSpecification,
+        writer: BinaryIO,
+        transformations: Optional[list[ProblemTransformation]] = None,
+    ) -> None:
+        """Exports a specification to its canonical representation
+
+        Args:
+            specification: The specification to export
+            transformations: Transformations to apply to the specification
+        """
+        sources = [s.text for s in specification.sources]
+
+        if transformations:
+            outline_generator = await ProblemOutlineGenerator.sources(
+                executor=self._executor, sources=sources
+            )
+            for tf in transformations or []:
+                outline_generator.add_transformation(tf)
+            _outline, transformation_data = await outline_generator.generate()
+        else:
+            transformation_data = []
+
+        async with self._executor.execute(
+            result_type=BinaryExecutorResult,
+            url="/sources/assemble",
+            method="POST",
+            json_data=json_dict(
+                sources=sources,
+                transformations=transformation_data,
+            ),
+        ) as res:
+            async for chunk in res.bytes():
+                writer.write(chunk)
+
     async def register_specification(
         self,
         specification: LocalSpecification,

src/opvious/executors/__init__.py

@@ -5,6 +5,7 @@
 from typing import Optional
 
 from .common import (
+    BinaryExecutorResult,
     Executor,
     ExecutorError,
     ExecutorResult,
@@ -24,6 +25,7 @@
     "Executor",
     "ExecutorError",
     "ExecutorResult",
+    "BinaryExecutorResult",
     "JsonExecutorResult",
     "JsonSeqExecutorResult",
     "PlainTextExecutorResult",

src/opvious/executors/aiohttp.py

@@ -5,6 +5,7 @@
 from typing import AsyncIterator, Optional
 
 from .common import (
+    BinaryExecutorResult,
     CONTENT_TYPE_HEADER,
     Executor,
     ExecutorError,
@@ -89,6 +90,12 @@ async def _send(
                             trace=trace,
                             reader=res.content,
                         )
+                    elif BinaryExecutorResult.is_eligible(ctype):
+                        yield BinaryExecutorResult(
+                            status=status,
+                            trace=trace,
+                            reader=res.content,
+                        )
                     else:
                         text = await res.text()
                         raise ExecutorError(

src/opvious/executors/common.py

@@ -153,6 +153,34 @@ def flush(self) -> str:
         return buf
 
 
+@dataclasses.dataclass
+class BinaryExecutorResult(ExecutorResult):
+    """Binary execution result"""
+
+    content_type = "application/octet-stream"
+    reader: Any = dataclasses.field(repr=False)
+
+    async def bytes(
+        self, assert_status: Optional[int] = 200
+    ) -> AsyncIterator[bytes]:
+        if assert_status:
+            self._assert_status(assert_status)
+
+        # Non-streaming
+        if isinstance(self.reader, bytes):
+            yield self.reader
+
+        # Streaming
+        if hasattr(self.reader, "__aiter__"):
+            async for chunk in self.reader:
+                yield chunk
+        elif hasattr(self.reader, "__iter__"):
+            for chunk in self.reader:
+                yield chunk
+        else:
+            raise Exception(f"Non-iterable reader: {self.reader}")
+
+
 @dataclasses.dataclass
 class JsonExecutorResult(ExecutorResult):
     """Unary JSON execution result"""
@@ -256,6 +284,8 @@ async def execute(
             accept = "application/json-seq;q=1, text/plain;q=0.1"
         elif result_type == PlainTextExecutorResult:
             accept = "text/plain"
+        elif result_type == BinaryExecutorResult:
+            accept = "application/octet-stream;q=1, text/plain;q=0.1"
         else:
             raise Exception(f"Unsupported result type: {result_type}")
         all_headers["accept"] = accept

src/opvious/executors/pyodide.py

@@ -4,6 +4,7 @@
 from typing import AsyncIterator, Optional
 
 from .common import (
+    BinaryExecutorResult,
     CONTENT_TYPE_HEADER,
     Executor,
     ExecutorError,
@@ -48,6 +49,9 @@ async def _send(
             yield PlainTextExecutorResult(
                 status=status, trace=trace, reader=text
             )
+        elif BinaryExecutorResult.is_eligible(ctype):
+            data = await res.js_response.bytes()
+            yield BinaryExecutorResult(status=status, trace=trace, reader=data)
         else:
             text = await res.js_response.text()
             raise ExecutorError(status=status, trace=trace, reason=text)

src/opvious/executors/urllib.py

@@ -5,6 +5,7 @@
 from typing import AsyncIterator, Optional
 
 from .common import (
+    BinaryExecutorResult,
     CONTENT_TYPE_HEADER,
     Headers,
     Executor,
@@ -54,6 +55,8 @@ async def _send(
             yield PlainTextExecutorResult(
                 status=status, trace=trace, reader=res
             )
+        elif BinaryExecutorResult.is_eligible(ctype):
+            yield BinaryExecutorResult(status=status, trace=trace, reader=res)
         else:
             raise ExecutorError(
                 status=status,