lang/py/avro/__main__.py - avro - Git at Google

 #!/usr/bin/env python3

 ##
 # 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."""
 import argparse
 import csv
 import functools
 import itertools
 import json
 import sys
 import warnings
 from pathlib import Path
 from typing import (
     IO,
     Any,
     AnyStr,
     Callable,
     Collection,
     Dict,
     Generator,
     Iterable,
     Mapping,
     Optional,
     Sequence,
     Tuple,
     Union,
     cast,
 )

 import avro
 import avro.datafile
 import avro.errors
 import avro.io
 import avro.schema


 def print_json(row: object) -> None:
     print(json.dumps(row))


 def print_json_pretty(row: object) -> None:
     """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: Mapping[str, object]) -> None:
     # 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_: str) -> Callable[[Mapping[str, object]], None]:
     return {"json": print_json, "json-pretty": print_json_pretty, "csv": print_csv}[format_]


 def record_match(expr: Any, record: Any) -> Callable[[Mapping[str, object]], bool]:
     warnings.warn(avro.errors.AvroWarning("There is no way to safely type check this."))
     return cast(Callable[[Mapping[str, object]], bool], eval(expr, None, {"r": record}))


 def _passthrough_keys_filter(obj: Mapping[str, object]) -> Dict[str, object]:
     return {**obj}


 def field_selector(fields: Collection[str]) -> Callable[[Mapping[str, object]], Dict[str, object]]:
     if fields:

         def keys_filter(obj: Mapping[str, object]) -> Dict[str, object]:
             return {k: obj[k] for k in obj.keys() & set(fields)}

         return keys_filter
     return _passthrough_keys_filter


 def print_avro(
     data_file_reader: avro.datafile.DataFileReader,
     header: bool,
     format_: str,
     filter_: str,
     skip: int = 0,
     count: int = 0,
     fields: Collection[str] = "",
 ) -> None:
     if header and format_ != "csv":
         raise avro.errors.UsageError("--header applies only to CSV format")
     predicate = functools.partial(record_match, filter_) if filter_ else None
     avro_filtered = filter(predicate, data_file_reader)
     avro_slice = itertools.islice(avro_filtered, skip, (skip + count) if count else None)
     fs = field_selector(fields)
     avro_fields = (fs(cast(Mapping[str, object], r)) for r in avro_slice)
     avro_enum = enumerate(avro_fields)
     printer = select_printer(format_)
     for i, record in avro_enum:
         if header and i == 0:
             _write_row(sorted(record))
         printer(record)


 def print_schema(data_file_reader: avro.datafile.DataFileReader) -> None:
     print(json.dumps(json.loads(data_file_reader.schema), indent=4))


 def cat(
     files: Sequence[IO[AnyStr]], print_schema_: bool, header: bool, format_: str, filter_: str, skip: int, count: int, fields: Collection[str]
 ) -> int:
     datum_reader = avro.io.DatumReader()
     for file_ in files:
         with avro.datafile.DataFileReader(file_, datum_reader) as data_file_reader:
             if print_schema_:
                 print_schema(data_file_reader)
                 continue
             print_avro(data_file_reader, header, format_, filter_, skip, count, fields)
     return 0


 def iter_json(info: Iterable[AnyStr], _: Any) -> Generator[object, None, None]:
     for i in info:
         row = (i if isinstance(i, str) else cast(bytes, i).decode()).strip()
         if row:
             yield json.loads(row)


 def convert(value: str, field: avro.schema.Field) -> Union[int, float, str, bytes, bool, None]:
     type_ = field.type.type
     if type_ in ("int", "long"):
         return int(value)
     if type_ in ("float", "double"):
         return float(value)
     if type_ == "string":
         return value
     if type_ == "bytes":
         return value.encode()
     if type_ == "boolean":
         return value.lower() in ("1", "t", "true")
     if type_ == "null":
         return None
     if type_ == "union":
         return convert_union(value, field)
     raise avro.errors.UsageError("No valid conversion type")


 def convert_union(value: str, field: avro.schema.Field) -> Union[int, float, str, bytes, bool, None]:
     for name in (s.name for s in field.type.schemas):
         try:
             return convert(value, name)
         except ValueError:
             continue
     raise avro.errors.UsageError("Exhausted Union Schema without finding a match")


 def iter_csv(info: IO[AnyStr], schema: avro.schema.RecordSchema) -> Generator[Dict[str, object], None, None]:
     header = [field.name for field in schema.fields]
     # If i is bytes, decode into a string.
     # If i is a string, no need to decode.
     csv_data = (cast(str, getattr(i, "decode", lambda: i)()) for i in info)
     for row in csv.reader(csv_data):
         values = [convert(v, f) for v, f in zip(row, schema.fields)]
         yield dict(zip(header, values))


 def guess_input_type(io_: IO[AnyStr]) -> str:
     ext = Path(io_.name).suffixes[0].lower()
     extensions = {".json": "json", ".js": "json", ".csv": "csv"}
     try:
         return extensions[ext]
     except KeyError as e:
         raise avro.errors.UsageError("Can't guess input file type (not .json or .csv)") from e


 def write(schema: avro.schema.RecordSchema, output: IO[AnyStr], input_files: Sequence[IO[AnyStr]], input_type: Optional[str] = None) -> int:
     input_type = input_type or guess_input_type(cast(IO[AnyStr], next(iter(input_files), sys.stdin)))
     iter_records = {"json": iter_json, "csv": iter_csv}[input_type]
     with avro.datafile.DataFileWriter(output, avro.io.DatumWriter(), schema) as writer:
         for file_ in input_files:
             for record in iter_records(file_, schema):
                 writer.append(record)
     return 0


 def csv_arg(arg: Optional[str]) -> Optional[Tuple[str, ...]]:
     return tuple(f.strip() for f in arg.split(",") if f.strip()) if (arg and arg.strip()) else None


 def _parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Display/write for Avro files")
     parser.add_argument("--version", action="version", version=avro.__version__)
     subparsers = parser.add_subparsers(required=True, dest="command") if sys.version_info >= (3, 7) else parser.add_subparsers(dest="command")
     subparser_cat = subparsers.add_parser("cat")
     subparser_cat.add_argument("--version", action="version", version=avro.__version__)
     subparser_cat.add_argument(
         "--count",
         "-n",
         default=None,
         type=int,
         help="number of records to print",
     )
     subparser_cat.add_argument(
         "--skip",
         "-s",
         type=int,
         default=0,
         help="number of records to skip",
     )
     subparser_cat.add_argument(
         "--format",
         "-f",
         default="json",
         choices=("json", "csv", "json-pretty"),
         help="record format",
     )
     subparser_cat.add_argument(
         "--header",
         "-H",
         default=False,
         action="store_true",
         help="print CSV header",
     )
     subparser_cat.add_argument(
         "--filter",
         "-F",
         default=None,
         help="filter records (e.g. r['age']>1)",
     )
     subparser_cat.add_argument(
         "--print-schema",
         "-p",
         action="store_true",
         default=False,
         help="print schema",
     )
     subparser_cat.add_argument("--fields", type=csv_arg, default=None, help="fields to show, comma separated (show all by default)")
     subparser_cat.add_argument("files", nargs="*", type=argparse.FileType("rb"), help="Files to read", default=(sys.stdin,))
     subparser_write = subparsers.add_parser("write")
     subparser_write.add_argument("--version", action="version", version=avro.__version__)
     subparser_write.add_argument("--schema", "-s", type=avro.schema.from_path, required=True, help="schema JSON file (required)")
     subparser_write.add_argument(
         "--input-type",
         "-f",
         choices=("json", "csv"),
         default=None,
         help="input file(s) type (json or csv)",
     )
     subparser_write.add_argument("--output", "-o", help="output file", type=argparse.FileType("wb"), default=sys.stdout)
     subparser_write.add_argument("input_files", nargs="*", type=argparse.FileType("r"), help="Files to read", default=(sys.stdin,))
     return parser.parse_args()


 def main() -> int:
     args = _parse_args()
     if args.command == "cat":
         return cat(args.files, args.print_schema, args.header, args.format, args.filter, args.skip, args.count, args.fields)
     if args.command == "write":
         return write(args.schema, args.output, args.input_files, args.input_type)
     raise avro.errors.UsageError(f"Unknown command - {args.command}")


 if __name__ == "__main__":
     raise SystemExit(main())
	#!/usr/bin/env python3

	##
	# 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."""
	import argparse
	import csv
	import functools
	import itertools
	import json
	import sys
	import warnings
	from pathlib import Path
	from typing import (
	IO,
	Any,
	AnyStr,
	Callable,
	Collection,
	Dict,
	Generator,
	Iterable,
	Mapping,
	Optional,
	Sequence,
	Tuple,
	Union,
	cast,
	)

	import avro
	import avro.datafile
	import avro.errors
	import avro.io
	import avro.schema


	def print_json(row: object) -> None:
	print(json.dumps(row))


	def print_json_pretty(row: object) -> None:
	"""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: Mapping[str, object]) -> None:
	# 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_: str) -> Callable[[Mapping[str, object]], None]:
	return {"json": print_json, "json-pretty": print_json_pretty, "csv": print_csv}[format_]


	def record_match(expr: Any, record: Any) -> Callable[[Mapping[str, object]], bool]:
	warnings.warn(avro.errors.AvroWarning("There is no way to safely type check this."))
	return cast(Callable[[Mapping[str, object]], bool], eval(expr, None, {"r": record}))


	def _passthrough_keys_filter(obj: Mapping[str, object]) -> Dict[str, object]:
	return {**obj}


	def field_selector(fields: Collection[str]) -> Callable[[Mapping[str, object]], Dict[str, object]]:
	if fields:

	def keys_filter(obj: Mapping[str, object]) -> Dict[str, object]:
	return {k: obj[k] for k in obj.keys() & set(fields)}

	return keys_filter
	return _passthrough_keys_filter


	def print_avro(
	data_file_reader: avro.datafile.DataFileReader,
	header: bool,
	format_: str,
	filter_: str,
	skip: int = 0,
	count: int = 0,
	fields: Collection[str] = "",
	) -> None:
	if header and format_ != "csv":
	raise avro.errors.UsageError("--header applies only to CSV format")
	predicate = functools.partial(record_match, filter_) if filter_ else None
	avro_filtered = filter(predicate, data_file_reader)
	avro_slice = itertools.islice(avro_filtered, skip, (skip + count) if count else None)
	fs = field_selector(fields)
	avro_fields = (fs(cast(Mapping[str, object], r)) for r in avro_slice)
	avro_enum = enumerate(avro_fields)
	printer = select_printer(format_)
	for i, record in avro_enum:
	if header and i == 0:
	_write_row(sorted(record))
	printer(record)


	def print_schema(data_file_reader: avro.datafile.DataFileReader) -> None:
	print(json.dumps(json.loads(data_file_reader.schema), indent=4))


	def cat(
	files: Sequence[IO[AnyStr]], print_schema_: bool, header: bool, format_: str, filter_: str, skip: int, count: int, fields: Collection[str]
	) -> int:
	datum_reader = avro.io.DatumReader()
	for file_ in files:
	with avro.datafile.DataFileReader(file_, datum_reader) as data_file_reader:
	if print_schema_:
	print_schema(data_file_reader)
	continue
	print_avro(data_file_reader, header, format_, filter_, skip, count, fields)
	return 0


	def iter_json(info: Iterable[AnyStr], _: Any) -> Generator[object, None, None]:
	for i in info:
	row = (i if isinstance(i, str) else cast(bytes, i).decode()).strip()
	if row:
	yield json.loads(row)


	def convert(value: str, field: avro.schema.Field) -> Union[int, float, str, bytes, bool, None]:
	type_ = field.type.type
	if type_ in ("int", "long"):
	return int(value)
	if type_ in ("float", "double"):
	return float(value)
	if type_ == "string":
	return value
	if type_ == "bytes":
	return value.encode()
	if type_ == "boolean":
	return value.lower() in ("1", "t", "true")
	if type_ == "null":
	return None
	if type_ == "union":
	return convert_union(value, field)
	raise avro.errors.UsageError("No valid conversion type")


	def convert_union(value: str, field: avro.schema.Field) -> Union[int, float, str, bytes, bool, None]:
	for name in (s.name for s in field.type.schemas):
	try:
	return convert(value, name)
	except ValueError:
	continue
	raise avro.errors.UsageError("Exhausted Union Schema without finding a match")


	def iter_csv(info: IO[AnyStr], schema: avro.schema.RecordSchema) -> Generator[Dict[str, object], None, None]:
	header = [field.name for field in schema.fields]
	# If i is bytes, decode into a string.
	# If i is a string, no need to decode.
	csv_data = (cast(str, getattr(i, "decode", lambda: i)()) for i in info)
	for row in csv.reader(csv_data):
	values = [convert(v, f) for v, f in zip(row, schema.fields)]
	yield dict(zip(header, values))


	def guess_input_type(io_: IO[AnyStr]) -> str:
	ext = Path(io_.name).suffixes[0].lower()
	extensions = {".json": "json", ".js": "json", ".csv": "csv"}
	try:
	return extensions[ext]
	except KeyError as e:
	raise avro.errors.UsageError("Can't guess input file type (not .json or .csv)") from e


	def write(schema: avro.schema.RecordSchema, output: IO[AnyStr], input_files: Sequence[IO[AnyStr]], input_type: Optional[str] = None) -> int:
	input_type = input_type or guess_input_type(cast(IO[AnyStr], next(iter(input_files), sys.stdin)))
	iter_records = {"json": iter_json, "csv": iter_csv}[input_type]
	with avro.datafile.DataFileWriter(output, avro.io.DatumWriter(), schema) as writer:
	for file_ in input_files:
	for record in iter_records(file_, schema):
	writer.append(record)
	return 0


	def csv_arg(arg: Optional[str]) -> Optional[Tuple[str, ...]]:
	return tuple(f.strip() for f in arg.split(",") if f.strip()) if (arg and arg.strip()) else None


	def _parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(description="Display/write for Avro files")
	parser.add_argument("--version", action="version", version=avro.__version__)
	subparsers = parser.add_subparsers(required=True, dest="command") if sys.version_info >= (3, 7) else parser.add_subparsers(dest="command")
	subparser_cat = subparsers.add_parser("cat")
	subparser_cat.add_argument("--version", action="version", version=avro.__version__)
	subparser_cat.add_argument(
	"--count",
	"-n",
	default=None,
	type=int,
	help="number of records to print",
	)
	subparser_cat.add_argument(
	"--skip",
	"-s",
	type=int,
	default=0,
	help="number of records to skip",
	)
	subparser_cat.add_argument(
	"--format",
	"-f",
	default="json",
	choices=("json", "csv", "json-pretty"),
	help="record format",
	)
	subparser_cat.add_argument(
	"--header",
	"-H",
	default=False,
	action="store_true",
	help="print CSV header",
	)
	subparser_cat.add_argument(
	"--filter",
	"-F",
	default=None,
	help="filter records (e.g. r['age']>1)",
	)
	subparser_cat.add_argument(
	"--print-schema",
	"-p",
	action="store_true",
	default=False,
	help="print schema",
	)
	subparser_cat.add_argument("--fields", type=csv_arg, default=None, help="fields to show, comma separated (show all by default)")
	subparser_cat.add_argument("files", nargs="*", type=argparse.FileType("rb"), help="Files to read", default=(sys.stdin,))
	subparser_write = subparsers.add_parser("write")
	subparser_write.add_argument("--version", action="version", version=avro.__version__)
	subparser_write.add_argument("--schema", "-s", type=avro.schema.from_path, required=True, help="schema JSON file (required)")
	subparser_write.add_argument(
	"--input-type",
	"-f",
	choices=("json", "csv"),
	default=None,
	help="input file(s) type (json or csv)",
	)
	subparser_write.add_argument("--output", "-o", help="output file", type=argparse.FileType("wb"), default=sys.stdout)
	subparser_write.add_argument("input_files", nargs="*", type=argparse.FileType("r"), help="Files to read", default=(sys.stdin,))
	return parser.parse_args()


	def main() -> int:
	args = _parse_args()
	if args.command == "cat":
	return cat(args.files, args.print_schema, args.header, args.format, args.filter, args.skip, args.count, args.fields)
	if args.command == "write":
	return write(args.schema, args.output, args.input_files, args.input_type)
	raise avro.errors.UsageError(f"Unknown command - {args.command}")


	if __name__ == "__main__":
	raise SystemExit(main())