Functions to Retrieve a Value from a Path, Retrieve Tokens from a Dict, and Replace Tokens in a Dict

tokens_in_dictionary.py

#! /usr/bin/env python3


import json
import re

from unittest import TestCase
import unittest.mock as mock


def get_property_from_dict(message: dict, path: list):
    """Get Property from Dictionary

    This function retrieves the value of a property at a
    specified location.

    Args:
        message (dict):
            The message the function traverses.
        path (list):
            An ordered list of nested properties or elements. This is
            the structure that identifies a property in a supplied
            message.

            path = ['romania', 'bar', 0, 'this']

            message = {
                "romania": {
                    "region": "transylvania"
                    "bar": [
                            { "this": "Cheia", "that": True},
                            { "this": "Brasov", "that": True}
                            { "this": "Paris", "that": Frue}
                        ]
                    }
                }
            }

    Returns:
        This function returns a the value of the property at the 
        specified location. The return value's data type can be any
        type.
    """
    for key in path:
        message = message[key]

    return message


def token_collect_from_dict(message: dict, message_path: list = None) -> dict:
    """Collect Tokens from Dictionary

    This function parses a dictionary marking token locations.
    The function returns a list of tokens and their locations
    in a given message.

    Args:
        message (dict):
            The message the function searches for tokens.
        message_path (dict):
            (Optional) An ordered list of nested properties or 
            elements. The path to the message's location in a
            larger dictionary.

            message = {
                "romania": {
                    "region": "transylvania",
                    "dateCreated": "<TOKEN.ANY>"
                    "bar": [
                            { "this": "Cheia", "that": True},
                            { "this": "Brasov", "that": True}
                            { "this": "<TOKEN.ff3c2f>", "that": Frue}
                        ]
                    }
                }
            }

            token_collect_from_dict(message['romania'], ['romania'])

    Returns:
        The function returns a list of tokens and their locations
        in a given message.

        [
            {'<TOKEN.a3bcf2>': ['romania', 'bar', 2, 'this']},
            {'<TOKEN.ANY>': ['romania', 'dateCreated']}
        ]
    """

    tokens = []
    message_path = [] if message_path is None else message_path

    for key in message.keys():

        path = message_path.copy()
        path.append(key)

        # Replace a token if found
        if type(message[key]) == str:
            res = re.findall('<TOKEN\..*>', message[key])

            if res:
                tokens.append({res[0]: path})

        # If the key is a dictionary, check it for tokens
        elif type(message[key]) == dict:
            res = token_collect_from_dict(message[key], path)

            if res:
                tokens += res

        # If the key is a list, check it for objects or tokens
        elif type(message[key]) == list:

            for element_position in range(len(message[key])):
                object_path = path.copy()

                if type(message[key][element_position]) == dict:
                    object_path.append(element_position)
                    res = token_collect_from_dict(
                        message[key][element_position], object_path)

                    if res:
                        tokens += res

    return tokens


def token_replace_in_dict(message: dict, tokens: dict) -> dict:
    """Replace Tokens

    This function replaces tokens in a given dictionary.

    Args:
        message (dict):
            The message the function searches for tokens.
        tokens (dict):
            The tokens the function searches for. The token dict
            includes tokens and values.

    Returns:
        The original dictinary with tokens replaced with
        token values. For example, the following message:

        {
            'kid': 'Sabrina',
            'incidents': [
                {'description': 'jump on desk', 'dateOccurred': '<TOKEN.ANY>'},
                {'description': 'paint on wall', 'dateOccurred': '<TOKEN.a3bcf2>'}
            ]
        }
        
        The following token bank:

        { "<TOKEN.ANY>": mock.ANY, "<TOKEN.a3bcf2>": '1234' }

        Results in the following output:

        {
            'kid': 'Sabrina',
            'incidents': [
                {'description': 'jump on desk', 'dateOccurred': <ANY>},
                {'description': 'paint on wall', 'dateOccurred': '1234'}
            ]
        }
    """

    for key in message.keys():
        # Replace a token if found
        if type(message[key]) == str and message[key] in tokens.keys():
            message[key] = tokens[message[key]]

        # If the key is a dictionary, check it for tokens
        elif type(message[key]) == dict:
            message[key] = \
                token_replace_in_dict(message[key], tokens)

        # If the key is a list, check it for objects or tokens
        elif type(message[key]) == list:
            for element_position in range(len(message[key])):
                if type(message[key][element_position]) == dict:
                    message[key][element_position] = \
                        token_replace_in_dict(message[key][element_position], tokens)

    return message

if __name__ == "__main__":
    
    test_dictionary_1 = {
        "name": "Ralphy",
        "record": [
            {"crime": "jump on desk", "occurrences": 2},
            {"crime": "<TOKEN.a3bcf2>", "occurrences": 1},
        ],
        "age": {
            "years": "<TOKEN.ANY>",
            "thinger": {
                "thang1": "<TOKEN.a3bcf2>",
                "thang2": "sabrina",
                "thang3": [
                    {"foo": "<TOKEN.a3bcf2>", "bar": True}
                ],
            }
        }
    }

    test_dictionary_2 = {
        "kid": "Sabrina",
        "incidents": [
            {"description": "jump on desk", "dateOccurred": "<TOKEN.ANY>"},
            {"description": "paint on wall", "dateOccurred": "<TOKEN.a3bcf2>"},
        ]
    }

    token_bank = {
        "<TOKEN.ANY>": mock.ANY,
        "<TOKEN.a3bcf2>": '1234',
    }
    print(test_dictionary_2)
    print(token_replace_in_dict(test_dictionary_2, token_bank))
    #print(get_property_from_dict(test_dictionary_1, ('record', 0, 'crime')))
    #print(token_collect_from_dict(test_dictionary_1))