Add pre-commit hook for custom_operator_name

.pre-commit-config.yaml

@@ -506,6 +506,12 @@ repos:
           (?x)
           ^airflow/providers/.*\.py$
         exclude: ^airflow/_vendor/
+      - id: check-decorated-operator-implements-custom-name
+        name: Check @task decorator implements custom_operator_name
+        language: python
+        entry: ./scripts/ci/pre_commit/pre_commit_decorator_operator_implements_custom_name.py
+        pass_filenames: true
+        files: ^airflow/.*\.py$
       - id: check-provide-create-sessions-imports
         language: pygrep
         name: Check provide_session and create_session imports

STATIC_CODE_CHECKS.rst

@@ -157,6 +157,8 @@ require Breeze Docker image to be build locally.
 +--------------------------------------------------------+------------------------------------------------------------------+---------+
 | check-daysago-import-from-utils                        | Make sure days_ago is imported from airflow.utils.dates          |         |
 +--------------------------------------------------------+------------------------------------------------------------------+---------+
+| check-decorated-operator-implements-custom-name        | Check @task decorator implements custom_operator_name            |         |
++--------------------------------------------------------+------------------------------------------------------------------+---------+
 | check-docstring-param-types                            | Check that docstrings do not specify param types                 |         |
 +--------------------------------------------------------+------------------------------------------------------------------+---------+
 | check-example-dags-urls                                | Check that example dags url include provider versions            |         |

airflow/decorators/branch_python.py

@@ -38,6 +38,8 @@ class _BranchPythonDecoratedOperator(DecoratedOperator, BranchPythonOperator):
         Defaults to False.
     """
 
+    custom_operator_name = "@task.branch"
+
     template_fields: Sequence[str] = ('op_args', 'op_kwargs')
     template_fields_renderers = {"op_args": "py", "op_kwargs": "py"}
 

airflow/providers/docker/decorators/docker.py

@@ -62,6 +62,8 @@ class _DockerDecoratedOperator(DecoratedOperator, DockerOperator):
         Defaults to False.
     """
 
+    custom_operator_name = "@task.docker"
+
     template_fields: Sequence[str] = ('op_args', 'op_kwargs')
 
     # since we won't mutate the arguments, we should just do the shallow copy

dev/breeze/src/airflow_breeze/pre_commit_ids.py

@@ -35,6 +35,7 @@
     'check-builtin-literals',
     'check-changelog-has-no-duplicates',
     'check-daysago-import-from-utils',
+    'check-decorated-operator-implements-custom-name',
     'check-docstring-param-types',
     'check-example-dags-urls',
     'check-executables-have-shebangs',

docs/apache-airflow/howto/create-custom-decorator.rst

@@ -31,6 +31,10 @@ tasks. The steps to create and register ``@task.foo`` are:
     ``airflow.decorators.base.DecoratedOperator``, Airflow will supply much of the needed functionality required
     to treat your new class as a taskflow native class.
 
+    You should also override the ``custom_operator_name`` attribute to provide a custom name for the task. For
+    example, ``_DockerDecoratedOperator`` in the ``apache-airflow-providers-docker`` provider sets this to
+    ``@task.docker`` to indicate the decorator name it implements.
+
 2. Create a ``foo_task`` function
 
     Once you have your decorated class, create a function like this, to convert

images/breeze/output-commands-hash.txt

@@ -54,7 +54,7 @@ setup:params:bc6b726d05240c22ac8522edc4a200ee
 setup:short_help:37a6259cc0c1dae299a7866489dff0bd
 shell:7baeb98c0103c594502ccf3d32ae24c6
 start-airflow:3e793b11dc2158c54bfc189bfe20d6f2
-static-checks:03e96245e60c225ed94514ad7b42ceb0
+static-checks:a86b25f9281c47abf362a221b316c7a2
 stop:8ebd8a42f1003495d37b884de5ac7ce6
 testing:chain:68934a3e9455fa72420237eb05902327
 testing:commands:470be9efe36272d49d4d43acb34b6fe5

scripts/ci/pre_commit/pre_commit_decorator_operator_implements_custom_name.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import ast
+import itertools
+import pathlib
+import sys
+from typing import Iterator
+
+
+def iter_decorated_operators(source: pathlib.Path) -> Iterator[ast.ClassDef]:
+    mod = ast.parse(source.read_text("utf-8"), str(source))
+    for node in ast.walk(mod):
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if not any(isinstance(base, ast.Name) and base.id == "DecoratedOperator" for base in node.bases):
+            continue  # Not a decorated operator.
+        yield node
+
+
+def check_missing_custom_operator_name(klass: ast.ClassDef) -> bool:
+    for node in ast.iter_child_nodes(klass):
+        if isinstance(node, ast.AnnAssign):
+            if isinstance(node.target, ast.Name) and node.target.id == "custom_operator_name":
+                return True
+        elif isinstance(node, ast.Assign):
+            if any(isinstance(t, ast.Name) and t.id == "custom_operator_name" for t in node.targets):
+                return True
+    return False
+
+
+def main(*args: str) -> int:
+    classes = itertools.chain.from_iterable(iter_decorated_operators(pathlib.Path(a)) for a in args[1:])
+    results = ((k.name, check_missing_custom_operator_name(k)) for k in classes)
+    failures = [name for name, success in results if not success]
+    for failure in failures:
+        print(f"Missing custom_operator_name in class: {failure}")
+    return len(failures)
+
+
+if __name__ == "__main__":
+    sys.exit(main(*sys.argv))