|
| 1 | +import os |
| 2 | +import glob |
| 3 | +import re |
| 4 | + |
| 5 | +def patch_noxfile(path): |
| 6 | + try: |
| 7 | + with open(path, 'r', encoding='utf-8') as f: |
| 8 | + content = f.read() |
| 9 | + except UnicodeDecodeError: |
| 10 | + return False |
| 11 | + |
| 12 | + if "import_profile" in content: |
| 13 | + return False |
| 14 | + |
| 15 | + # 1. Add "import_profile" to nox.options.sessions |
| 16 | + # It looks like: |
| 17 | + # nox.options.sessions = [ |
| 18 | + # "unit", |
| 19 | + # ... |
| 20 | + # "docs", |
| 21 | + # ] |
| 22 | + # We want to add it right before the closing bracket. |
| 23 | + if "nox.options.sessions =" in content: |
| 24 | + content = re.sub( |
| 25 | + r'(\n\s*"docs",\n)(\s*\])', |
| 26 | + r'\1 "import_profile",\n\2', |
| 27 | + content |
| 28 | + ) |
| 29 | + # If "docs", is not the last one, maybe just add it before the closing bracket |
| 30 | + if "import_profile" not in content: |
| 31 | + content = re.sub( |
| 32 | + r'(\n)(\s*\]\n)', |
| 33 | + r'\1 "import_profile",\n\2', |
| 34 | + content, |
| 35 | + count=1, |
| 36 | + flags=re.DOTALL |
| 37 | + ) |
| 38 | + # Make sure it only replaced the first matching closing bracket for sessions |
| 39 | + # wait, it's safer to just replace: |
| 40 | + # content = content.replace('"docs",\n', '"docs",\n "import_profile",\n') |
| 41 | + |
| 42 | + # 2. Append the import_profile session at the end |
| 43 | + # We need to extract the module namespace. It's usually "google". |
| 44 | + # Let's extract python-gapic-namespace or python-gapic-name from somewhere, |
| 45 | + # or just hardcode "google" like the previous agent, which worked for CI. |
| 46 | + |
| 47 | + session_text = """ |
| 48 | +
|
| 49 | +@nox.session(python="3.15") |
| 50 | +def import_profile(session): |
| 51 | + \"\"\"Ensure import times remain below defined thresholds.\"\"\" |
| 52 | + profiler_script = ( |
| 53 | + CURRENT_DIRECTORY.parent.parent / "scripts" / "import_profiler" / "profiler.py" |
| 54 | + ) |
| 55 | + if not profiler_script.exists(): |
| 56 | + session.skip("The import profiler script was not found.") |
| 57 | +
|
| 58 | + session.install(".") |
| 59 | + session.run( |
| 60 | + "python", |
| 61 | + str(profiler_script), |
| 62 | + "--package", |
| 63 | + "{pkg}", |
| 64 | + "--iterations", |
| 65 | + "10", |
| 66 | + "--fail-threshold", |
| 67 | + "5000", |
| 68 | + ) |
| 69 | +""" |
| 70 | + |
| 71 | + package_name = path.split('/')[1] |
| 72 | + session_text = session_text.replace("{pkg}", package_name) |
| 73 | + |
| 74 | + content += session_text |
| 75 | + |
| 76 | + with open(path, 'w', encoding='utf-8') as f: |
| 77 | + f.write(content) |
| 78 | + return True |
| 79 | + |
| 80 | +if __name__ == "__main__": |
| 81 | + count = 0 |
| 82 | + for path in glob.glob('packages/*/noxfile.py'): |
| 83 | + # don't modify the goldens here, we use bazel for them |
| 84 | + if "goldens" in path: |
| 85 | + continue |
| 86 | + if patch_noxfile(path): |
| 87 | + count += 1 |
| 88 | + |
| 89 | + print(f"Patched {count} noxfiles.") |
0 commit comments