remove_experimental_decorators.py
1 """ 2 Script to automatically remove @experimental decorators from functions 3 that have been experimental for more than a configurable cutoff period (default: 6 months). 4 """ 5 6 import argparse 7 import ast 8 import json 9 import subprocess 10 from dataclasses import dataclass 11 from datetime import datetime, timedelta, timezone 12 from pathlib import Path 13 from urllib.request import urlopen 14 15 16 @dataclass 17 class ExperimentalDecorator: 18 version: str 19 line_number: int 20 end_line_number: int 21 column: int 22 age_days: int 23 content: str 24 skip: bool = False 25 26 27 def get_tracked_python_files() -> list[Path]: 28 """Get all tracked Python files in the repository.""" 29 result = subprocess.check_output(["git", "ls-files", "*.py"], text=True) 30 return [Path(f) for f in result.strip().split("\n") if f] 31 32 33 def get_mlflow_release_dates() -> dict[str, datetime]: 34 """Fetch MLflow release dates from PyPI API.""" 35 with urlopen("https://pypi.org/pypi/mlflow/json") as response: 36 data = json.loads(response.read().decode()) 37 38 release_dates: dict[str, datetime] = {} 39 for version, releases in data["releases"].items(): 40 if releases: # Some versions might have empty release lists 41 # Get the earliest release date for this version 42 upload_times: list[str] = [r["upload_time"] for r in releases if "upload_time" in r] 43 if upload_times: 44 earliest_time = min(upload_times) 45 # Parse ISO format datetime and convert to UTC 46 release_date = datetime.fromisoformat(earliest_time.replace("Z", "+00:00")) 47 if release_date.tzinfo is None: 48 release_date = release_date.replace(tzinfo=timezone.utc) 49 release_dates[version] = release_date 50 51 return release_dates 52 53 54 def find_experimental_decorators( 55 file_path: Path, release_dates: dict[str, datetime], now: datetime 56 ) -> list[ExperimentalDecorator]: 57 """ 58 Find all @experimental decorators in a Python file using AST and return their information 59 with computed age. 60 """ 61 content = file_path.read_text() 62 tree = ast.parse(content) 63 decorators: list[ExperimentalDecorator] = [] 64 65 for node in ast.walk(tree): 66 if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): 67 continue 68 69 for decorator in node.decorator_list: 70 if not isinstance(decorator, ast.Call): 71 continue 72 73 if not (isinstance(decorator.func, ast.Name) and decorator.func.id == "experimental"): 74 continue 75 76 version = _extract_version_from_ast_decorator(decorator) 77 if not version or version not in release_dates: 78 continue 79 80 release_date = release_dates[version] 81 age_days = (now - release_date).days 82 83 decorator_info = ExperimentalDecorator( 84 version=version, 85 line_number=decorator.lineno, 86 end_line_number=decorator.end_lineno or decorator.lineno, 87 column=decorator.col_offset + 1, # 1-indexed 88 age_days=age_days, 89 content=ast.unparse(decorator), 90 skip=_is_skip(decorator), 91 ) 92 decorators.append(decorator_info) 93 94 return decorators 95 96 97 def _extract_version_from_ast_decorator(decorator: ast.Call) -> str | None: 98 """Extract version string from AST decorator node.""" 99 for keyword in decorator.keywords: 100 if keyword.arg == "version" and isinstance(keyword.value, ast.Constant): 101 return str(keyword.value.value) 102 return None 103 104 105 def _is_skip(decorator: ast.Call) -> bool: 106 """Check if the decorator has skip=True.""" 107 return any( 108 kw.arg == "skip" and isinstance(kw.value, ast.Constant) and kw.value.value is True 109 for kw in decorator.keywords 110 ) 111 112 113 def remove_decorators_from_file( 114 file_path: Path, 115 decorators_to_remove: list[ExperimentalDecorator], 116 dry_run: bool, 117 ) -> list[ExperimentalDecorator]: 118 if not decorators_to_remove: 119 return [] 120 121 lines = file_path.read_text().splitlines(keepends=True) 122 # Create a set of line numbers to remove for quick lookup (handle ranges) 123 lines_to_remove: set[int] = set() 124 for decorator in decorators_to_remove: 125 lines_to_remove.update(range(decorator.line_number, decorator.end_line_number + 1)) 126 127 new_lines: list[str] = [] 128 129 for line_num, line in enumerate(lines, 1): 130 if line_num not in lines_to_remove: 131 new_lines.append(line) 132 133 if not dry_run: 134 file_path.write_text("".join(new_lines)) 135 136 return decorators_to_remove 137 138 139 def main() -> None: 140 """Main entry point.""" 141 parser = argparse.ArgumentParser( 142 description="Remove @experimental decorators older than specified cutoff period" 143 ) 144 parser.add_argument( 145 "--dry-run", action="store_true", help="Show what would be removed without making changes" 146 ) 147 parser.add_argument( 148 "--cutoff-days", 149 type=int, 150 default=180, 151 help="Number of days after which to remove decorators (default: 180)", 152 ) 153 parser.add_argument( 154 "files", nargs="*", help="Python files to process (defaults to all tracked Python files)" 155 ) 156 157 args = parser.parse_args() 158 release_dates = get_mlflow_release_dates() 159 # Calculate cutoff date using configurable cutoff days 160 now = datetime.now(timezone.utc) 161 cutoff_date = now - timedelta(days=args.cutoff_days) 162 print(f"Cutoff date: {cutoff_date.strftime('%Y-%m-%d %H:%M:%S UTC')}") 163 164 python_files = [Path(f) for f in args.files] if args.files else get_tracked_python_files() 165 for file_path in python_files: 166 if not file_path.exists(): 167 continue 168 169 # First, find all experimental decorators in the file with computed ages 170 decorators = find_experimental_decorators(file_path, release_dates, now) 171 if not decorators: 172 continue 173 174 # Filter to only decorators that should be removed (older than cutoff days) 175 old_decorators = [d for d in decorators if d.age_days > args.cutoff_days and not d.skip] 176 177 # Log skipped decorators (those with skip=True and past the cutoff) 178 if args.dry_run: 179 for decorator in decorators: 180 if decorator.skip and decorator.age_days > args.cutoff_days: 181 print( 182 f"{file_path}:{decorator.line_number}:{decorator.column}: " 183 f"Skipped (skip=True) {decorator.content}" 184 ) 185 if not old_decorators: 186 continue 187 188 # Remove old decorators 189 if removed := remove_decorators_from_file(file_path, old_decorators, args.dry_run): 190 for decorator in removed: 191 action = "Would remove" if args.dry_run else "Removed" 192 print( 193 f"{file_path}:{decorator.line_number}:{decorator.column}: " 194 f"{action} {decorator.content} (age: {decorator.age_days} days)" 195 ) 196 197 198 if __name__ == "__main__": 199 main()