@@ -839,6 +839,52 @@ def imports_on_separate_lines(logical_line):
839839 yield found , "E401 multiple imports on one line"
840840
841841
842+ def imports_on_top_of_file (logical_line , indent_level , checker_state , noqa ):
843+ r"""Imports are always put at the top of the file, just after any module
844+ comments and docstrings, and before module globals and constants.
845+
846+ Okay: import os
847+ Okay: # this is a comment\nimport os
848+ Okay: '''this is a module docstring'''\nimport os
849+ Okay: r'''this is a module docstring'''\nimport os
850+ Okay: __version__ = "123"\nimport os
851+ E402: a=1\nimport os
852+ E402: 'One string'\n"Two string"\nimport os
853+ E402: a=1\nfrom sys import x
854+
855+ Okay: if x:\n import os
856+ """
857+ def is_string_literal (line ):
858+ if line [0 ] in 'uUbB' :
859+ line = line [1 :]
860+ if line and line [0 ] in 'rR' :
861+ line = line [1 :]
862+ return line and (line [0 ] == '"' or line [0 ] == "'" )
863+
864+ if indent_level : # Allow imports in conditional statements or functions
865+ return
866+ if not logical_line : # Allow empty lines or comments
867+ return
868+ if noqa :
869+ return
870+ line = logical_line
871+ if line .startswith ('import ' ) or line .startswith ('from ' ):
872+ if checker_state .get ('seen_non_imports' , False ):
873+ yield 0 , "E402 import not at top of file"
874+ elif line .startswith ('__version__ ' ):
875+ # These lines should be included after the module's docstring, before
876+ # any other code, separated by a blank line above and below.
877+ return
878+ elif is_string_literal (line ):
879+ # The first literal is a docstring, allow it. Otherwise, report error.
880+ if checker_state .get ('seen_docstring' , False ):
881+ checker_state ['seen_non_imports' ] = True
882+ else :
883+ checker_state ['seen_docstring' ] = True
884+ else :
885+ checker_state ['seen_non_imports' ] = True
886+
887+
842888def compound_statements (logical_line ):
843889 r"""Compound statements (on the same line) are generally discouraged.
844890
@@ -1251,6 +1297,8 @@ def __init__(self, filename=None, lines=None,
12511297 self .hang_closing = options .hang_closing
12521298 self .verbose = options .verbose
12531299 self .filename = filename
1300+ # Dictionary where a checker can store its custom state.
1301+ self ._checker_states = {}
12541302 if filename is None :
12551303 self .filename = 'stdin'
12561304 self .lines = lines or []
@@ -1306,10 +1354,16 @@ def run_check(self, check, argument_names):
13061354 arguments .append (getattr (self , name ))
13071355 return check (* arguments )
13081356
1357+ def init_checker_state (self , name , argument_names ):
1358+ """ Prepares a custom state for the specific checker plugin."""
1359+ if 'checker_state' in argument_names :
1360+ self .checker_state = self ._checker_states .setdefault (name , {})
1361+
13091362 def check_physical (self , line ):
13101363 """Run all physical checks on a raw input line."""
13111364 self .physical_line = line
13121365 for name , check , argument_names in self ._physical_checks :
1366+ self .init_checker_state (name , argument_names )
13131367 result = self .run_check (check , argument_names )
13141368 if result is not None :
13151369 (offset , text ) = result
@@ -1368,6 +1422,7 @@ def check_logical(self):
13681422 for name , check , argument_names in self ._logical_checks :
13691423 if self .verbose >= 4 :
13701424 print (' ' + name )
1425+ self .init_checker_state (name , argument_names )
13711426 for offset , text in self .run_check (check , argument_names ) or ():
13721427 if not isinstance (offset , tuple ):
13731428 for token_offset , pos in mapping :
0 commit comments