|
41 | 41 | from collections.abc import Callable, Generator, Iterable, Sequence
|
42 | 42 | from getopt import getopt
|
43 | 43 | from io import BufferedIOBase, BufferedReader, BytesIO
|
44 |
| -from itertools import chain, groupby |
| 44 | +from itertools import chain |
45 | 45 | from typing import (
|
46 | 46 | TYPE_CHECKING,
|
47 | 47 | Any,
|
@@ -598,17 +598,10 @@ def stripped_lines(
|
598 | 598 | if ignore_imports or ignore_signatures:
|
599 | 599 | tree = astroid.parse("".join(lines))
|
600 | 600 | if ignore_imports:
|
601 |
| - node_is_import_by_lineno = ( |
602 |
| - (node.lineno, isinstance(node, (nodes.Import, nodes.ImportFrom))) |
603 |
| - for node in tree.body |
604 |
| - ) |
605 |
| - line_begins_import = { |
606 |
| - lineno: all(is_import for _, is_import in node_is_import_group) |
607 |
| - for lineno, node_is_import_group in groupby( |
608 |
| - node_is_import_by_lineno, key=lambda x: x[0] |
609 |
| - ) |
610 |
| - } |
611 |
| - current_line_is_import = False |
| 601 | + import_lines = {} |
| 602 | + for node in tree.nodes_of_class((nodes.Import, nodes.ImportFrom)): |
| 603 | + for lineno in range(node.lineno, (node.end_lineno or node.lineno) + 1): |
| 604 | + import_lines[lineno] = True |
612 | 605 | if ignore_signatures:
|
613 | 606 |
|
614 | 607 | def _get_functions(
|
@@ -664,9 +657,7 @@ def _get_functions(
|
664 | 657 | docstring = None
|
665 | 658 | line = ""
|
666 | 659 | if ignore_imports:
|
667 |
| - current_line_is_import = line_begins_import.get( |
668 |
| - lineno, current_line_is_import |
669 |
| - ) |
| 660 | + current_line_is_import = import_lines.get(lineno, False) |
670 | 661 | if current_line_is_import:
|
671 | 662 | line = ""
|
672 | 663 | if ignore_comments:
|
|
0 commit comments