@@ -839,6 +839,53 @@ def imports_on_separate_lines(logical_line):
839839 yield found , "E401 multiple imports on one line"
840840
841841
842+ def module_imports_on_top_of_file (
843+ logical_line , indent_level , checker_state , noqa ):
844+ r"""Imports are always put at the top of the file, just after any module
845+ comments and docstrings, and before module globals and constants.
846+
847+ Okay: import os
848+ Okay: # this is a comment\nimport os
849+ Okay: '''this is a module docstring'''\nimport os
850+ Okay: r'''this is a module docstring'''\nimport os
851+ Okay: __version__ = "123"\nimport os
852+ E402: a=1\nimport os
853+ E402: 'One string'\n"Two string"\nimport os
854+ E402: a=1\nfrom sys import x
855+
856+ Okay: if x:\n import os
857+ """
858+ def is_string_literal (line ):
859+ if line [0 ] in 'uUbB' :
860+ line = line [1 :]
861+ if line and line [0 ] in 'rR' :
862+ line = line [1 :]
863+ return line and (line [0 ] == '"' or line [0 ] == "'" )
864+
865+ if indent_level : # Allow imports in conditional statements or functions
866+ return
867+ if not logical_line : # Allow empty lines or comments
868+ return
869+ if noqa :
870+ return
871+ line = logical_line
872+ if line .startswith ('import ' ) or line .startswith ('from ' ):
873+ if checker_state .get ('seen_non_imports' , False ):
874+ yield 0 , "E402 import not at top of file"
875+ elif line .startswith ('__version__ ' ):
876+ # These lines should be included after the module's docstring, before
877+ # any other code, separated by a blank line above and below.
878+ return
879+ elif is_string_literal (line ):
880+ # The first literal is a docstring, allow it. Otherwise, report error.
881+ if checker_state .get ('seen_docstring' , False ):
882+ checker_state ['seen_non_imports' ] = True
883+ else :
884+ checker_state ['seen_docstring' ] = True
885+ else :
886+ checker_state ['seen_non_imports' ] = True
887+
888+
842889def compound_statements (logical_line ):
843890 r"""Compound statements (on the same line) are generally discouraged.
844891
@@ -1251,6 +1298,8 @@ def __init__(self, filename=None, lines=None,
12511298 self .hang_closing = options .hang_closing
12521299 self .verbose = options .verbose
12531300 self .filename = filename
1301+ # Dictionary where a checker can store its custom state.
1302+ self ._checker_states = {}
12541303 if filename is None :
12551304 self .filename = 'stdin'
12561305 self .lines = lines or []
@@ -1306,10 +1355,16 @@ def run_check(self, check, argument_names):
13061355 arguments .append (getattr (self , name ))
13071356 return check (* arguments )
13081357
1358+ def init_checker_state (self , name , argument_names ):
1359+ """ Prepares a custom state for the specific checker plugin."""
1360+ if 'checker_state' in argument_names :
1361+ self .checker_state = self ._checker_states .setdefault (name , {})
1362+
13091363 def check_physical (self , line ):
13101364 """Run all physical checks on a raw input line."""
13111365 self .physical_line = line
13121366 for name , check , argument_names in self ._physical_checks :
1367+ self .init_checker_state (name , argument_names )
13131368 result = self .run_check (check , argument_names )
13141369 if result is not None :
13151370 (offset , text ) = result
@@ -1368,6 +1423,7 @@ def check_logical(self):
13681423 for name , check , argument_names in self ._logical_checks :
13691424 if self .verbose >= 4 :
13701425 print (' ' + name )
1426+ self .init_checker_state (name , argument_names )
13711427 for offset , text in self .run_check (check , argument_names ) or ():
13721428 if not isinstance (offset , tuple ):
13731429 for token_offset , pos in mapping :
0 commit comments