Support module and package names in the codemod context (#662)

* Support module and package names in the codemod context

* PR feedback

* Reorganize module name and relative name logic to libcst.helpers.module

* Force rebuild

libcst/codemod/_cli.py

@@ -16,7 +16,7 @@
 import traceback
 from dataclasses import dataclass, replace
 from multiprocessing import cpu_count, Pool
-from pathlib import Path, PurePath
+from pathlib import Path
 from typing import Any, AnyStr, cast, Dict, List, Optional, Sequence, Union
 
 from libcst import parse_module, PartialParserConfig
@@ -32,6 +32,7 @@
     TransformSkip,
     TransformSuccess,
 )
+from libcst.helpers import calculate_module_and_package
 from libcst.metadata import FullRepoManager
 
 _DEFAULT_GENERATED_CODE_MARKER: str = f"@gen{''}erated"
@@ -184,30 +185,6 @@ def exec_transform_with_prettyprint(
     return maybe_code
 
 
-def _calculate_module(repo_root: Optional[str], filename: str) -> Optional[str]:
-    # Given an absolute repo_root and an absolute filename, calculate the
-    # python module name for the file.
-    if repo_root is None:
-        # We don't have a repo root, so this is impossible to calculate.
-        return None
-
-    try:
-        relative_filename = PurePath(filename).relative_to(repo_root)
-    except ValueError:
-        # This file seems to be out of the repo root.
-        return None
-
-    # get rid of extension
-    relative_filename = relative_filename.with_suffix("")
-
-    # get rid of any special cases
-    if relative_filename.stem in ["__init__", "__main__"]:
-        relative_filename = relative_filename.parent
-
-    # Now, convert to dots to represent the python module.
-    return ".".join(relative_filename.parts)
-
-
 @dataclass(frozen=True)
 class ExecutionResult:
     # File we have results for
@@ -271,10 +248,20 @@ def _execute_transform(  # noqa: C901
         transformer.context = replace(
             transformer.context,
             filename=filename,
-            full_module_name=_calculate_module(config.repo_root, filename),
             scratch={},
         )
 
+        # attempt to work out the module and package name for this file
+        module_name_and_package = calculate_module_and_package(
+            config.repo_root, filename
+        )
+        if module_name_and_package is not None:
+            transformer.context = replace(
+                transformer.context,
+                full_module_name=module_name_and_package.name,
+                full_package_name=module_name_and_package.package,
+            )
+
         # Run the transform, bail if we failed or if we aren't formatting code
         try:
             input_tree = parse_module(

libcst/codemod/_context.py

@@ -44,6 +44,12 @@ class CodemodContext:
     #: in the repo named ``foo/bar/baz.py``.
     full_module_name: Optional[str] = None
 
+    #: The current package if a codemod is being executed against a file that
+    #: lives on disk, and the repository root is correctly configured. This
+    #: Will take the form of a dotted name such as ``foo.bar`` for a file
+    #: in the repo named ``foo/bar/baz.py``
+    full_package_name: Optional[str] = None
+
     #: The current top level metadata wrapper for the module being modified.
     #: To access computed metadata when inside an actively running codemod, use
     #: the :meth:`~libcst.MetadataDependent.get_metadata` method on

libcst/helpers/__init__.py

@@ -4,11 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 #
 
-from libcst.helpers._statement import (
-    get_absolute_module,
-    get_absolute_module_for_import,
-    get_absolute_module_for_import_or_raise,
-)
 from libcst.helpers._template import (
     parse_template_expression,
     parse_template_module,
@@ -19,9 +14,17 @@
     get_full_name_for_node,
     get_full_name_for_node_or_raise,
 )
-from libcst.helpers.module import insert_header_comments
+from libcst.helpers.module import (
+    calculate_module_and_package,
+    get_absolute_module,
+    get_absolute_module_for_import,
+    get_absolute_module_for_import_or_raise,
+    insert_header_comments,
+    ModuleNameAndPackage,
+)
 
 __all__ = [
+    "calculate_module_and_package",
     "get_absolute_module",
     "get_absolute_module_for_import",
     "get_absolute_module_for_import_or_raise",
@@ -32,4 +35,5 @@
     "parse_template_module",
     "parse_template_statement",
     "parse_template_expression",
+    "ModuleNameAndPackage",
 ]

libcst/helpers/module.py

@@ -3,13 +3,16 @@
 # 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 typing import List
+from pathlib import PurePath
+from typing import List, Optional
 
-import libcst
+from libcst import Comment, EmptyLine, ImportFrom, Module
+from libcst.helpers.expression import get_full_name_for_node
 
 
-def insert_header_comments(node: libcst.Module, comments: List[str]) -> libcst.Module:
+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
@@ -25,9 +28,91 @@ def insert_header_comments(node: libcst.Module, comments: List[str]) -> libcst.M
 
     comment_lines = islice(node.header, last_comment_index + 1)
     empty_lines = islice(node.header, last_comment_index + 1, None)
-    inserted_lines = [
-        libcst.EmptyLine(comment=libcst.Comment(value=comment)) for comment in comments
-    ]
+    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
+
+
+@dataclass(frozen=True)
+class ModuleNameAndPackage:
+    name: str
+    package: str
+
+
+def calculate_module_and_package(
+    repo_root: Optional[str], filename: str
+) -> Optional[ModuleNameAndPackage]:
+    # Given an absolute repo_root and an absolute filename, calculate the
+    # python module name for the file.
+    if repo_root is None:
+        # We don't have a repo root, so this is impossible to calculate.
+        return None
+
+    try:
+        relative_filename = PurePath(filename).relative_to(repo_root)
+    except ValueError:
+        # This file seems to be out of the repo root.
+        return None
+
+    # get rid of extension
+    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)