Merge pull request #834 from EricCousineau-TRI/issue/833

E402: Add "with" statement to allowed keywords

2 files changed, 20 insertions, 4 deletions

diff --git a/pycodestyle.py b/pycodestyle.py
index 4ade087..6ad5456 100755
--- a/pycodestyle.py
+++ b/pycodestyle.py
@@ -1073,7 +1073,8 @@ def module_imports_on_top_of_file(
             line = line[1:]
         return line and (line[0] == '"' or line[0] == "'")
 
-    allowed_try_keywords = ('try', 'except', 'else', 'finally')
+    allowed_keywords = (
+        'try', 'except', 'else', 'finally', 'with', 'if', 'elif')
 
     if indent_level:  # Allow imports in conditional statement/function
         return
@@ -1087,9 +1088,9 @@ def module_imports_on_top_of_file(
             yield 0, "E402 module level import not at top of file"
     elif re.match(DUNDER_REGEX, line):
         return
-    elif any(line.startswith(kw) for kw in allowed_try_keywords):
-        # Allow try, except, else, finally keywords intermixed with
-        # imports in order to support conditional importing
+    elif any(line.startswith(kw) for kw in allowed_keywords):
+        # Allow certain keywords intermixed with imports in order to
+        # support conditional or filtered importing
         return
     elif is_string_literal(line):
         # The first literal is a docstring, allow it. Otherwise, report
diff --git a/testsuite/E40.py b/testsuite/E40.py
index f9a18fc..6c71fa2 100644
--- a/testsuite/E40.py
+++ b/testsuite/E40.py
@@ -34,6 +34,21 @@ finally:
     print('made attempt to import foo')
 
 import bar
+#: Okay
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", DeprecationWarning)
+    import foo
+
+import bar
+#: Okay
+if False:
+    import foo
+elif not True:
+    import bar
+else:
+    import mwahaha
+
+import bar
 #: E402
 VERSION = '1.2.3'