• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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, Mapping, Sequence, Tuple
19
20from pw_cli.file_filter import FileFilter
21from pw_cli import find_config
22from pw_presubmit.format.core import (
23    FileFormatter,
24    FormattedFileContents,
25    FormatFixStatus,
26)
27
28
29DEFAULT_PYTHON_FILE_PATTERNS = FileFilter(endswith=['.py'])
30
31# Pigweed prescribes a .black.toml naming pattern for black configs.
32_PIGWEED_BLACK_CONFIG_PATTERN = '.black.toml'
33
34
35class BlackFormatter(FileFormatter):
36    """A formatter that runs ``black`` on files."""
37
38    def __init__(self, config_file: Path | bool = True, **kwargs):
39        """Creates a formatter for Python that uses black.
40
41        Args:
42            config_file: The black config file to use to configure black. This
43                defaults to ``True``, which formats files using the nearest
44                ``.black.toml`` file in the parent directory of the file being
45                formatted. ``False`` disables this behavior entirely.
46        """
47        kwargs.setdefault('mnemonic', 'Python (black)')
48        kwargs.setdefault('file_patterns', DEFAULT_PYTHON_FILE_PATTERNS)
49        super().__init__(**kwargs)
50        self._config_file_override: Path | None = (
51            config_file if isinstance(config_file, Path) else None
52        )
53        self._enable_config_lookup = config_file is True
54
55    def _config_file_for(self, file_path: Path) -> Path | None:
56        if self._config_file_override:
57            return self._config_file_override
58        if not self._enable_config_lookup:
59            return None
60
61        # Search for Pigweed's `.black.toml`
62        configs = find_config.configs_in_parents(
63            _PIGWEED_BLACK_CONFIG_PATTERN, file_path
64        )
65        return next(configs, None)
66
67    def _config_file_args(self, file_path: Path) -> Sequence[str]:
68        config = self._config_file_for(file_path)
69        if config:
70            return ('--config', str(config))
71
72        return ()
73
74    def format_file_in_memory(
75        self, file_path: Path, file_contents: bytes
76    ) -> FormattedFileContents:
77        """Uses ``black`` to check the formatting of the requested file.
78
79        The file at ``file_path`` is NOT modified by this check.
80
81        Returns:
82            A populated
83            :py:class:`pw_presubmit.format.core.FormattedFileContents` that
84            contains either the result of formatting the file, or an error
85            message.
86        """
87        proc = self.run_tool(
88            'black',
89            [*self._config_file_args(file_path), '-q', '-'],
90            input=file_contents,
91        )
92        ok = proc.returncode == 0
93        formatted_file_contents = proc.stdout if ok else b''
94
95        # On Windows, Black's stdout always has CRLF line endings.
96        if os.name == 'nt':
97            formatted_file_contents = formatted_file_contents.replace(
98                b'\r\n', b'\n'
99            )
100
101        return FormattedFileContents(
102            ok=ok,
103            formatted_file_contents=formatted_file_contents,
104            error_message=None if ok else proc.stderr.decode(),
105        )
106
107    def format_file(self, file_path: Path) -> FormatFixStatus:
108        """Formats the provided file in-place using ``black``.
109
110        Returns:
111            A FormatFixStatus that contains relevant errors/warnings.
112        """
113        proc = self.run_tool(
114            'black',
115            [*self._config_file_args(file_path), '-q', file_path],
116        )
117        ok = proc.returncode == 0
118        return FormatFixStatus(
119            ok=ok,
120            error_message=None if ok else proc.stderr.decode(),
121        )
122
123    def format_files(
124        self, paths: Iterable[Path], keep_warnings: bool = True
125    ) -> Iterator[Tuple[Path, FormatFixStatus]]:
126        """Uses ``black`` to format the specified files in-place.
127
128        Returns:
129            An iterator of ``Path`` and
130            :py:class:`pw_presubmit.format.core.FormatFixStatus` pairs for each
131            file that was not successfully formatted. If ``keep_warnings`` is
132            ``True``, any successful format operations with warnings will also
133            be returned.
134        """
135        paths_by_config: Mapping[Path | None, Iterable[Path]] = {}
136        if self._config_file_override:
137            paths_by_config = {self._config_file_override: paths}
138        elif self._enable_config_lookup:
139            paths_by_config = find_config.paths_by_nearest_config(
140                _PIGWEED_BLACK_CONFIG_PATTERN, paths
141            )
142        else:
143            paths_by_config = {None: paths}
144
145        for config, group_paths in paths_by_config.items():
146            config_file_args = ('--config', str(config)) if config else ()
147            proc = self.run_tool(
148                'black',
149                [*config_file_args, '-q', *group_paths],
150            )
151
152            # If there's an error, fall back to per-file formatting to figure
153            # out which file has problems.
154            if proc.returncode != 0:
155                yield from super().format_files(group_paths, keep_warnings)
156