From 932288c179cb705825c45545a9b205d8980f8aeb Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Fri, 1 Mar 2024 12:51:43 -0800 Subject: [PATCH] fix: make to_obj() more efficient (#20) --- .github/workflows/python-package.yml | 2 +- avrogen/avrojson.py | 30 ++++++++++++++++++---------- avrogen/dict_wrapper.py | 12 +++++------ avrogen/protocol.py | 3 ++- avrogen/schema.py | 3 ++- setup.py | 2 +- 6 files changed, 32 insertions(+), 20 deletions(-) diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 4e12ab8..42e17f8 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -15,7 +15,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] + python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] avro-version: ["1.10.2", "1.11.0"] steps: diff --git a/avrogen/avrojson.py b/avrogen/avrojson.py index 2f8a720..79ced34 100644 --- a/avrogen/avrojson.py +++ b/avrogen/avrojson.py @@ -16,22 +16,32 @@ _PRIMITIVE_TYPES = set(schema.PRIMITIVE_TYPES) +_json_converter = None +_json_converter_tuples = None + +def set_global_json_converter(json_converter: "AvroJsonConverter") -> None: + global _json_converter + _json_converter = json_converter + global _json_converter_tuples + _json_converter_tuples = json_converter.with_tuple_union(True) + +def get_global_json_converter(tuples: bool = False) -> "AvroJsonConverter": + if tuples: + assert _json_converter_tuples + return _json_converter_tuples + assert _json_converter + return _json_converter + class AvroJsonConverter(object): - def __init__(self, use_logical_types=False, logical_types=logical.DEFAULT_LOGICAL_TYPES, schema_types=None): + def __init__(self, use_logical_types=False, logical_types=logical.DEFAULT_LOGICAL_TYPES, fastavro: bool = False, schema_types=None): self.use_logical_types = use_logical_types self.logical_types = logical_types or {} self.schema_types = schema_types or {} - self.fastavro = False - - # Register self with all the schema objects. - for klass in self.schema_types.values(): - klass._json_converter = self + self.fastavro = fastavro - def with_tuple_union(self, enable=True) -> 'AvroJsonConverter': - ret = AvroJsonConverter(self.use_logical_types, self.logical_types, self.schema_types) - ret.fastavro = enable - return ret + def with_tuple_union(self, tuples=True) -> 'AvroJsonConverter': + return AvroJsonConverter(self.use_logical_types, self.logical_types, tuples, self.schema_types) def validate(self, expected_schema, datum, skip_logical_types=False) -> bool: if self.use_logical_types and expected_schema.props.get('logicalType') and not skip_logical_types \ diff --git a/avrogen/dict_wrapper.py b/avrogen/dict_wrapper.py index d8f5b6a..daf9d91 100644 --- a/avrogen/dict_wrapper.py +++ b/avrogen/dict_wrapper.py @@ -12,7 +12,6 @@ class DictWrapper: _inner_dict: dict RECORD_SCHEMA: ClassVar[RecordSchema] - _json_converter: ClassVar["AvroJsonConverter"] def __init__(self): self._inner_dict = {} @@ -43,17 +42,18 @@ def _construct_with_defaults(cls: Type[TC]) -> TC: return cls._construct({}) @classmethod - def _get_json_converter(cls) -> "AvroJsonConverter": - # This attribute will be set by the AvroJsonConverter's init method. - return cls._json_converter + def _get_json_converter(cls, tuples: bool = False) -> "AvroJsonConverter": + import avrogen.avrojson + + return avrogen.avrojson.get_global_json_converter(tuples) @classmethod def from_obj(cls: Type[TC], obj: dict, tuples: bool = False) -> TC: - conv = cls._get_json_converter().with_tuple_union(tuples) + conv = cls._get_json_converter(tuples=tuples) return conv.from_json_object(obj, cls.RECORD_SCHEMA) def to_obj(self, tuples: bool = False) -> dict: - conv = self._get_json_converter().with_tuple_union(tuples) + conv = self._get_json_converter(tuples=tuples) return conv.to_json_object(self, self.RECORD_SCHEMA) def to_avro_writable(self, fastavro: bool = False) -> dict: diff --git a/avrogen/protocol.py b/avrogen/protocol.py index a56ef6f..94c613a 100644 --- a/avrogen/protocol.py +++ b/avrogen/protocol.py @@ -160,7 +160,8 @@ def generate_protocol(protocol_json, use_logical_types=False, custom_imports=Non writer.untab() writer.write('\n}\n') - writer.write('_json_converter = %s\n\n' % avro_json_converter) + writer.write('_json_converter = %s\n' % avro_json_converter) + writer.write('avrojson.set_global_json_converter(_json_converter)\n') value = main_out.getvalue() main_out.close() return value, schema_names, request_names diff --git a/avrogen/schema.py b/avrogen/schema.py index 6789b1d..8db5679 100644 --- a/avrogen/schema.py +++ b/avrogen/schema.py @@ -83,7 +83,8 @@ def generate_schema(schema_json, use_logical_types=False, custom_imports=None, a writer.untab() writer.write('\n}\n\n') - writer.write(f'_json_converter = {avro_json_converter}\n\n') + writer.write(f'_json_converter = {avro_json_converter}\n') + writer.write('avrojson.set_global_json_converter(_json_converter)\n\n') value = main_out.getvalue() main_out.close() diff --git a/setup.py b/setup.py index b268b6b..833812e 100644 --- a/setup.py +++ b/setup.py @@ -20,7 +20,7 @@ setup( name="avro-gen3", - version="0.7.11", + version="0.7.12", description="Avro record class and specific record reader generator", long_description=long_description, long_description_content_type="text/markdown",