2222# "rich>=13.6.0",
2323# ]
2424# ///
25+ """Check that no new ``airflow.sdk`` imports are introduced in ``airflow-core``.
26+
27+ All *existing* imports are recorded in ``generated/known_sdk_imports_in_core.txt``
28+ as ``relative/path::N`` entries (one per file), where ``N`` is the maximum
29+ number of ``airflow.sdk`` import statements allowed in that file. A file whose
30+ current count exceeds the recorded limit is treated as a violation -- core
31+ should not gain new runtime dependencies on the Task SDK. Two escape hatches
32+ exist: ``# noqa: SDK001`` on a single import line, which keeps that line out of
33+ the count and out of the allowlist entirely, or raising the file's recorded
34+ count via ``--generate FILE`` when a file genuinely needs several.
35+
36+ ``airflow-core/src/airflow/example_dags`` is skipped entirely: those files are
37+ Dags authored with the Task SDK, not core runtime code.
38+
39+ Modes
40+ -----
41+ Default (files passed by prek/pre-commit):
42+ Check only the supplied files; fail if any file's count exceeds the limit.
43+ When a file's count has *decreased*, the allowlist entry is tightened
44+ automatically and the hook exits with a non-zero code so that pre-commit
45+ reports the modified allowlist -- just stage
46+ ``generated/known_sdk_imports_in_core.txt`` and re-run.
47+
48+ ``--all-files``:
49+ Walk all of ``airflow-core/src/airflow`` and check every ``.py`` file.
50+
51+ ``--cleanup``:
52+ Remove entries for files that no longer exist. Safe to run at any time;
53+ does not add new entries or raise limits.
54+
55+ ``--generate [FILE ...]``:
56+ With files, record the current counts for just those files and leave every
57+ other entry alone -- the usual way to approve a deliberate new import.
58+ With no files, *rebuild* the allowlist from scratch; that also approves any
59+ unrelated drift, so it is meant for the initial setup or a clean-up sprint.
60+ """
61+
2562from __future__ import annotations
2663
2764import argparse
28- import sys
65+ import ast
66+ from collections .abc import Iterable
2967from pathlib import Path
3068
31- from common_prek_utils import find_import_violations , report_import_violations
69+ from common_prek_utils import AIRFLOW_CORE_ROOT_PATH , AllowlistManager , find_import_violations
70+ from rich .console import Console
71+
72+ console = Console (color_system = "standard" , width = 200 )
73+
74+ REPO_ROOT = AIRFLOW_CORE_ROOT_PATH .parent
75+ CORE_SRC_ROOT = AIRFLOW_CORE_ROOT_PATH / "src" / "airflow"
3276
3377NOCHECK_CODE = "SDK001"
3478
79+ # Dags are authored *with* the task SDK, so example Dags shipped in core are not core code
80+ # for the purpose of this check and never take part in the ratchet.
81+ EXCLUDED_DIR_NAMES = frozenset ({"example_dags" })
82+
3583
3684def check_file_for_sdk_imports (file_path : Path ) -> list [tuple [int , str ]]:
3785 """Check file for airflow.sdk imports. Returns list of (line_num, import_statement)."""
@@ -42,23 +90,153 @@ def check_file_for_sdk_imports(file_path: Path) -> list[tuple[int, str]]:
4290 )
4391
4492
45- def main ():
46- parser = argparse .ArgumentParser (description = "Check for SDK imports in airflow-core files" )
47- parser .add_argument ("files" , nargs = "*" , help = "Files to check" )
48- args = parser .parse_args ()
93+ def _can_parse (path : Path ) -> bool :
94+ """Whether *path* is valid enough Python for its import count to be trustworthy.
4995
50- if not args .files :
51- return
96+ ``find_import_violations`` swallows parse failures and reports zero violations for
97+ them, which is indistinguishable from a file that genuinely has none. Callers must
98+ filter those files out *before* counting, so a file mid-edit with a syntax error
99+ never looks like a legitimate drop in its allowlist entry.
100+ """
101+ try :
102+ ast .parse (path .read_text (encoding = "utf-8" ), filename = str (path ))
103+ except (OSError , UnicodeDecodeError , SyntaxError ):
104+ return False
105+ return True
52106
53- report_import_violations (
54- args .files ,
55- check_func = check_file_for_sdk_imports ,
56- violation_label = "SDK import(s) in core files" ,
57- nocheck_code = NOCHECK_CODE ,
58- only_python_files = True ,
107+
108+ class SdkImportsAllowlistManager (AllowlistManager ):
109+ def __init__ (self , allowlist_file : Path ) -> None :
110+ super ().__init__ (allowlist_file , repo_root = REPO_ROOT )
111+
112+ def is_excluded (self , path : Path ) -> bool :
113+ return not EXCLUDED_DIR_NAMES .isdisjoint (path .parts )
114+
115+ def iter_files (self ) -> Iterable [Path ]:
116+ return (
117+ path for path in CORE_SRC_ROOT .rglob ("*.py" ) if not self .is_excluded (path ) and _can_parse (path )
118+ )
119+
120+ def check (self , files : list [Path ], allowlist : dict [str , int ]) -> int :
121+ checkable = []
122+ for path in files :
123+ if self .is_excluded (path ):
124+ continue
125+ if path .exists () and path .suffix == ".py" and not _can_parse (path ):
126+ console .print (
127+ f"[yellow]Skipping unparsable file (not counted, allowlist entry left "
128+ f"untouched):[/yellow] { path }
129+ )
130+ continue
131+ checkable .append (path )
132+ return super ().check (checkable , allowlist )
133+
134+ def count_occurrences (self , path : Path ) -> int :
135+ return len (check_file_for_sdk_imports (path ))
136+
137+ def generate_for (self , files : list [Path ]) -> int :
138+ """Rewrite the allowlist entries for *files* only, leaving every other entry untouched."""
139+ allowlist = self .load ()
140+ updated : list [tuple [str , int , int ]] = []
141+
142+ for path in files :
143+ if self .is_excluded (path ) or not path .exists () or path .suffix != ".py" :
144+ continue
145+ if not _can_parse (path ):
146+ console .print (f"[yellow]Skipping unparsable file:[/yellow] { path } )
147+ continue
148+ try :
149+ rel = str (path .resolve ().relative_to (self .repo_root .resolve ()))
150+ except ValueError :
151+ console .print (f"[yellow]Skipping file outside the repository:[/yellow] { path } )
152+ continue
153+ before = allowlist .get (rel , 0 )
154+ after = self .count_occurrences (path )
155+ if before == after :
156+ continue
157+ if after :
158+ allowlist [rel ] = after
159+ else :
160+ del allowlist [rel ]
161+ updated .append ((rel , before , after ))
162+
163+ if not updated :
164+ console .print ("[green]Allowlist already up to date for the given files.[/green]" )
165+ return 0
166+
167+ self .save (allowlist )
168+ console .print (
169+ f"[green]Updated[/green] [cyan]{ self .allowlist_file .relative_to (self .repo_root )}
170+ )
171+ for rel , before , after in updated :
172+ console .print (f" [cyan]{ rel } { before } { after } )
173+ return 0
174+
175+ def violation_panel_text (self ) -> str :
176+ return (
177+ "New [bold]airflow.sdk[/bold] import detected in airflow-core.\n "
178+ "Core (scheduler/API server) should not gain new runtime dependencies "
179+ "on the Task SDK.\n \n "
180+ "The preferred fix is to drop the import. If it has to stay, pick one:\n \n "
181+ "[bold]1. A single deliberate import line[/bold] -- append `# noqa: SDK001` to it.\n "
182+ " The line is then invisible to this check and never reaches the allowlist.\n \n "
183+ "[bold]2. A file that genuinely needs several[/bold] -- record its new count:\n \n "
184+ " [cyan]uv run ./scripts/ci/prek/check_sdk_imports_in_core.py --generate FILE ...[/cyan]\n \n "
185+ " then commit the updated [cyan]generated/known_sdk_imports_in_core.txt[/cyan].\n "
186+ " Passing no FILE rebuilds the whole allowlist and approves any unrelated\n "
187+ " drift with it, so prefer the scoped form."
188+ )
189+
190+
191+ def main (argv : list [str ] | None = None ) -> int :
192+ parser = argparse .ArgumentParser (
193+ description = "Prevent new airflow.sdk imports in airflow-core." ,
194+ formatter_class = argparse .RawDescriptionHelpFormatter ,
195+ epilog = __doc__ ,
196+ )
197+ parser .add_argument ("files" , nargs = "*" , metavar = "FILE" , help = "Files to check (provided by prek)" )
198+ parser .add_argument (
199+ "--all-files" ,
200+ action = "store_true" ,
201+ help = "Check every Python file under airflow-core/src/airflow" ,
202+ )
203+ parser .add_argument (
204+ "--cleanup" ,
205+ action = "store_true" ,
206+ help = "Remove stale entries from the allowlist and exit" ,
59207 )
208+ parser .add_argument (
209+ "--generate" ,
210+ action = "store_true" ,
211+ help = (
212+ "Record the current counts for the given FILEs and exit; "
213+ "with no FILE, rebuild the whole allowlist from the current codebase"
214+ ),
215+ )
216+ args = parser .parse_args (argv )
217+
218+ manager = SdkImportsAllowlistManager (REPO_ROOT / "generated" / "known_sdk_imports_in_core.txt" )
219+ files = [Path (f ).resolve () for f in args .files ]
220+
221+ if args .generate :
222+ return manager .generate_for (files ) if files else manager .generate ()
223+
224+ if args .cleanup :
225+ return manager .cleanup ()
226+
227+ allowlist = manager .load ()
228+
229+ if args .all_files :
230+ return manager .check (list (manager .iter_files ()), allowlist )
231+
232+ if not files :
233+ console .print (
234+ "[yellow]No files provided. Pass filenames or use --all-files to scan the whole repo.[/yellow]"
235+ )
236+ return 0
237+
238+ return manager .check (files , allowlist )
60239
61240
62241if __name__ == "__main__" :
63- main ()
64- sys .exit (0 )
242+ raise SystemExit (main ())
0 commit comments