python YAML，JSON等......

Posted 2021-05-09

#! /usr/bin/env python
from __future__ import print_function

import argparse
import contextlib
import functools
import json
import os
import sys

import yaml
try:
    from yaml import CSafeDumper as yaml_dumper
except ImportError:
    from yaml import SafeDumper as yaml_dumper
yaml_dump = functools.partial(
    yaml.dump,
    Dumper=yaml_dumper,
    default_flow_style=False,
    explicit_start=True,
    explicit_end=True,
)

try:
    basestring
except NameError:
    basestring = (str, bytes)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'json_filename',
        type=argparse.FileType('r'),
        default='-',
        nargs='*',
        help='json file to parse',
    )
    parser.add_argument(
        '--save',
        '-s',
        type=bool,
        nargs='?',
        const=True,
        default=False,
        help='save output to .yaml files',
    )
    return parser.parse_args()


@contextlib.contextmanager
def open_input(input_file=None, mode='rt'):
    """ yield a self-closing file_object
        if input_file supports .read(), just pass it through
        fallback to sys.stdin
    """
    if not hasattr(input_file, 'read'):
        try:
            input_file = open(input_file, mode)
        except IOError:
            input_file = sys.stdin

    if getattr(input_file, 'fileno', lambda: -1)() in {0, 1, 2}:
        yield input_file
    else:
        try:
            with contextlib.closing(input_file) as file_object:
                yield file_object
        except AttributeError:
            # apparently, this wasn't actually a file
            pass


@contextlib.contextmanager
def open_output_from_input(input_file=None, ext='yaml', mode='wt'):
    """ yield a self-closing writable file object
        file name matches input_file with extension defaulting to 'yaml'
        if input_file is stdin or missing, yield stdout
    """
    input_file_name = getattr(input_file, 'name', None)
    if input_file_name in {None, '-', '<stdin>'}:
        yield sys.stdout
    else:
        new_file = os.path.splitext(input_file_name)[0]
        output_file = os.path.extsep.join((new_file, ext))
        with open(output_file, mode) as output_fp:
            yield output_fp


def parse_json_file(json_file, save):
    with open_input(json_file) as json_fp:
        obj = json.load(json_fp)
    yaml = yaml_dump(obj)
    print(yaml)
    if save:
        with open_output_from_input(json_file) as output_fp:
            print(yaml, file=output_fp)


if __name__ == '__main__':
    args = parse_args()
    for json_file in args.json_filename:
        parse_json_file(json_file, args.save)