Coverage for databooks/metadata.py: 94%

1"""Metadata wrapper functions for cleaning notebook metadata.""" 

2 

3from pathlib import Path 

4from typing import Any, Callable, List, Optional, Sequence 

5 

6from databooks import JupyterNotebook 

7from databooks.data_models.cell import BaseCell 

8from databooks.logging import get_logger, set_verbose 

9 

10logger = get_logger(__file__) 

11 

12 

13def clear( 

14 read_path: Path, 

15 write_path: Optional[Path] = None, 

16 notebook_metadata_keep: Sequence[str] = (), 

17 cell_metadata_keep: Sequence[str] = (), 

18 cell_fields_keep: Sequence[str] = (), 

19 check: bool = False, 

20 verbose: bool = False, 

21 overwrite: bool = False, 

22 **kwargs: Any, 

23) -> bool: 

24 """ 

25 Clear Jupyter Notebook metadata. 

26 

27 Clear metadata (at notebook and cell level) and write clean 

28 notebook. By default, remove all metadata. 

29 :param read_path: Path of notebook file with metadata to be cleaned 

30 :param write_path: Path of notebook file with metadata to be cleaned 

31 :param notebook_metadata_keep: Notebook metadata fields to keep 

32 :param cell_metadata_keep: Cell metadata fields to keep 

33 :param cell_fields_keep: Cell fields to keep 

34 :param check: Don't write any files, check whether there is unwanted metadata 

35 :param verbose: Log written files 

36 :param overwrite: Whether to overwrite files (if exists) 

37 :param kwargs: Additional keyword arguments to pass to 

38 `databooks.data_models.JupyterNotebook.clear_metadata` 

39 :return: Whether notebooks are equal 

40 """ 

41 if verbose: 

42 set_verbose(logger) 

43 

44 if write_path is None: 

45 write_path = read_path 

46 notebook = JupyterNotebook.parse_file(read_path) 

47 

48 # Get fields to remove from cells and keep notebook schema 

49 cell_fields = {field for cell in notebook.cells for field, _ in cell if field} 

50 cell_fields_keep = list(cell_fields_keep) + list(BaseCell.model_fields) 

51 

52 cell_remove_fields = [ 

53 field for field in cell_fields if field not in cell_fields_keep 

54 ] 

55 

56 notebook.clear_metadata( 

57 notebook_metadata_keep=notebook_metadata_keep, 

58 cell_metadata_keep=cell_metadata_keep, 

59 cell_remove_fields=cell_remove_fields, 

60 **kwargs, 

61 ) 

62 nb_equals = notebook == JupyterNotebook.parse_file(read_path) 

63 

64 if nb_equals or check: 

65 msg = ( 

66 "no metadata to remove." 

67 if nb_equals 

68 else "only check (unwanted metadata found)." 

69 ) 

70 logger.debug(f"No action taken for {read_path} - {msg}") 

71 else: 

72 notebook.write(path=write_path, overwrite=overwrite) 

73 logger.debug(f"Removed metadata from {read_path}, saved as {write_path}") 

74 

75 return nb_equals 

76 

77 

78def clear_all( 

79 read_paths: List[Path], 

80 write_paths: List[Path], 

81 *, 

82 progress_callback: Callable[[], None] = lambda: None, 

83 **clear_kwargs: Any, 

84) -> List[bool]: 

85 """ 

86 Clear metadata for multiple notebooks at notebooks and cell level. 

87 

88 :param read_paths: Paths of notebook to remove metadata 

89 :param write_paths: Paths of where to write cleaned notebooks 

90 :param progress_callback: Callback function to report progress 

91 :param clear_kwargs: Keyword arguments to be passed to `databooks.metadata.clear` 

92 :return: Whether the notebooks contained or not unwanted metadata 

93 """ 

94 if len(read_paths) != len(write_paths): 

95 raise ValueError( 

96 "Read and write paths must have same length." 

97 f" Got {len(read_paths)} and {len(write_paths)}" 

98 ) 

99 checks = [] 

100 for nb_path, write_path in zip(read_paths, write_paths): 

101 checks.append(clear(read_path=nb_path, write_path=write_path, **clear_kwargs)) 

102 progress_callback() 

103 return checks