NOTICE

The upcoming release of Featuretools 1.0.0 contains several breaking changes. Users are encouraged to test this version prior to release by installing from GitHub:

pip install https://github.com/alteryx/featuretools/archive/woodwork-integration.zip

For details on migrating to the new version, refer to Transitioning to Featuretools Version 1.0. Please report any issues in the Featuretools GitHub repo or by messaging in Alteryx Open Source Slack.

Source code for featuretools.variable_types.utils

import pandas as pd

from featuretools.utils.gen_utils import find_descendents
from featuretools.utils.plot_utils import (
    check_graphviz,
    get_graphviz_format,
    save_graph
)
from featuretools.variable_types.variable import Text, Variable


def find_variable_types():
    """
    Retrieves all Variable types as a dictionary where key is type_string
        of Variable, and value is a Variable object.

    Args:
        None

    Returns:
        variable_types (dict):

    """
    return {vtype.type_string: vtype for vtype in find_descendents(Variable)
            if vtype != Variable and vtype != Text}


def list_variable_types():
    """
    Retrieves all Variable types as a dataframe, with the column headers
        of name, type_string, and description.

    Args:
        None

    Returns:
        variable_types (pd.DataFrame): a DataFrame with column headers of
            name, type_strings, and description.
    """
    v_types = list(find_variable_types().values())
    v_type_strings = list(find_variable_types().keys())
    v_names = [x.__name__ for x in v_types]
    descriptions = [v.__doc__ for v in v_types]
    return pd.DataFrame({'name': v_names,
                         'type_string': v_type_strings,
                         'description': descriptions})


def graph_variable_types(to_file=None):
    """
    Create a UML diagram-ish graph of all the Variables.

    Args:
        to_file (str, optional) : Path to where the plot should be saved.
            If set to None (as by default), the plot will not be saved.

    Returns:
        graphviz.Digraph : Graph object that can directly be displayed in
            Jupyter notebooks.
    """
    graphviz = check_graphviz()
    format_ = get_graphviz_format(graphviz=graphviz,
                                  to_file=to_file)

    # Initialize a new directed graph
    graph = graphviz.Digraph('variables', format=format_)
    graph.attr(rankdir="LR")
    graph.node(Variable.__name__, shape='Mdiamond')

    all_variables_types = list(find_variable_types().values())
    all_variables_types.sort(key=lambda x: x.__name__)

    for node in all_variables_types:
        for parent in node.__bases__:
            graph.edge(parent.__name__, node.__name__)

    if to_file:
        save_graph(graph, to_file, format_)
    return graph