Refactor: Consolidate import textwrap

airflow/decorators/base.py

@@ -18,9 +18,9 @@
 
 import inspect
 import itertools
+import textwrap
 import warnings
 from functools import cached_property
-from textwrap import dedent
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -293,7 +293,7 @@ def _hook_apply_defaults(self, *args, **kwargs):
 
     def get_python_source(self):
         raw_source = inspect.getsource(self.python_callable)
-        res = dedent(raw_source)
+        res = textwrap.dedent(raw_source)
         res = remove_task_decorator(res, self.custom_operator_name)
         return res
 

airflow/example_dags/example_passing_params_via_test_command.py

@@ -20,7 +20,7 @@
 
 import datetime
 import os
-from textwrap import dedent
+import textwrap
 
 import pendulum
 
@@ -67,7 +67,7 @@ def print_env_vars(test_mode=None):
 ) as dag:
     run_this = my_py_command(params={"miff": "agg"})
 
-    my_command = dedent(
+    my_command = textwrap.dedent(
         """
         echo "'foo' was passed in via Airflow CLI Test command with value '$FOO'"
         echo "'miff' was passed in via BashOperator with value '$MIFF'"

airflow/example_dags/tutorial.py

@@ -24,8 +24,8 @@
 
 # [START tutorial]
 # [START import_module]
+import textwrap
 from datetime import datetime, timedelta
-from textwrap import dedent
 
 # The DAG object; we'll need this to instantiate a DAG
 from airflow import DAG
@@ -87,7 +87,7 @@
     # [END basic_task]
 
     # [START documentation]
-    t1.doc_md = dedent(
+    t1.doc_md = textwrap.dedent(
         """\
     #### Task Documentation
     You can document your task using the attributes `doc_md` (markdown),
@@ -105,7 +105,7 @@
     # [END documentation]
 
     # [START jinja_template]
-    templated_command = dedent(
+    templated_command = textwrap.dedent(
         """
     {% for i in range(5) %}
         echo "{{ ds }}"

airflow/example_dags/tutorial_dag.py

@@ -24,7 +24,7 @@
 # [START tutorial]
 # [START import_module]
 import json
-from textwrap import dedent
+import textwrap
 
 import pendulum
 
@@ -94,7 +94,7 @@ def load(**kwargs):
         task_id="extract",
         python_callable=extract,
     )
-    extract_task.doc_md = dedent(
+    extract_task.doc_md = textwrap.dedent(
         """\
     #### Extract task
     A simple Extract task to get data ready for the rest of the data pipeline.
@@ -107,7 +107,7 @@ def load(**kwargs):
         task_id="transform",
         python_callable=transform,
     )
-    transform_task.doc_md = dedent(
+    transform_task.doc_md = textwrap.dedent(
         """\
     #### Transform task
     A simple Transform task which takes in the collection of order data from xcom
@@ -120,7 +120,7 @@ def load(**kwargs):
         task_id="load",
         python_callable=load,
     )
-    load_task.doc_md = dedent(
+    load_task.doc_md = textwrap.dedent(
         """\
     #### Load task
     A simple Load task which takes in the result of the Transform task, by reading it

airflow/operators/python.py

@@ -25,13 +25,13 @@
 import shutil
 import subprocess
 import sys
+import textwrap
 import types
 import warnings
 from abc import ABCMeta, abstractmethod
 from collections.abc import Container
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from textwrap import dedent
 from typing import TYPE_CHECKING, Any, Callable, Collection, Iterable, Mapping, Sequence, cast
 
 import dill
@@ -398,7 +398,7 @@ def execute(self, context: Context) -> Any:
 
     def get_python_source(self):
         """Return the source of self.python_callable."""
-        return dedent(inspect.getsource(self.python_callable))
+        return textwrap.dedent(inspect.getsource(self.python_callable))
 
     def _write_args(self, file: Path):
         if self.op_args or self.op_kwargs:

scripts/ci/pre_commit/pre_commit_insert_extras.py

@@ -19,8 +19,8 @@
 
 import os
 import sys
+import textwrap
 from pathlib import Path
-from textwrap import wrap
 
 AIRFLOW_SOURCES_DIR = Path(__file__).parents[3].resolve()
 
@@ -57,7 +57,7 @@
     global_constants_file_path = (
         AIRFLOW_SOURCES_DIR / "dev" / "breeze" / "src" / "airflow_breeze" / "global_constants.py"
     )
-    extras_list = wrap(", ".join(EXTRAS_DEPENDENCIES.keys()), 100)
+    extras_list = textwrap.wrap(", ".join(EXTRAS_DEPENDENCIES.keys()), 100)
     extras_list = [line + "\n" for line in extras_list]
     extras_code = [f"    {extra}\n" for extra in EXTRAS_DEPENDENCIES.keys()]
     insert_documentation(install_file_path, extras_list, INSTALL_HEADER, INSTALL_FOOTER)

scripts/in_container/run_migration_reference.py

@@ -23,8 +23,8 @@
 
 import os
 import re
+import textwrap
 from pathlib import Path
-from textwrap import wrap
 from typing import TYPE_CHECKING, Iterable
 
 from alembic.script import ScriptDirectory
@@ -125,7 +125,7 @@ def update_docs(revisions: Iterable[Script]):
                 revision=wrap_backticks(rev.revision) + revision_suffix(rev),
                 down_revision=wrap_backticks(rev.down_revision),
                 version=wrap_backticks(rev.module.airflow_version),  # type: ignore
-                description="\n".join(wrap(rev.doc, width=60)),
+                description="\n".join(textwrap.wrap(rev.doc, width=60)),
             )
         )
 

setup.py

@@ -28,10 +28,10 @@
 import os
 import subprocess
 import sys
+import textwrap
 import unittest
 from copy import deepcopy
 from pathlib import Path
-from textwrap import wrap
 from typing import Iterable
 
 from setuptools import Command, Distribution, find_namespace_packages, setup
@@ -206,7 +206,7 @@ def finalize_options(self) -> None:
 
     def run(self) -> None:
         """List extras."""
-        print("\n".join(wrap(", ".join(EXTRAS_DEPENDENCIES.keys()), 100)))
+        print("\n".join(textwrap.wrap(", ".join(EXTRAS_DEPENDENCIES.keys()), 100)))
 
 
 def git_version() -> str:

tests/dag_processing/test_job_runner.py

@@ -27,11 +27,11 @@
 import random
 import socket
 import sys
+import textwrap
 import threading
 import time
 from datetime import datetime, timedelta
 from logging.config import dictConfig
-from textwrap import dedent
 from unittest import mock
 from unittest.mock import MagicMock, PropertyMock
 
@@ -892,7 +892,7 @@ def fake_processor_(*args, **kwargs):
     @mock.patch("airflow.dag_processing.manager.Stats.timing")
     def test_send_file_processing_statsd_timing(self, statsd_timing_mock, tmp_path):
         path_to_parse = tmp_path / "temp_dag.py"
-        dag_code = dedent(
+        dag_code = textwrap.dedent(
             """
         from airflow import DAG
         dag = DAG(dag_id='temp_dag', schedule='0 0 * * *')

tests/dags/test_default_impersonation.py

@@ -17,8 +17,8 @@
 # under the License.
 from __future__ import annotations
 
+import textwrap
 from datetime import datetime
-from textwrap import dedent
 
 from airflow.models import DAG
 from airflow.operators.bash import BashOperator
@@ -34,7 +34,7 @@
 
 deelevated_user = "airflow_test_user"
 
-test_command = dedent(
+test_command = textwrap.dedent(
     f"""\
     if [ '{deelevated_user}' != "$(whoami)" ]; then
         echo current user $(whoami) is not {deelevated_user}!

tests/dags/test_impersonation.py

@@ -17,8 +17,8 @@
 # under the License.
 from __future__ import annotations
 
+import textwrap
 from datetime import datetime
-from textwrap import dedent
 
 from airflow.models import DAG
 from airflow.operators.bash import BashOperator
@@ -34,7 +34,7 @@
 
 run_as_user = "airflow_test_user"
 
-test_command = dedent(
+test_command = textwrap.dedent(
     f"""\
     if [ '{run_as_user}' != "$(whoami)" ]; then
         echo current user is not {run_as_user}!

tests/dags/test_no_impersonation.py

@@ -17,8 +17,8 @@
 # under the License.
 from __future__ import annotations
 
+import textwrap
 from datetime import datetime
-from textwrap import dedent
 
 from airflow.models import DAG
 from airflow.operators.bash import BashOperator
@@ -32,7 +32,7 @@
 
 dag = DAG(dag_id="test_no_impersonation", default_args=args)
 
-test_command = dedent(
+test_command = textwrap.dedent(
     """\
     sudo ls
     if [ $? -ne 0 ]; then