1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
from dataclasses import dataclass
from itertools import islice
from pathlib import PurePath
from typing import List, Optional
from libcst import Comment, EmptyLine, ImportFrom, Module
from libcst._types import StrPath
from libcst.helpers.expression import get_full_name_for_node
def insert_header_comments(node: Module, comments: List[str]) -> Module:
"""
Insert comments after last non-empty line in header. Use this to insert one or more
comments after any copyright preamble in a :class:`~libcst.Module`. Each comment in
the list of ``comments`` must start with a ``#`` and will be placed on its own line
in the appropriate location.
"""
# Split the lines up into a contiguous comment-containing section and
# the empty whitespace section that follows
last_comment_index = -1
for i, line in enumerate(node.header):
if line.comment is not None:
last_comment_index = i
comment_lines = islice(node.header, last_comment_index + 1)
empty_lines = islice(node.header, last_comment_index + 1, None)
inserted_lines = [EmptyLine(comment=Comment(value=comment)) for comment in comments]
# pyre-fixme[60]: Concatenation not yet support for multiple variadic tuples:
# `*comment_lines, *inserted_lines, *empty_lines`.
return node.with_changes(header=(*comment_lines, *inserted_lines, *empty_lines))
def get_absolute_module(
current_module: Optional[str], module_name: Optional[str], num_dots: int
) -> Optional[str]:
if num_dots == 0:
# This is an absolute import, so the module is correct.
return module_name
if current_module is None:
# We don't actually have the current module available, so we can't compute
# the absolute module from relative.
return None
# We have the current module, as well as the relative, let's compute the base.
modules = current_module.split(".")
if len(modules) < num_dots:
# This relative import goes past the base of the repository, so we can't calculate it.
return None
base_module = ".".join(modules[:-num_dots])
# Finally, if the module name was supplied, append it to the end.
if module_name is not None:
# If we went all the way to the top, the base module should be empty, so we
# should return the relative bit as absolute. Otherwise, combine the base
# module and module name using a dot separator.
base_module = (
f"{base_module}.{module_name}" if len(base_module) > 0 else module_name
)
# If they tried to import all the way to the root, return None. Otherwise,
# return the module itself.
return base_module if len(base_module) > 0 else None
def get_absolute_module_for_import(
current_module: Optional[str], import_node: ImportFrom
) -> Optional[str]:
# First, let's try to grab the module name, regardless of relative status.
module = import_node.module
module_name = get_full_name_for_node(module) if module is not None else None
# Now, get the relative import location if it exists.
num_dots = len(import_node.relative)
return get_absolute_module(current_module, module_name, num_dots)
def get_absolute_module_for_import_or_raise(
current_module: Optional[str], import_node: ImportFrom
) -> str:
module = get_absolute_module_for_import(current_module, import_node)
if module is None:
raise Exception(f"Unable to compute absolute module for {import_node}")
return module
def get_absolute_module_from_package(
current_package: Optional[str], module_name: Optional[str], num_dots: int
) -> Optional[str]:
if num_dots == 0:
# This is an absolute import, so the module is correct.
return module_name
if current_package is None or current_package == "":
# We don't actually have the current module available, so we can't compute
# the absolute module from relative.
return None
# see importlib._bootstrap._resolve_name
# https://github.com/python/cpython/blob/3.10/Lib/importlib/_bootstrap.py#L902
bits = current_package.rsplit(".", num_dots - 1)
if len(bits) < num_dots:
return None
base = bits[0]
return "{}.{}".format(base, module_name) if module_name else base
def get_absolute_module_from_package_for_import(
current_package: Optional[str], import_node: ImportFrom
) -> Optional[str]:
# First, let's try to grab the module name, regardless of relative status.
module = import_node.module
module_name = get_full_name_for_node(module) if module is not None else None
# Now, get the relative import location if it exists.
num_dots = len(import_node.relative)
return get_absolute_module_from_package(current_package, module_name, num_dots)
def get_absolute_module_from_package_for_import_or_raise(
current_package: Optional[str], import_node: ImportFrom
) -> str:
module = get_absolute_module_from_package_for_import(current_package, import_node)
if module is None:
raise Exception(f"Unable to compute absolute module for {import_node}")
return module
@dataclass(frozen=True)
class ModuleNameAndPackage:
name: str
package: str
def calculate_module_and_package(
repo_root: StrPath, filename: StrPath
) -> ModuleNameAndPackage:
# Given an absolute repo_root and an absolute filename, calculate the
# python module name for the file.
relative_filename = PurePath(filename).relative_to(repo_root)
relative_filename = relative_filename.with_suffix("")
# handle special cases
if relative_filename.stem in ["__init__", "__main__"]:
relative_filename = relative_filename.parent
package = name = ".".join(relative_filename.parts)
else:
name = ".".join(relative_filename.parts)
package = ".".join(relative_filename.parts[:-1])
return ModuleNameAndPackage(name, package)
|