Derive nested value from dict by providing sequence of keys
"""Extract a nested value from a json-like dictionary object."""
from collections.abc import Sequence
from typing import Optional, SupportsIndex, TypeGuard, get_args

Key = SupportsIndex
# `Optional` stands for `null` values
JSONValue = Optional[str | int | float | bool]
NestedJSON = dict | list
KEY_ERROR_MSG = 'Keys sequence "{keys}" do not exists in object "{json_obj}"'


class StopExtractionError(Exception):
    """Extraction reached the end of nesting, but keys are not over yet."""


def extract_nested_value(
    json_obj: NestedJSON,
    keys: Sequence[Key],
    safe=True,
) -> JSONValue | NestedJSON:
    """
    Extract a nested value from a json-like dictionary object.

    Args:
        json_obj: The json-like dictionary to extract the nested value from.
        keys: A sequence of keys to traverse the nested structure.
        safe: When set to True, returns None
            in case of absence the sequence of keys in a nested structure
            When set to False, raises a KeyError in such a case

    Returns:
        Any: The nested value, or None if it could not be found.

    Raises:
        KeyError: when `safe` is False
            and the sequence of keys is absent in a nested structure
    """
    nested_obj: JSONValue | NestedJSON = json_obj

    for key in keys:
        try:
            nested_obj = _extract_or_stop(nested_obj, key)
        except StopExtractionError:
            if safe:
                return None
            raise KeyError(
                KEY_ERROR_MSG.format(keys=keys, json_obj=json_obj),
            )

    return nested_obj


def _extract_or_stop(
    nested_obj: JSONValue | NestedJSON,
    key: Key,
) -> JSONValue | NestedJSON:
    """Extract element from `nested_obj` by `key` index.

    Args:
        nested_obj: JSON object (nested or final)
            to get the nested element by indexing
        key: key for indexing the `nested_obj`

    Returns:
        Either another nested JSON or final JSON value

    Raises:
        StopExtractionError: when element is not found

    """
    if _is_nested(nested_obj):
        try:
            return nested_obj[key]
        except (KeyError, TypeError):
            raise StopExtractionError()

    raise StopExtractionError()


def _is_nested(json_obj: JSONValue | NestedJSON) -> TypeGuard[NestedJSON]:
    return isinstance(json_obj, get_args(NestedJSON))

src/nested_value_extractor.py
skatromb/nested-value-extractor
"""Tests for nested_value_extractor."""
from collections.abc import Sequence

import pytest

from src.nested_value_extractor import (
    JSONValue,
    Key,
    NestedJSON,
    extract_nested_value,
)

NON_HAPPY_ARGS = (
    ({1: "string shouldn't be indexed"}, [1, 2]),
    ({1: 2}, [1, 'end of nesting reached']),
    ({1: 2}, ['non existent key']),
    (1, ['non-indexable object']),
    ({1: 'nested key fail'}, [[1]]),
)


@pytest.mark.parametrize('safe', [True, False])
@pytest.mark.parametrize(
    'json_obj, keys, expected',
    [
        ({1: '1-st level value'}, [1], '1-st level value'),
        ({1: {2: {3: '3-rd level value'}}}, [1, 2, 3], '3-rd level value'),
        ({'no keys in': 'same dict out'}, [], {'no keys in': 'same dict out'}),
    ],
)
def test_happy_ways(
    json_obj: dict,
    keys: Sequence[Key],
    safe: bool,
    expected: NestedJSON | JSONValue,
) -> None:
    """Tests happy ways."""
    assert extract_nested_value(json_obj, keys, safe) == expected


@pytest.mark.parametrize(
    'json_obj, keys, expected',
    [(*args, None) for args in NON_HAPPY_ARGS],
)
def test_not_found_safe(
    json_obj: dict,
    keys: Sequence[Key],
    expected: None,
) -> None:
    """Tests cases when keys sequence is not found."""
    assert extract_nested_value(json_obj, keys) is expected


@pytest.mark.parametrize('safe', [False])
@pytest.mark.parametrize(
    'json_obj, keys, expected_exception',
    [(*args, pytest.raises(KeyError)) for args in NON_HAPPY_ARGS],
)
def test_exceptions(
    json_obj: dict,
    keys: Sequence[Key],
    safe: bool,
    expected_exception: pytest.ExceptionInfo,
):
    """Tests cases when keys sequence is not found and exception is thrown."""
    with expected_exception:
        extract_nested_value(json_obj, keys, safe)

tests/test_nested_value_extractor.py
skatromb/nested-value-extractor