Gift Exchange gift-chain generator

exchange.py

"""
This code is for generating a single, circular "loop" of names of family
members for gift giving as part of a Secret Santa tradition.

The input is a file named "family.json". The JSON file should be of a single
mapping where the keys are the family members, and the values would be a
that contains the key "group".

The purpose of the "group" is to define households of people so that  this code
can prevent people within the same household getting one another as recipients.

Example "family.json" file:

``` json
{
    "Abrigale": {"group": 1},
    "Brian": {"group": 1},
    "Carrie": {"group": 1},

    "Derrick": {"group": 2},
    "Emily": {"group": 2},

    "Frank": {"group": 3},

    "Hariette": {"group": 4},
    "Issac": {"group": 4},
    "Jessica": {"group": 4},
    "Ken": {"group": 4}
}

In this example JSON, there are 4 households. Frank lives alone.

One possible putput would be:

Hariette => Brian => Emily => Jessica => Frank => Issac => Derrick => Carrie
=> Ken => Abrigale => Hariette.
```
"""
from argparse import ArgumentParser
from copy import deepcopy
import json
import logging
from pathlib import Path
import random
from typing import Any, Dict, Generator, Optional, Union

Graph = Dict[str, Dict[str, Any]]

logger = logging.getLogger(__name__)


def count_components(nodes_dict: Graph,
                     target: str = "target") -> int:
    """
    Counts the number of components in the given graph.

    Counts the number of components (cycles) in a directed graph represented as
    a dictionary of nodes, where each node has a `target` property pointing to
    another node.

    Parameters
    ----------
    nodes_dict : Dict[str, Dict[str, Any]]
        The dictionary representing the graph. Each key is a node name with a
        dictionary of properties, including `target` which points to the key of
        another node.

    target : str, default "target"
        The key of a node's properties where it stores the node it directs to.

    Returns
    -------
    int
        The number of components in the graph.
    """
    visited = set()  # Keeps track of visited nodes
    components = 0

    # Helper function for depth-first search to find cycles
    def dfs(node: str):
        nonlocal components
        stack = []
        while node not in visited:
            visited.add(node)
            stack.append(node)
            node = nodes_dict[node].get(target, "target")
        # If the last node's target is not in stack, it's a chain pointing to
        # a cycle.
        if node in stack:
            # Found a new component (cycle)
            components += 1

    # Iterate over all nodes and perform DFS from unvisited nodes
    for node in nodes_dict.keys():
        if node not in visited:
            dfs(node)

    return components


def create_random_graph(nodes: Graph, target: str = "target",
                        group_by: Optional[str] = None) -> Graph:
    """
    Create a single component of nodes, each conneced to the next once.

    Creates a randomly generated graph that connects each node in a directed
    manner to another node such that there is a single circle within the graph.

    If `group_by` is provided also ensures that each node targets someone with
    a different value for the provided key in the `group_by` parameter.

    Parameters
    ----------
    nodes : Dict[str, Dict[str, Any]]
        A dictionary of named nodes to be connected in a graph and
        their properties.

    target : str, default "target"
        The key of a node's properties where it stores the node it directs to.

    group_by : str, default None
        An optional str to group nodes by. By providing this this function will
        PREVENT a node being directed into another node of the same group.

    Returns
    -------
    Dict[str, Dict[str, Any]]
        A copy of the input dictionary as the input, except the key "target" is
        added to value dictionary of each key.
    """
    # Make a copy of the given nodes to play nicely.
    nodes = deepcopy(nodes)
    attempts = 1

    while attempts < 1000:
        node_list = [(key, value) for key, value in nodes.items()]
        already_recipient = []

        for giver, _ in node_list:
            recipient_nodes = [
                r for r, _ in node_list
                if (
                    r not in already_recipient
                    and giver != r
                    and (
                        not group_by
                        or nodes[giver][group_by] != nodes[r][group_by]
                    )
                )
            ]
            if recipient_nodes:
                recipient = random.choice(recipient_nodes)
                nodes[giver][target] = recipient
                already_recipient.append(recipient)
            else:
                break
        else:
            if count_components(nodes) == 1:
                break
        attempts += 1
    else:
        raise Exception("Could not find a suitable graph in 1000 attempts.")

    logger.info(f"Graph computed in {attempts} attempts.")
    return nodes


def gen_component_nodes(graph: Graph, target: str = "target"
                        ) -> Generator[None, None, str]:
    """
    Yields the nodes of a component of circularly directed nodes.

    Parameters
    ----------
    graph : Dictionary of nodes (keys) and their propertied (values)
        One of the properties should be "target" which references the name
        of the next node in the directed graph.

    target : str, default "target"
        The key of the node's properties which points to the next node in the
        directed graph.

    Returns
    -------
    Generator
        A generator of nodes, including the first as the target of the last.
    """
    graph = deepcopy(graph)
    left = random.choice(list(graph.keys()))
    props = graph.pop(left)
    last_right = left
    yield left
    while graph:
        right = props[target]
        yield right
        left = right
        props = graph.pop(right)
    yield last_right


def process(json_file_path: Union[Path, str] = "family.json"):
    """
    Given a path to a JSON file, print a single chain of gift-givers.

    Parameters
    ----------
    json_file_path : Path or str, default "family.json"
    """
    file_path = Path(json_file_path)
    if file_path.exists():
        with open(file_path) as json_file:
            people = json.load(json_file)
        graph = create_random_graph(people, target="target", group_by="group")
        print(" => ".join(gen_component_nodes(graph)))
    else:
        logger.error(f"{json_file_path} does not seem to exist.")


if __name__ == "__main__":
    parser = ArgumentParser("Generate a single chain of give-exchange pairs")
    parser.add_argument(
        '-f', '--json-file', default='family.json', dest="json_file_path",
        help='The path to the input JSON file (default is "family.json").'
    )

    options = parser.parse_args()
    process(**vars(options))