Tools to work with hierarchical data

hierarchy.py

"""
Copyright 2024 Dustin Oprea

MIT LICENSE

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the “Software”), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

def translate_hierarchy_with_mappings(
        data, mappings, outer_parent_tuple=None, level=0, do_require_all=True):
    """Translate names in a hierarchy according to a mappings dictionary. Deeper
    elements are represented as tuples (though the top-level ones can be as
    well). the `outer` name is the human name and the `inner` name is the
    internal/storage name [that we're trying to obscure via this functionality].

    Typically we want all attributes to be present if we're using this
    functionality at all. `do_require_all` should be set to False to support any
    preexisting data that is not covered.
    """

    # Allows us to translate lists of blobs
    if isinstance(data, list) is True:

        results = []
        for i, x in enumerate(data):

            y = translate_hierarchy_with_mappings(
                    x,
                    mappings,
                    outer_parent_tuple=outer_parent_tuple,
                    level=level + 1,
                    do_require_all=do_require_all)

            results.append(y)

        return results

    elif isinstance(data, dict) is False:

        # Some scalar (or, less likely, an unsupported type, which wouldn't
        # make sense in this context since we're writing data). Just pass it
        # through.

        return data


    # It's a dictionary

    translated = {}
    for outer_name, value in data.items():

        # This is what we're look for the translation under. It can be a string
        # or a tuple.
        if outer_parent_tuple is None:
            reference_outer_tuple = outer_name
        else:
            reference_outer_tuple = outer_parent_tuple + (outer_name,)

        try:
            inner_name = mappings[reference_outer_tuple]

        except KeyError:

            # If not enforcing then just set into the output dictionary verbatim
            if do_require_all is False:
                translated[outer_name] = value
                continue

            raise


        # Descend

        # This is the parent-tuple that we'll be passing forward. It will
        # always be a tuple.
        if outer_parent_tuple is None:
            outer_parent_tuple2 = (outer_name,)
        else:
            outer_parent_tuple2 = outer_parent_tuple + (outer_name,)


        updated_value = \
            translate_hierarchy_with_mappings(
                value,
                mappings,
                outer_parent_tuple=outer_parent_tuple2,
                level=level + 1,
                do_require_all=do_require_all)

        translated[inner_name] = updated_value


    return translated


def _invert_mappings_tuple(mappings, input_key_tuple, cache):

    # If it's a 1-tuple, then try looking it up as a string. If not found, the
    # mapping is not complete and it's an error.

    if len(input_key_tuple) == 1:
        input_name = input_key_tuple[0]
        output_name = mappings[input_name]
        return (output_name,)


    # Prevent rerecursions

    try:
        return cache[input_key_tuple]
    except KeyError:
        pass


    # We have a tuple that isn't in the mapping. It must be an intermediate
    # form.

    # This retains the tuple type
    prefix_input_key_tuple = input_key_tuple[:-1]

    output_tuple = \
        _invert_mappings_tuple(
            mappings,
            prefix_input_key_tuple,
            cache)

    output_key_name = mappings[input_key_tuple]
    output_key_tuple = output_tuple + (output_key_name,)


    # Memoize so siblings caller contexts can benefit

    cache[input_key_tuple] = output_key_tuple


    return output_key_tuple


def get_reverse_hierarchy_mappings(mappings):
    """Returns a set of reverse mappings that can be provided to
    `translate_hierarchy_with_mappings` alongside already-mapped data in order
    to restore the pre-mapped data.
    """

    rmappings = {}
    cache = {}
    for outer_key, inner_member_name in mappings.items():
        if outer_key.__class__ is str:
            rmappings[inner_member_name] = outer_key

        elif outer_key.__class__ is tuple:
            inner_member_key = _invert_mappings_tuple(mappings, outer_key, cache)
            outer_member_name = outer_key[-1]

            rmappings[inner_member_key] = outer_member_name

        else:
            raise \
                Exception(
                    "We don't handle keys of type [{}]: [{}] [{}]".format(
                    key.__class__.__name__, key, value))


    return rmappings


def get_value_from_hierarchy_with_string_reference(
        record, reference, separator='.'):

# TODO(dustin): Add test

    parts = reference.split(separator)
    ptr = record
    while parts:

        # Allow us to specify nonexistent parts
        if ptr is None:
            ptr = {}

        part, parts = parts[0], parts[1:]

        # This will raise KeyError if the reference is invalid
        ptr = ptr[part]

    return ptr


def pick_from_hierarchical_records_gen(
        records, attribute_mappings, default_comma_separate_lists=False):
    """Given an iterable and a dictionayr of mappings (the keys of which can be
    hierarchical), yield a dictionary of mapped values for each record.
    """

    assert \
        attribute_mappings.__class__ is dict, \
        "Attribute mappings must be a dictionary: [{}]".format(
            attribute_mappings.__class__.__name__)

    assert \
        attribute_mappings, \
        "Attribute mappings is empty."

    for record in records:
        assembled = {}
        for from_, to in attribute_mappings.items():

            try:
                ptr = get_value_from_hierarchy_with_string_reference(
                        record,
                        from_)

            except KeyError:

                # Supports nonexistent parts
                ptr = None


            # If None, we couldn't find this attribute
            if ptr is None:
                ptr = ''

            if ptr.__class__ is list and default_comma_separate_lists is True:
                ptr = ','.join([str(x) for x in ptr])

            assembled[to] = ptr


        yield assembled