jsonファイルの構造を任意深度で可視化する(python)

何をするコード？

pythonを使ってjsonの構造を可視化します。
ファイルのパスと表示したい深度を実行時引数で指定できます。
表示するのはjsonのkeyです。

何に使ったの？

jsonで書かれた深層学習用データセットの構造を知りたいと思った時に使いました。

コード

json-visualizeKey.py

import json
import sys
from pathlib import Path
import re


def open_json(path):
    with open(path) as data:
        return json.load(data)


def print_json(json_data, depth=1, max_depth=None):
    if max_depth is not None and depth > max_depth:
        return

    indentation = "  " * (depth - 1)

    if isinstance(json_data, dict):
        print(f"{indentation}" + "{")
        for key, value in json_data.items():
            print(f"{indentation}" + "  " + f"{key}:")
            print_json(value, depth + 1, max_depth)
        print(f"{indentation}" + "}")

    elif isinstance(json_data, list):
        print(f"{indentation}[")
        for item in json_data:
            print_json(item, depth + 1, max_depth)
        print(f"{indentation}]")


if __name__ == "__main__":
    arg_len = len(sys.argv)
    # 調査したいJsonファイルのパスを書く(実行時引数で上書きされる)
    path_string = (
        "~/Desktop/Dataset.json"
    )
    # 調査したいJsonの深さを書く(実行時引数で上書きされる)
    depth = 3

    if arg_len == 2:
        path_string = sys.argv[1]
    elif arg_len == 3:
        path_string, depth_string = sys.argv[1:3]
        pattern = r"\d+"
        matches = re.findall(pattern, depth_string)
        if len(matches) != 1:
            print("depthは整数である必要があります")
            exit()
        depth = int(depth_string)
    elif arg_len > 3:
        print("\n予期しない引数")
        print("python json-visualizeKey.py json-file-name view-depth")
        print("の形式で実行してください")
        exit()

    path = Path(path_string)
    if not path.exists():
        print(f"{path} は存在しません")
        exit()

    json_data = open_json(path)
    print_json(json_data, max_depth=depth)