xref: /aosp_15_r20/external/pigweed/pw_presubmit/py/pw_presubmit/format/python.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"""Code formatter plugins for Python."""
15
16import os
17from pathlib import Path
18from typing import Iterable, Iterator, Optional, Sequence, Tuple, Union
19
20from pw_presubmit.format.core import (
21    FileFormatter,
22    FormattedFileContents,
23    FormatFixStatus,
24)
25
26
27class BlackFormatter(FileFormatter):
28    """A formatter that runs ``black`` on files."""
29
30    def __init__(self, config_file: Optional[Path], **kwargs):
31        super().__init__(**kwargs)
32        self.config_file = config_file
33
34    def _config_file_args(self) -> Sequence[Union[str, Path]]:
35        if self.config_file is not None:
36            return ('--config', Path(self.config_file))
37
38        return ()
39
40    def format_file_in_memory(
41        self, file_path: Path, file_contents: bytes
42    ) -> FormattedFileContents:
43        """Uses ``black`` to check the formatting of the requested file.
44
45        The file at ``file_path`` is NOT modified by this check.
46
47        Returns:
48            A populated
49            :py:class:`pw_presubmit.format.core.FormattedFileContents` that
50            contains either the result of formatting the file, or an error
51            message.
52        """
53        proc = self.run_tool(
54            'black',
55            [*self._config_file_args(), '-q', '-'],
56            input=file_contents,
57        )
58        ok = proc.returncode == 0
59        formatted_file_contents = proc.stdout if ok else b''
60
61        # On Windows, Black's stdout always has CRLF line endings.
62        if os.name == 'nt':
63            formatted_file_contents = formatted_file_contents.replace(
64                b'\r\n', b'\n'
65            )
66
67        return FormattedFileContents(
68            ok=ok,
69            formatted_file_contents=formatted_file_contents,
70            error_message=None if ok else proc.stderr.decode(),
71        )
72
73    def format_file(self, file_path: Path) -> FormatFixStatus:
74        """Formats the provided file in-place using ``black``.
75
76        Returns:
77            A FormatFixStatus that contains relevant errors/warnings.
78        """
79        proc = self.run_tool(
80            'black',
81            [*self._config_file_args(), '-q', file_path],
82        )
83        ok = proc.returncode == 0
84        return FormatFixStatus(
85            ok=ok,
86            error_message=None if ok else proc.stderr.decode(),
87        )
88
89    def format_files(
90        self, paths: Iterable[Path], keep_warnings: bool = True
91    ) -> Iterator[Tuple[Path, FormatFixStatus]]:
92        """Uses ``black`` to format the specified files in-place.
93
94        Returns:
95            An iterator of ``Path`` and
96            :py:class:`pw_presubmit.format.core.FormatFixStatus` pairs for each
97            file that was not successfully formatted. If ``keep_warnings`` is
98            ``True``, any successful format operations with warnings will also
99            be returned.
100        """
101        proc = self.run_tool(
102            'black',
103            [*self._config_file_args(), '-q', *paths],
104        )
105
106        # If there's an error, fall back to per-file formatting to figure out
107        # which file has problems.
108        if proc.returncode != 0:
109            yield from super().format_files(paths, keep_warnings)
110