Refactor consolidate import from io (apache#34377)

airflow/cli/commands/config_command.py

@@ -17,7 +17,7 @@
 """Config sub-commands."""
 from __future__ import annotations
 
-import io
+from io import StringIO
 
 import pygments
 from pygments.lexers.configs import IniLexer
@@ -31,7 +31,7 @@
 @providers_configuration_loaded
 def show_config(args):
     """Show current application configuration."""
-    with io.StringIO() as output:
+    with StringIO() as output:
         conf.write(
             output,
             section=args.section,
@@ -44,9 +44,9 @@ def show_config(args):
             only_defaults=args.defaults,
         )
         code = output.getvalue()
-        if should_use_colors(args):
-            code = pygments.highlight(code=code, formatter=get_terminal_formatter(), lexer=IniLexer())
-        print(code)
+    if should_use_colors(args):
+        code = pygments.highlight(code=code, formatter=get_terminal_formatter(), lexer=IniLexer())
+    print(code)
 
 
 @providers_configuration_loaded

airflow/cli/utils.py

@@ -21,7 +21,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    import io
+    from io import IOBase
 
 
 class CliConflictError(Exception):
@@ -30,7 +30,7 @@ class CliConflictError(Exception):
     pass
 
 
-def is_stdout(fileio: io.IOBase) -> bool:
+def is_stdout(fileio: IOBase) -> bool:
     """Check whether a file IO is stdout.
 
     The intended use case for this helper is to check whether an argument parsed

airflow/configuration.py

@@ -19,7 +19,6 @@
 import contextlib
 import datetime
 import functools
-import io
 import itertools
 import json
 import logging
@@ -35,6 +34,7 @@
 from configparser import ConfigParser, NoOptionError, NoSectionError
 from contextlib import contextmanager
 from copy import deepcopy
+from io import StringIO
 from json.decoder import JSONDecodeError
 from typing import IO, TYPE_CHECKING, Any, Dict, Generator, Iterable, Pattern, Set, Tuple, Union
 from urllib.parse import urlsplit
@@ -1792,7 +1792,7 @@ def load_test_config(self):
         unit_test_config_file = pathlib.Path(__file__).parent / "config_templates" / "unit_tests.cfg"
         unit_test_config = unit_test_config_file.read_text()
         self.remove_all_read_configurations()
-        with io.StringIO(unit_test_config) as test_config_file:
+        with StringIO(unit_test_config) as test_config_file:
             self.read_file(test_config_file)
         # set fernet key to a random value
         global FERNET_KEY

airflow/serialization/serializers/pandas.py

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+from io import BytesIO
 from typing import TYPE_CHECKING
 
 from airflow.utils.module_loading import qualname
@@ -57,14 +58,12 @@ def deserialize(classname: str, version: int, data: object) -> pd.DataFrame:
     if version > __version__:
         raise TypeError(f"serialized {version} of {classname} > {__version__}")
 
-    import io
-
     from pyarrow import parquet as pq
 
     if not isinstance(data, str):
         raise TypeError(f"serialized {classname} has wrong data type {type(data)}")
 
-    buf = io.BytesIO(bytes.fromhex(data))
-    df = pq.read_table(buf).to_pandas()
+    with BytesIO(bytes.fromhex(data)) as buf:
+        df = pq.read_table(buf).to_pandas()
 
     return df

airflow/utils/file.py

@@ -18,10 +18,10 @@
 from __future__ import annotations
 
 import ast
-import io
 import logging
 import os
 import zipfile
+from io import TextIOWrapper
 from pathlib import Path
 from typing import Generator, NamedTuple, Pattern, Protocol, overload
 
@@ -189,7 +189,7 @@ def open_maybe_zipped(fileloc, mode="r"):
     """
     _, archive, filename = ZIP_REGEX.search(fileloc).groups()
     if archive and zipfile.is_zipfile(archive):
-        return io.TextIOWrapper(zipfile.ZipFile(archive, mode=mode).open(filename))
+        return TextIOWrapper(zipfile.ZipFile(archive, mode=mode).open(filename))
     else:
         return open(fileloc, mode=mode)
 

airflow/www/decorators.py

@@ -22,7 +22,7 @@
 import itertools
 import json
 import logging
-from io import BytesIO as IO
+from io import BytesIO
 from typing import Callable, TypeVar, cast
 
 import pendulum
@@ -156,12 +156,10 @@ def zipper(response):
                 or "Content-Encoding" in response.headers
             ):
                 return response
