Added removed import handling (#5)

* Recognize if import other than `typing` was removed from main import block
* Add test cases
* Add `--concurrent-files` option

Author: cdce8p

README.md

@@ -66,6 +66,9 @@ Always print verbose logging.
 Max number of files that should be changed. No performance improvements,
 since the limit is only applied **after** all files have been processed.
 
+**`--concurrent-files`**  
+Number of files to process concurrently during initial load.
+
 **`--full-reorder`**  
 Use additional options from [python-reorder-imports][pri] to rewrite
 - `--py38-plus` (default): Imports from `mypy_extensions` and `typing_extensions` when possible.

python_typing_update/__main__.py

@@ -13,9 +13,21 @@
 logger = logging.getLogger(__name__)
 
 
+class CustomHelpFormatter(argparse.HelpFormatter):
+    def __init__(
+        self, prog: str, indent_increment: int = 2,
+        max_help_position: int = 24, width: int | None = None,
+    ) -> None:
+        max_help_position = 40
+        super().__init__(
+            prog, indent_increment=indent_increment,
+            max_help_position=max_help_position, width=width)
+
+
 async def async_main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(
         description="Tool to update Python typing syntax.",
+        formatter_class=CustomHelpFormatter,
     )
     parser.add_argument(
         '-v', '--verbose',
@@ -30,6 +42,10 @@ async def async_main(argv: list[str] | None = None) -> int:
         '--limit', type=int, default=0,
         help="Max number of files that should be changed. No performance improvement!",
     )
+    parser.add_argument(
+        '--concurrent-files', metavar="NUM", type=int, default=100,
+        help="Number of files to process concurrently during initial load. (default: %(default)s)"
+    )
     parser.add_argument(
         '--full-reorder',
         action='store_true',

python_typing_update/const.py

@@ -1,7 +1,10 @@
 # ---------------------------------------------------------------------------
 # Licensed under the MIT License. See LICENSE file for license information.
 # ---------------------------------------------------------------------------
+from __future__ import annotations
+
 from enum import Flag, auto
+from typing import NamedTuple
 
 version = (0, 2, 0)
 dev_version = 1
@@ -11,6 +14,11 @@
     version_str += f'-dev{dev_version}'
 
 
+class FileAttributes(NamedTuple):
+    status: FileStatus
+    imports: set[str]
+
+
 class FileStatus(Flag):
     CLEAR = 0
     COMMENT = auto()

python_typing_update/main.py

@@ -6,19 +6,22 @@
 import argparse
 import asyncio
 import builtins
+import io
 import logging
+from collections.abc import Iterable
 from functools import partial
 from io import StringIO
 
+import aiofiles
 import reorder_python_imports
 from autoflake import _main as autoflake_main
 from isort.main import main as isort_main
 from pyupgrade._main import main as pyupgrade_main
 
-from .const import FileStatus
+from .const import FileAttributes, FileStatus
 from .utils import (
     async_check_uncommitted_changes, async_restore_files,
-    check_comment_between_imports, check_files_exist)
+    check_comment_between_imports, check_files_exist, extract_imports)
 
 logger = logging.getLogger("typing-update")
 
@@ -107,6 +110,32 @@ async def typing_update(
     return 0, filename
 
 
+async def async_load_files(
+    args: argparse.Namespace,
+    filenames: Iterable[str], *,
+    check_comments: bool,
+) -> dict[str, FileAttributes]:
+    """Process files from file list."""
+    active_tasks: int = 0
+
+    async def async_load_file(filename: str) -> tuple[str, FileAttributes]:
+        """Load file into memory and perform token analysis."""
+        nonlocal active_tasks
+        while active_tasks > args.concurrent_files:
+            await asyncio.sleep(0)
+        active_tasks += 1
+        async with aiofiles.open(filename, encoding="utf-8") as fp:
+            data = await fp.read()
+        file_status = check_comment_between_imports(io.StringIO(data)) \
+            if check_comments is True else FileStatus.CLEAR
+        imports_set = extract_imports(io.StringIO(data))
+        active_tasks -= 1
+        return filename, FileAttributes(file_status, imports_set)
+
+    results = await asyncio.gather(*[async_load_file(file_) for file_ in filenames])
+    return dict(results)
+
+
 async def async_run(args: argparse.Namespace) -> int:
     """Update Python typing syntax.
 
@@ -129,15 +158,11 @@ async def async_run(args: argparse.Namespace) -> int:
         print("Abort! Commit all changes to '.py' files before running again.")
         return 11
 
-    filenames: dict[str, FileStatus] = {}
-    for filename in args.filenames:
-        with open(filename) as fp:
-            result = check_comment_between_imports(fp)
-            filenames[filename] = result
+    filenames: dict[str, FileAttributes] = await async_load_files(args, args.filenames, check_comments=True)
 
     if args.only_force:
-        filenames = {filename: file_status for filename, file_status in filenames.items()
-                     if file_status != FileStatus.CLEAR}
+        filenames = {filename: attrs for filename, attrs in filenames.items()
+                     if attrs.status != FileStatus.CLEAR}
 
     loop = asyncio.get_running_loop()
     files_updated: list[str] = []
@@ -148,7 +173,7 @@ async def async_run(args: argparse.Namespace) -> int:
     builtins.print = lambda *args, **kwargs: None
 
     return_values = await asyncio.gather(
-        *[typing_update(loop, filename, args, file_status) for filename, file_status in filenames.items()])
+        *[typing_update(loop, filename, args, attrs.status) for filename, attrs in filenames.items()])
     for status, filename in return_values:
         if status == 0:
             files_updated.append(filename)
@@ -176,35 +201,53 @@ async def async_run(args: argparse.Namespace) -> int:
 
     files_updated_set: set[str] = set(files_updated)
     files_with_comments = sorted([
-        filename for filename, file_status in filenames.items()
-        if FileStatus.COMMENT in file_status and filename in files_updated_set
+        filename for filename, attrs in filenames.items()
+        if FileStatus.COMMENT in attrs.status and filename in files_updated_set
     ])
-    if files_with_comments:
+    files_imports_changed: list[str] = []
+    for file_, attrs in (await async_load_files(args, files_updated_set, check_comments=False)).items():
+        import_diff = filenames[file_].imports.difference(attrs.imports)
+        for import_ in import_diff:
+            if not import_.startswith('typing'):
+                files_imports_changed.append(file_)
+                break
+    files_imports_changed = sorted(files_imports_changed)
+    files_no_automatic_update = set(files_with_comments + files_imports_changed)
+
+    if files_no_automatic_update:
         if args.force or args.only_force:
             print("Force mode selected!")
             print("Make sure to double check:")
             for file_ in files_with_comments:
                 print(f" - {file_}")
+            if files_with_comments and files_imports_changed:
+                print(" --")
+            for file_ in files_imports_changed:
+                print(f" - {file_}")
         else:
             print("Could not update all files, check:")
             for file_ in files_with_comments:
                 print(f" - {file_}")
-            await async_restore_files(files_with_comments)
+            if files_with_comments and files_imports_changed:
+                print(" --")
+            for file_ in files_imports_changed:
+                print(f" - {file_}")
+            await async_restore_files(files_no_automatic_update)
 
     print("---")
     print(f"All files: {len(filenames)}")
     print(f"No changes: {len(files_no_changes)}")
-    print(f"Files updated: {len(files_updated) - len(files_with_comments)}")
-    print(f"Files (no automatic update): {len(files_with_comments)}")
+    print(f"Files updated: {len(files_updated) - len(files_no_automatic_update)}")
+    print(f"Files (no automatic update): {len(files_no_automatic_update)}")
 
     if (
-        not files_with_comments
+        not files_no_automatic_update
         and not args.force
         and not args.only_force
         and args.verbose == 0
     ):
         return 0
-    if files_with_comments:
+    if files_no_automatic_update:
         return 2
     if args.verbose > 0:
         return 12

python_typing_update/utils.py

@@ -113,3 +113,58 @@ def check_comment_between_imports(fp: TextIO) -> FileStatus:
             # Report all comments in the main import block
             return FileStatus.COMMENT
     return FileStatus.CLEAR
+
+
+def extract_imports(fp: TextIO) -> set[str]:
+    """Create set of all imports in main import block."""
+    flag_in_import_block: bool = False
+    flag_relative_import: bool | None = None
+    flag_imports: bool = False
+    flag_last_token_name: bool = False
+    current_package: str = ''
+    imports: set[str] = set()
+
+    tokens = tokenize.generate_tokens(fp.readline)
+    while True:
+        try:
+            t = next(tokens)
+            if flag_in_import_block is True:
+                if t.type == token.NEWLINE:
+                    if flag_relative_import is False:
+                        imports.add(current_package)
+                    flag_in_import_block = False
+                    flag_relative_import = None
+                    flag_imports = False
+                    flag_last_token_name = False
+                    current_package = ''
+                elif t.type == token.NAME and t.string == 'import':
+                    flag_imports = True
+                elif t.type == token.NAME and flag_last_token_name is False:
+                    if (
+                        flag_relative_import is False
+                        or flag_relative_import is True
+                        and flag_imports is False
+                    ):
+                        current_package += t.string
+                    elif flag_relative_import is True and flag_imports is True:
+                        imports.add(f"{current_package}.{t.string}")
+                elif t.type == token.OP and t.string == '.':
+                    current_package += '.'
+                elif t.type == token.OP and t.string == ',' and flag_relative_import is False:
+                    imports.add(current_package)
+                    current_package = ''
+
+                flag_last_token_name = (t.type == token.NAME and t.string != 'import')
+                continue
+            if t.type == token.NAME:
+                if t.string == 'import':
+                    flag_in_import_block = True
+                    flag_relative_import = False
+                elif t.string == 'from':
+                    flag_in_import_block = True
+                    flag_relative_import = True
+                else:
+                    break
+        except StopIteration:
+            break
+    return imports

requirements.txt

@@ -1,3 +1,4 @@
+aiofiles==0.6.0
 autoflake==1.4
 black==20.8b1
 isort==5.7.0

requirements_test.txt

@@ -1,5 +1,4 @@
 -r requirements_test_pre_commit.txt
-aiofiles==0.6.0
 mypy==0.812
 pre-commit==2.11.1
 pylint==2.7.2

tests/fixtures/unused_import_1.py

@@ -0,0 +1,7 @@
+"""Test unused import retention."""
+from typing import Any, List
+
+import const
+
+var1: List[str]
+var2: Any

tests/fixtures/unused_import_1_fixed.py

@@ -0,0 +1,9 @@
+"""Test unused import retention."""
+from __future__ import annotations
+
+from typing import Any
+
+import const
+
+var1: list[str]
+var2: Any

tests/fixtures/unused_import_2.py

@@ -0,0 +1,7 @@
+"""Test unused import retention."""
+from typing import Any, List
+
+from const import xyz
+
+var1: List[str]
+var2: Any