Coverage for databooks/conflicts.py: 92%

1"""Functions to resolve any git conflicts between notebooks.""" 

2 

3from __future__ import annotations 

4 

5from pathlib import Path 

6from typing import Any, Callable, Dict, List, Optional, Tuple, cast 

7 

8from git import Repo 

9 

10from databooks.common import get_logger, set_verbose, write_notebook 

11from databooks.data_models.base import BaseCells, DiffModel 

12from databooks.data_models.notebook import Cell, Cells, JupyterNotebook 

13from databooks.git_utils import ConflictFile, get_conflict_blobs, get_repo 

14 

15logger = get_logger(__file__) 

16 

17 

18class DiffJupyterNotebook(DiffModel): 

19 """Protocol for mypy static type checking.""" 

20 

21 nbformat: int 

22 nbformat_minor: int 

23 metadata: Dict[str, Any] 

24 cells: BaseCells[Any] 

25 

26 

27def path2conflicts( 

28 nb_paths: List[Path], repo: Optional[Repo] = None 

29) -> List[ConflictFile]: 

30 """ 

31 Get the difference model from the path based on the git conflict information. 

32 

33 :param nb_path: Path to file with conflicts (must be notebook paths) 

34 :return: Generator of `DiffModel`s, to be resolved 

35 """ 

36 if any(nb_path.suffix not in ("", ".ipynb") for nb_path in nb_paths): 

37 raise ValueError( 

38 "Expected either notebook files, a directory or glob expression." 

39 ) 

40 common_parent = max(set.intersection(*[set(p.parents) for p in nb_paths])) 

41 repo = get_repo(common_parent) if repo is None else repo 

42 return [ 

43 file 

44 for file in get_conflict_blobs(repo=repo) 

45 if any(file.filename.match(str(p.name)) for p in nb_paths) 

46 ] 

47 

48 

49def conflict2nb( 

50 conflict_file: ConflictFile, 

51 *, 

52 keep_first: bool = True, 

53 cells_first: Optional[bool] = None, 

54 ignore_none: bool = True, 

55 verbose: bool = False, 

56) -> JupyterNotebook: 

57 """ 

58 Merge diffs from conflicts and return valid a notebook. 

59 

60 :param conflict_file: A `databooks.git_utils.ConflictFile` with conflicts 

61 :param keep_first: Whether to keep the metadata of the first or last notebook 

62 :param cells_first: Whether to keep the cells of the first or last notebook 

63 :param ignore_none: Keep all metadata fields even if it's included in only one 

64 notebook 

65 :param verbose: Log written files and metadata conflicts 

66 :return: Resolved conflicts as a `databooks.data_models.notebook.JupyterNotebook` 

67 model 

68 """ 

69 if verbose: 

70 set_verbose(logger) 

71 

72 nb_1 = JupyterNotebook.parse_raw(conflict_file.first_contents) 

73 nb_2 = JupyterNotebook.parse_raw(conflict_file.last_contents) 

74 if nb_1.metadata != nb_2.metadata: 

75 msg = ( 

76 f"Notebook metadata conflict for {conflict_file.filename}. Keeping " 

77 + "first." 

78 if keep_first 

79 else "last." 

80 ) 

81 logger.debug(msg) 

82 

83 diff_nb = cast(DiffModel, nb_1 - nb_2) 

84 nb = cast( 

85 JupyterNotebook, 

86 diff_nb.resolve( 

87 ignore_none=ignore_none, 

88 keep_first=keep_first, 

89 keep_first_cells=cells_first, 

90 first_id=conflict_file.first_log, 

91 last_id=conflict_file.last_log, 

92 ), 

93 ) 

94 logger.debug(f"Resolved conflicts in {conflict_file.filename}.") 

95 return nb 

96 

97 

98def conflicts2nbs( 

99 conflict_files: List[ConflictFile], 

100 *, 

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

102 **conflict2nb_kwargs: Any, 

103) -> None: 

104 """ 

105 Get notebooks from conflicts. 

106 

107 Wrap `databooks.conflicts.conflict2nb` to write notebooks to list of 

108 `databooks.git_utils.ConflictFile`. 

109 :param conflict_files: Files with source conflict files and one-liner git logs 

110 :param progress_callback: Callback function to report progress 

111 :param conflict2nb_kwargs: Keyword arguments to be passed to 

112 `databooks.conflicts.conflict2nb` 

113 :return: 

114 """ 

115 for conflict in conflict_files: 

116 nb = conflict2nb(conflict, **conflict2nb_kwargs) 

117 write_notebook(nb=nb, path=conflict.filename) 

118 progress_callback() 

119 

120 

121def cells_equals(diff_cells: Cells[Tuple[List[Cell], List[Cell]]]) -> List[bool]: 

122 """Return if cells in `DiffCells` are equal.""" 

123 return [cell_first == cell_last for cell_first, cell_last in diff_cells]