-            gzip_buffer = IO()
-            gzip_file = gzip.GzipFile(mode="wb", fileobj=gzip_buffer)
-            gzip_file.write(response.data)
-            gzip_file.close()
-
-            response.data = gzip_buffer.getvalue()
+            with BytesIO() as gzip_buffer:
+                with gzip.GzipFile(mode="wb", fileobj=gzip_buffer) as gzip_file:
+                    gzip_file.write(response.data)
+                response.data = gzip_buffer.getvalue()
             response.headers["Content-Encoding"] = "gzip"
             response.headers["Vary"] = "Accept-Encoding"
             response.headers["Content-Length"] = len(response.data)

tests/auth/managers/fab/cli_commands/test_role_command.py

@@ -17,9 +17,9 @@
 # under the License.
 from __future__ import annotations
 
-import io
 import json
 from contextlib import redirect_stdout
+from io import StringIO
 from typing import TYPE_CHECKING
 
 import pytest
@@ -96,7 +96,7 @@ def test_cli_list_roles(self):
         self.appbuilder.sm.add_role("FakeTeamA")
         self.appbuilder.sm.add_role("FakeTeamB")
 
-        with redirect_stdout(io.StringIO()) as stdout:
+        with redirect_stdout(StringIO()) as stdout:
             role_command.roles_list(self.parser.parse_args(["roles", "list"]))
             stdout = stdout.getvalue()
 

tests/auth/managers/fab/cli_commands/test_user_command.py

@@ -16,12 +16,12 @@
 # under the License.
 from __future__ import annotations
 
-import io
 import json
 import os
 import re
 import tempfile
 from contextlib import redirect_stdout
+from io import StringIO
 
 import pytest
 
@@ -121,7 +121,7 @@ def test_cli_delete_user(self):
                 "test3",
             ]
         )
-        with redirect_stdout(io.StringIO()) as stdout:
+        with redirect_stdout(StringIO()) as stdout:
             user_command.users_delete(args)
         assert 'User "test3" deleted' in stdout.getvalue()
 
@@ -152,7 +152,7 @@ def test_cli_delete_user_by_email(self):
                 "[email protected]",
             ]
         )
-        with redirect_stdout(io.StringIO()) as stdout:
+        with redirect_stdout(StringIO()) as stdout:
             user_command.users_delete(args)
         assert 'User "test4" deleted' in stdout.getvalue()
 
@@ -225,7 +225,7 @@ def test_cli_list_users(self):
                 ]
             )
             user_command.users_create(args)
-        with redirect_stdout(io.StringIO()) as stdout:
+        with redirect_stdout(StringIO()) as stdout:
             user_command.users_list(self.parser.parse_args(["users", "list"]))
             stdout = stdout.getvalue()
         for i in range(3):
@@ -368,7 +368,7 @@ def test_cli_add_user_role(self, create_user_test4):
         ), "User should not yet be a member of role 'Op'"
 
         args = self.parser.parse_args(["users", "add-role", "--username", "test4", "--role", "Op"])
-        with redirect_stdout(io.StringIO()) as stdout:
+        with redirect_stdout(StringIO()) as stdout:
             user_command.users_manage_role(args, remove=False)
         assert 'User "test4" added to role "Op"' in stdout.getvalue()
 
@@ -382,7 +382,7 @@ def test_cli_remove_user_role(self, create_user_test4):
         ), "User should have been created with role 'Viewer'"
 
         args = self.parser.parse_args(["users", "remove-role", "--username", "test4", "--role", "Viewer"])
-        with redirect_stdout(io.StringIO()) as stdout:
+        with redirect_stdout(StringIO()) as stdout:
             user_command.users_manage_role(args, remove=True)
         assert 'User "test4" removed from role "Viewer"' in stdout.getvalue()
 

tests/cli/commands/test_cheat_sheet_command.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import contextlib
-import io
+from io import StringIO
 from unittest import mock
 
 from airflow.cli import cli_parser
@@ -95,7 +95,7 @@ def setup_class(cls):
 
     @mock.patch("airflow.cli.cli_parser.airflow_commands", MOCK_COMMANDS)
     def test_should_display_index(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             args = self.parser.parse_args(["cheat-sheet"])
             args.func(args)
         output = temp_stdout.getvalue()

tests/cli/commands/test_config_command.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import contextlib
-import io
+from io import StringIO
 from unittest import mock
 
 import pytest
@@ -32,7 +32,7 @@ class TestCliConfigList:
     def setup_class(cls):
         cls.parser = cli_parser.get_parser()
 
-    @mock.patch("airflow.cli.commands.config_command.io.StringIO")
+    @mock.patch("airflow.cli.commands.config_command.StringIO")
     @mock.patch("airflow.cli.commands.config_command.conf")
     def test_cli_show_config_should_write_data(self, mock_conf, mock_stringio):
         config_command.show_config(self.parser.parse_args(["config", "list", "--color", "off"]))
@@ -48,7 +48,7 @@ def test_cli_show_config_should_write_data(self, mock_conf, mock_stringio):
             only_defaults=False,
         )
 
