Upgrade command (#2886)

* Add initial version of upgrade command

* Add autocompletion

* Initial attempt at a codemod

* Better codemod

* Support for passing a path

* Update strawberry/codemods/annotated_unions.py

Co-authored-by: Marcelo Trylesinski <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update tests/codemods/test_annotated_unions.py

Co-authored-by: Marcelo Trylesinski <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Support for `union`

* Proper support for `union` thanks to @Kludex

* Update strawberry/codemods/annotated_unions.py

Co-authored-by: Marcelo Trylesinski <[email protected]>

* Add support for both union syntax

* Add support for multiple paths

* Add support for passing old syntax

* Allow to specify python target and if to use typing extensions

* Tests and fixes

* Add docs and release notes

* Lint

* Type fixes

* Fix runtime types

---------

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

.pre-commit-config.yaml

@@ -3,13 +3,13 @@ repos:
     rev: 23.3.0
     hooks:
       - id: black
-        exclude: ^tests/codegen/snapshots/python/
+        exclude: ^tests/\w+/snapshots/
 
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     rev: v0.0.272
     hooks:
       - id: ruff
-        exclude: ^tests/codegen/snapshots/python/
+        exclude: ^tests/\w+/snapshots/
 
   - repo: https://github.com/patrick91/pre-commit-alex
     rev: aa5da9e54b92ab7284feddeaf52edf14b1690de3

RELEASE.md

@@ -0,0 +1,13 @@
+Release type: minor
+
+This release introduces a new command called `upgrade`, this command can be used
+to run codemods on your codebase to upgrade to the latest version of Strawberry.
+
+At the moment we only support upgrading unions to use the new syntax with
+annotated, but in future we plan to add more commands to help with upgrading.
+
+Here's how you can use the command to upgrade your codebase:
+
+```shell
+strawberry upgrade annotated-union .
+```

docs/README.md

@@ -10,6 +10,7 @@
 - [Subscriptions](./general/subscriptions.md)
 - [Why](./general/why.md)
 - [Breaking changes](./breaking-changes.md)
+- [Breaking changes](./general/upgrades.md)
 - [FAQ](./faq.md)
 
 ## Types

docs/general/upgrades.md

@@ -0,0 +1,24 @@
+---
+title: Upgrading Strawberry
+---
+
+# Upgrading Strawberry
+
+<!--alex ignore-->
+
+We try to keep Strawberry as backwards compatible as possible, but sometimes we
+need to make updates to the public API. While we try to deprecate APIs before
+removing them, we also want to make it as easy as possible to upgrade to the
+latest version of Strawberry.
+
+For this reason we provide a CLI command that can automatically upgrade your
+codebase to use the updated APIs.
+
+At the moment we only support updating unions to use the new syntax with
+annotated, but in future we plan to add more commands to help with upgrading.
+
+Here's how you can use the command to upgrade your codebase:
+
+```shell
+strawberry upgrade annotated-union .
+```

pyproject.toml

@@ -139,6 +139,7 @@ strawberry = "strawberry.cli:run"
 line-length = 88
 extend-exclude = '''
 tests/codegen/snapshots/
+tests/cli/snapshots/
 '''
 
 [tool.pytest.ini_options]
@@ -338,6 +339,7 @@ src = ["strawberry", "tests"]
 "tests/federation/printer/*" = ["E501"]
 "tests/test_printer/test_basic.py" = ["E501"]
 "tests/pyright/test_federation.py" = ["E501"]
+"tests/codemods/*" = ["E501"]
 "tests/test_printer/test_schema_directives.py" = ["E501"]
 "tests/*" = ["RSE102", "SLF001", "TCH001", "TCH002", "TCH003", "ANN001", "ANN201", "PLW0603", "PLC1901", "S603", "S607", "B018"]
 "strawberry/extensions/tracing/__init__.py" = ["TCH004"]

strawberry/cli/__init__.py

@@ -1,7 +1,7 @@
-from .commands.codegen import codegen  # noqa
-from .commands.export_schema import export_schema  # noqa
-from .commands.server import server  # noqa
-
+from .commands.codegen import codegen as codegen  # noqa
+from .commands.export_schema import export_schema as export_schema  # noqa
+from .commands.server import server as server  # noqa
+from .commands.upgrade import upgrade as upgrade  # noqa
 
 from .app import app
 

strawberry/cli/commands/upgrade/__init__.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import glob
+import pathlib  # noqa: TCH003
+import sys
+from typing import List
+
+import rich
+import typer
+from libcst.codemod import CodemodContext
+
+from strawberry.cli.app import app
+from strawberry.codemods.annotated_unions import ConvertUnionToAnnotatedUnion
+
+from ._run_codemod import run_codemod
+
+codemods = {
+    "annotated-union": ConvertUnionToAnnotatedUnion,
+}
+
+
+# TODO: add support for running all of them
+@app.command(help="Upgrades a Strawberry project to the latest version")
+def upgrade(
+    codemod: str = typer.Argument(
+        ...,
+        autocompletion=lambda: list(codemods.keys()),
+        help="Name of the upgrade to run",
+    ),
+    paths: List[pathlib.Path] = typer.Argument(file_okay=True, dir_okay=True),
+    python_target: str = typer.Option(
+        ".".join(str(x) for x in sys.version_info[:2]),
+        "--python-target",
+        help="Python version to target",
+    ),
+    use_typing_extensions: bool = typer.Option(
+        False,
+        "--use-typing-extensions",
+        help="Use typing_extensions instead of typing for newer features",
+    ),
+) -> None:
+    if codemod not in codemods:
+        rich.print(f'[red]Upgrade named "{codemod}" does not exist')
+
+        raise typer.Exit(2)
+
+    python_target_version = tuple(int(x) for x in python_target.split("."))
+
+    transformer = ConvertUnionToAnnotatedUnion(
+        CodemodContext(),
+        use_pipe_syntax=python_target_version >= (3, 10),
+        use_typing_extensions=use_typing_extensions,
+    )
+
+    files: list[str] = []
+
+    for path in paths:
+        if path.is_dir():
+            glob_path = str(path / "**/*.py")
+            files.extend(glob.glob(glob_path, recursive=True))
+        else:
+            files.append(str(path))
+
+    files = list(set(files))
+
+    results = list(run_codemod(transformer, files))
+    changed = [result for result in results if result.changed]
+
+    rich.print()
+    rich.print("[green]Upgrade completed successfully, here's a summary:")
+    rich.print(f"  - {len(changed)} files changed")
+    rich.print(f"  - {len(results) - len(changed)} files skipped")
+
+    if changed:
+        raise typer.Exit(1)

strawberry/cli/commands/upgrade/_fake_progress.py

@@ -0,0 +1,21 @@
+from typing import Any
+
+from rich.progress import TaskID
+
+
+class FakeProgress:
+    """A fake progress bar that does nothing.
+
+    This is used when the user has only one file to process."""
+
+    def advance(self, task_id: TaskID) -> None:
+        pass
+
+    def add_task(self, *args: Any, **kwargs: Any) -> TaskID:
+        return TaskID(0)
+
+    def __enter__(self) -> "FakeProgress":
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        pass

strawberry/cli/commands/upgrade/_run_codemod.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import contextlib
+import os
+from multiprocessing import Pool, cpu_count
+from typing import TYPE_CHECKING, Any, Dict, Generator, Sequence, Type, Union
+
+from libcst.codemod._cli import ExecutionConfig, ExecutionResult, _execute_transform
+from libcst.codemod._dummy_pool import DummyPool
+from rich.progress import Progress
+
+from ._fake_progress import FakeProgress
+
+if TYPE_CHECKING:
+    from libcst.codemod import Codemod
+
+ProgressType = Union[Type[Progress], Type[FakeProgress]]
+PoolType = Union[Type[Pool], Type[DummyPool]]  # type: ignore
+
+
+def _execute_transform_wrap(
+    job: Dict[str, Any],
+) -> ExecutionResult:
+    # TODO: maybe capture warnings?
+    with open(os.devnull, "w") as null:  # noqa: PTH123
+        with contextlib.redirect_stderr(null):
+            return _execute_transform(**job)
+
+
+def _get_progress_and_pool(
+    total_files: int, jobs: int
+) -> tuple[PoolType, ProgressType]:
+    poll_impl: PoolType = Pool  # type: ignore
+    progress_impl: ProgressType = Progress
+
+    if total_files == 1 or jobs == 1:
+        poll_impl = DummyPool
+
+    if total_files == 1:
+        progress_impl = FakeProgress
+
+    return poll_impl, progress_impl
+
+
+def run_codemod(
+    codemod: Codemod,
+    files: Sequence[str],
+) -> Generator[ExecutionResult, None, None]:
+    chunk_size = 4
+    total = len(files)
+    jobs = min(cpu_count(), (total + chunk_size - 1) // chunk_size)
+
+    config = ExecutionConfig()
+
+    pool_impl, progress_impl = _get_progress_and_pool(total, jobs)
+
+    tasks = [
+        {
+            "transformer": codemod,
+            "filename": filename,
+            "config": config,
+        }
+        for filename in files
+    ]
+
+    with pool_impl(processes=jobs) as p, progress_impl() as progress:  # type: ignore
+        task_id = progress.add_task("[cyan]Updating...", total=len(tasks))
+
+        for result in p.imap_unordered(
+            _execute_transform_wrap, tasks, chunksize=chunk_size
+        ):
+            progress.advance(task_id)
+
+            yield result