Coverage for e2xgrader/preprocessors/extractattachments.py: 81%
53 statements
« prev ^ index » next coverage.py v7.4.2, created at 2024-03-14 13:22 +0100
« prev ^ index » next coverage.py v7.4.2, created at 2024-03-14 13:22 +0100
1import os
2import sys
3from binascii import a2b_base64
5from nbgrader.preprocessors import NbGraderPreprocessor
6from traitlets import Set, Unicode
8from ..utils.extra_cells import is_attachment_cell, is_diagram
11class ExtractAttachments(NbGraderPreprocessor):
12 subdirectory = Unicode("attachments").tag(config=True)
14 output_filename_template = Unicode("attach_{cell_index}_{name}").tag(config=True)
16 extract_output_types = Set(
17 {"image/png", "image/jpeg", "image/svg+xml", "application/pdf"}
18 ).tag(config=True)
20 def preprocess(self, nb, resources):
21 # Get files directory if it has been specified
22 self.path = os.path.join(resources["metadata"]["path"], self.subdirectory)
23 os.makedirs(self.path, exist_ok=True)
25 for cell_index, cell in enumerate(nb.cells):
26 cell, resources = self.preprocess_cell(cell, resources, cell_index)
28 return nb, resources
30 def preprocess_cell(self, cell, resources, cell_index):
31 """
32 Apply a transformation on each cell,
34 Parameters
35 ----------
36 cell : NotebookNode cell
37 Notebook cell being processed
38 resources : dictionary
39 Additional resources used in the conversion process. Allows
40 preprocessors to pass variables into the Jinja engine.
41 cell_index : int
42 Index of the cell being processed (see base.py)
43 """
45 if not (is_attachment_cell(cell) or is_diagram(cell)):
46 return cell, resources
48 if is_diagram(cell):
49 # Remove the attachment line from source
50 cell.source = cell.source.replace("", "")
52 # Get files directory if it has been specified
53 output_files_dir = resources.get("output_files_dir", None)
55 # Make sure outputs key exists
56 if not isinstance(resources["outputs"], dict):
57 resources["outputs"] = {}
59 to_delete = []
61 # Loop through all of the attachments in the cell
62 for name, attach in cell.get("attachments", {}).items():
63 for mime, data in attach.items():
64 if mime not in self.extract_output_types:
65 continue
67 # Binary files are base64-encoded, SVG is already XML
68 if mime in {"image/png", "image/jpeg", "application/pdf"}:
69 # data is b64-encoded as text (str, unicode),
70 # we want the original bytes
71 data = a2b_base64(data)
72 elif sys.platform == "win32":
73 data = data.replace("\n", "\r\n").encode("UTF-8")
74 else:
75 data = data.encode("UTF-8")
77 filename = self.output_filename_template.format(
78 cell_index=cell_index,
79 name=name,
80 )
81 link = os.path.join(self.subdirectory, filename)
83 if output_files_dir is not None:
84 filename = os.path.join(self.path, filename)
86 if name.endswith(".gif") and mime == "image/png":
87 filename = filename.replace(".gif", ".png")
89 # Write the attachments to a file and add a link
91 with open(filename, "wb") as f:
92 f.write(data)
94 if "image" in mime:
95 cell.source += f'\n<a href="{link}"><img src="{link}"></a>\n'
96 else:
97 cell.source += f'\n<a href="{link}">{name}</a>\n'
98 if not is_diagram(cell):
99 to_delete.append(name)
101 # now we need to change the cell source so that it links to the
102 # filename instead of `attachment:`
103 attach_str = "attachment:" + name
104 if attach_str in cell.source:
105 cell.source = cell.source.replace(attach_str, filename)
107 for name in to_delete:
108 del cell.attachments[name]
110 return cell, resources