| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262 |
- #!/usr/bin/env python
- # Licensed to the Apache Software Foundation (ASF) under one
- # or more contributor license agreements. See the NOTICE file
- # distributed with this work for additional information
- # regarding copyright ownership. The ASF licenses this file
- # to you under the Apache License, Version 2.0 (the
- # "License"); you may not use this file except in compliance
- # with the License. You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- """Command line utlity for reading and writing Avro files."""
- from avro.io import DatumReader, DatumWriter
- from avro.datafile import DataFileReader, DataFileWriter
- import avro.schema
- try:
- import json
- except ImportError:
- import simplejson as json
- import csv
- from sys import stdout, stdin
- from itertools import ifilter, imap
- from functools import partial
- from os.path import splitext
- class AvroError(Exception):
- pass
- def print_json(row):
- print(json.dumps(row))
- def print_json_pretty(row):
- print(json.dumps(row, indent=4))
- _write_row = csv.writer(stdout).writerow
- _encoding = stdout.encoding or "UTF-8"
- def _encode(v, encoding=_encoding):
- if not isinstance(v, basestring):
- return v
- return v.encode(_encoding)
- def print_csv(row):
- # We sort the keys to the fields will be in the same place
- # FIXME: Do we want to do it in schema order?
- _write_row([_encode(row[key]) for key in sorted(row)])
- def select_printer(format):
- return {
- "json" : print_json,
- "json-pretty" : print_json_pretty,
- "csv" : print_csv
- }[format]
- def record_match(expr, record):
- return eval(expr, None, {"r" : record})
- def parse_fields(fields):
- fields = fields or ''
- if not fields.strip():
- return None
- return [field.strip() for field in fields.split(',') if field.strip()]
- def field_selector(fields):
- fields = set(fields)
- def keys_filter(obj):
- return dict((k, obj[k]) for k in (set(obj) & fields))
- return keys_filter
- def print_avro(avro, opts):
- if opts.header and (opts.format != "csv"):
- raise AvroError("--header applies only to CSV format")
- # Apply filter first
- if opts.filter:
- avro = ifilter(partial(record_match, opts.filter), avro)
- for i in xrange(opts.skip):
- try:
- next(avro)
- except StopIteration:
- return
- fields = parse_fields(opts.fields)
- if fields:
- avro = imap(field_selector(fields), avro)
- printer = select_printer(opts.format)
- for i, record in enumerate(avro):
- if i == 0 and opts.header:
- _write_row(sorted(record.keys()))
- if i >= opts.count:
- break
- printer(record)
- def print_schema(avro):
- schema = avro.meta["avro.schema"]
- # Pretty print
- print json.dumps(json.loads(schema), indent=4)
- def cat(opts, args):
- if not args:
- raise AvroError("No files to show")
- for filename in args:
- try:
- fo = open(filename, "rb")
- except (OSError, IOError), e:
- raise AvroError("Can't open %s - %s" % (filename, e))
- avro = DataFileReader(fo, DatumReader())
- if opts.print_schema:
- print_schema(avro)
- continue
- print_avro(avro, opts)
- def _open(filename, mode):
- if filename == "-":
- return {
- "rb" : stdin,
- "wb" : stdout
- }[mode]
- return open(filename, mode)
- def iter_json(info, _):
- return imap(json.loads, info)
- def convert(value, field):
- type = field.type.type
- if type == "union":
- return convert_union(value, field)
- return {
- "int" : int,
- "long" : long,
- "float" : float,
- "double" : float,
- "string" : str,
- "bytes" : str,
- "boolean" : bool,
- "null" : lambda _: None,
- "union" : lambda v: convert_union(v, field),
- }[type](value)
- def convert_union(value, field):
- for name in [s.name for s in field.type.schemas]:
- try:
- return convert(name)(value)
- except ValueError:
- continue
- def iter_csv(info, schema):
- header = [field.name for field in schema.fields]
- for row in csv.reader(info):
- values = [convert(v, f) for v, f in zip(row, schema.fields)]
- yield dict(zip(header, values))
- def guess_input_type(files):
- if not files:
- return None
- ext = splitext(files[0])[1].lower()
- if ext in (".json", ".js"):
- return "json"
- elif ext in (".csv",):
- return "csv"
- return None
- def write(opts, files):
- if not opts.schema:
- raise AvroError("No schema specified")
- input_type = opts.input_type or guess_input_type(files)
- if not input_type:
- raise AvroError("Can't guess input file type (not .json or .csv)")
- try:
- schema = avro.schema.parse(open(opts.schema, "rb").read())
- out = _open(opts.output, "wb")
- except (IOError, OSError), e:
- raise AvroError("Can't open file - %s" % e)
- writer = DataFileWriter(out, DatumWriter(), schema)
- iter_records = {"json" : iter_json, "csv" : iter_csv}[input_type]
- for filename in (files or ["-"]):
- info = _open(filename, "rb")
- for record in iter_records(info, schema):
- writer.append(record)
- writer.close()
- def main(argv=None):
- import sys
- from optparse import OptionParser, OptionGroup
- argv = argv or sys.argv
- parser = OptionParser(description="Display/write for Avro files",
- version="1.8.2",
- usage="usage: %prog cat|write [options] FILE [FILE...]")
- # cat options
- cat_options = OptionGroup(parser, "cat options")
- cat_options.add_option("-n", "--count", default=float("Infinity"),
- help="number of records to print", type=int)
- cat_options.add_option("-s", "--skip", help="number of records to skip",
- type=int, default=0)
- cat_options.add_option("-f", "--format", help="record format",
- default="json",
- choices=["json", "csv", "json-pretty"])
- cat_options.add_option("--header", help="print CSV header", default=False,
- action="store_true")
- cat_options.add_option("--filter", help="filter records (e.g. r['age']>1)",
- default=None)
- cat_options.add_option("--print-schema", help="print schema",
- action="store_true", default=False)
- cat_options.add_option('--fields', default=None,
- help='fields to show, comma separated (show all by default)')
- parser.add_option_group(cat_options)
- # write options
- write_options = OptionGroup(parser, "write options")
- write_options.add_option("--schema", help="schema file (required)")
- write_options.add_option("--input-type",
- help="input file(s) type (json or csv)",
- choices=["json", "csv"], default=None)
- write_options.add_option("-o", "--output", help="output file", default="-")
- parser.add_option_group(write_options)
- opts, args = parser.parse_args(argv[1:])
- if len(args) < 1:
- parser.error("You much specify `cat` or `write`") # Will exit
- command = args.pop(0)
- try:
- if command == "cat":
- cat(opts, args)
- elif command == "write":
- write(opts, args)
- else:
- raise AvroError("Unknown command - %s" % command)
- except AvroError, e:
- parser.error("%s" % e) # Will exit
- except Exception, e:
- raise SystemExit("panic: %s" % e)
- if __name__ == "__main__":
- main()
|