Coverage for src/marimushka/notebook.py: 97%

1"""Notebook module for handling marimo notebooks. 

2 

3This module provides the Notebook class for representing and exporting marimo notebooks. 

4""" 

5 

6import dataclasses 

7import subprocess 

8from enum import Enum 

9from pathlib import Path 

10 

11from loguru import logger 

12 

13 

14class Kind(Enum): 

15 """Kind of notebook.""" 

16 

17 NB = "notebook" 

18 NB_WASM = "notebook_wasm" 

19 APP = "app" 

20 

21 @classmethod 

22 def from_str(cls, value: str) -> "Kind": 

23 """Represent a factory method to parse a string into a Kind enumeration instance. 

24 

25 This method attempts to match the input string to an existing kind defined 

26 in the Kind enumeration. If the input string does not match any valid kind, 

27 an error is raised detailing the invalid value and listing acceptable kinds. 

28 

29 Args: 

30 value (str): A string representing the kind to parse into a Kind instance. 

31 

32 Returns: 

33 Kind: An instance of the Kind enumeration corresponding to the input string. 

34 

35 Raises: 

36 ValueError: If the input string does not match any valid Kind value. 

37 

38 """ 

39 try: 

40 return Kind(value) 

41 except ValueError as e: 

42 raise ValueError(f"Invalid Kind: {value!r}. Must be one of {[k.value for k in Kind]}") from e 

43 

44 @property 

45 def command(self) -> list[str]: 

46 """Get the command list associated with a specific Kind instance. 

47 

48 The command property returns a list of command strings that correspond 

49 to different kinds of operations based on the Kind instance. 

50 

51 Attributes: 

52 command: A list of strings representing the command. 

53 

54 Returns: 

55 list[str]: A list of command strings for the corresponding Kind instance. 

56 

57 """ 

58 commands = { 

59 Kind.NB: ["uvx", "marimo", "export", "html", "--sandbox"], 

60 Kind.NB_WASM: ["uvx", "marimo", "export", "html-wasm", "--sandbox", "--mode", "edit"], 

61 Kind.APP: ["uvx", "marimo", "export", "html-wasm", "--sandbox", "--mode", "run", "--no-show-code"], 

62 } 

63 return commands[self] 

64 

65 @property 

66 def html_path(self) -> Path: 

67 """Provide a property to determine the HTML path for different kinds of objects. 

68 

69 This property computes the corresponding directory path based on the kind 

70 of the object, such as notebooks, notebooks_wasm, or apps. 

71 

72 @return: A Path object representing the relevant directory path for the 

73 current kind. 

74 

75 @rtype: Path 

76 """ 

77 paths = { 

78 Kind.NB: Path("notebooks"), 

79 Kind.NB_WASM: Path("notebooks_wasm"), 

80 Kind.APP: Path("apps"), 

81 } 

82 return paths[self] 

83 

84 

85@dataclasses.dataclass(frozen=True) 

86class Notebook: 

87 """Represents a marimo notebook. 

88 

89 This class encapsulates a marimo notebook (.py file) and provides methods 

90 for exporting it to HTML/WebAssembly format. 

91 

92 Attributes: 

93 path (Path): Path to the marimo notebook (.py file) 

94 kind (Kind): How the notebook ts treated 

95 

96 """ 

97 

98 path: Path 

99 kind: Kind = Kind.NB 

100 

101 def __post_init__(self): 

102 """Validate the notebook path after initialization. 

103 

104 Raises: 

105 FileNotFoundError: If the file does not exist 

106 ValueError: If the path is not a file or not a Python file 

107 

108 """ 

109 if not self.path.exists(): 

110 raise FileNotFoundError(f"File not found: {self.path}") 

111 if not self.path.is_file(): 

112 raise ValueError(f"Path is not a file: {self.path}") 

113 if not self.path.suffix == ".py": 

114 raise ValueError(f"File is not a Python file: {self.path}") 

115 

116 def export(self, output_dir: Path) -> bool: 

117 """Export the notebook to HTML/WebAssembly format. 

118 

119 This method exports the marimo notebook to HTML/WebAssembly format. 

120 If is_app is True, the notebook is exported in "run" mode with code hidden, 

121 suitable for applications. Otherwise, it's exported in "edit" mode, 

122 suitable for interactive notebooks. 

123 

124 Args: 

125 output_dir (Path): Directory where the exported HTML file will be saved 

126 

127 Returns: 

128 bool: True if export succeeded, False otherwise 

129 

130 """ 

131 cmd = self.kind.command 

132 

133 try: 

134 # Create the full output path and ensure the directory exists 

135 output_file: Path = output_dir / f"{self.path.stem}.html" 

136 output_file.parent.mkdir(parents=True, exist_ok=True) 

137 

138 # Add the notebook path and output file to command 

139 cmd.extend([str(self.path), "-o", str(output_file)]) 

140 

141 # Run marimo export command 

142 logger.debug(f"Running command: {cmd}") 

143 subprocess.run(cmd, capture_output=True, text=True, check=True) 

144 return True 

145 except subprocess.CalledProcessError as e: 

146 # Handle marimo export errors 

147 logger.error(f"Error exporting {self.path}:") 

148 logger.error(f"Command output: {e.stderr}") 

149 return False 

150 except Exception as e: 

151 # Handle unexpected errors 

152 logger.error(f"Unexpected error exporting {self.path}: {e}") 

153 return False 

154 

155 @property 

156 def display_name(self) -> str: 

157 """Return the display name for the notebook.""" 

158 return self.path.stem.replace("_", " ") 

159 

160 @property 

161 def html_path(self) -> Path: 

162 """Return the path to the exported HTML file.""" 

163 return self.kind.html_path / f"{self.path.stem}.html" 

164 

165 

166def folder2notebooks(folder: Path | str | None, kind: Kind = Kind.NB) -> list[Notebook]: 

167 """Find all marimo notebooks in a directory.""" 

168 if folder is None or folder == "": 

169 return [] 

170 

171 # which files are included here? 

172 notebooks = list(Path(folder).glob("*.py")) 

173 

174 # uvx marimo export html-wasm / html --sandbox (--mode edit/run) ( 

175 return [Notebook(path=nb, kind=kind) for nb in notebooks]