Merge pull request #717 from taion/annot-missing-space

Add E252 on missing whitespace for annotated parameter defaults

2 files changed, 23 insertions, 5 deletions

diff --git a/pycodestyle.py b/pycodestyle.py
index d8dfcf4..7131f2a 100755
--- a/pycodestyle.py
+++ b/pycodestyle.py
@@ -869,7 +869,8 @@ def whitespace_around_named_parameter_equals(logical_line, tokens):
     r"""Don't use spaces around the '=' sign in function arguments.
 
     Don't use spaces around the '=' sign when used to indicate a
-    keyword argument or a default parameter value.
+    keyword argument or a default parameter value, except when using a type
+    annotation.
 
     Okay: def complex(real, imag=0.0):
     Okay: return magic(r=real, i=imag)
@@ -882,13 +883,18 @@ def whitespace_around_named_parameter_equals(logical_line, tokens):
 
     E251: def complex(real, imag = 0.0):
     E251: return magic(r = real, i = imag)
+    E252: def complex(real, image: float=0.0):
     """
     parens = 0
     no_space = False
+    require_space = False
     prev_end = None
     annotated_func_arg = False
     in_def = bool(STARTSWITH_DEF_REGEX.match(logical_line))
+
     message = "E251 unexpected spaces around keyword / parameter equals"
+    missing_message = "E252 missing whitespace around parameter equals"
+
     for token_type, text, start, end, line in tokens:
         if token_type == tokenize.NL:
             continue
@@ -896,6 +902,10 @@ def whitespace_around_named_parameter_equals(logical_line, tokens):
             no_space = False
             if start != prev_end:
                 yield (prev_end, message)
+        if require_space:
+            require_space = False
+            if start == prev_end:
+                yield (prev_end, missing_message)
         if token_type == tokenize.OP:
             if text in '([':
                 parens += 1
@@ -905,10 +915,15 @@ def whitespace_around_named_parameter_equals(logical_line, tokens):
                 annotated_func_arg = True
             elif parens and text == ',' and parens == 1:
                 annotated_func_arg = False
-            elif parens and text == '=' and not annotated_func_arg:
-                no_space = True
-                if start != prev_end:
-                    yield (prev_end, message)
+            elif parens and text == '=':
+                if not annotated_func_arg:
+                    no_space = True
+                    if start != prev_end:
+                        yield (prev_end, message)
+                else:
+                    require_space = True
+                    if start == prev_end:
+                        yield (prev_end, missing_message)
             if not parens:
                 annotated_func_arg = False
 
diff --git a/testsuite/E25.py b/testsuite/E25.py
index dde95b8..71d3f80 100644
--- a/testsuite/E25.py
+++ b/testsuite/E25.py
@@ -42,3 +42,6 @@ async def add(a: int = 0, b: int = 0) -> int:
 #: E272:1:6
 async  def add(a: int = 0, b: int = 0) -> int:
     return a + b
+#: E252:1:15 E252:1:16 E252:1:27 E252:1:36
+def add(a: int=0, b: int =0, c: int= 0) -> int:
+    return a + b + c