Fixed doctest cases and improved binary helper function

2 files changed, 24 insertions, 18 deletions

diff --git a/pycodestyle.py b/pycodestyle.py
index 355a947..f721fb0 100755
--- a/pycodestyle.py
+++ b/pycodestyle.py
@@ -1064,7 +1064,12 @@ def _break_around_binary_operators(tokens):
         if ('\n' in text or '\r' in text) and token_type != tokenize.STRING:
             line_break = True
         else:
-            yield (token_type, text, previous_token_type, previous_text,
+            prev_is_binary = _is_binary_operator(
+                previous_token_type, previous_text
+            )
+            curr_is_binary = _is_binary_operator(token_type, text)
+
+            yield (curr_is_binary, prev_is_binary,
                    line_break, unary_context, start)
             unary_context = text in '([{,;'
             line_break = False
@@ -1082,23 +1087,23 @@ def break_before_binary_operator(logical_line, tokens):
     W503: (width == 0\n + height == 0)
     W503: (width == 0\n and height == 0)
 
-    Okay: (width == 0 +\n height == 0)
+    W504: (width == 0 +\n height == 0)
+    W504: var = (1 &\n       ~2)
     Okay: foo(\n    -x)
     Okay: foo(x\n    [])
     Okay: x = '''\n''' + ''
     Okay: foo(x,\n    -y)
     Okay: foo(x,  # comment\n    -y)
-    Okay: var = (1 &\n       ~2)
     Okay: var = (1 /\n       -2)
     Okay: var = (1 +\n       -1 +\n       -2)
     """
     for context in _break_around_binary_operators(tokens):
-        (token_type, text, previous_token_type, previous_text,
+        (curr_is_binary, prev_is_binary,
          line_break, unary_context, start) = context
-        if (_is_binary_operator(token_type, text) and line_break and
-                not unary_context and
-                not _is_binary_operator(previous_token_type,
-                                        previous_text)):
+
+        binary_break = line_break and curr_is_binary and not prev_is_binary
+
+        if (binary_break and not unary_context):
             yield start, "W503 line break before binary operator"
 
 
@@ -1112,24 +1117,25 @@ def break_after_binary_operator(logical_line, tokens):
     W504: (width == 0 +\n height == 0)
     W504: (width == 0 and\n height == 0)
 
-    Okay: (width == 0\n + height == 0)
+    W503: (width == 0\n + height == 0)
+    W503: var = (1\n       & ~2)
+    W503: var = (1\n       / -2)
+    W503: var = (1\n       + -1\n       + -2)
+
     Okay: foo(\n    -x)
     Okay: foo(x\n    [])
     Okay: x = '''\n''' + ''
     Okay: x = '' + '''\n'''
     Okay: foo(x,\n    -y)
     Okay: foo(x,  # comment\n    -y)
-    Okay: var = (1\n       & ~2)
-    Okay: var = (1\n       / -2)
-    Okay: var = (1\n       + -1\n       + -2)
     """
     for context in _break_around_binary_operators(tokens):
-        (token_type, text, previous_token_type, previous_text,
+        (curr_is_binary, prev_is_binary,
          line_break, unary_context, start) = context
-        if (_is_binary_operator(previous_token_type, previous_text)
-                and line_break
-                and not unary_context
-                and not _is_binary_operator(token_type, text)):
+
+        binary_break = line_break and prev_is_binary and not curr_is_binary
+
+        if (binary_break and not unary_context):
             error_pos = (start[0] - 1, start[1])
             yield error_pos, "W504 line break after binary operator"
 
diff --git a/setup.cfg b/setup.cfg
index 803bc10..15b3968 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -3,5 +3,5 @@ universal = 1
 
 [pycodestyle]
 select =
-ignore = E226,E24
+ignore = E226,E24,W504
 max_line_length = 79