Coverage for src / marimushka / notebook.py: 100%
144 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-28 17:41 +0000
1"""Notebook module for handling marimo notebooks.
3This module provides the core abstractions for representing and exporting marimo
4notebooks to various HTML and WebAssembly formats. It defines the Notebook class
5which encapsulates a marimo notebook file and handles the export process via
6subprocess calls to the marimo CLI.
8Key Components:
9 Kind: Enumeration of notebook export types (static HTML, interactive WASM, app mode).
10 Notebook: Dataclass representing a single marimo notebook with export capabilities.
11 folder2notebooks: Utility function to discover notebooks in a directory.
13Export Modes:
14 The module supports three export modes through the Kind enum:
16 - NB (notebook): Exports to static HTML using `marimo export html --sandbox`.
17 Best for documentation and read-only sharing.
19 - NB_WASM (notebook_wasm): Exports to interactive WebAssembly using
20 `marimo export html-wasm --mode edit --sandbox`. Users can edit and run
21 code in the browser.
23 - APP (app): Exports to WebAssembly in run mode using
24 `marimo export html-wasm --mode run --no-show-code --sandbox`. Presents
25 a clean application interface with code hidden.
27Example::
29 from marimushka.notebook import Notebook, Kind, folder2notebooks
30 from pathlib import Path
32 # Create a notebook and export it
33 nb = Notebook(Path("my_notebook.py"), kind=Kind.APP)
34 result = nb.export(Path("_site/apps"))
35 if result.success:
36 print(f"Exported to {result.output_path}")
38 # Discover all notebooks in a directory
39 notebooks = folder2notebooks(Path("notebooks"), kind=Kind.NB)
40"""
42import dataclasses
43import os
44import shutil
45import subprocess # nosec B404
46from enum import Enum
47from pathlib import Path
49from loguru import logger
51from .audit import AuditLogger, get_audit_logger
52from .exceptions import (
53 ExportError,
54 ExportExecutableNotFoundError,
55 ExportSubprocessError,
56 NotebookExportResult,
57 NotebookInvalidError,
58 NotebookNotFoundError,
59)
60from .security import sanitize_error_message, set_secure_file_permissions, validate_bin_path, validate_path_traversal
63class Kind(Enum):
64 """Enumeration of notebook export types.
66 This enum defines the three ways a marimo notebook can be exported,
67 each with different capabilities and use cases. The choice of Kind
68 affects both the marimo export command used and the output directory.
70 Attributes:
71 NB: Static HTML export. Creates a read-only HTML representation
72 of the notebook. Output goes to the 'notebooks/' subdirectory.
73 Uses command: `marimo export html --sandbox`
74 NB_WASM: Interactive WebAssembly export in edit mode. Creates a
75 fully interactive notebook that runs in the browser with code
76 editing capabilities. Output goes to the 'notebooks_wasm/'
77 subdirectory. Uses command: `marimo export html-wasm --mode edit --sandbox`
78 APP: WebAssembly export in run/app mode. Creates an application
79 interface with code hidden from users. Ideal for dashboards
80 and user-facing tools. Output goes to the 'apps/' subdirectory.
81 Uses command: `marimo export html-wasm --mode run --no-show-code --sandbox`
83 """
85 NB = "notebook"
86 NB_WASM = "notebook_wasm"
87 APP = "app"
89 @classmethod
90 def from_str(cls, value: str) -> "Kind":
91 """Represent a factory method to parse a string into a Kind enumeration instance.
93 This method attempts to match the input string to an existing kind defined
94 in the Kind enumeration. If the input string does not match any valid kind,
95 an error is raised detailing the invalid value and listing acceptable kinds.
97 Args:
98 value (str): A string representing the kind to parse into a Kind instance.
100 Returns:
101 Kind: An instance of the Kind enumeration corresponding to the input string.
103 Raises:
104 ValueError: If the input string does not match any valid Kind value.
106 Examples:
107 >>> from marimushka.notebook import Kind
108 >>> Kind.from_str("notebook")
109 <Kind.NB: 'notebook'>
110 >>> Kind.from_str("app")
111 <Kind.APP: 'app'>
112 >>> Kind.from_str("invalid")
113 Traceback (most recent call last):
114 ...
115 ValueError: Invalid Kind: 'invalid'. Must be one of ['notebook', 'notebook_wasm', 'app']
117 """
118 try:
119 return Kind(value)
120 except ValueError as e:
121 msg = f"Invalid Kind: {value!r}. Must be one of {[k.value for k in Kind]}"
122 raise ValueError(msg) from e
124 @property
125 def command(self) -> list[str]:
126 """Get the command list associated with a specific Kind instance.
128 The command property returns a list of command strings that correspond
129 to different kinds of operations based on the Kind instance.
131 Attributes:
132 command: A list of strings representing the command.
134 Returns:
135 list[str]: A list of command strings for the corresponding Kind instance.
137 Examples:
138 >>> from marimushka.notebook import Kind
139 >>> Kind.NB.command
140 ['marimo', 'export', 'html']
141 >>> Kind.APP.command
142 ['marimo', 'export', 'html-wasm', '--mode', 'run', '--no-show-code']
144 """
145 commands = {
146 Kind.NB: ["marimo", "export", "html"],
147 Kind.NB_WASM: ["marimo", "export", "html-wasm", "--mode", "edit"],
148 Kind.APP: ["marimo", "export", "html-wasm", "--mode", "run", "--no-show-code"],
149 }
150 return commands[self]
152 @property
153 def html_path(self) -> Path:
154 """Provide a property to determine the HTML path for different kinds of objects.
156 This property computes the corresponding directory path based on the kind
157 of the object, such as notebooks, notebooks_wasm, or apps.
159 @return: A Path object representing the relevant directory path for the
160 current kind.
162 @rtype: Path
164 Examples:
165 >>> from marimushka.notebook import Kind
166 >>> str(Kind.NB.html_path)
167 'notebooks'
168 >>> str(Kind.APP.html_path)
169 'apps'
171 """
172 paths = {
173 Kind.NB: Path("notebooks"),
174 Kind.NB_WASM: Path("notebooks_wasm"),
175 Kind.APP: Path("apps"),
176 }
177 return paths[self]
180@dataclasses.dataclass(frozen=True)
181class Notebook:
182 """Represents a marimo notebook.
184 This class encapsulates a marimo notebook (.py file) and provides methods
185 for exporting it to HTML/WebAssembly format.
187 Attributes:
188 path (Path): Path to the marimo notebook (.py file)
189 kind (Kind): How the notebook ts treated
191 """
193 path: Path
194 kind: Kind = Kind.NB
196 def __post_init__(self) -> None:
197 """Validate the notebook path after initialization.
199 Raises:
200 NotebookNotFoundError: If the file does not exist.
201 NotebookInvalidError: If the path is not a file or not a Python file.
203 """
204 if not self.path.exists():
205 raise NotebookNotFoundError(self.path)
206 if not self.path.is_file():
207 raise NotebookInvalidError(self.path, reason="path is not a file")
208 if not self.path.suffix == ".py":
209 raise NotebookInvalidError(self.path, reason="not a Python file")
211 def export(
212 self,
213 output_dir: Path,
214 sandbox: bool = True,
215 bin_path: Path | None = None,
216 timeout: int = 300,
217 audit_logger: AuditLogger | None = None,
218 ) -> NotebookExportResult:
219 """Export the notebook to HTML/WebAssembly format.
221 This method exports the marimo notebook to HTML/WebAssembly format.
222 If is_app is True, the notebook is exported in "run" mode with code hidden,
223 suitable for applications. Otherwise, it's exported in "edit" mode,
224 suitable for interactive notebooks.
226 Args:
227 output_dir: Directory where the exported HTML file will be saved.
228 sandbox: Whether to run the notebook in a sandbox. Defaults to True.
229 bin_path: The directory where the executable is located. Defaults to None.
230 timeout: Maximum time in seconds for the export process. Defaults to 300.
231 audit_logger: Logger for audit events. If None, uses default logger.
233 Returns:
234 NotebookExportResult indicating success or failure with details.
236 """
237 if audit_logger is None:
238 audit_logger = get_audit_logger()
240 # Resolve executable
241 exe = self._resolve_executable(bin_path, audit_logger)
242 if isinstance(exe, NotebookExportResult):
243 return exe
245 # Prepare output path
246 output_file_or_error = self._prepare_output_path(output_dir, audit_logger)
247 if isinstance(output_file_or_error, NotebookExportResult):
248 return output_file_or_error
249 output_file = output_file_or_error
251 # Build and run command
252 cmd = self._build_command(exe, sandbox, output_file)
253 return self._run_export_subprocess(cmd, output_file, timeout, audit_logger)
255 def _resolve_executable(self, bin_path: Path | None, audit_logger: AuditLogger) -> str | NotebookExportResult:
256 """Resolve the executable path.
258 Args:
259 bin_path: Optional directory where the executable is located.
260 audit_logger: Audit logger for security logging.
262 Returns:
263 Executable string on success, or NotebookExportResult on error.
265 """
266 executable = "uvx"
268 if not bin_path:
269 return executable
271 # Validate bin_path for security
272 try:
273 validated_bin_path = validate_bin_path(bin_path)
274 audit_logger.log_path_validation(bin_path, "bin_path", True)
275 except ValueError as e:
276 err: ExportError = ExportExecutableNotFoundError(executable, bin_path)
277 sanitized_error = sanitize_error_message(str(e))
278 logger.error(f"Invalid bin_path: {sanitized_error}")
279 audit_logger.log_path_validation(bin_path, "bin_path", False, sanitized_error)
280 return NotebookExportResult.failed(self.path, err)
282 # Construct the full executable path
283 # Use shutil.which to find it with platform-specific extensions (like .exe on Windows)
284 exe = shutil.which(executable, path=str(validated_bin_path))
285 if not exe:
286 # Fallback: try constructing the path directly
287 exe_path = validated_bin_path / executable
288 if exe_path.is_file() and os.access(exe_path, os.X_OK):
289 return str(exe_path)
291 err = ExportExecutableNotFoundError(executable, validated_bin_path)
292 logger.error(str(err))
293 audit_logger.log_export(self.path, None, False, str(err))
294 return NotebookExportResult.failed(self.path, err)
296 return exe
298 def _prepare_output_path(self, output_dir: Path, audit_logger: AuditLogger) -> Path | NotebookExportResult:
299 """Validate and prepare the output path.
301 Args:
302 output_dir: Directory where the exported HTML file will be saved.
303 audit_logger: Audit logger for security logging.
305 Returns:
306 Output file Path on success, or NotebookExportResult on error.
308 """
309 output_file = output_dir / f"{self.path.stem}.html"
311 # Validate output path to prevent path traversal
312 try:
313 validate_path_traversal(output_file)
314 audit_logger.log_path_validation(output_file, "output_path", True)
315 except ValueError as e:
316 sanitized_error = sanitize_error_message(str(e))
317 err = ExportSubprocessError(
318 notebook_path=self.path,
319 command=[], # Command not yet built
320 return_code=-1,
321 stderr=f"Invalid output path: {sanitized_error}",
322 )
323 logger.error(str(err))
324 audit_logger.log_path_validation(output_file, "output_path", False, sanitized_error)
325 return NotebookExportResult.failed(self.path, err)
327 # Create output directory
328 try:
329 output_file.parent.mkdir(parents=True, exist_ok=True)
330 except OSError as e: # pragma: no cover
331 sanitized_error = sanitize_error_message(str(e))
332 err = ExportSubprocessError(
333 notebook_path=self.path,
334 command=[], # Command not yet built
335 return_code=-1,
336 stderr=f"Failed to create output directory: {sanitized_error}",
337 )
338 logger.error(str(err))
339 audit_logger.log_export(self.path, None, False, sanitized_error)
340 return NotebookExportResult.failed(self.path, err)
342 return output_file
344 def _build_command(self, exe: str, sandbox: bool, output_file: Path) -> list[str]:
345 """Build the export command.
347 Args:
348 exe: Executable to use (e.g., 'uvx' or full path).
349 sandbox: Whether to run the notebook in a sandbox.
350 output_file: Path where the exported HTML file will be saved.
352 Returns:
353 Command list ready for subprocess execution.
355 """
356 cmd = [exe, *self.kind.command]
357 if sandbox:
358 cmd.append("--sandbox")
359 else:
360 cmd.append("--no-sandbox")
361 cmd.extend([str(self.path), "-o", str(output_file)])
362 return cmd
364 def _run_export_subprocess(
365 self, cmd: list[str], output_file: Path, timeout: int, audit_logger: AuditLogger
366 ) -> NotebookExportResult:
367 """Run the export subprocess and handle results.
369 Args:
370 cmd: Command list to execute.
371 output_file: Path where the exported HTML file will be saved.
372 timeout: Maximum time in seconds for the export process.
373 audit_logger: Audit logger for security logging.
375 Returns:
376 NotebookExportResult indicating success or failure.
378 """
379 try:
380 # Run marimo export command with timeout
381 logger.debug(f"Running command: {cmd}")
382 result = subprocess.run(cmd, capture_output=True, text=True, check=False, timeout=timeout) # nosec B603 # noqa: S603
384 nb_logger = logger.bind(subprocess=f"[{self.path.name}] ")
386 if result.stdout:
387 nb_logger.info(f"stdout:\n{result.stdout.strip()}")
389 if result.stderr:
390 nb_logger.warning(f"stderr:\n{result.stderr.strip()}")
392 if result.returncode != 0:
393 sanitized_stderr = sanitize_error_message(result.stderr)
394 err = ExportSubprocessError(
395 notebook_path=self.path,
396 command=cmd,
397 return_code=result.returncode,
398 stdout=result.stdout,
399 stderr=sanitized_stderr,
400 )
401 nb_logger.error(str(err))
402 audit_logger.log_export(self.path, None, False, sanitized_stderr)
403 return NotebookExportResult.failed(self.path, err)
405 # Set secure permissions on output file
406 try:
407 set_secure_file_permissions(output_file, mode=0o644)
408 except ValueError as e: # pragma: no cover
409 logger.warning(f"Could not set secure permissions on {output_file}: {e}")
411 audit_logger.log_export(self.path, output_file, True)
412 return NotebookExportResult.succeeded(self.path, output_file)
414 except subprocess.TimeoutExpired:
415 err = ExportSubprocessError(
416 notebook_path=self.path,
417 command=cmd,
418 return_code=-1,
419 stderr=f"Export timed out after {timeout} seconds",
420 )
421 logger.error(str(err))
422 audit_logger.log_export(self.path, None, False, f"timeout after {timeout}s")
423 return NotebookExportResult.failed(self.path, err)
424 except FileNotFoundError as e:
425 # Executable not found in PATH
426 exec_err = ExportExecutableNotFoundError(cmd[0])
427 sanitized_error = sanitize_error_message(str(e))
428 logger.error(f"{exec_err}: {sanitized_error}")
429 audit_logger.log_export(self.path, None, False, sanitized_error)
430 return NotebookExportResult.failed(self.path, exec_err)
431 except subprocess.SubprocessError as e:
432 sanitized_error = sanitize_error_message(str(e))
433 err = ExportSubprocessError(
434 notebook_path=self.path,
435 command=cmd,
436 return_code=-1,
437 stderr=sanitized_error,
438 )
439 logger.error(str(err))
440 audit_logger.log_export(self.path, None, False, sanitized_error)
441 return NotebookExportResult.failed(self.path, err)
443 @property
444 def display_name(self) -> str:
445 """Return the display name for the notebook.
447 The display name is derived from the notebook filename by replacing
448 underscores with spaces, making it more human-readable.
450 Returns:
451 str: Human-friendly display name with underscores replaced by spaces.
453 Examples:
454 >>> # Demonstrating the transformation logic
455 >>> filename = "my_cool_notebook"
456 >>> display_name = filename.replace("_", " ")
457 >>> display_name
458 'my cool notebook'
460 """
461 return self.path.stem.replace("_", " ")
463 @property
464 def html_path(self) -> Path:
465 """Return the path to the exported HTML file."""
466 return self.kind.html_path / f"{self.path.stem}.html"
469def folder2notebooks(folder: Path | str | None, kind: Kind = Kind.NB) -> list[Notebook]:
470 """Discover and create Notebook instances for all Python files in a directory.
472 This function scans a directory for Python files (*.py) and creates Notebook
473 instances for each one. It assumes all Python files in the directory are
474 valid marimo notebooks. The resulting list is sorted alphabetically by
475 filename to ensure consistent ordering across runs.
477 Args:
478 folder: Path to the directory to scan for notebooks. Can be a Path
479 object, a string path, or None. If None or empty string, returns
480 an empty list.
481 kind: The export type for all discovered notebooks. Defaults to Kind.NB
482 (static HTML export). All notebooks in the folder will be assigned
483 this kind.
485 Returns:
486 A list of Notebook instances, one for each Python file found in the
487 directory, sorted alphabetically by filename. Returns an empty list
488 if the folder is None, empty, or contains no Python files.
490 Raises:
491 NotebookNotFoundError: If a discovered file does not exist (unlikely
492 in normal usage but possible in race conditions).
493 NotebookInvalidError: If a discovered path is not a valid file.
495 Example::
497 from pathlib import Path
498 from marimushka.notebook import folder2notebooks, Kind
500 # Get all notebooks from a directory as static HTML exports
501 notebooks = folder2notebooks(Path("notebooks"), Kind.NB)
503 # Get all notebooks as interactive apps
504 apps = folder2notebooks("apps", Kind.APP)
506 # Handle empty or missing directories gracefully
507 empty = folder2notebooks(None) # Returns []
508 empty = folder2notebooks("") # Returns []
510 Examples:
511 >>> from marimushka.notebook import folder2notebooks
512 >>> # When folder is None, returns empty list
513 >>> folder2notebooks(None)
514 []
515 >>> # When folder is empty string, returns empty list
516 >>> folder2notebooks("")
517 []
519 """
520 if folder is None or folder == "":
521 return []
523 notebooks = list(Path(folder).glob("*.py"))
524 notebooks.sort()
526 return [Notebook(path=nb, kind=kind) for nb in notebooks]