CLI: add replace package name (#372)

* add replace package name

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Borda

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased] - 2025-MM-DD
 
+### Added
+
+- CLI: replace package name in requirements ([#372](https://github.com/Lightning-AI/utilities/pull/372))
+
+
 ### Changed
 
 -

src/lightning_utilities/cli/__main__.py

@@ -4,7 +4,11 @@
 #
 
 import lightning_utilities
-from lightning_utilities.cli.dependencies import prune_pkgs_in_requirements, replace_oldest_ver
+from lightning_utilities.cli.dependencies import (
+    prune_packages_in_requirements,
+    replace_oldest_version,
+    replace_package_in_requirements,
+)
 
 
 def main() -> None:
@@ -13,8 +17,9 @@ def main() -> None:
 
     Fire({
         "requirements": {
-            "prune-pkgs": prune_pkgs_in_requirements,
-            "set-oldest": replace_oldest_ver,
+            "prune-pkgs": prune_packages_in_requirements,
+            "set-oldest": replace_oldest_version,
+            "replace-pkg": replace_package_in_requirements,
         },
         "version": lambda: print(lightning_utilities.__version__),
     })

src/lightning_utilities/cli/dependencies.py

@@ -4,6 +4,7 @@
 #
 import glob
 import os.path
+import re
 from collections.abc import Sequence
 from pprint import pprint
 from typing import Union
@@ -15,7 +16,7 @@
     REQUIREMENT_FILES_ALL += [REQUIREMENT_ROOT]
 
 
-def prune_pkgs_in_requirements(
+def prune_packages_in_requirements(
     packages: Union[str, Sequence[str]], req_files: Union[str, Sequence[str]] = REQUIREMENT_FILES_ALL
 ) -> None:
     """Remove some packages from given requirement files."""
@@ -49,9 +50,35 @@ def _replace_min(fname: str) -> None:
         fw.write(req)
 
 
-def replace_oldest_ver(req_files: Union[str, Sequence[str]] = REQUIREMENT_FILES_ALL) -> None:
+def replace_oldest_version(req_files: Union[str, Sequence[str]] = REQUIREMENT_FILES_ALL) -> None:
     """Replace the min package version by fixed one."""
     if isinstance(req_files, str):
         req_files = [req_files]
     for fname in req_files:
         _replace_min(fname)
+
+
+def _replace_package_name(requirements: list[str], old_package: str, new_package: str) -> list[str]:
+    """Replace one package by another with same version in given requirement file.
+
+    >>> _replace_package_name(["torch>=1.0 # comment", "torchvision>=0.2", "torchtext <0.3"], "torch", "pytorch")
+    ['pytorch>=1.0 # comment', 'torchvision>=0.2', 'torchtext <0.3']
+
+    """
+    for i, req in enumerate(requirements):
+        requirements[i] = re.sub(r"^" + re.escape(old_package) + r"(?=[ <=>#]|$)", new_package, req)
+    return requirements
+
+
+def replace_package_in_requirements(
+    old_package: str, new_package: str, req_files: Union[str, Sequence[str]] = REQUIREMENT_FILES_ALL
+) -> None:
+    """Replace one package by another with same version in given requirement files."""
+    if isinstance(req_files, str):
+        req_files = [req_files]
+    for fname in req_files:
+        with open(fname) as fopen:
+            reqs = fopen.readlines()
+        reqs = _replace_package_name(reqs, old_package, new_package)
+        with open(fname, "w") as fw:
+            fw.writelines(reqs)

tests/unittests/cli/test_dependencies.py

@@ -1,6 +1,10 @@
 from pathlib import Path
 
-from lightning_utilities.cli.dependencies import prune_pkgs_in_requirements, replace_oldest_ver
+from lightning_utilities.cli.dependencies import (
+    prune_packages_in_requirements,
+    replace_oldest_version,
+    replace_package_in_requirements,
+)
 
 _PATH_ROOT = Path(__file__).parent.parent.parent
 
@@ -9,7 +13,7 @@ def test_prune_packages(tmpdir):
     req_file = tmpdir / "requirements.txt"
     with open(req_file, "w") as fp:
         fp.writelines(["fire\n", "abc>=0.1\n"])
-    prune_pkgs_in_requirements("abc", req_files=[str(req_file)])
+    prune_packages_in_requirements("abc", req_files=[str(req_file)])
     with open(req_file) as fp:
         lines = fp.readlines()
     assert lines == ["fire\n"]
@@ -19,7 +23,17 @@ def test_oldest_packages(tmpdir):
     req_file = tmpdir / "requirements.txt"
     with open(req_file, "w") as fp:
         fp.writelines(["fire>0.2\n", "abc>=0.1\n"])
-    replace_oldest_ver(req_files=[str(req_file)])
+    replace_oldest_version(req_files=[str(req_file)])
     with open(req_file) as fp:
         lines = fp.readlines()
     assert lines == ["fire>0.2\n", "abc==0.1\n"]
+
+
+def test_replace_packages(tmpdir):
+    req_file = tmpdir / "requirements.txt"
+    with open(req_file, "w") as fp:
+        fp.writelines(["torchvision>=0.2\n", "torch>=1.0 # comment\n", "torchtext <0.3\n"])
+    replace_package_in_requirements(old_package="torch", new_package="pytorch", req_files=[str(req_file)])
+    with open(req_file) as fp:
+        lines = fp.readlines()
+    assert lines == ["torchvision>=0.2\n", "pytorch>=1.0 # comment\n", "torchtext <0.3\n"]