#!/usr/bin/env python
# -*- coding:utf-8 -*-

# Copyright 2015 Pablo Santiago Blum de Aguiar <scorphus@gmail.com>. All rights
# reserved. Use of this source code is governed by Apache License, Version 2.0,
# that can be found on https://opensource.org/licenses/Apache-2.0

import os
import sys
import yaml


def main(argv):
    '''Usage: extract_yaml.py <yam-file> <key>[.<key>...]'''

    def usage(yaml_dict=None):
        print main.__doc__
        if yaml_dict:
            print 'Available keys:\n  {}'.format(
                '\n  '.join(sorted(yaml_dict.keys()))
            )

    if len(argv) < 2 or not os.path.isfile(argv[1]):
        usage()
        return 1

    with open(sys.argv[1]) as yaml_file:
        try:
            yaml_dict = yaml.load(yaml_file)
        except Exception as e:
            print 'Could not read {}: {}'.format(sys.argv[1], e)
            return 2

        keys = [x for x in argv[2].split('.')] if len(argv) > 2 else []

        if not keys:
            usage(yaml_dict)
            return 3

        for key in keys:
            if key not in yaml_dict:
                usage(yaml_dict)
                return 4

            yaml_dict = yaml_dict[key]

        print yaml_dict


if __name__ == '__main__':
    status = main(sys.argv)
    sys.exit(status)

@@ -1068,130 +1068,221 @@
     
-if key not in yaml_dict:%0A                usage(yaml_dict)%0A                return 4%0A%0A            yaml_dict = yaml_dict%5Bkey%5D
+try:%0A                idx = int(key)%0A                key = idx%0A            except:%0A                pass%0A            try:%0A                yaml_dict = yaml_dict%5Bkey%5D%0A            except:%0A                yaml_dict = ''
 %0A%0A  

support integer keys