pylint_type_inference.py

"""Plugin for pylint to allow gcloud type inference.

Supports:
- Specifying the return type of an instance methods in the case
  that no arguments are passed in.
"""

import importlib

import astroid


BUILDER = astroid.builder.AstroidBuilder()
INFER_CALL_RESULT = 'infer_call_result'
NO_ARGS_RETURN_MAPPING = {
    'gcloud.datastore.entity': {
        'Entity': {
            'key': ('gcloud.datastore.key', 'Key'),
        },
    },
}
METHOD_RETURN_MAPPING = {
    'pkg_resources': {
        'get_distribution': ('pkg_resources', 'Distribution'),
    },
}


def register(linter):
    """Dummy function needed to make this module a valid PyLint plugin."""
    return linter


def get_astroid_class(module, class_name):
    """Gets the astroid.scoped_nodes.Class object for a local class.

    NOTE: This caches results after creation.
    """
    key = (module, class_name)
    if key not in get_astroid_class.cache:
        module_obj = importlib.import_module(module)
        # NOTE: We don't check uniqueness of results here.
        get_astroid_class.cache[key] = BUILDER.module_build(
            module_obj).getattr(class_name)[0]
    return get_astroid_class.cache[key]
get_astroid_class.cache = {}


class FunctionWrapper(astroid.bases.Proxy):
    """Proxy class to allow inferred type of function return value."""

    def __init__(self, function_obj, return_class, only_no_args=True):
        if not isinstance(function_obj,
                          astroid.scoped_nodes.Function):
            raise TypeError(function_obj)
        super(FunctionWrapper, self).__init__(function_obj)
        # Using "private" variables to avoid polluting namespace of
        # proxied object.
        self.__function_obj = function_obj
        self.__return_class = return_class
        self.__only_no_args = only_no_args

    def __getattribute__(self, name):
        # We let astroid.bases.Proxy handle all attributes except
        # for the method we wish to over-ride and define below.
        if name == INFER_CALL_RESULT:
            # NOTE: We have to use object.__getattribute__ to
            #       avoid an infinite loop caused by the invocation
            #       of __getattribute__ when using `.` lookup.
            return object.__getattribute__(self, INFER_CALL_RESULT)
        return super(FunctionWrapper, self).__getattribute__(name)

    def infer_call_result(self, caller, context=None):
        """Replacement method to describe a function return value.

        This return value "description" is a prescribed type inference for
        values returned by our current proxied astroid.scoped_nodes.Function.
        """
        replace_result = True
        if self.__only_no_args:
            callcontext = getattr(context, 'callcontext', None)
            if not (getattr(callcontext, 'args', []) == [] and
                    getattr(callcontext, 'nargs', {}) == {}):
                replace_result = False

        if replace_result:
            return_instance = astroid.bases.Instance(
                proxied=self.__return_class)
            return iter([return_instance, astroid.bases.YES])
        else:
            return self.__function_obj.infer_call_result(
                caller, context=context)


def make_explicit_inference(function_obj, module, class_name,
                            only_no_args=True):
    """Method factory to allow explicit type inference in astroid.

    :type function_obj: :class:`astroid.scoped_nodes.Function`
    :param function_obj: An function with a return type specified.

    :type module: string
    :param module: The module of the return type.

    :type class_name: string
    :param class_name: The class name of the return type.

    :type only_no_args: bool
    :param only_no_args: Value that determines whether the return type will be
                         used in the case of no arguments or in all cases.

    NOTE: We use a factory to avoid re-use of scope in a for loop.
    """
    def explicit_inference(unused_self, unused_context):
        """Method to allow explicit type inference in astroid."""
        # NOTE: unused_self should always be `function_obj`.
        stub_class = get_astroid_class(module, class_name)
        return iter([FunctionWrapper(function_obj,
                                     return_class=stub_class,
                                     only_no_args=only_no_args)])

    return explicit_inference


def transform_class_type_inference(class_node):
    """Class transformer to allow class methods to specify return type.

    Checks for the module of the current class in the mapping
    of methods which have a return type specified when no arguments
    are passed (NO_ARGS_RETURN_MAPPING).

    For each method in the current class, patches the corresponding
    astroid Function object.
    """
    module_mapping = NO_ARGS_RETURN_MAPPING.get(class_node.root().name)
    if module_mapping is None:
        return

    class_mapping = module_mapping.get(class_node.name)
    if class_mapping is None:
        return

    for function_name, new_type_info in class_mapping.iteritems():
        function_matches = class_node.getattr(function_name)
        if len(function_matches) > 1:
            raise ValueError('Function %r has multiple matches.' % (
                function_name,))

        function_obj = function_matches[0]
        if function_obj._explicit_inference is not None:
            raise ValueError('Expected explicit_inference to be '
                             'unset on %r.%r function.' % (
                                 class_node.name, function_name))

        # NOTE: We don't bind this to `function_obj`.
        function_obj._explicit_inference = make_explicit_inference(
            function_obj, new_type_info[0], new_type_info[1])


def transform_fn_type_inference(function_node):
    """Function transformer to allow specifying return type.

    This works on methods owned directly by a module with a return
    type which never changes (i.e. doesn't depend on branching
    within method).
    """
    if not isinstance(function_node.parent, astroid.scoped_nodes.Module):
        return

    module_name = function_node.parent.name
    module_mapping = METHOD_RETURN_MAPPING.get(module_name)
    if module_mapping is None:
        return

    new_type_info = module_mapping.get(function_node.name)
    if new_type_info is None:
        return

    if function_node._explicit_inference is not None:
        raise ValueError('Expected explicit_inference to be '
                         'unset on %r.%r function.' % (
                             module_name, function_node.name))

    # NOTE: We don't bind this to `function_node`.
    function_node._explicit_inference = make_explicit_inference(
        function_node, new_type_info[0], new_type_info[1],
        only_no_args=False)


# Add transforms for type inference.
astroid.MANAGER.register_transform(
    astroid.scoped_nodes.Class, transform_class_type_inference)
astroid.MANAGER.register_transform(
    astroid.scoped_nodes.Function, transform_fn_type_inference)