Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/libcst/codemod/_cli.py: 25%
225 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-25 06:43 +0000
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2#
3# This source code is licensed under the MIT license found in the
4# LICENSE file in the root directory of this source tree.
5#
6"""
7Provides helpers for CLI interaction.
8"""
10import difflib
11import os.path
12import re
13import subprocess
14import sys
15import time
16import traceback
17from dataclasses import dataclass, replace
18from multiprocessing import cpu_count, Pool
19from pathlib import Path
20from typing import Any, AnyStr, cast, Dict, List, Optional, Sequence, Union
22from libcst import parse_module, PartialParserConfig
23from libcst.codemod._codemod import Codemod
24from libcst.codemod._dummy_pool import DummyPool
25from libcst.codemod._runner import (
26 SkipFile,
27 SkipReason,
28 transform_module,
29 TransformExit,
30 TransformFailure,
31 TransformResult,
32 TransformSkip,
33 TransformSuccess,
34)
35from libcst.helpers import calculate_module_and_package
36from libcst.metadata import FullRepoManager
38_DEFAULT_GENERATED_CODE_MARKER: str = f"@gen{''}erated"
41def invoke_formatter(formatter_args: Sequence[str], code: AnyStr) -> AnyStr:
42 """
43 Given a code string, run an external formatter on the code and return new
44 formatted code.
45 """
47 # Make sure there is something to run
48 if len(formatter_args) == 0:
49 raise Exception("No formatter configured but code formatting requested.")
51 # Invoke the formatter, giving it the code as stdin and assuming the formatted
52 # code comes from stdout.
53 work_with_bytes = isinstance(code, bytes)
54 return cast(
55 AnyStr,
56 subprocess.check_output(
57 formatter_args,
58 input=code,
59 universal_newlines=not work_with_bytes,
60 encoding=None if work_with_bytes else "utf-8",
61 ),
62 )
65def print_execution_result(result: TransformResult) -> None:
66 for warning in result.warning_messages:
67 print(f"WARNING: {warning}", file=sys.stderr)
69 if isinstance(result, TransformFailure):
70 error = result.error
71 if isinstance(error, subprocess.CalledProcessError):
72 print(error.output.decode("utf-8"), file=sys.stderr)
73 print(result.traceback_str, file=sys.stderr)
76def gather_files(
77 files_or_dirs: Sequence[str], *, include_stubs: bool = False
78) -> List[str]:
79 """
80 Given a list of files or directories (can be intermingled), return a list of
81 all python files that exist at those locations. If ``include_stubs`` is ``True``,
82 this will include ``.py`` and ``.pyi`` stub files. If it is ``False``, only
83 ``.py`` files will be included in the returned list.
84 """
85 ret: List[str] = []
86 for fd in files_or_dirs:
87 if os.path.isfile(fd):
88 ret.append(fd)
89 elif os.path.isdir(fd):
90 ret.extend(
91 str(p)
92 for p in Path(fd).rglob("*.py*")
93 if Path.is_file(p)
94 and (
95 str(p).endswith("py") or (include_stubs and str(p).endswith("pyi"))
96 )
97 )
98 return sorted(ret)
101def diff_code(
102 oldcode: str, newcode: str, context: int, *, filename: Optional[str] = None
103) -> str:
104 """
105 Given two strings representing a module before and after a codemod, produce
106 a unified diff of the changes with ``context`` lines of context. Optionally,
107 assign the ``filename`` to the change, and if it is not available, assume
108 that the change was performed on stdin/stdout. If no change is detected,
109 return an empty string instead of returning an empty unified diff. This is
110 comparable to revision control software which only shows differences for
111 files that have changed.
112 """
114 if oldcode == newcode:
115 return ""
117 if filename:
118 difflines = difflib.unified_diff(
119 oldcode.split("\n"),
120 newcode.split("\n"),
121 fromfile=filename,
122 tofile=filename,
123 lineterm="",
124 n=context,
125 )
126 else:
127 difflines = difflib.unified_diff(
128 oldcode.split("\n"), newcode.split("\n"), lineterm="", n=context
129 )
130 return "\n".join(difflines)
133def exec_transform_with_prettyprint(
134 transform: Codemod,
135 code: str,
136 *,
137 include_generated: bool = False,
138 generated_code_marker: str = _DEFAULT_GENERATED_CODE_MARKER,
139 format_code: bool = False,
140 formatter_args: Sequence[str] = (),
141 python_version: Optional[str] = None,
142) -> Optional[str]:
143 """
144 Given an instantiated codemod and a string representing a module, transform that
145 code by executing the transform, optionally invoking the formatter and finally
146 printing any generated warnings to stderr. If the code includes the generated
147 marker at any spot and ``include_generated`` is not set to ``True``, the code
148 will not be modified. If ``format_code`` is set to ``False`` or the instantiated
149 codemod does not modify the code, the code will not be formatted. If a
150 ``python_version`` is provided, then we will parse the module using
151 this version. Otherwise, we will use the version of the currently executing python
152 binary.
154 In all cases a module will be returned. Whether it is changed depends on the
155 input parameters as well as the codemod itself.
156 """
158 if not include_generated and generated_code_marker in code:
159 print(
160 "WARNING: Code is generated and we are set to ignore generated code, "
161 + "skipping!",
162 file=sys.stderr,
163 )
164 return code
166 result = transform_module(transform, code, python_version=python_version)
167 maybe_code: Optional[str] = (
168 None
169 if isinstance(result, (TransformFailure, TransformExit, TransformSkip))
170 else result.code
171 )
173 if maybe_code is not None and format_code:
174 try:
175 maybe_code = invoke_formatter(formatter_args, maybe_code)
176 except Exception as ex:
177 # Failed to format code, treat as a failure and make sure that
178 # we print the exception for debugging.
179 maybe_code = None
180 result = TransformFailure(
181 error=ex,
182 traceback_str=traceback.format_exc(),
183 warning_messages=result.warning_messages,
184 )
186 # Finally, print the output, regardless of what happened
187 print_execution_result(result)
188 return maybe_code
191@dataclass(frozen=True)
192class ExecutionResult:
193 # File we have results for
194 filename: str
195 # Whether we actually changed the code for the file or not
196 changed: bool
197 # The actual result
198 transform_result: TransformResult
201@dataclass(frozen=True)
202class ExecutionConfig:
203 blacklist_patterns: Sequence[str] = ()
204 format_code: bool = False
205 formatter_args: Sequence[str] = ()
206 generated_code_marker: str = _DEFAULT_GENERATED_CODE_MARKER
207 include_generated: bool = False
208 python_version: Optional[str] = None
209 repo_root: Optional[str] = None
210 unified_diff: Optional[int] = None
213def _execute_transform( # noqa: C901
214 transformer: Codemod,
215 filename: str,
216 config: ExecutionConfig,
217) -> ExecutionResult:
218 for pattern in config.blacklist_patterns:
219 if re.fullmatch(pattern, filename):
220 return ExecutionResult(
221 filename=filename,
222 changed=False,
223 transform_result=TransformSkip(
224 skip_reason=SkipReason.BLACKLISTED,
225 skip_description=f"Blacklisted by pattern {pattern}.",
226 ),
227 )
229 try:
230 with open(filename, "rb") as fp:
231 oldcode = fp.read()
233 # Skip generated files
234 if (
235 not config.include_generated
236 and config.generated_code_marker.encode("utf-8") in oldcode
237 ):
238 return ExecutionResult(
239 filename=filename,
240 changed=False,
241 transform_result=TransformSkip(
242 skip_reason=SkipReason.GENERATED,
243 skip_description="Generated file.",
244 ),
245 )
247 # Somewhat gross hack to provide the filename in the transform's context.
248 # We do this after the fork so that a context that was initialized with
249 # some defaults before calling parallel_exec_transform_with_prettyprint
250 # will be updated per-file.
251 transformer.context = replace(
252 transformer.context,
253 filename=filename,
254 scratch={},
255 )
257 # determine the module and package name for this file
258 try:
259 module_name_and_package = calculate_module_and_package(
260 config.repo_root or ".", filename
261 )
262 transformer.context = replace(
263 transformer.context,
264 full_module_name=module_name_and_package.name,
265 full_package_name=module_name_and_package.package,
266 )
267 except ValueError as ex:
268 print(
269 f"Failed to determine module name for {filename}: {ex}", file=sys.stderr
270 )
272 # Run the transform, bail if we failed or if we aren't formatting code
273 try:
274 input_tree = parse_module(
275 oldcode,
276 config=(
277 PartialParserConfig(python_version=str(config.python_version))
278 if config.python_version is not None
279 else PartialParserConfig()
280 ),
281 )
282 output_tree = transformer.transform_module(input_tree)
283 newcode = output_tree.bytes
284 encoding = output_tree.encoding
285 except KeyboardInterrupt:
286 return ExecutionResult(
287 filename=filename, changed=False, transform_result=TransformExit()
288 )
289 except SkipFile as ex:
290 return ExecutionResult(
291 filename=filename,
292 changed=False,
293 transform_result=TransformSkip(
294 skip_reason=SkipReason.OTHER,
295 skip_description=str(ex),
296 warning_messages=transformer.context.warnings,
297 ),
298 )
299 except Exception as ex:
300 return ExecutionResult(
301 filename=filename,
302 changed=False,
303 transform_result=TransformFailure(
304 error=ex,
305 traceback_str=traceback.format_exc(),
306 warning_messages=transformer.context.warnings,
307 ),
308 )
310 # Call formatter if needed, but only if we actually changed something in this
311 # file
312 if config.format_code and newcode != oldcode:
313 try:
314 newcode = invoke_formatter(config.formatter_args, newcode)
315 except KeyboardInterrupt:
316 return ExecutionResult(
317 filename=filename,
318 changed=False,
319 transform_result=TransformExit(),
320 )
321 except Exception as ex:
322 return ExecutionResult(
323 filename=filename,
324 changed=False,
325 transform_result=TransformFailure(
326 error=ex,
327 traceback_str=traceback.format_exc(),
328 warning_messages=transformer.context.warnings,
329 ),
330 )
332 # Format as unified diff if needed, otherwise save it back
333 changed = oldcode != newcode
334 if config.unified_diff:
335 newcode = diff_code(
336 oldcode.decode(encoding),
337 newcode.decode(encoding),
338 config.unified_diff,
339 filename=filename,
340 )
341 else:
342 # Write back if we changed
343 if changed:
344 with open(filename, "wb") as fp:
345 fp.write(newcode)
346 # Not strictly necessary, but saves space in pickle since we won't use it
347 newcode = ""
349 # Inform success
350 return ExecutionResult(
351 filename=filename,
352 changed=changed,
353 transform_result=TransformSuccess(
354 warning_messages=transformer.context.warnings, code=newcode
355 ),
356 )
357 except KeyboardInterrupt:
358 return ExecutionResult(
359 filename=filename, changed=False, transform_result=TransformExit()
360 )
361 except Exception as ex:
362 return ExecutionResult(
363 filename=filename,
364 changed=False,
365 transform_result=TransformFailure(
366 error=ex,
367 traceback_str=traceback.format_exc(),
368 warning_messages=transformer.context.warnings,
369 ),
370 )
373class Progress:
374 ERASE_CURRENT_LINE: str = "\r\033[2K"
376 def __init__(self, *, enabled: bool, total: int) -> None:
377 self.enabled = enabled
378 self.total = total
379 # 1/100 = 0, len("0") = 1, precision = 0, more digits for more files
380 self.pretty_precision: int = len(str(self.total // 100)) - 1
381 # Pretend we start processing immediately. This is not true, but it's
382 # close enough to true.
383 self.started_at: float = time.time()
385 def print(self, finished: int) -> None:
386 if not self.enabled:
387 return
388 left = self.total - finished
389 percent = 100.0 * (float(finished) / float(self.total))
390 elapsed_time = max(time.time() - self.started_at, 0)
392 print(
393 f"{self.ERASE_CURRENT_LINE}{self._human_seconds(elapsed_time)} {percent:.{self.pretty_precision}f}% complete, {self.estimate_completion(elapsed_time, finished, left)} estimated for {left} files to go...",
394 end="",
395 file=sys.stderr,
396 )
398 def _human_seconds(self, seconds: Union[int, float]) -> str:
399 """
400 This returns a string which is a human-ish readable elapsed time such
401 as 30.42s or 10m 31s
402 """
404 minutes, seconds = divmod(seconds, 60)
405 hours, minutes = divmod(minutes, 60)
406 if hours > 0:
407 return f"{hours:.0f}h {minutes:02.0f}m {seconds:02.0f}s"
408 elif minutes > 0:
409 return f"{minutes:02.0f}m {seconds:02.0f}s"
410 else:
411 return f"{seconds:02.2f}s"
413 def estimate_completion(
414 self, elapsed_seconds: float, files_finished: int, files_left: int
415 ) -> str:
416 """
417 Computes a really basic estimated completion given a number of
418 operations still to do.
419 """
421 if files_finished <= 0:
422 # Technically infinite but calculating sounds better.
423 return "[calculating]"
425 fps = files_finished / elapsed_seconds
426 estimated_seconds_left = files_left / fps
427 return self._human_seconds(estimated_seconds_left)
429 def clear(self) -> None:
430 if not self.enabled:
431 return
432 print(self.ERASE_CURRENT_LINE, end="", file=sys.stderr)
435def _print_parallel_result(
436 exec_result: ExecutionResult,
437 progress: Progress,
438 *,
439 unified_diff: bool,
440 show_successes: bool,
441 hide_generated: bool,
442 hide_blacklisted: bool,
443) -> None:
444 filename = exec_result.filename
445 result = exec_result.transform_result
447 if isinstance(result, TransformSkip):
448 # Skipped file, print message and don't write back since not changed.
449 if not (
450 (result.skip_reason is SkipReason.BLACKLISTED and hide_blacklisted)
451 or (result.skip_reason is SkipReason.GENERATED and hide_generated)
452 ):
453 progress.clear()
454 print(f"Codemodding {filename}", file=sys.stderr)
455 print_execution_result(result)
456 print(
457 f"Skipped codemodding {filename}: {result.skip_description}\n",
458 file=sys.stderr,
459 )
460 elif isinstance(result, TransformFailure):
461 # Print any exception, don't write the file back.
462 progress.clear()
463 print(f"Codemodding {filename}", file=sys.stderr)
464 print_execution_result(result)
465 print(f"Failed to codemod {filename}\n", file=sys.stderr)
466 elif isinstance(result, TransformSuccess):
467 if show_successes or result.warning_messages:
468 # Print any warnings, save the changes if there were any.
469 progress.clear()
470 print(f"Codemodding {filename}", file=sys.stderr)
471 print_execution_result(result)
472 print(
473 f"Successfully codemodded {filename}"
474 + (" with warnings\n" if result.warning_messages else "\n"),
475 file=sys.stderr,
476 )
478 # In unified diff mode, the code is a diff we must print.
479 if unified_diff and result.code:
480 print(result.code)
483@dataclass(frozen=True)
484class ParallelTransformResult:
485 """
486 The result of running
487 :func:`~libcst.codemod.parallel_exec_transform_with_prettyprint` against
488 a series of files. This is a simple summary, with counts for number of
489 successfully codemodded files, number of files that we failed to codemod,
490 number of warnings generated when running the codemod across the files, and
491 the number of files that we skipped when running the codemod.
492 """
494 #: Number of files that we successfully transformed.
495 successes: int
496 #: Number of files that we failed to transform.
497 failures: int
498 #: Number of warnings generated when running transform across files.
499 warnings: int
500 #: Number of files skipped because they were blacklisted, generated
501 #: or the codemod requested to skip.
502 skips: int
505# Unfortunate wrapper required since there is no `istarmap_unordered`...
506def _execute_transform_wrap(
507 job: Dict[str, Any],
508) -> ExecutionResult:
509 return _execute_transform(**job)
512def parallel_exec_transform_with_prettyprint( # noqa: C901
513 transform: Codemod,
514 files: Sequence[str],
515 *,
516 jobs: Optional[int] = None,
517 unified_diff: Optional[int] = None,
518 include_generated: bool = False,
519 generated_code_marker: str = _DEFAULT_GENERATED_CODE_MARKER,
520 format_code: bool = False,
521 formatter_args: Sequence[str] = (),
522 show_successes: bool = False,
523 hide_generated: bool = False,
524 hide_blacklisted: bool = False,
525 hide_progress: bool = False,
526 blacklist_patterns: Sequence[str] = (),
527 python_version: Optional[str] = None,
528 repo_root: Optional[str] = None,
529) -> ParallelTransformResult:
530 """
531 Given a list of files and an instantiated codemod we should apply to them,
532 fork and apply the codemod in parallel to all of the files, including any
533 configured formatter. The ``jobs`` parameter controls the maximum number of
534 in-flight transforms, and needs to be at least 1. If not included, the number
535 of jobs will automatically be set to the number of CPU cores. If ``unified_diff``
536 is set to a number, changes to files will be printed to stdout with
537 ``unified_diff`` lines of context. If it is set to ``None`` or left out, files
538 themselves will be updated with changes and formatting. If a
539 ``python_version`` is provided, then we will parse each source file using
540 this version. Otherwise, we will use the version of the currently executing python
541 binary.
543 A progress indicator as well as any generated warnings will be printed to stderr.
544 To supress the interactive progress indicator, set ``hide_progress`` to ``True``.
545 Files that include the generated code marker will be skipped unless the
546 ``include_generated`` parameter is set to ``True``. Similarly, files that match
547 a supplied blacklist of regex patterns will be skipped. Warnings for skipping
548 both blacklisted and generated files will be printed to stderr along with
549 warnings generated by the codemod unless ``hide_blacklisted`` and
550 ``hide_generated`` are set to ``True``. Files that were successfully codemodded
551 will not be printed to stderr unless ``show_successes`` is set to ``True``.
553 To make this API possible, we take an instantiated transform. This is due to
554 the fact that lambdas are not pickleable and pickling functions is undefined.
555 This means we're implicitly relying on fork behavior on UNIX-like systems, and
556 this function will not work on Windows systems. To create a command-line utility
557 that runs on Windows, please instead see
558 :func:`~libcst.codemod.exec_transform_with_prettyprint`.
559 """
561 # Ensure that we have no duplicates, otherwise we might get race conditions
562 # on write.
563 files = sorted({os.path.abspath(f) for f in files})
564 total = len(files)
565 progress = Progress(enabled=not hide_progress, total=total)
567 chunksize = 4
568 # Grab number of cores if we need to
569 jobs = min(
570 jobs if jobs is not None else cpu_count(),
571 (len(files) + chunksize - 1) // chunksize,
572 )
574 if jobs < 1:
575 raise Exception("Must have at least one job to process!")
577 if total == 0:
578 return ParallelTransformResult(successes=0, failures=0, skips=0, warnings=0)
580 if repo_root is not None:
581 # Make sure if there is a root that we have the absolute path to it.
582 repo_root = os.path.abspath(repo_root)
583 # Spin up a full repo metadata manager so that we can provide metadata
584 # like type inference to individual forked processes.
585 print("Calculating full-repo metadata...", file=sys.stderr)
586 metadata_manager = FullRepoManager(
587 repo_root,
588 files,
589 transform.get_inherited_dependencies(),
590 )
591 metadata_manager.resolve_cache()
592 transform.context = replace(
593 transform.context,
594 metadata_manager=metadata_manager,
595 )
596 print("Executing codemod...", file=sys.stderr)
598 config = ExecutionConfig(
599 repo_root=repo_root,
600 unified_diff=unified_diff,
601 include_generated=include_generated,
602 generated_code_marker=generated_code_marker,
603 format_code=format_code,
604 formatter_args=formatter_args,
605 blacklist_patterns=blacklist_patterns,
606 python_version=python_version,
607 )
609 if total == 1 or jobs == 1:
610 # Simple case, we should not pay for process overhead.
611 # Let's just use a dummy synchronous pool.
612 jobs = 1
613 pool_impl = DummyPool
614 else:
615 pool_impl = Pool
616 # Warm the parser, pre-fork.
617 parse_module(
618 "",
619 config=(
620 PartialParserConfig(python_version=python_version)
621 if python_version is not None
622 else PartialParserConfig()
623 ),
624 )
626 successes: int = 0
627 failures: int = 0
628 warnings: int = 0
629 skips: int = 0
631 with pool_impl(processes=jobs) as p: # type: ignore
632 args = [
633 {
634 "transformer": transform,
635 "filename": filename,
636 "config": config,
637 }
638 for filename in files
639 ]
640 try:
641 for result in p.imap_unordered(
642 _execute_transform_wrap, args, chunksize=chunksize
643 ):
644 # Print an execution result, keep track of failures
645 _print_parallel_result(
646 result,
647 progress,
648 unified_diff=bool(unified_diff),
649 show_successes=show_successes,
650 hide_generated=hide_generated,
651 hide_blacklisted=hide_blacklisted,
652 )
653 progress.print(successes + failures + skips)
655 if isinstance(result.transform_result, TransformFailure):
656 failures += 1
657 elif isinstance(result.transform_result, TransformSuccess):
658 successes += 1
659 elif isinstance(
660 result.transform_result, (TransformExit, TransformSkip)
661 ):
662 skips += 1
664 warnings += len(result.transform_result.warning_messages)
665 finally:
666 progress.clear()
668 # Return whether there was one or more failure.
669 return ParallelTransformResult(
670 successes=successes, failures=failures, skips=skips, warnings=warnings
671 )