Fix argument check for variable arguments.

2 files changed, 13 insertions, 1 deletions

diff --git a/checkers/classes.py b/checkers/classes.py
index 7ced271..fd76146 100644
--- a/checkers/classes.py
+++ b/checkers/classes.py
@@ -353,7 +353,10 @@ a metaclass class method.'}
 
     def _check_exit(self, node):
         positional = sum(1 for arg in node.args.args if arg.name != 'self')
-        if positional != 3 and not node.args.vararg:
+        if positional < 3 and not node.args.vararg:
+            self.add_message('bad-context-manager',
+                             node=node)  
+        elif positional > 3:
             self.add_message('bad-context-manager',
                              node=node)    
 
diff --git a/test/input/func_bad_context_manager.py b/test/input/func_bad_context_manager.py
index 620190d..2439c47 100644
--- a/test/input/func_bad_context_manager.py
+++ b/test/input/func_bad_context_manager.py
@@ -49,5 +49,14 @@ class SecondBadContextManager(object):
     def __exit__(self, exc_type, value, tb, stack):
         pass
 
+class ThirdBadContextManager(object):
+    """ Too many arguments and variable arguments """
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, value, tb, stack, *args):
+        pass
+