jedcunningham commented on code in PR #72861:
URL: https://github.com/apache/airflow/pull/72861#discussion_r3981414444
##########
airflow-core/src/airflow/utils/dag_version_inflation_checker.py:
##########
@@ -312,24 +312,35 @@ def __init__(self, from_imports: dict[str, tuple[str,
str]]):
def is_dag_constructor(self, node: ast.Call) -> bool:
"""Check if a call is a Dag constructor."""
- # to handle use case "from airflow import sdk" and "with sdk.DAG()"
- if isinstance(node.func, ast.Attribute) and
isinstance(node.func.value, ast.Name):
- if node.func.value.id in self.from_imports:
- module, original = self.from_imports[node.func.value.id]
- if (module == "airflow" or module.startswith("airflow.")) and
node.func.attr in (
- "DAG",
- "dag",
- ):
+ # The Dag file is not imported yet, so there is no object to test —
only names and its imports.
+ # A *DAG/*Dag suffix counts on its own, so a subclass named anything
else is missed.
+ # Lowercase "dag" is an ordinary word and needs an import — and only a
plain name can
+ # be an imported one, since "from airflow import DAG as D" says
nothing about config.D().
+ func = node.func
+
+ # DAG(...), TeamDAG(...), an alias like D(...), or the @dag(...)
decorator
+ if isinstance(func, ast.Name):
+ if func.id in self.from_imports:
+ module, original = self.from_imports[func.id]
+ if self._is_airflow_module(module) and original in ("DAG",
"dag"):
return True
+ return func.id.endswith(("DAG", "Dag"))
- # to handle use case "from airflow import DAG" form or "from
airflow.decorator import dag"
- if isinstance(node.func, ast.Name) and node.func.id in
self.from_imports:
- module, original = self.from_imports[node.func.id]
- if (module == "airflow" or module.startswith("airflow.")) and
original in ("DAG", "dag"):
- return True
+ # sdk.DAG(...), airflow.sdk.DAG(...), or the @sdk.dag(...) decorator
+ if isinstance(func, ast.Attribute):
+ if func.attr == "dag" and isinstance(func.value, ast.Name):
Review Comment:
We could, but I'd rather do it in a follow up since the mechanism to do it
is a bit different. I'll look at that soon.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
