xref: /aosp_15_r20/external/pigweed/pw_cli/py/pw_cli/diff.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2024 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Print git-style diffs."""
15
16
17import difflib
18import sys
19from typing import Iterable, Sequence, Union
20
21from pw_cli.color import colors
22
23_COLOR = colors()
24
25
26def colorize_diff_line(line: str) -> str:
27    if line.startswith('--- ') or line.startswith('+++ '):
28        return _COLOR.bold_white(line)
29    if line.startswith('-'):
30        return _COLOR.red(line)
31    if line.startswith('+'):
32        return _COLOR.green(line)
33    if line.startswith('@@ '):
34        return _COLOR.cyan(line)
35    return line
36
37
38def colorize_diff(lines: Union[str, Iterable[str]]) -> str:
39    """Takes a diff str or list of str lines and returns a colorized version."""
40    if isinstance(lines, str):
41        lines = lines.splitlines(True)
42
43    return ''.join(colorize_diff_line(line) for line in lines)
44
45
46def print_diff(
47    a: Sequence[str], b: Sequence[str], fromfile="", tofile="", indent=0
48) -> None:
49    """Print a git-style diff of two string sequences."""
50    indent_str = ' ' * indent
51    diff = colorize_diff(difflib.unified_diff(a, b, fromfile, tofile))
52
53    for line in diff.splitlines(True):
54        sys.stdout.write(indent_str + line)
55
56    sys.stdout.write("\n")
57