amoghrajesh commented on code in PR #72861:
URL: https://github.com/apache/airflow/pull/72861#discussion_r3976485129
##########
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:
Do we want to also support `airflow.sdk.dag(dag_id=......)` with no alias
like` import airflow.sdk as sdk`?
(rare style but worth putting it)
--
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]