[contextmanager-generator-missing-cleanup] Warn about context manager without try/finally in generator functions (#9133)

Author: rhyn0

doc/data/messages/c/contextmanager-generator-missing-cleanup/bad.py

@@ -0,0 +1,14 @@
+import contextlib
+
+
+@contextlib.contextmanager
+def cm():
+    contextvar = "acquired context"
+    print("cm enter")
+    yield contextvar
+    print("cm exit")
+
+
+def genfunc_with_cm():  # [contextmanager-generator-missing-cleanup]
+    with cm() as context:
+        yield context * 2

doc/data/messages/c/contextmanager-generator-missing-cleanup/details.rst

@@ -0,0 +1,10 @@
+Instantiating and using a contextmanager inside a generator function can
+result in unexpected behavior if there is an expectation that the context is only
+available for the generator function. In the case that the generator is not closed or destroyed
+then the context manager is held suspended as is.
+
+This message warns on the generator function instead of the contextmanager function
+because the ways to use a contextmanager are many.
+A contextmanager can be used as a decorator (which immediately has ``__enter__``/``__exit__`` applied)
+and the use of ``as ...`` or discard of the return value also implies whether the context needs cleanup or not.
+So for this message, warning the invoker of the contextmanager is important.

doc/data/messages/c/contextmanager-generator-missing-cleanup/good.py

@@ -0,0 +1,49 @@
+import contextlib
+
+
+@contextlib.contextmanager
+def good_cm_except():
+    contextvar = "acquired context"
+    print("good cm enter")
+    try:
+        yield contextvar
+    except GeneratorExit:
+        print("good cm exit")
+
+
+def genfunc_with_cm():
+    with good_cm_except() as context:
+        yield context * 2
+
+
+def genfunc_with_discard():
+    with good_cm_except():
+        yield "discarded"
+
+
+@contextlib.contextmanager
+def good_cm_yield_none():
+    print("good cm enter")
+    yield
+    print("good cm exit")
+
+
+def genfunc_with_none_yield():
+    with good_cm_yield_none() as var:
+        print(var)
+        yield "constant yield"
+
+
+@contextlib.contextmanager
+def good_cm_finally():
+    contextvar = "acquired context"
+    print("good cm enter")
+    try:
+        yield contextvar
+    finally:
+        print("good cm exit")
+
+
+def good_cm_finally_genfunc():
+    with good_cm_finally() as context:
+        yield context * 2

doc/data/messages/c/contextmanager-generator-missing-cleanup/related.rst

@@ -0,0 +1,2 @@
+- `Rationale <https://discuss.python.org/t/preventing-yield-inside-certain-context-managers/1091>`_
+- `CPython Issue <https://github.com/python/cpython/issues/81924#issuecomment-1093830682>`_

doc/user_guide/checkers/features.rst

@@ -166,6 +166,9 @@ Basic checker Messages
   This is a particular case of W0104 with its own message so you can easily
   disable it if you're using those strings as documentation, instead of
   comments.
+:contextmanager-generator-missing-cleanup (W0135): *The context used in function %r will not be exited.*
+  Used when a contextmanager is used inside a generator function and the
+  cleanup is not handled.
 :unnecessary-pass (W0107): *Unnecessary pass statement*
   Used when a "pass" statement can be removed without affecting the behaviour
   of the code.

doc/user_guide/messages/messages_overview.rst

@@ -228,6 +228,7 @@ All messages in the warning category:
    warning/comparison-with-callable
    warning/confusing-with-statement
    warning/consider-ternary-expression
+   warning/contextmanager-generator-missing-cleanup
    warning/dangerous-default-value
    warning/deprecated-argument
    warning/deprecated-attribute

doc/whatsnew/fragments/2832.new_check

@@ -0,0 +1,6 @@
+Checks for generators that use contextmanagers that don't handle cleanup properly.
+Is meant to raise visibilty on the case that a generator is not fully exhausted and the contextmanager is not cleaned up properly.
+A contextmanager must yield a non-constant value and not handle cleanup for GeneratorExit.
+The using generator must attempt to use the yielded context value `with x() as y` and not just `with x()`.
+
+Closes #2832

pylint/checkers/base/__init__.py

@@ -23,6 +23,7 @@
 from pylint.checkers.base.basic_error_checker import BasicErrorChecker
 from pylint.checkers.base.comparison_checker import ComparisonChecker
 from pylint.checkers.base.docstring_checker import DocStringChecker
+from pylint.checkers.base.function_checker import FunctionChecker
 from pylint.checkers.base.name_checker import (
     KNOWN_NAME_TYPES_WITH_STYLE,
     AnyStyle,
@@ -46,3 +47,4 @@ def register(linter: PyLinter) -> None:
     linter.register_checker(DocStringChecker(linter))
     linter.register_checker(PassChecker(linter))
     linter.register_checker(ComparisonChecker(linter))
+    linter.register_checker(FunctionChecker(linter))

pylint/checkers/base/function_checker.py

@@ -0,0 +1,135 @@
+# Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
+# For details: https://github.com/pylint-dev/pylint/blob/main/LICENSE
+# Copyright (c) https://github.com/pylint-dev/pylint/blob/main/CONTRIBUTORS.txt
+
+"""Function checker for Python code."""
+
+from __future__ import annotations
+
+from itertools import chain
+
+from astroid import nodes
+
+from pylint.checkers import utils
+from pylint.checkers.base.basic_checker import _BasicChecker
+
+
+class FunctionChecker(_BasicChecker):
+    """Check if a function definition handles possible side effects."""
+
+    msgs = {
+        "W0135": (
+            "The context used in function %r will not be exited.",
+            "contextmanager-generator-missing-cleanup",
+            "Used when a contextmanager is used inside a generator function"
+            " and the cleanup is not handled.",
+        )
+    }
+
+    @utils.only_required_for_messages("contextmanager-generator-missing-cleanup")
+    def visit_functiondef(self, node: nodes.FunctionDef) -> None:
+        self._check_contextmanager_generator_missing_cleanup(node)
+
+    @utils.only_required_for_messages("contextmanager-generator-missing-cleanup")
+    def visit_asyncfunctiondef(self, node: nodes.AsyncFunctionDef) -> None:
+        self._check_contextmanager_generator_missing_cleanup(node)
+
+    def _check_contextmanager_generator_missing_cleanup(
+        self, node: nodes.FunctionDef
+    ) -> None:
+        """Check a FunctionDef to find if it is a generator
+        that uses a contextmanager internally.
+
+        If it is, check if the contextmanager is properly cleaned up. Otherwise, add message.
+
+        :param node: FunctionDef node to check
+        :type node: nodes.FunctionDef
+        """
+        # if function does not use a Yield statement, it cant be a generator
+        with_nodes = list(node.nodes_of_class(nodes.With))
+        if not with_nodes:
+            return
+        # check for Yield inside the With statement
+        yield_nodes = list(
+            chain.from_iterable(
+                with_node.nodes_of_class(nodes.Yield) for with_node in with_nodes
+            )
+        )
+        if not yield_nodes:
+            return
+
+        # infer the call that yields a value, and check if it is a contextmanager
+        for with_node in with_nodes:
+            for call, held in with_node.items:
+                if held is None:
+                    # if we discard the value, then we can skip checking it
+                    continue
+
+                # safe infer is a generator
+                inferred_node = getattr(utils.safe_infer(call), "parent", None)
+                if not isinstance(inferred_node, nodes.FunctionDef):
+                    continue
+                if self._node_fails_contextmanager_cleanup(inferred_node, yield_nodes):
+                    self.add_message(
+                        "contextmanager-generator-missing-cleanup",
+                        node=node,
+                        args=(node.name,),
+                    )
+
+    @staticmethod
+    def _node_fails_contextmanager_cleanup(
+        node: nodes.FunctionDef, yield_nodes: list[nodes.Yield]
+    ) -> bool:
+        """Check if a node fails contextmanager cleanup.
+
+        Current checks for a contextmanager:
+            - only if the context manager yields a non-constant value
+            - only if the context manager lacks a finally, or does not catch GeneratorExit
+
+        :param node: Node to check
+        :type node: nodes.FunctionDef
+        :return: True if fails, False otherwise
+        :param yield_nodes: List of Yield nodes in the function body
+        :type yield_nodes: list[nodes.Yield]
+        :rtype: bool
+        """
+
+        def check_handles_generator_exceptions(try_node: nodes.Try) -> bool:
+            # needs to handle either GeneratorExit, Exception, or bare except
+            for handler in try_node.handlers:
+                if handler.type is None:
+                    # handles all exceptions (bare except)
+                    return True
+                inferred = utils.safe_infer(handler.type)
+                if inferred and inferred.qname() in {
+                    "builtins.GeneratorExit",
+                    "builtins.Exception",
+                }:
+                    return True
+            return False
+
+        # if context manager yields a non-constant value, then continue checking
+        if any(
+            yield_node.value is None or isinstance(yield_node.value, nodes.Const)
+            for yield_node in yield_nodes
+        ):
+            return False
+        # if function body has multiple Try, filter down to the ones that have a yield node
+        try_with_yield_nodes = [
+            try_node
+            for try_node in node.nodes_of_class(nodes.Try)
+            if any(try_node.nodes_of_class(nodes.Yield))
+        ]
+        if not try_with_yield_nodes:
+            # no try blocks at all, so checks after this line do not apply
+            return True
+        # if the contextmanager has a finally block, then it is fine
+        if all(try_node.finalbody for try_node in try_with_yield_nodes):
+            return False
+        # if the contextmanager catches GeneratorExit, then it is fine
+        if all(
+            check_handles_generator_exceptions(try_node)
+            for try_node in try_with_yield_nodes
+        ):
+            return False
+        return True

tests/functional/c/consider/consider_using_with.py

@@ -186,9 +186,7 @@ def test_futures():
     pass
 
 
-global_pool = (
-    multiprocessing.Pool()
-)  # must not trigger, will be used in nested scope
+global_pool = multiprocessing.Pool()  # must not trigger, will be used in nested scope
 
 
 def my_nested_function():