Improve add_rule.py and add_plugin.py scripts (#3725)

.github/workflows/ci.yaml

@@ -110,14 +110,16 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - name: "Install Rust toolchain"
-        run: rustup show
+        run: rustup component add rustfmt
       - uses: Swatinem/rust-cache@v2
-      - run: ./scripts/add_rule.py --name DoTheThing --code PLC999 --linter pylint
+      - run: ./scripts/add_rule.py --name DoTheThing --prefix PL --code C0999 --linter pylint
       - run: cargo check
+      - run: cargo fmt --all --check
       - run: |
           ./scripts/add_plugin.py test --url https://pypi.org/project/-test/0.1.0/ --prefix TST
-          ./scripts/add_rule.py --name FirstRule --code TST001 --linter test
+          ./scripts/add_rule.py --name FirstRule --prefix TST --code 001 --linter test
       - run: cargo check
+      - run: cargo fmt --all --check
 
   typos:
     name: "spell check"

scripts/_utils.py

@@ -13,5 +13,12 @@ def pascal_case(linter_name: str) -> str:
     return "".join(word.title() for word in linter_name.split("-"))
 
 
+def snake_case(name: str) -> str:
+    """Convert from PascalCase to snake_case."""
+    return "".join(
+        f"_{word.lower()}" if word.isupper() else word for word in name
+    ).lstrip("_")
+
+
 def get_indent(line: str) -> str:
     return re.match(r"^\s*", line).group()  # type: ignore[union-attr]

scripts/add_plugin.py

@@ -40,8 +40,8 @@ mod tests {
     use test_case::test_case;
 
     use crate::registry::Rule;
-    use crate::test::test_path;
     use crate::settings;
+    use crate::test::test_path;
 
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.as_ref(), path.to_string_lossy());
@@ -61,15 +61,17 @@ mod tests {
     rules_dir = plugin_dir / "rules"
     rules_dir.mkdir(exist_ok=True)
 
-    with (rules_dir / "mod.rs").open("w+") as fp:
+    (rules_dir / "mod.rs").touch()
-        fp.write("\n\n")
 
     # Create the snapshots subdirectory
     (plugin_dir / "snapshots").mkdir(exist_ok=True)
 
     # Add the plugin to `rules/mod.rs`.
-    with (ROOT_DIR / "crates/ruff/src/rules/mod.rs").open("a") as fp:
+    rules_mod_path = ROOT_DIR / "crates/ruff/src/rules/mod.rs"
-        fp.write(f"pub mod {dir_name(plugin)};")
+    lines = rules_mod_path.read_text().strip().splitlines()
+    lines.append(f"pub mod {dir_name(plugin)};")
+    lines.sort()
+    rules_mod_path.write_text("\n".join(lines) + "\n")
 
     # Add the relevant sections to `src/registry.rs`.
     content = (ROOT_DIR / "crates/ruff/src/registry.rs").read_text()
@@ -95,7 +97,7 @@ mod tests {
     with (ROOT_DIR / "crates/ruff/src/codes.rs").open("r") as fp:
         while (line := next(fp)).strip() != "// ruff":
             text += line
-        text += " " * 8 + f"// {plugin}\n"
+        text += " " * 8 + f"// {plugin}\n\n"
         text += line
         text += fp.read()
 

scripts/add_rule.py

@@ -5,30 +5,24 @@ Example usage:
 
     python scripts/add_rule.py \
         --name PreferListBuiltin \
-        --code PIE807 \
+        --prefix PIE \
+        --code 807 \
         --linter flake8-pie
 """
 
 import argparse
 
-from _utils import ROOT_DIR, dir_name, get_indent, pascal_case
+from _utils import ROOT_DIR, dir_name, get_indent, pascal_case, snake_case
 
 
-def snake_case(name: str) -> str:
+def main(*, name: str, prefix: str, code: str, linter: str) -> None:
-    """Convert from PascalCase to snake_case."""
-    return "".join(
-        f"_{word.lower()}" if word.isupper() else word for word in name
-    ).lstrip("_")
-
-
-def main(*, name: str, code: str, linter: str) -> None:
     """Generate boilerplate for a new rule."""
     # Create a test fixture.
     with (
         ROOT_DIR
         / "crates/ruff/resources/test/fixtures"
         / dir_name(linter)
-        / f"{code}.py"
+        / f"{prefix}{code}.py"
     ).open(
         "a",
     ):
@@ -42,17 +36,37 @@ def main(*, name: str, code: str, linter: str) -> None:
     content = mod_rs.read_text()
 
     with mod_rs.open("w") as fp:
+        has_added_testcase = False
+        lines = []
         for line in content.splitlines():
-            if line.strip() == "fn rules(rule_code: Rule, path: &Path) -> Result<()> {":
+            if not has_added_testcase and (
+                line.strip() == "fn rules(rule_code: Rule, path: &Path) -> Result<()> {"
+            ):
                 indent = get_indent(line)
-                fp.write(
+                filestem = f"{prefix}{code}" if linter != "pylint" else snake_case(name)
-                    f'{indent}#[test_case(Rule::{name}, Path::new("{code}.py"); '
+                lines.append(
-                    f'"{code}")]',
+                    f'{indent}#[test_case(Rule::{name}, Path::new("{filestem}.py");'
+                    f' "{prefix}{code}")]',
                 )
+                lines.sort(
+                    key=lambda line: line.split('Path::new("')[1]
+                    if linter != "pylint"
+                    else line.split(");")[1],
+                )
+                fp.write("\n".join(lines))
                 fp.write("\n")
+                lines.clear()
+                has_added_testcase = True
 
+            if has_added_testcase:
                 fp.write(line)
                 fp.write("\n")
+            elif line.strip() == "":
+                fp.write("\n".join(lines))
+                fp.write("\n\n")
+                lines.clear()
+            else:
+                lines.append(line)
 
     # Add the exports
     rules_dir = plugin_module / "rules"
@@ -60,27 +74,38 @@ def main(*, name: str, code: str, linter: str) -> None:
 
     contents = rules_mod.read_text()
     parts = contents.split("\n\n")
+
+    new_pub_use = f"pub use {rule_name_snake}::{{{rule_name_snake}, {name}}}"
+    new_mod = f"mod {rule_name_snake};"
+
     if len(parts) == 2:
-        new_contents = parts[0] + "\n"
+        pub_use_contents = parts[0].split(";\n")
-        new_contents += f"pub use {rule_name_snake}::{{{rule_name_snake}, {name}}};"
+        pub_use_contents.append(new_pub_use)
-        new_contents += "\n"
+        pub_use_contents.sort()
-        new_contents += "\n"
+
-        new_contents += parts[1]
+        mod_contents = parts[1].splitlines()
-        new_contents += f"mod {rule_name_snake};"
+        mod_contents.append(new_mod)
+        mod_contents.sort()
+
+        new_contents = ";\n".join(pub_use_contents)
+        new_contents += "\n\n"
+        new_contents += "\n".join(mod_contents)
         new_contents += "\n"
+
         rules_mod.write_text(new_contents)
     else:
         with rules_mod.open("a") as fp:
-            fp.write(f"pub use {rule_name_snake}::{{{rule_name_snake}, {name}}};")
+            fp.write(f"{new_pub_use};")
-            fp.write("\n")
+            fp.write("\n\n")
-            fp.write(f"mod {rule_name_snake};")
+            fp.write(f"{new_mod}")
             fp.write("\n")
 
     # Add the relevant rule function.
     with (rules_dir / f"{rule_name_snake}.rs").open("w") as fp:
         fp.write(
-            """use ruff_macros::{derive_message_formats, violation};
+            """\
 use ruff_diagnostics::Violation;
+use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
 
@@ -96,38 +121,58 @@ impl Violation for %s {
 """
             % (name, name),
         )
-        fp.write("\n")
         fp.write(
             f"""
-/// {code}
+/// {prefix}{code}
 pub fn {rule_name_snake}(checker: &mut Checker) {{}}
 """,
         )
-        fp.write("\n")
 
     # Add the relevant code-to-violation pair to `src/registry.rs`.
     content = (ROOT_DIR / "crates/ruff/src/registry.rs").read_text()
 
     seen_macro = False
     has_written = False
+    has_seen_linter = False
     with (ROOT_DIR / "crates/ruff/src/registry.rs").open("w") as fp:
+        lines = []
         for line in content.splitlines():
+            if has_written:
                 fp.write(line)
                 fp.write("\n")
-
-            if has_written:
                 continue
 
             if line.startswith("ruff_macros::register_rules!"):
                 seen_macro = True
+                fp.write(line)
+                fp.write("\n")
                 continue
 
             if not seen_macro:
+                fp.write(line)
+                fp.write("\n")
                 continue
 
             if line.strip() == f"// {linter}":
                 indent = get_indent(line)
-                fp.write(f"{indent}rules::{dir_name(linter)}::rules::{name},")
+                lines.append(f"{indent}rules::{dir_name(linter)}::rules::{name},")
+                has_seen_linter = True
+                fp.write(line)
+                fp.write("\n")
+                continue
+
+            if not has_seen_linter:
+                fp.write(line)
+                fp.write("\n")
+                continue
+
+            if not line.strip().startswith("// "):
+                lines.append(line)
+            else:
+                lines.sort()
+                fp.write("\n".join(lines))
+                fp.write("\n")
+                fp.write(line)
                 fp.write("\n")
                 has_written = True
 
@@ -138,8 +183,20 @@ pub fn {rule_name_snake}(checker: &mut Checker) {{}}
         while (line := next(fp)).strip() != f"// {linter}":
             text += line
         text += line
+
+        lines = []
+        while (line := next(fp)).strip() != "":
+            lines.append(line)
+
         linter_variant = pascal_case(linter)
-        text += " " * 8 + f'({linter_variant}, "{code}") => Rule::{name},\n'
+        lines.append(
+            " " * 8 + f"""({linter_variant}, "{code}") => Rule::{name},\n""",
+        )
+        lines.sort()
+
+        text += "".join(lines)
+        text += "\n"
+
         text += fp.read()
 
     with (ROOT_DIR / "crates/ruff/src/codes.rs").open("w") as fp:
@@ -158,20 +215,29 @@ if __name__ == "__main__":
         "--name",
         type=str,
         required=True,
-        help="The name of the check to generate, in PascalCase (e.g., 'LineTooLong').",
+        help=(
+            "The name of the check to generate, in PascalCase "
+            "(e.g., 'PreferListBuiltin')."
+        ),
+    )
+    parser.add_argument(
+        "--prefix",
+        type=str,
+        required=True,
+        help="Prefix code for the plugin (e.g. 'PIE').",
     )
     parser.add_argument(
         "--code",
         type=str,
         required=True,
-        help="The code of the check to generate (e.g., 'A001').",
+        help="The code of the check to generate (e.g., '807').",
     )
     parser.add_argument(
         "--linter",
         type=str,
         required=True,
-        help="The source with which the check originated (e.g., 'flake8-builtins').",
+        help="The source with which the check originated (e.g., 'flake8-pie').",
     )
     args = parser.parse_args()
 
-    main(name=args.name, code=args.code, linter=args.linter)
+    main(name=args.name, prefix=args.prefix, code=args.code, linter=args.linter)