|
| 1 | +from typing import Union, Mapping, Optional, Any, Iterable |
| 2 | + |
| 3 | +from io import BytesIO |
| 4 | + |
| 5 | +from fastavro import schemaless_reader, schemaless_writer, parse_schema |
| 6 | +from fastavro.types import Schema |
| 7 | + |
| 8 | +from .base import Serializer, Deserializer, SerializationContext |
| 9 | +from .exceptions import SerializationError |
| 10 | + |
| 11 | +__all__ = ("AvroSerializer", "AvroDeserializer") |
| 12 | + |
| 13 | + |
| 14 | +class AvroSerializer(Serializer): |
| 15 | + def __init__( |
| 16 | + self, |
| 17 | + schema: Schema, |
| 18 | + strict: bool = False, |
| 19 | + strict_allow_default: bool = False, |
| 20 | + disable_tuple_notation: bool = False, |
| 21 | + ): |
| 22 | + """ |
| 23 | + Serializer that returns data in Avro format. |
| 24 | +
|
| 25 | + For more information see fastavro [schemaless_writer](https://fastavro.readthedocs.io/en/latest/writer.html#fastavro._write_py.schemaless_writer) method. |
| 26 | +
|
| 27 | + :param schema: The avro schema. |
| 28 | + :param strict: If set to True, an error will be raised if records do not contain exactly the same fields that the schema states. |
| 29 | + Default - `False` |
| 30 | + :param strict_allow_default: If set to True, an error will be raised if records do not contain exactly the same fields that the schema states unless it is a missing field that has a default value in the schema. |
| 31 | + Default - `False` |
| 32 | + :param disable_tuple_notation: If set to True, tuples will not be treated as a special case. Therefore, using a tuple to indicate the type of a record will not work. |
| 33 | + Default - `False` |
| 34 | + """ |
| 35 | + self._schema = parse_schema(schema) |
| 36 | + self._strict = strict |
| 37 | + self._strict_allow_default = strict_allow_default |
| 38 | + self._disable_tuple_notation = disable_tuple_notation |
| 39 | + |
| 40 | + def __call__(self, value: Any, ctx: SerializationContext) -> bytes: |
| 41 | + data = BytesIO() |
| 42 | + |
| 43 | + with BytesIO() as data: |
| 44 | + try: |
| 45 | + schemaless_writer( |
| 46 | + data, |
| 47 | + self._schema, |
| 48 | + value, |
| 49 | + strict=self._strict, |
| 50 | + strict_allow_default=self._strict_allow_default, |
| 51 | + disable_tuple_notation=self._disable_tuple_notation, |
| 52 | + ) |
| 53 | + except ValueError as exc: |
| 54 | + raise SerializationError(str(exc)) from exc |
| 55 | + |
| 56 | + return data.getvalue() |
| 57 | + |
| 58 | + |
| 59 | +class AvroDeserializer(Deserializer): |
| 60 | + def __init__( |
| 61 | + self, |
| 62 | + schema: Schema, |
| 63 | + reader_schema: Optional[Schema] = None, |
| 64 | + return_record_name: bool = False, |
| 65 | + return_record_name_override: bool = False, |
| 66 | + return_named_type: bool = False, |
| 67 | + return_named_type_override: bool = False, |
| 68 | + handle_unicode_errors: str = "strict", |
| 69 | + ): |
| 70 | + """ |
| 71 | + Deserializer that parses data from Avro. |
| 72 | +
|
| 73 | + For more information see fastavro [schemaless_reader](https://fastavro.readthedocs.io/en/latest/reader.html#fastavro._read_py.schemaless_reader) method. |
| 74 | +
|
| 75 | + :param schema: The Avro schema. |
| 76 | + :param reader_schema: If the schema has changed since being written then the new schema can be given to allow for schema migration. |
| 77 | + Default - `None` |
| 78 | + :param return_record_name: If true, when reading a union of records, the result will be a tuple where the first value is the name of the record and the second value is the record itself. |
| 79 | + Default - `False` |
| 80 | + :param return_record_name_override: If true, this will modify the behavior of return_record_name so that the record name is only returned for unions where there is more than one record. For unions that only have one record, this option will make it so that the record is returned by itself, not a tuple with the name. |
| 81 | + Default - `False` |
| 82 | + :param return_named_type: If true, when reading a union of named types, the result will be a tuple where the first value is the name of the type and the second value is the record itself NOTE: Using this option will ignore return_record_name and return_record_name_override. |
| 83 | + Default - `False` |
| 84 | + :param return_named_type_override: If true, this will modify the behavior of return_named_type so that the named type is only returned for unions where there is more than one named type. For unions that only have one named type, this option will make it so that the named type is returned by itself, not a tuple with the name. |
| 85 | + Default - `False` |
| 86 | + :param handle_unicode_errors: Should be set to a valid string that can be used in the errors argument of the string decode() function. |
| 87 | + Default - `"strict"` |
| 88 | + """ |
| 89 | + super().__init__() |
| 90 | + self._schema = parse_schema(schema) |
| 91 | + self._reader_schema = parse_schema(reader_schema) if reader_schema else None |
| 92 | + self._return_record_name = return_record_name |
| 93 | + self._return_record_name_override = return_record_name_override |
| 94 | + self._return_named_type = return_named_type |
| 95 | + self._return_named_type_override = return_named_type_override |
| 96 | + self._handle_unicode_errors = handle_unicode_errors |
| 97 | + |
| 98 | + def __call__( |
| 99 | + self, value: bytes, ctx: SerializationContext |
| 100 | + ) -> Union[Iterable[Mapping], Mapping]: |
| 101 | + try: |
| 102 | + return schemaless_reader( |
| 103 | + BytesIO(value), |
| 104 | + self._schema, |
| 105 | + reader_schema=self._reader_schema, |
| 106 | + return_record_name=self._return_record_name, |
| 107 | + return_record_name_override=self._return_record_name_override, |
| 108 | + return_named_type=self._return_named_type, |
| 109 | + return_named_type_override=self._return_named_type_override, |
| 110 | + handle_unicode_errors=self._handle_unicode_errors, |
| 111 | + ) |
| 112 | + except EOFError as exc: |
| 113 | + raise SerializationError(str(exc)) from exc |
0 commit comments