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

1import os 

2import sys 

3from binascii import a2b_base64 

4 

5from nbgrader.preprocessors import NbGraderPreprocessor 

6from traitlets import Set, Unicode 

7 

8from ..utils.extra_cells import is_attachment_cell, is_diagram 

9 

10 

11class ExtractAttachments(NbGraderPreprocessor): 

12 subdirectory = Unicode("attachments").tag(config=True) 

13 

14 output_filename_template = Unicode("attach_{cell_index}_{name}").tag(config=True) 

15 

16 extract_output_types = Set( 

17 {"image/png", "image/jpeg", "image/svg+xml", "application/pdf"} 

18 ).tag(config=True) 

19 

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) 

24 

25 for cell_index, cell in enumerate(nb.cells): 

26 cell, resources = self.preprocess_cell(cell, resources, cell_index) 

27 

28 return nb, resources 

29 

30 def preprocess_cell(self, cell, resources, cell_index): 

31 """ 

32 Apply a transformation on each cell, 

33 

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 """ 

44 

45 if not (is_attachment_cell(cell) or is_diagram(cell)): 

46 return cell, resources 

47 

48 if is_diagram(cell): 

49 # Remove the attachment line from source 

50 cell.source = cell.source.replace("![diagram](attachment:diagram.png)", "") 

51 

52 # Get files directory if it has been specified 

53 output_files_dir = resources.get("output_files_dir", None) 

54 

55 # Make sure outputs key exists 

56 if not isinstance(resources["outputs"], dict): 

57 resources["outputs"] = {} 

58 

59 to_delete = [] 

60 

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 

66 

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") 

76 

77 filename = self.output_filename_template.format( 

78 cell_index=cell_index, 

79 name=name, 

80 ) 

81 link = os.path.join(self.subdirectory, filename) 

82 

83 if output_files_dir is not None: 

84 filename = os.path.join(self.path, filename) 

85 

86 if name.endswith(".gif") and mime == "image/png": 

87 filename = filename.replace(".gif", ".png") 

88 

89 # Write the attachments to a file and add a link 

90 

91 with open(filename, "wb") as f: 

92 f.write(data) 

93 

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) 

100 

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) 

106 

107 for name in to_delete: 

108 del cell.attachments[name] 

109 

110 return cell, resources