-    @mock.patch("airflow.cli.commands.config_command.io.StringIO")
+    @mock.patch("airflow.cli.commands.config_command.StringIO")
     @mock.patch("airflow.cli.commands.config_command.conf")
     def test_cli_show_config_should_write_data_specific_section(self, mock_conf, mock_stringio):
         config_command.show_config(
@@ -68,21 +68,21 @@ def test_cli_show_config_should_write_data_specific_section(self, mock_conf, moc
 
     @conf_vars({("core", "testkey"): "test_value"})
     def test_cli_show_config_should_display_key(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(self.parser.parse_args(["config", "list", "--color", "off"]))
         output = temp_stdout.getvalue()
         assert "[core]" in output
         assert "testkey = test_value" in temp_stdout.getvalue()
 
     def test_cli_show_config_should_only_show_comments_when_no_defaults(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(self.parser.parse_args(["config", "list", "--color", "off"]))
         output = temp_stdout.getvalue()
         lines = output.splitlines()
         assert all(not line.startswith("#") or line.endswith("= ") for line in lines if line)
 
     def test_cli_show_config_shows_descriptions(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--include-descriptions"])
             )
@@ -95,7 +95,7 @@ def test_cli_show_config_shows_descriptions(self):
         assert all(not line.startswith("# Variable:") for line in lines if line)
 
     def test_cli_show_config_shows_examples(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--include-examples"])
             )
@@ -107,7 +107,7 @@ def test_cli_show_config_shows_examples(self):
         assert all(not line.startswith("# Variable:") for line in lines if line)
 
     def test_cli_show_config_shows_variables(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--include-env-vars"])
             )
@@ -119,7 +119,7 @@ def test_cli_show_config_shows_variables(self):
         assert any(line.startswith("# Variable:") for line in lines if line)
 
     def test_cli_show_config_shows_sources(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--include-sources"])
             )
@@ -131,7 +131,7 @@ def test_cli_show_config_shows_sources(self):
         assert all(not line.startswith("# Variable:") for line in lines if line)
 
     def test_cli_show_config_defaults(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--defaults"])
             )
@@ -146,7 +146,7 @@ def test_cli_show_config_defaults(self):
 
     @conf_vars({("core", "task_runner"): "test-runner"})
     def test_cli_show_config_defaults_not_show_conf_changes(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--defaults"])
             )
@@ -156,7 +156,7 @@ def test_cli_show_config_defaults_not_show_conf_changes(self):
 
     @mock.patch("os.environ", {"AIRFLOW__CORE__TASK_RUNNER": "test-env-runner"})
     def test_cli_show_config_defaults_do_not_show_env_changes(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--defaults"])
             )
@@ -166,29 +166,29 @@ def test_cli_show_config_defaults_do_not_show_env_changes(self):
 
     @conf_vars({("core", "task_runner"): "test-runner"})
     def test_cli_show_changed_defaults_when_overridden_in_conf(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(self.parser.parse_args(["config", "list", "--color", "off"]))
         output = temp_stdout.getvalue()
         lines = output.splitlines()
         assert any(line.startswith("task_runner = test-runner") for line in lines if line)
 
     @mock.patch("os.environ", {"AIRFLOW__CORE__TASK_RUNNER": "test-env-runner"})
     def test_cli_show_changed_defaults_when_overridden_in_env(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(self.parser.parse_args(["config", "list", "--color", "off"]))
         output = temp_stdout.getvalue()
         lines = output.splitlines()
         assert any(line.startswith("task_runner = test-env-runner") for line in lines if line)
 
     def test_cli_has_providers(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(self.parser.parse_args(["config", "list", "--color", "off"]))
         output = temp_stdout.getvalue()
         lines = output.splitlines()
         assert any(line.startswith("celery_config_options") for line in lines if line)
 
     def test_cli_comment_out_everything(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.show_config(
                 self.parser.parse_args(["config", "list", "--color", "off", "--comment-out-everything"])
             )
@@ -204,7 +204,7 @@ def setup_class(cls):
 
     @conf_vars({("core", "test_key"): "test_value"})
     def test_should_display_value(self):
-        with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
             config_command.get_value(self.parser.parse_args(["config", "get-value", "core", "test_key"]))
 
         assert "test_value" == temp_stdout.getvalue().strip()