Prospective submodule for automatic higher order version of the janitor.functions functions to use `pandas.DataFrame.pip` and get autocompletion with jedi.

higher_order.py

#! /usr/bin/env python

# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright https://github.com/VPerrollaz
#
# Distributed under terms of the %LICENSE% license.

"""
Automatically generate higher order version of the janitor.functions functions.
"""
import sys
import inspect
import typing
import copy
import pandas as pd
import janitor


def _validate(obj):
    """Check if the object is a function from the API"""
    return inspect.getsource(obj).startswith("@pf.register_dataframe_method")


def _get_all_functions():
    """Get the list of functions to highroderize."""
    return [
        function
        for _, function in inspect.getmembers(janitor.functions, inspect.isfunction)
        if _validate(function)
    ]


def _reference_function() -> typing.Callable[[pd.DataFrame], pd.DataFrame]:
    """Function to be used for its signature and annotations."""

    def interne(df: pd.DataFrame) -> pd.DataFrame:
        return df

    return interne


def _adjust_signature(original_function, new_function):
    """Change the signature."""
    original_signature = inspect.signature(original_function)
    new_parameters = list()
    for clef, val in original_signature.parameters.items():
        if clef != "df":
            new_parameters.append(val)

    new_return_annotation = inspect.signature(_reference_function).return_annotation
    new_signature = inspect.Signature(
        parameters=new_parameters, return_annotation=new_return_annotation
    )
    new_function.__signature__ = new_signature


def _adjust_annotations(original_function, new_function):
    annotations = copy.deepcopy(original_function.__annotations__)
    try:
        annotations.pop("df")
    except KeyError:
        pass
    annotations["return"] = _reference_function.__annotations__["return"]
    new_function.__annotations__ = annotations


def _helper(original_function):
    """Create a higher order function which does not take a DataFrame argument"""

    def new_function(*args, **kwargs):
        def interne(df: pd.DataFrame) -> pd.DataFrame:
            return original_function(df, *args, **kwargs)

        return interne

    new_function.__name__ = original_function.__name__
    new_function.__doc__ = original_function.__doc__
    _adjust_signature(original_function, new_function)
    _adjust_annotations(original_function, new_function)
    return new_function


for function in _get_all_functions():
    setattr(sys.modules[__name__], function.__name__, _helper(function))

janitor_completion.ipynb