1"""Base classes and utilities for readers and writers.
2
3Authors:
4
5* Brian Granger
6"""
7
8# -----------------------------------------------------------------------------
9# Copyright (C) 2008-2011 The IPython Development Team
10#
11# Distributed under the terms of the BSD License. The full license is in
12# the file LICENSE, distributed as part of this software.
13# -----------------------------------------------------------------------------
14
15# -----------------------------------------------------------------------------
16# Imports
17# -----------------------------------------------------------------------------
18from __future__ import annotations
19
20from base64 import decodebytes, encodebytes
21
22# -----------------------------------------------------------------------------
23# Code
24# -----------------------------------------------------------------------------
25
26
27def restore_bytes(nb):
28 """Restore bytes of image data from unicode-only formats.
29
30 Base64 encoding is handled elsewhere. Bytes objects in the notebook are
31 always b64-encoded. We DO NOT encode/decode around file formats.
32 """
33 for ws in nb.worksheets:
34 for cell in ws.cells:
35 if cell.cell_type == "code":
36 for output in cell.outputs:
37 if "png" in output:
38 output.png = output.png.encode("ascii")
39 if "jpeg" in output:
40 output.jpeg = output.jpeg.encode("ascii")
41 return nb
42
43
44# output keys that are likely to have multiline values
45_multiline_outputs = ["text", "html", "svg", "latex", "javascript", "json"]
46
47
48def rejoin_lines(nb):
49 """rejoin multiline text into strings
50
51 For reversing effects of ``split_lines(nb)``.
52
53 This only rejoins lines that have been split, so if text objects were not split
54 they will pass through unchanged.
55
56 Used when reading JSON files that may have been passed through split_lines.
57 """
58 for ws in nb.worksheets:
59 for cell in ws.cells:
60 if cell.cell_type == "code":
61 if "input" in cell and isinstance(cell.input, list):
62 cell.input = "\n".join(cell.input)
63 for output in cell.outputs:
64 for key in _multiline_outputs:
65 item = output.get(key, None)
66 if isinstance(item, list):
67 output[key] = "\n".join(item)
68 else: # text cell
69 for key in ["source", "rendered"]:
70 item = cell.get(key, None)
71 if isinstance(item, list):
72 cell[key] = "\n".join(item)
73 return nb
74
75
76def split_lines(nb):
77 """split likely multiline text into lists of strings
78
79 For file output more friendly to line-based VCS. ``rejoin_lines(nb)`` will
80 reverse the effects of ``split_lines(nb)``.
81
82 Used when writing JSON files.
83 """
84 for ws in nb.worksheets:
85 for cell in ws.cells:
86 if cell.cell_type == "code":
87 if "input" in cell and isinstance(cell.input, str):
88 cell.input = cell.input.splitlines()
89 for output in cell.outputs:
90 for key in _multiline_outputs:
91 item = output.get(key, None)
92 if isinstance(item, str):
93 output[key] = item.splitlines()
94 else: # text cell
95 for key in ["source", "rendered"]:
96 item = cell.get(key, None)
97 if isinstance(item, str):
98 cell[key] = item.splitlines()
99 return nb
100
101
102# b64 encode/decode are never actually used, because all bytes objects in
103# the notebook are already b64-encoded, and we don't need/want to double-encode
104
105
106def base64_decode(nb):
107 """Restore all bytes objects in the notebook from base64-encoded strings.
108
109 Note: This is never used
110 """
111 for ws in nb.worksheets:
112 for cell in ws.cells:
113 if cell.cell_type == "code":
114 for output in cell.outputs:
115 if "png" in output:
116 if isinstance(output.png, str):
117 output.png = output.png.encode("ascii")
118 output.png = decodebytes(output.png)
119 if "jpeg" in output:
120 if isinstance(output.jpeg, str):
121 output.jpeg = output.jpeg.encode("ascii")
122 output.jpeg = decodebytes(output.jpeg)
123 return nb
124
125
126def base64_encode(nb):
127 """Base64 encode all bytes objects in the notebook.
128
129 These will be b64-encoded unicode strings
130
131 Note: This is never used
132 """
133 for ws in nb.worksheets:
134 for cell in ws.cells:
135 if cell.cell_type == "code":
136 for output in cell.outputs:
137 if "png" in output:
138 output.png = encodebytes(output.png).decode("ascii")
139 if "jpeg" in output:
140 output.jpeg = encodebytes(output.jpeg).decode("ascii")
141 return nb
142
143
144class NotebookReader:
145 """A class for reading notebooks."""
146
147 def reads(self, s, **kwargs):
148 """Read a notebook from a string."""
149 msg = "loads must be implemented in a subclass"
150 raise NotImplementedError(msg)
151
152 def read(self, fp, **kwargs):
153 """Read a notebook from a file like object"""
154 return self.read(fp.read(), **kwargs)
155
156
157class NotebookWriter:
158 """A class for writing notebooks."""
159
160 def writes(self, nb, **kwargs):
161 """Write a notebook to a string."""
162 msg = "loads must be implemented in a subclass"
163 raise NotImplementedError(msg)
164
165 def write(self, nb, fp, **kwargs):
166 """Write a notebook to a file like object"""
167 return fp.write(self.writes(nb, **kwargs))