pydantic-avro 0.8.1__tar.gz → 0.9.1__tar.gz

{pydantic_avro-0.8.1 → pydantic_avro-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pydantic-avro
-Version: 0.8.1
+Version: 0.9.1
 Summary: Converting pydantic classes to avro schemas
 Home-page: https://github.com/godatadriven/pydantic-avro
 License: MIT

{pydantic_avro-0.8.1 → pydantic_avro-0.9.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pydantic-avro"
-version = "0.8.1"
+version = "0.9.1"
 description = "Converting pydantic classes to avro schemas"
 authors = ["Peter van 't Hof' <peter.vanthof@godatadriven.com>"]
 
@@ -18,7 +18,7 @@ python = ">=3.8.1,<4.0"
 pydantic = ">=1.4,<3.0"
 
 [tool.poetry.dev-dependencies]
-coverage = { version = "^7.2.2", extras = ["toml"] }
+coverage = { version = "^7.6.1", extras = ["toml"] }
 pytest = "^8.3.5"
 pytest-mock = "^3.10.0"
 pyproject-flake8 = "^7.0.0"
@@ -28,6 +28,7 @@ pytest-cov = "^5.0.0"
 mypy = "^1.1.1"
 avro = "^1.12.0"
 fastavro = "^1.8.1"
+typing-extensions = "^4.13.2"
 
 [tool.poetry.scripts]
 pydantic-avro = "pydantic_avro.__main__:root_main"

{pydantic_avro-0.8.1 → pydantic_avro-0.9.1}/setup.py

@@ -18,7 +18,7 @@ entry_points = \
 
 setup_kwargs = {
     'name': 'pydantic-avro',
-    'version': '0.8.1',
+    'version': '0.9.1',
     'description': 'Converting pydantic classes to avro schemas',
     'long_description': '[![Python package](https://github.com/godatadriven/pydantic-avro/actions/workflows/python-package.yml/badge.svg)](https://github.com/godatadriven/pydantic-avro/actions/workflows/python-package.yml)\n[![codecov](https://codecov.io/gh/godatadriven/pydantic-avro/branch/main/graph/badge.svg?token=5L08GOERAW)](https://codecov.io/gh/godatadriven/pydantic-avro)\n[![PyPI version](https://badge.fury.io/py/pydantic-avro.svg)](https://badge.fury.io/py/pydantic-avro)\n[![CodeQL](https://github.com/godatadriven/pydantic-avro/actions/workflows/codeql-analysis.yml/badge.svg)](https://github.com/godatadriven/pydantic-avro/actions/workflows/codeql-analysis.yml)\n\n# pydantic-avro\n\nThis library can convert a pydantic class to a avro schema or generate python code from a avro schema.\n\n### Install\n\n```bash\npip install pydantic-avro\n```\n\n### Pydantic class to avro schema\n\n```python\nimport json\nfrom typing import Optional\n\nfrom pydantic_avro.base import AvroBase\n\n\nclass TestModel(AvroBase):\n    key1: str\n    key2: int\n    key2: Optional[str]\n\n\nschema_dict: dict = TestModel.avro_schema()\nprint(json.dumps(schema_dict))\n\n```\n\n### Avro schema to pydantic\n\n```shell\n# Print to stdout\npydantic-avro avro_to_pydantic --asvc /path/to/schema.asvc\n\n# Save it to a file\npydantic-avro avro_to_pydantic --asvc /path/to/schema.asvc --output /path/to/output.py\n```\n\n### Specify expected Avro type\n\n```python\nfrom datetime import datetime\nfrom pydantic import Field\nfrom pydantic_avro.base import AvroBase \n\nclass ExampleModel(AvroBase):\n    field1: int = Field(..., avro_type="long")  # Explicitly set Avro type to "long"\n    field2: datetime = Field(..., avro_type="timestamp-millis")  # Explicitly set Avro type to "timestamp-millis"\n```\n\n### Install for developers\n\n###### Install package\n\n- Requirement: Poetry 1.*\n\n```shell\npoetry install\n```\n\n###### Run unit tests\n```shell\npytest\ncoverage run -m pytest  # with coverage\n# or (depends on your local env) \npoetry run pytest\npoetry run coverage run -m pytest  # with coverage\n```\n\n##### Run linting\n\nThe linting is checked in the github workflow. To fix and review issues run this:\n```shell\nblack .   # Auto fix all issues\nisort .   # Auto fix all issues\npflake .  # Only display issues, fixing is manual\n```\n',
     'author': "Peter van 't Hof'",

{pydantic_avro-0.8.1 → pydantic_avro-0.9.1}/src/pydantic_avro/to_avro/types.py

@@ -61,6 +61,8 @@ AVRO_TYPE_MAPPING = {
     },
 }
 
+PRIMITVE_TYPES = ["int", "long", "float", "double", "boolean", "null"]
+
 
 def get_definition(ref: str, schema: dict):
     """Reading definition of base schema for nested structs"""
@@ -133,10 +135,16 @@ class AvroTypeConverter:
         at = field_props.get("avro_type")
         if "allOf" in field_props and len(field_props["allOf"]) == 1:
             r = field_props["allOf"][0]["$ref"]
+        if ("prefixItems" in field_props or ("minItems" in field_props and "maxItems" in field_props)) and t == "array":
+            t = "tuple"
         u = field_props.get("anyOf")
+        o = field_props.get("oneOf")
+        discriminator = field_props.get("discriminator")
 
         if u is not None:
-            return self._union_to_avro(u, avro_type_dict)
+            return self._union_to_avro(u, avro_type_dict, discriminator)
+        elif o is not None:
+            return self._union_to_avro(o, avro_type_dict, discriminator)
         elif r is not None:
             return self._handle_references(r, avro_type_dict)
         elif t is None:
@@ -162,6 +170,8 @@ class AvroTypeConverter:
             avro_type_dict["type"] = "boolean"
         elif t == "null":
             avro_type_dict["type"] = "null"
+        elif t == "tuple":
+            return self._tuple_to_avro(field_props, avro_type_dict)
         else:
             raise NotImplementedError(
                 f"Type '{t}' not support yet, "
@@ -216,34 +226,98 @@ class AvroTypeConverter:
     def _object_to_avro(self, field_props: dict) -> dict:
         """Returns a type of an object field"""
         a = field_props.get("additionalProperties")
-        value_type = "string" if a is None else self._get_avro_type_dict(a)["type"]
+        if not a or isinstance(a, bool):
+            value_type = "string"
+        else:
+            value_type = self._get_avro_type_dict(a)["type"]
         return {"type": "map", "values": value_type}
 
-    def _union_to_avro(self, field_props: list, avro_type_dict: dict) -> dict:
-        """Returns a type of a union field"""
+    def _union_to_avro(self, field_props: list, avro_type_dict: dict, discriminator: Optional[dict] = None) -> dict:
+        """Returns a type of a union field, including discriminated unions"""
+        # Handle discriminated unions
+        if discriminator is not None:
+            return self._discriminated_union_to_avro(field_props, avro_type_dict, discriminator)
+
+        # Standard union handling (unchanged)
         avro_type_dict["type"] = []
         for union_element in field_props:
             t = self._get_avro_type_dict(union_element)
             avro_type_dict["type"].append(t["type"])
         return avro_type_dict
 
+    def _tuple_to_avro(self, field_props: dict, avro_type_dict: dict) -> dict:
+        """Returns a type of a tuple field"""
+        prefix_items = field_props.get("prefixItems")
+        if not prefix_items:
+            # Pydantic v1 there is no prefixItems, but minItems and maxItems and items is a list.
+            prefix_items = field_props.get("items")
+            if not prefix_items:
+                raise ValueError(f"Tuple Field '{field_props}' does not have any items .")
+        possible_types = []
+        for prefix_item in prefix_items:
+            item_type = self._get_avro_type(prefix_item, avro_type_dict).get("type")
+            if not item_type:
+                raise ValueError(f"Field '{avro_type_dict}' does not have a defined type.")
+            if isinstance(item_type, list):
+                possible_types.extend([x for x in item_type if x not in possible_types])
+            elif item_type not in possible_types:
+                possible_types.append(item_type)
+        avro_type_dict["type"] = {"type": "array", "items": possible_types}
+        return avro_type_dict
+
     def _array_to_avro(self, field_props: dict, avro_type_dict: dict) -> dict:
         """Returns a type of an array field"""
-        items = field_props["items"]
+        items = field_props.get("items")
+        # In pydantic v1. items is a list, we need to handle this case first
+        # E.g. [{'type': 'number'}, {'type': 'number'}]}
+        if isinstance(items, list):
+            if not len(items):
+                raise ValueError(f"Field '{field_props}' does not have valid items.")
+            items = items[0]
+        if not isinstance(items, dict):
+            raise ValueError(f"Field '{field_props}' does not have valid items.")
         tn = self._get_avro_type_dict(items)
         # If items in array are an object:
         if "$ref" in items:
             tn = tn["type"]
         # If items in array are a logicalType
-        if (
-            isinstance(tn, dict)
-            and isinstance(tn.get("type", {}), dict)
-            and tn.get("type", {}).get("logicalType") is not None
-        ):
+        if isinstance(tn, dict) and isinstance(tn.get("type", None), (dict, list)):
             tn = tn["type"]
-        # If items in array are an array, the structure must be corrected
-        if isinstance(tn, dict) and isinstance(tn.get("type", {}), dict) and tn.get("type", {}).get("type") == "array":
-            items = tn["type"]["items"]
-            tn = {"type": "array", "items": items}
         avro_type_dict["type"] = {"type": "array", "items": tn}
         return avro_type_dict
+
+    def _discriminated_union_to_avro(self, variants: list, avro_type_dict: dict, discriminator: dict) -> dict:
+        """Handles discriminated unions with a discriminator field"""
+        discriminator_property = discriminator.get("propertyName", "type")
+        variant_types = []
+
+        # Process each variant in the union
+        for variant in variants:
+            # Get the discriminator value for this variant
+            disc_value = None
+            if "properties" in variant and discriminator_property in variant.get("properties", {}):
+                prop = variant["properties"][discriminator_property]
+                if "const" in prop:
+                    disc_value = prop["const"]
+                elif "enum" in prop and len(prop["enum"]) == 1:
+                    disc_value = prop["enum"][0]
+
+            # If we can't determine the discriminator value, generate a regular record
+            if disc_value is None:
+                variant_type = self._get_avro_type_dict(variant)
+            else:
+                # Create a named record for this variant
+                variant_name = f"{disc_value}Variant"
+                if "$ref" in variant:
+                    ref_schema = get_definition(variant["$ref"], self.root_schema)
+                    variant_fields = self.fields_to_avro_dicts(ref_schema)
+                else:
+                    variant_fields = self.fields_to_avro_dicts(variant)
+
+                variant_type = {"type": {"type": "record", "name": variant_name, "fields": variant_fields}}
+
+            variant_types.append(variant_type["type"])
+
+        # Set the union of all variant types
+        avro_type_dict["type"] = variant_types
+        return avro_type_dict