Fix a not-context-manager false positive.

If the parent of the generator is not the context manager itself,
that means that it could have been returned from another
function which was the real context manager.
The approach we took is more of a hack rather than a real
solution: use an inference context when inferring the context manager
value in order to store all the inferred statements. Walk all the
inferred statements for the given *ctx_mgr* and if we find one
function scope which is decorated, consider it to be the real
manager and give up, otherwise emit not-context-manager.
The problem is that we can't say for sure what the context
manager actually is, because the inference will just return
the final result (which in the case of using contextlib.contextmanager,
it can be a generator). See the test file for not_context_manager for a couple
of self explaining tests. Closes issue #577.

3 files changed, 62 insertions, 5 deletions

diff --git a/pylint/checkers/typecheck.py b/pylint/checkers/typecheck.py
index d56e66b..12c77c6 100644
--- a/pylint/checkers/typecheck.py
+++ b/pylint/checkers/typecheck.py
@@ -16,6 +16,7 @@
 """try to find more bugs in the code using astroid inference capabilities
 """
 
+import collections
 import re
 import shlex
 import sys
@@ -28,6 +29,7 @@ from astroid import (
 from astroid.bases import BUILTINS
 from astroid import objects
 from astroid import helpers
+import six
 
 from pylint.interfaces import IAstroidChecker, INFERENCE, INFERENCE_FAILURE
 from pylint.checkers import BaseChecker
@@ -35,6 +37,16 @@ from pylint.checkers.utils import (
     is_super, check_messages, decorated_with_property,
     decorated_with, node_ignores_exception)
 
+def _unflatten(iterable):
+    for elem in iterable:
+        elem = elem[0]
+        if isinstance(elem, collections.Sequence):
+            for subelem in _unflatten(elem):
+                yield subelem
+        else:
+            yield elem
+
+
 MSGS = {
     'E1101': ('%s %r has no %r member',
               'no-member',
@@ -687,18 +699,37 @@ accessed. Python regular expressions are accepted.'}
             # Anything else is an error
             self.add_message('invalid-slice-index', node=node)
 
+
     @check_messages('not-context-manager')
     def visit_with(self, node):
         for ctx_mgr, _ in node.items:
-            infered = helpers.safe_infer(ctx_mgr)
+            context = astroid.bases.InferenceContext()
+            infered = helpers.safe_infer(ctx_mgr, context=context)
             if infered is None or infered is astroid.YES:
                 continue
 
-            # Check if we are dealing with a function decorated
-            # with contextlib.contextmanager.
             if isinstance(infered, astroid.bases.Generator):
-                func = infered.parent
-                if not decorated_with(func, ['contextlib.contextmanager']):
+                # Check if we are dealing with a function decorated
+                # with contextlib.contextmanager.
+                if decorated_with(infered.parent, ['contextlib.contextmanager']):
+                    continue
+                # If the parent of the generator is not the context manager itself,
+                # that means that it could have been returned from another
+                # function which was the real context manager.
+                # The following approach is more of a hack rather than a real
+                # solution: walk all the inferred statements for the
+                # given *ctx_mgr* and if you find one function scope
+                # which is decorated, consider it to be the real
+                # manager and give up, otherwise emit not-context-manager.
+                # See the test file for not_context_manager for a couple
+                # of self explaining tests.
+                for path in six.moves.filter(None, _unflatten(context.path)):
+                    scope = path.scope()
+                    if not isinstance(scope, astroid.Function):
+                        continue 
+                    if decorated_with(scope, ['contextlib.contextmanager']):
+                        break
+                else:
                     self.add_message('not-context-manager',
                                      node=node, args=(infered.name, ))
             else:
diff --git a/pylint/test/functional/not_context_manager.py b/pylint/test/functional/not_context_manager.py
index 00af36c..97f6e97 100644
--- a/pylint/test/functional/not_context_manager.py
+++ b/pylint/test/functional/not_context_manager.py
@@ -108,3 +108,28 @@ class FullContextManager(ManagerMixin):
         return self
     def __exit__(self, *args):
         pass
+
+# Test a false positive with returning a generator
+# from a context manager.
+def generator():
+    yield 42
+
+@contextmanager
+def context_manager_returning_generator():
+    return generator()
+
+with context_manager_returning_generator():
+    pass
+
+FIRST = [context_manager_returning_generator()]
+with FIRST[0]:
+    pass
+
+def other_indirect_func():
+    return generator()
+
+def not_context_manager():
+    return other_indirect_func()
+
+with not_context_manager(): # [not-context-manager]
+    pass
diff --git a/pylint/test/functional/not_context_manager.txt b/pylint/test/functional/not_context_manager.txt
index 4fd8ce8..7fad3af 100644
--- a/pylint/test/functional/not_context_manager.txt
+++ b/pylint/test/functional/not_context_manager.txt
@@ -2,3 +2,4 @@ not-context-manager:23::Context manager 'NotAManager' doesn't implement __enter_
 not-context-manager:37::Context manager 'dec' doesn't implement __enter__ and __exit__.
 not-context-manager:55::Context manager 'int' doesn't implement __enter__ and __exit__.
 not-context-manager:90::Context manager 'int' doesn't implement __enter__ and __exit__.
+not-context-manager:134::Context manager 'generator' doesn't implement __enter__ and __exit__.