Simplify the not-context-manager check, by doing a simple parent lookup on a Generator.

3 files changed, 36 insertions, 9 deletions

diff --git a/pylint/checkers/typecheck.py b/pylint/checkers/typecheck.py
index c86f8b9..58ce4be 100644
--- a/pylint/checkers/typecheck.py
+++ b/pylint/checkers/typecheck.py
@@ -634,16 +634,10 @@ accessed. Python regular expressions are accepted.'}
             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):
-                if (isinstance(ctx_mgr, astroid.Name) and
-                        infered.parent and
-                        isinstance(infered.parent, astroid.Function)):
-                    func = infered.parent
-                else:
-                    func = safe_infer(ctx_mgr.func)
-
-                if func is None and func is astroid.YES:
-                    continue
+                func = infered.parent
                 if not decorated_with(func, ['contextlib.contextmanager']):
                     self.add_message('not-context-manager',
                                      node=node, args=(infered.name, ))
diff --git a/pylint/test/functional/not_context_manager.py b/pylint/test/functional/not_context_manager.py
index 5c0a587..f63b5a2 100644
--- a/pylint/test/functional/not_context_manager.py
+++ b/pylint/test/functional/not_context_manager.py
@@ -57,3 +57,35 @@ with tropa: # [not-context-manager]
 
 with hopa:
     pass
+
+
+# Tests that no messages are emitted for function calls
+# which return managers
+
+def wrapper():
+    return dec()
+
+with wrapper():
+    pass
+
+# Tests for properties returning managers.
+
+class Property(object):
+
+    @property
+    def ctx(self):
+        return dec()
+
+    @property
+    def not_ctx(self):
+        return 42
+
+
+lala = Property()
+with lala.ctx:
+    # Don't emit when the context manager is the
+    # result of accessing a property.
+    pass
+
+with lala.not_ctx: # [not-context-manager]
+    pass
diff --git a/pylint/test/functional/not_context_manager.txt b/pylint/test/functional/not_context_manager.txt
index a43387d..4fd8ce8 100644
--- a/pylint/test/functional/not_context_manager.txt
+++ b/pylint/test/functional/not_context_manager.txt
@@ -1,3 +1,4 @@
 not-context-manager:23::Context manager 'NotAManager' doesn't implement __enter__ and __exit__.
 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__.