Helper library for (somewhat) safely extracting data from JSON in Python

Well, to be fair this actually doesn't work with JSON, but with dictionaries. But it's meant for possibly unknown dictionaries, so the most common application is HTTP request data turned into a dictionary.

There's a few solutions to the problem of safe(ish) data extraction in Python.

Wild West

response_dict = response.json()
machine_id = response_dict["servers"][0]["machine_id"]
# Yee, haw!

Now to be fair, this can work fine a lot of the time, but some APIs will change, or don't do proper HTTP error handling. They may suddenly change the JSON structure on you to indicate an error which causes some confusion. And this angers type checkers that you may or may not be using.

Pydantic

I'm not going to take the time to write this out, but you can make a Pydantic model for everything and try to load it into the model. This is pretty solid, but cumbersome.

Other libraries

I've used other libraries that attempt to map dictioaries to dataclasses. This also works well, but if you're calling one endpoint in one function, and need just a small bit of it, maybe there's an easier way?

Good, but long

response_dict = response.json()
assert isinstance(dict, response_dict)
assert "servers" in response_dict
assert isinstance(list, response_dict["servers"])
assert len(response_dict["servers"]) > 0
assert isinstance(dict, response_dict["servers"][0])
assert "machine_id" in response_dict["servers"][0]
machine_id = response_dict["servers"][0]["machine_id"]
assert isisntance(str, machine_id)

This may be overkill for many, but it can be somewhat necessary if you use strict type checking.

Better

response_dict = response.json()
servers = extract.dictionary(servers)
assert len(servers) > 0
extract.string(servers[0], "machine_id")

Now it's actually in this example that I realized my extract helper library does not have list index support, which would be useful for this particular case.

Where this really shines is actually in cases like this.

machine_id = extract.string(response.json(), "really", "nested", "structures", "machine_id")

In case it isn't obvious, here's the JSON that would map to.

{
  "really": {
    "nested": {
      "structures": {
        "machine_id": "there it is!"
      }
    }
  }
}

I wish this example was far-fetched, but sometimes it isn't!

Extract helper

So here's the helper, itself. I may write in support for list indexes at some point. It's released into the public domain.

from typing import Any


def _extract(data: dict[Any, Any], *key_path: str) -> Any:
    nested_data = data
    key_path_list = list(key_path)
    final_key = key_path_list.pop()
    for key in key_path_list:
        if key not in nested_data:
            raise ValueError(f"Key {key} not in data")
        if not isinstance(nested_data[key], dict):
            raise TypeError("Not a dictionary")
        nested_data = nested_data[key]
    if final_key not in nested_data:
        raise ValueError(f"Key {final_key} not in data")
    return nested_data[final_key]


def integer(data: dict[Any, Any], *key_path: str) -> int:
    value = _extract(data, *key_path)
    if not isinstance(value, int):
        raise TypeError("Value is not an int.")
    return value


def string(data: dict[Any, Any], *key_path: str) -> str:
    value = _extract(data, *key_path)
    if not isinstance(value, str):
        raise TypeError(
            f"Value is not a str. It's a {type(value)} with a value of: {value}"
        )
    return value


def a_list(data: dict[Any, Any], *key_path: str) -> list[Any]:
    value = _extract(data, *key_path)
    if not isinstance(value, list):
        raise TypeError("Value is not a list.")
    return value


def dictionary(data: dict[Any, Any], *key_path: str) -> dict[Any, Any]:
    value = _extract(data, *key_path)
    if not isinstance(value, dict):
        raise TypeError("Value is not a dict.")
    return value

And a couple of test cases that I've been using with pytest.

def test_integer() -> None:
    assert extract.integer({"an int": 5}, "an int") == 5
    assert extract.integer({"a dict": {"an int": 5}}, "a dict", "an int") == 5

    with pytest.raises(TypeError):
        extract.integer({"a str": "five"}, "a str")

    with pytest.raises(ValueError):
        extract.integer({"an int": 5}, "nope")

    with pytest.raises(TypeError):
        extract.integer({"a dict": {"an int": 5}}, "a dict", "an int", "nope")

    with pytest.raises(ValueError):
        extract.integer({"a dict": {"an int": 5}}, "a dict", "nope")

    with pytest.raises(ValueError):
        extract.integer({"a dict": {"an int": 5}}, "nope", "an int")


def test_string() -> None:
    assert extract.string({"a str": "five"}, "a str") == "five"
    assert extract.string({"a dict": {"a str": "five"}}, "a dict", "a str") == "five"

    with pytest.raises(TypeError):
        extract.string({"a str": 5}, "a str")

Assertions...

Some may complain that my examples include the assert keyword for non-test code. And this is a valid complaint, AssertionErrors aren't pretty, and you can ignore assertions with a certain Python runtime flag. In practice, I never find myself using that flag. The extract helper library avoids that, thankfully. But I can't say that I never use assert in production code, generally though just to appease mypy, my type checker.