Support module and package names in the codemod context

libcst/codemod/_cli.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass, replace
 from multiprocessing import cpu_count, Pool
 from pathlib import Path, PurePath
-from typing import Any, AnyStr, cast, Dict, List, Optional, Sequence, Union
+from typing import Any, AnyStr, cast, Dict, List, Optional, Sequence, Union, Tuple
 
 from libcst import parse_module, PartialParserConfig
 from libcst.codemod._codemod import Codemod
@@ -184,28 +184,31 @@ def exec_transform_with_prettyprint(
     return maybe_code
 
 
-def _calculate_module(repo_root: Optional[str], filename: str) -> Optional[str]:
+def _calculate_module_and_package(repo_root: Optional[str], filename: str) -> Tuple[Optional[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
+        return None, None
 
     try:
         relative_filename = PurePath(filename).relative_to(repo_root)
     except ValueError:
         # This file seems to be out of the repo root.
-        return None
+        return None, None
 
     # get rid of extension
     relative_filename = relative_filename.with_suffix("")
 
-    # get rid of any special cases
+    # handle special cases
     if relative_filename.stem in ["__init__", "__main__"]:
         relative_filename = relative_filename.parent
+        package = module = ".".join(relative_filename.parts)
+    else:
+        module = ".".join(relative_filename.parts)
+        package = ".".join(relative_filename.parts[:-1])
 
-    # Now, convert to dots to represent the python module.
-    return ".".join(relative_filename.parts)
+    return module, package
 
 
 @dataclass(frozen=True)
@@ -264,14 +267,20 @@ def _execute_transform(  # noqa: C901
                 ),
             )
 
+        # attempt to work out the module and package name for this file
+        full_module_name, full_package_name = _calculate_module_and_package(
+            config.repo_root, filename
+        )
+
         # Somewhat gross hack to provide the filename in the transform's context.
         # We do this after the fork so that a context that was initialized with
         # some defaults before calling parallel_exec_transform_with_prettyprint
         # will be updated per-file.
         transformer.context = replace(
             transformer.context,
             filename=filename,
-            full_module_name=_calculate_module(config.repo_root, filename),
+            full_module_name=full_module_name,
+            full_package_name=full_package_name,
             scratch={},
         )
 

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/codemod/tests/test_cli.py

@@ -5,62 +5,69 @@
 #
 from typing import Optional
 
-from libcst.codemod._cli import _calculate_module
+from libcst.codemod._cli import _calculate_module_and_package
 from libcst.testing.utils import data_provider, UnitTest
 
 
 class TestPackageCalculation(UnitTest):
     @data_provider(
         (
             # Providing no root should give back no module.
-            (None, "/some/dummy/file.py", None),
+            (None, "/some/dummy/file.py", None, None),
             # Providing a file outside the root should give back no module.
-            ("/home/username/root", "/some/dummy/file.py", None),
-            ("/home/username/root/", "/some/dummy/file.py", None),
-            ("/home/username/root", "/home/username/file.py", None),
+            ("/home/username/root", "/some/dummy/file.py", None, None),
+            ("/home/username/root/", "/some/dummy/file.py", None, None),
+            ("/home/username/root", "/home/username/file.py", None, None),
             # Various files inside the root should give back valid modules.
-            ("/home/username/root", "/home/username/root/file.py", "file"),
-            ("/home/username/root/", "/home/username/root/file.py", "file"),
+            ("/home/username/root", "/home/username/root/file.py", "file", ""),
+            ("/home/username/root/", "/home/username/root/file.py", "file", ""),
             (
                 "/home/username/root/",
                 "/home/username/root/some/dir/file.py",
                 "some.dir.file",
+                "some.dir",
             ),
             # Various special files inside the root should give back valid modules.
             (
                 "/home/username/root/",
                 "/home/username/root/some/dir/__init__.py",
                 "some.dir",
+                "some.dir",
             ),
             (
                 "/home/username/root/",
                 "/home/username/root/some/dir/__main__.py",
                 "some.dir",
+                "some.dir",
             ),
             # some windows tests
             (
                 "c:/Program Files/",
                 "d:/Program Files/some/dir/file.py",
                 None,
+                None,
             ),
             (
                 "c:/Program Files/other/",
                 "c:/Program Files/some/dir/file.py",
                 None,
+                None,
             ),
             (
                 "c:/Program Files/",
                 "c:/Program Files/some/dir/file.py",
                 "some.dir.file",
+                "some.dir",
             ),
             (
                 "c:/Program Files/",
                 "c:/Program Files/some/dir/__main__.py",
                 "some.dir",
+                "some.dir",
             ),
         ),
     )
-    def test_calculate_module(
-        self, repo_root: Optional[str], filename: str, module: str
+    def test_calculate_module_and_package(
+        self, repo_root: Optional[str], filename: str, module: Optional[str], package: Optional[str]
     ) -> None:
-        self.assertEqual(_calculate_module(repo_root, filename), module)
+        self.assertEqual(_calculate_module_and_package(repo_root, filename), (module, package))