| #!/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 |
| # |
| # https://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 utility for reading and writing Avro files.""" |
| |
| from __future__ import absolute_import, division, print_function |
| |
| import csv |
| import itertools |
| import json |
| import os.path |
| import sys |
| from functools import partial |
| from optparse import OptionGroup, OptionParser |
| |
| import avro |
| import avro.schema |
| from avro.datafile import DataFileReader, DataFileWriter |
| from avro.io import DatumReader, DatumWriter |
| |
| _AVRO_DIR = os.path.abspath(os.path.dirname(avro.__file__)) |
| |
| ifilter = getattr(itertools, 'ifilter', filter) |
| imap = getattr(itertools, 'imap', map) |
| |
| try: |
| unicode |
| except NameError: |
| unicode = str |
| |
| try: |
| long |
| except NameError: |
| long = int |
| |
| |
| def _version(): |
| with open(os.path.join(_AVRO_DIR, 'VERSION.txt')) as v: |
| return v.read() |
| |
| |
| _AVRO_VERSION = _version() |
| |
| |
| class AvroError(Exception): |
| pass |
| |
| |
| def print_json(row): |
| print(json.dumps(row)) |
| |
| |
| def print_json_pretty(row): |
| """Pretty print JSON""" |
| # Need to work around https://bugs.python.org/issue16333 |
| # where json.dumps leaves trailing spaces. |
| result = json.dumps(row, sort_keys=True, indent=4).replace(' \n', '\n') |
| print(result) |
| |
| |
| _write_row = csv.writer(sys.stdout).writerow |
| |
| |
| 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([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 range(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.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: |
| with DataFileReader(open(filename, 'rb'), DatumReader()) as avro: |
| if opts.print_schema: |
| print_schema(avro) |
| continue |
| print_avro(avro, opts) |
| |
| |
| def _open(filename, mode): |
| if filename == "-": |
| return { |
| "rb": sys.stdin, |
| "wb": sys.stdout |
| }[mode] |
| |
| return open(filename, mode) |
| |
| |
| def iter_json(info, _): |
| for i in info: |
| try: |
| s = i.decode() |
| except AttributeError: |
| s = i |
| yield json.loads(s) |
| |
| |
| 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": unicode, |
| "bytes": bytes, |
| "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((getattr(i, "decode", lambda: i)() for i in 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 = os.path.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)") |
| iter_records = {"json": iter_json, "csv": iter_csv}[input_type] |
| |
| try: |
| with open(opts.schema) as schema_file: |
| schema = avro.schema.parse(schema_file.read()) |
| out = _open(opts.output, "wb") |
| except (IOError, OSError) as e: |
| raise AvroError("Can't open file - %s" % e) |
| |
| writer = DataFileWriter(getattr(out, 'buffer', out), DatumWriter(), schema) |
| |
| 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): |
| argv = argv or sys.argv |
| |
| parser = OptionParser(description="Display/write for Avro files", |
| version=_AVRO_VERSION, |
| 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 as e: |
| parser.error("%s" % e) # Will exit |
| |
| |
| if __name__ == "__main__": |
| main() |
