Coverage for databooks/git_utils.py: 95%

1"""Git helper functions.""" 

2 

3from dataclasses import dataclass 

4from enum import Enum 

5from pathlib import Path 

6from typing import Dict, List, Optional, Sequence, Union, cast, overload 

7 

8from git import Git, InvalidGitRepositoryError, NoSuchPathError 

9from git.index import IndexFile 

10from git.objects.blob import Blob 

11from git.objects.commit import Commit 

12from git.objects.tree import Tree 

13from git.repo import Repo 

14 

15from databooks.common import find_common_parent 

16from databooks.logging import get_logger, set_verbose 

17 

18logger = get_logger(name=__file__) 

19 

20 

21class ChangeType(Enum): 

22 """Git change kinds, retaining the original public enum values.""" 

23 

24 A = 1 

25 C = 2 

26 D = 3 

27 R = 4 

28 M = 5 

29 T = 6 

30 U = 7 

31 

32 

33@dataclass 

34class UnmergedBlob: 

35 """Container for git unmerged blobs.""" 

36 

37 filename: Path 

38 stage: Dict[int, Blob] 

39 

40 

41@dataclass 

42class ConflictFile: 

43 """Container for path and different versions of conflicted notebooks.""" 

44 

45 filename: Path 

46 first_log: str 

47 last_log: str 

48 first_contents: str | bytes 

49 last_contents: str | bytes 

50 

51 

52@dataclass 

53class Contents: 

54 """Container for path of file versions.""" 

55 

56 path: Optional[Path] 

57 contents: Optional[str | bytes] 

58 

59 

60@dataclass 

61class DiffContents: 

62 """Container for path and different versions of conflicted notebooks.""" 

63 

64 a: Contents 

65 b: Contents 

66 change_type: ChangeType 

67 

68 

69@overload 

70def blob2str(blob: None) -> None: ... 

71 

72 

73@overload 

74def blob2str(blob: Blob) -> bytes: ... 

75 

76 

77def blob2str(blob: Optional[Blob]) -> Optional[bytes]: 

78 """Get the blob contents if they exist (otherwise return `None`).""" 

79 return blob.data_stream.read() if blob is not None else None 

80 

81 

82def blob2commit(blob: Blob, repo: Repo) -> str: 

83 """Get the short commit message from blob hash.""" 

84 _git = Git(working_dir=repo.working_dir) 

85 commit_id = _git.log(find_object=blob, max_count=1, all=True, oneline=True) 

86 return ( 

87 commit_id 

88 if len(commit_id) > 0 

89 else _git.stash("list", "--oneline", "--max-count", "1", "--find-object", blob) 

90 ) 

91 

92 

93def diff2contents( 

94 blob: Blob, 

95 ref: Optional[Union[IndexFile, Tree, Commit, str]], 

96 path: Path, 

97 not_exists: bool = False, 

98) -> Optional[str | bytes]: 

99 """ 

100 Get the blob contents from the diff. 

101 

102 Depends on whether we are diffing against current working tree and if object exists 

103 at diff time (added or deleted objects only exist at one side). If comparing 

104 against working tree (`ref=None`) we return the current file contents. 

105 :param blob: git diff blob 

106 :param ref: git reference 

107 :param path: path to object 

108 :param not_exists: whether object exists at 'diff time' (added or removed objects 

109 do not exist) 

110 :return: blob contents as a string (if exists) 

111 """ 

112 if not_exists: 

113 return None 

114 elif ref is None: 

115 return path.read_text(encoding="utf-8") 

116 else: 

117 return blob2str(blob) 

118 

119 

120def get_repo(path: Path) -> Optional[Repo]: 

121 """Find git repo in current or parent directories.""" 

122 try: 

123 repo = Repo(path=path, search_parent_directories=True) 

124 except (InvalidGitRepositoryError, NoSuchPathError): 

125 logger.debug(f"No repo found at {path}.") 

126 return None 

127 logger.debug(f"Repo found at: {repo.working_dir}.") 

128 return repo 

129 

130 

131def get_conflict_blobs(repo: Repo) -> List[ConflictFile]: 

132 """Get the source files for conflicts.""" 

133 unmerged_blobs = repo.index.unmerged_blobs() 

134 blobs = ( 

135 UnmergedBlob(filename=Path(k), stage=dict(v)) 

136 for k, v in unmerged_blobs.items() 

137 if 0 not in dict(v).keys() # only get blobs that could not be merged 

138 ) 

139 

140 if not isinstance(repo.working_dir, (Path, str)): 

141 raise RuntimeError( 

142 "Expected `repo` to be `pathlib.Path` or `str`, got" 

143 f" {type(repo.working_dir)}." 

144 ) 

145 return [ 

146 ConflictFile( 

147 filename=repo.working_dir / blob.filename, 

148 first_log=blob2commit(blob=blob.stage[2], repo=repo), 

149 last_log=blob2commit(blob=blob.stage[3], repo=repo), 

150 first_contents=blob2str(blob.stage[2]), 

151 last_contents=blob2str(blob.stage[3]), 

152 ) 

153 for blob in blobs 

154 ] 

155 

156 

157def get_nb_diffs( 

158 ref_base: Optional[str] = None, 

159 ref_remote: Optional[str] = None, 

160 paths: Sequence[Path] = (), 

161 *, 

162 repo: Optional[Repo] = None, 

163 verbose: bool = False, 

164) -> List[DiffContents]: 

165 """ 

166 Get the noteebook(s) git diff(s). 

167 

168 By default, diffs are compared with the current working directory. That is, staged 

169 files will still show up in the diffs. Only return the diffs for notebook files. 

170 """ 

171 if verbose: 

172 set_verbose(logger) 

173 

174 common_path = find_common_parent(paths or [Path.cwd()]) 

175 repo = get_repo(path=common_path) if repo is None else repo 

176 if repo is None or repo.working_dir is None: 

177 raise ValueError("No repo found - cannot compute diffs.") 

178 

179 base = repo.index if ref_base is None else repo.tree(ref_base) 

180 remote = None if ref_remote is None else repo.tree(ref_remote) 

181 

182 logger.debug( 

183 f"Looking for diffs on path(s) {[p.resolve() for p in paths]}.\n" 

184 f"Comparing `{base}` and `{remote}`." 

185 ) 

186 repo_root_dir = Path(repo.working_dir) 

187 return [ 

188 DiffContents( 

189 a=Contents( 

190 path=Path(cast(str, d.a_path)), 

191 contents=diff2contents( 

192 blob=cast(Blob, d.a_blob), 

193 ref=base, 

194 path=repo_root_dir / cast(str, d.a_path), 

195 not_exists=d.change_type is ChangeType.A, 

196 ), 

197 ), 

198 b=Contents( 

199 path=Path(cast(str, d.b_path)), 

200 contents=diff2contents( 

201 blob=cast(Blob, d.b_blob), 

202 ref=remote, 

203 path=repo_root_dir / cast(str, d.b_path), 

204 not_exists=d.change_type is ChangeType.D, 

205 ), 

206 ), 

207 change_type=ChangeType[cast(str, d.change_type)], 

208 ) 

209 for d in base.diff( 

210 other=remote, paths=list(paths) or list(repo_root_dir.rglob("*.ipynb")) 

211 ) 

212 ]