1#!/usr/bin/env python3
2
3import os
4import subprocess
5import sys
6from subprocess import CompletedProcess
7
8
9def usage():
10    print("""Usage: jspecify_update.py <dir>
11This script updates all projects in a directory to use JSpecify annotations instead of AndroidX
12nullness annotations. If no directory is provided, uses the root frameworks/support dir.
13""")
14    sys.exit(1)
15
16
17def run_gradle_task(gradle_path: str, task: str, args: list[str] | None = None) -> \
18        CompletedProcess[bytes]:
19    """Runs a Gradle task ans returns the result."""
20    cmd = [gradle_path, task] + (args or [])
21    return subprocess.run(cmd, capture_output=True)
22
23
24def update_annotations(dir_path: str, gradle_path: str) -> bool:
25    """If the project in dir_path has a lint task, runs lintFix to shift all AndroidX nullness
26    annotations to type-use position. Returns whether any updates were applied."""
27    tasks_result = run_gradle_task(gradle_path, "tasks")
28    if "lintFix" not in str(tasks_result.stdout):
29        print(f"Lint task does not exist for {dir_path}")
30        return False
31
32    fix_result = run_gradle_task(gradle_path, "lintFix", ["-Pandroidx.useJSpecifyAnnotations=true"])
33    return fix_result.returncode != 0
34
35
36def update_imports_for_file(filepath: str) -> bool:
37    """Replaces AndroidX nullness annotation imports with JSpecify imports. Returns whether any
38    updates were needed."""
39    with open(filepath, "r") as f:
40        lines = f.readlines()
41
42    replacements = {
43        "import androidx.annotation.NonNull": "import org.jspecify.annotations.NonNull;\n",
44        "import androidx.annotation.Nullable": "import org.jspecify.annotations.Nullable;\n"
45    }
46    updated_count = 0
47    for i, line in enumerate(lines):
48        for target, replacement in replacements.items():
49            if target in line:
50                lines[i] = replacement
51                updated_count += 1
52                break
53
54        if updated_count == len(replacements):
55            break
56
57    replaced = updated_count > 0
58    if replaced:
59        with open(filepath, "w") as f:
60            f.writelines(lines)
61    return replaced
62
63
64def reformat_files(gradle_path: str) -> None:
65    """Runs the java formatter to fix imports for the specified files."""
66    run_gradle_task(gradle_path, "javaFormat", ["--fix-imports-only"]),
67
68
69def update_imports(dir_path: str, gradle_path: str) -> bool:
70    """For each java file in the directory, replaces the AndroidX nullness imports with JSpecify
71    imports and runs the java formatter to correct the new import order."""
72    java_files = []
73    for sub_dir_path, _, filenames in os.walk(dir_path):
74        for filename in filenames:
75            (_, ext) = os.path.splitext(filename)
76            if ext == ".java":
77                file_path = os.path.join(sub_dir_path, filename)
78                if update_imports_for_file(file_path):
79                    java_files.append(file_path)
80
81    if java_files:
82        reformat_files(gradle_path)
83    return java_files
84
85
86def add_jspecify_dependency(build_gradle_path: str) -> None:
87    """Adds a JSpecify dependency to the build file."""
88    with open(build_gradle_path, "r") as f:
89        lines = f.readlines()
90
91    jspecify_dependency = "    api(libs.jspecify)\n"
92    if jspecify_dependency in lines:
93        print(f"JSpecify dependency already present for {build_gradle_path}")
94        return
95
96    dependencies_start = None
97    for i in range(len(lines)):
98        line = lines[i]
99        if line.startswith("dependencies {"):
100            dependencies_start = i
101            break
102
103    if not dependencies_start:
104        print(f"No dependencies block found for {build_gradle_path}")
105        return
106
107    lines.insert(dependencies_start + 1, "    api(libs.jspecify)\n")
108    with open(build_gradle_path, "w") as f:
109        f.writelines(lines)
110
111
112def process_dir(dir_path: str, root_dir_path: str) -> None:
113    """Updates the directory to use JSpecify annotations."""
114    print(f"Processing {dir_path}")
115    os.chdir(dir_path)
116    gradle_path = os.path.join(root_dir_path, "gradlew")
117    if update_annotations(dir_path, gradle_path):
118        print(f"Lint fixes applied in {dir_path}")
119        if update_imports(dir_path, gradle_path):
120            add_jspecify_dependency(os.path.join(dir_path, "build.gradle"))
121
122
123def main(start_dir: str | None) -> None:
124    # Location of this script: under support/development
125    script_path = os.path.realpath(__file__)
126    # Move up to the support dir
127    support_dir = os.path.dirname(os.path.dirname(script_path))
128
129    # Search the specified directory, or the support dir if there wasn't one
130    if not start_dir:
131        start_dir = support_dir
132    else:
133        start_dir = os.path.abspath(start_dir)
134    for dir_path, _, filenames in os.walk(start_dir):
135        if dir_path == support_dir:
136            continue
137        if "build.gradle" in filenames:
138            process_dir(dir_path, support_dir)
139
140
141if __name__ == '__main__':
142    if len(sys.argv) == 1:
143        main(None)
144    elif len(sys.argv) == 2:
145        main(sys.argv[1])
146    else:
147        usage()
148