pylint-missing-await.py

# http://pylint.pycqa.org/en/latest/how_tos/custom_checkers.html

import astroid
from pylint.checkers import BaseChecker
from pylint.interfaces import IAstroidChecker


def register(linter):
    linter.register_checker(MissingAwaitChecker(linter))


class MissingAwaitChecker(BaseChecker):
    __implements__ = IAstroidChecker

    name = "missing-await"
    priority = -1
    msgs = {
        'E????': (
            "Called an async function, but forgot to write 'await'.",
            "missing-await",
            "Checks that you remember to use 'await' when calling an async "
            + "function.",
        ),
    }
    options = ()

    def __init__(self, *args, **kwargs):
        super(MissingAwaitChecker, self).__init__(*args, **kwargs)
        self._last_awaited_node = None

    def visit_await(self, node):
        self._last_awaited_node = node

    def visit_call(self, node):
        if node is self._last_awaited_node:
            return

        nodes_to_resolve = [node]
        resolved_nodes = set()
        while nodes_to_resolve:
            node_to_resolve = nodes_to_resolve.pop()
            if node_to_resolve is astroid.Uninferable():
                continue
            elif isinstance(
                    node_to_resolve,
                    (astroid.FunctionDef, astroid.AsyncFunctionDef)
            ):
                resolved_nodes.add(node_to_resolve)
            elif isinstance(node_to_resolve, astroid.Call):
                try:
                    nodes_to_resolve += node_to_resolve.func.inferred()
                except astroid.InferenceError:
                    continue
            elif isinstance(node_to_resolve, astroid.BoundMethod):
                try:
                    nodes_to_resolve += node_to_resolve.inferred()
                except astroid.InferenceError:
                    continue

        for resolved_node in resolved_nodes:
            if isinstance(resolved_node, astroid.AsyncFunctionDef):
                self.add_message(
                    'missing-await', node=node,
                )
                return