Add support for quick identification of just the top-level imports, before functions and classes.

2 files changed, 10 insertions, 4 deletions

diff --git a/isort/identify.py b/isort/identify.py
index 156d866a..46e3df5f 100644
--- a/isort/identify.py
+++ b/isort/identify.py
@@ -3,13 +3,15 @@ Eventually this will likely replace parse.py
 """
 from functools import partial
 from pathlib import Path
-from typing import Iterator, NamedTuple, Optional, TextIO
+from typing import Iterator, NamedTuple, Optional, TextIO, Tuple
 
 from isort.parse import _normalize_line, _strip_syntax, skip_line
 
 from .comments import parse as parse_comments
 from .settings import DEFAULT_CONFIG, Config
 
+STATEMENT_DECLARATIONS: Tuple[str, ...] = ("def ", "cdef ", "cpdef ", "class ", "@", "async def")
+
 
 class Import(NamedTuple):
     line_number: int
@@ -36,7 +38,10 @@ class Import(NamedTuple):
 
 
 def imports(
-    input_stream: TextIO, config: Config = DEFAULT_CONFIG, file_path: Optional[Path] = None
+    input_stream: TextIO,
+    config: Config = DEFAULT_CONFIG,
+    file_path: Optional[Path] = None,
+    top_only: bool = False,
 ) -> Iterator[Import]:
     """Parses a python file taking out and categorizing imports."""
     in_quote = ""
@@ -48,6 +53,8 @@ def imports(
         )
 
         if skipping_line:
+            if top_only and not in_quote and line.startswith(STATEMENT_DECLARATIONS):
+                break
             continue
 
         line, *end_of_line_comment = line.split("#", 1)
diff --git a/isort/output.py b/isort/output.py
index 1d64785b..e0855de6 100644
--- a/isort/output.py
+++ b/isort/output.py
@@ -7,10 +7,9 @@ from isort.format import format_simplified
 
 from . import parse, sorting, wrap
 from .comments import add_to_line as with_comments
+from .identify import STATEMENT_DECLARATIONS
 from .settings import DEFAULT_CONFIG, Config
 
-STATEMENT_DECLARATIONS: Tuple[str, ...] = ("def ", "cdef ", "cpdef ", "class ", "@", "async def")
-
 
 def sorted_imports(
     parsed: parse.ParsedContent,