structurize 2.16.2__py3-none-any.whl

avrotize/prototoavro.py

@@ -0,0 +1,382 @@
+"""
+Module to convert Protobuf .proto files to Avro schema.
+"""
+
+import json
+import os
+import re
+from typing import Dict, List, Tuple
+from avrotize.common import pascal
+from avrotize.dependency_resolver import sort_messages_by_dependencies, inline_dependencies_of
+from . import proto2parser
+from . import proto3parser
+
+AvroSchema = Dict[str, 'AvroSchema'] | List['AvroSchema'] | str | None
+
+class ProtoToAvroConverter:
+    """Class to convert Protobuf .proto files to Avro schema."""
+
+    isomorphic_types = ['float', 'double', 'bytes', 'string']
+
+    def __init__(self, proto_root: str = None):
+        """Initialize ProtoToAvroConverter.
+        
+        Args:
+            proto_root (str): Optional root directory for resolving proto imports. 
+                             When provided, imports are resolved relative to this directory.
+        """
+        self.imported_types: Dict[str, str] = {}
+        self.generated_types: Dict[str, str] = {}
+        self.forward_references: Dict[str, str] = {} # table for resolvbing forward references
+        self.proto_root: str = proto_root
+
+    def proto_type_to_avro_primitive(self, proto_type: str)-> Tuple[bool, str]:
+        """
+        Map Protobuf types to Avro primitive types.
+
+        Args:
+            proto_type (str): Protobuf type to convert.
+
+        Returns:
+            str or dict: Corresponding Avro type.
+        """
+        mapping = {
+            'google.protobuf.Empty': 'null',  # Special handling may be required
+            'bool': 'boolean',
+            'int32': 'int',
+            'uint32': 'int',
+            'sint32': 'int',
+            'int64': 'long',
+            'uint64': 'long',
+            'sint64': 'long',
+            'fixed32': 'int',
+            'fixed64': 'long',
+            'sfixed32': 'int',
+            'sfixed64': 'long',
+            'google.protobuf.Timestamp': {
+                "type": "long",
+                "logicalType": "timestamp-micros"
+            }
+        }
+        if proto_type in self.isomorphic_types:
+            return True, proto_type
+        mapped = mapping.get(proto_type, None)
+        if mapped:
+            return True, mapped
+        return False, proto_type
+
+    def build_forward_references_from_message(self, proto_message_type: proto2parser.Message | proto3parser.Message, avro_namespace: str):
+        """
+        Build forward references from a Protobuf message.
+
+        Args:
+            proto_message_type: The message type from the parsed proto file.
+            avro_namespace (str): The namespace for the message.
+        """
+        for _, nested_message in proto_message_type.messages.items():
+            nested_namespace = avro_namespace + '.' + proto_message_type.name + '_types'
+            self.build_forward_references_from_message(nested_message, nested_namespace)
+        for _, enum_type in proto_message_type.enums.items():
+            nested_namespace = avro_namespace + '.' + proto_message_type.name + '_types'
+            self.forward_references[nested_namespace+'.'+enum_type.name] = "enum"
+        self.forward_references[avro_namespace+'.'+proto_message_type.name] = "record"
+
+    def build_forward_references_from_file(self, proto_file:  proto3parser.ProtoFile| proto2parser.ProtoFile, avro_namespace: str):
+        """
+        Build forward references from a Protobuf file.
+
+        Args:
+            proto_file: The parsed proto file.
+            avro_namespace (str): The namespace for the message.
+        """
+        for _, enum_type in proto_file.enums.items():
+             self.forward_references[avro_namespace+'.'+enum_type.name] = "enum"
+        for _, message in proto_file.messages.items():
+            self.build_forward_references_from_message(message, avro_namespace)
+
+    def convert_proto_to_avro_schema(self, proto_file_path: str, avro_namespace: str, message_type: str) -> list:
+        """
+        Convert .proto file to Avro schema.
+
+        Args:
+            proto_file_path (str): Path to the Protobuf .proto file.
+
+        Returns:
+            list: Avro schema as a list of dictionaries.
+        """
+        with open(proto_file_path, 'r', encoding='utf-8') as proto_file:
+            proto_schema = proto_file.read()
+
+        # Determine whether we have proto3 or proto2 and parse the data
+        if re.search(r'syntax\s*=\s*"proto3"', proto_schema):
+            data: proto3parser.ProtoFile = proto3parser.parse(proto_schema)
+        else:
+            data: proto2parser.ProtoFile = proto2parser.parse(proto_schema)
+
+        # Build forward references
+        self.build_forward_references_from_file(data, avro_namespace)
+        # Avro schema header
+        avro_schema = []
+        for import_ in data.imports:
+            # Handle protobuf imports
+            if import_.startswith('google/protobuf/'):
+                script_path = os.path.dirname(os.path.abspath(__file__))
+                avsc_dir = os.path.join(script_path, 'prototypes')
+                # Load the corresponding avsc file from ./prototypes at this script's path into avro_schema
+                avsc = f'{avsc_dir}/{import_.replace("google/protobuf/", "").replace(".proto", ".avsc")}'
+                with open(avsc, 'r', encoding='utf-8') as avsc_file:
+                    types = json.load(avsc_file)
+                    for t in types:
+                        qualified_name = t["namespace"] + "." + t["name"]
+                        self.imported_types[qualified_name] = t
+            else:
+                # Resolve import path: try proto_root first, then fall back to file-relative path
+                import_path = None
+                
+                if self.proto_root:
+                    # Try resolving relative to proto_root
+                    candidate_path = os.path.join(self.proto_root, import_)
+                    if os.path.exists(candidate_path):
+                        import_path = candidate_path
+                
+                if not import_path:
+                    # Fall back to resolving relative to the directory of the current proto file
+                    cwd = os.path.join(os.getcwd(), os.path.dirname(proto_file_path))
+                    candidate_path = os.path.join(cwd, import_)
+                    if os.path.exists(candidate_path):
+                        import_path = candidate_path
+                
+                # Raise an exception if the imported file does not exist
+                if not import_path:
+                    raise FileNotFoundError(f'Import file {import_} does not exist. Searched in proto_root: {self.proto_root}, and relative to: {os.path.dirname(proto_file_path)}')
+                
+                package_name = pascal(import_.replace('.proto', ''))
+                import_namespace = (avro_namespace + '.' + package_name) if avro_namespace else package_name
+                avro_schema.extend(self.convert_proto_to_avro_schema(import_path, import_namespace, message_type))
+
+
+        # Convert enum fields
+        for _, enum_type in data.enums.items():
+            self.handle_enum(enum_type, avro_schema, avro_namespace)
+
+        # Convert message fields
+        for _, m in data.messages.items():
+            self.handle_message(m, avro_schema, avro_namespace)
+
+
+        # Sort the messages in avro_schema by dependencies
+        if message_type:
+            message_schema = next(
+                (message for message in avro_schema if message['type'] == "record" and message['name'] == message_type), None)
+            if not message_schema:
+                raise ValueError(f'Message type {message_type} not found in the Avro schema.')
+            else:
+                inline_dependencies_of(avro_schema, message_schema)
+                return message_schema
+        else:
+            avro_schema = sort_messages_by_dependencies(avro_schema)
+        return avro_schema
+
+    @staticmethod
+    def clean_comment(comment: str):
+        """
+        Clean comments by stripping slashes, newlines, linefeeds, and extra whitespace.
+
+        Args:
+            comment (str): The comment to clean.
+
+        Returns:
+            str: Cleaned comment.
+        """
+        if comment:
+            return comment.replace('//', '').replace('\n', '').lstrip().rstrip()
+        return None
+
+    def handle_enum(self, enum_type: proto2parser.Enum | proto3parser.Enum, avro_schema: AvroSchema, avro_namespace: str) -> AvroSchema:
+        """
+        Convert enum fields to avro schema.
+
+        Args:
+            enum_type: The enum type from the parsed proto file.
+            avro_schema (list): The list to append the converted enum schema.
+            namespace (str): The namespace for the enum.
+        """
+        comment = self.clean_comment(
+            enum_type.comment.content if enum_type.comment and enum_type.comment.content else None)
+
+        # Create avro schema
+        avro_enum: AvroSchema = {
+            'name': enum_type.name,
+            'type': 'enum',
+            'namespace': avro_namespace,
+            'symbols': [],
+            'ordinals': {}
+        }
+
+        if comment:
+            avro_enum['doc'] = comment
+        for value in enum_type.fields:
+            avro_enum['symbols'].append(value.name)
+            avro_enum['ordinals'][value.name] = int(value.number)
+        avro_schema.append(avro_enum)
+        self.generated_types[avro_enum['namespace']+'.'+avro_enum['name']] = "enum"
+        return avro_enum
+
+    def handle_message(self, proto_message_type: proto2parser.Message | proto3parser.Message, avro_schema: AvroSchema, avro_namespace: str)-> AvroSchema:
+        """
+        Convert protobuf messages to avro records.
+
+        Args:
+            m: The message type from the parsed proto file.
+            avro_schema (list): The list to append the converted message schema.
+            namespace (str): The namespace for the message.
+        """
+        dependencies = []
+
+        comment = self.clean_comment(proto_message_type.comment.content if proto_message_type.comment and proto_message_type.comment.content else None)
+        avro_record: AvroSchema = {
+            'type': 'record',
+            'name': proto_message_type.name,
+            'namespace': avro_namespace,
+            'fields': []
+        }
+        if comment:
+            avro_record['doc'] = comment
+        for proto_field in proto_message_type.fields:
+            avro_type = self.get_avro_type_for_field(proto_message_type, avro_namespace, avro_schema, dependencies, proto_field)
+            comment = self.clean_comment(proto_field.comment.content if proto_field.comment and proto_field.comment.content else None)
+
+            avro_field = {
+                'name': proto_field.name,
+                'type': avro_type,
+            }
+
+            if comment:
+                avro_field['doc'] = comment
+
+            avro_record['fields'].append(avro_field)
+
+        for proto_field in proto_message_type.oneofs:
+            avro_oneof: AvroSchema = {
+                'name': proto_field.name,
+                'type': []
+            }
+            comment = self.clean_comment(proto_field.comment.content if proto_field.comment and proto_field.comment.content else None)
+            if comment:
+                avro_oneof['doc'] = comment
+            for oneof_field in proto_field.fields:
+                avro_type = self.get_avro_type_for_field(proto_message_type, avro_namespace, avro_schema, dependencies, oneof_field)
+                comment = self.clean_comment(oneof_field.comment.content if oneof_field.comment and oneof_field.comment.content else None)
+                if comment:
+                    oneof_field['doc'] = comment
+                avro_oneof['type'].append(avro_type)
+            avro_record['fields'].append(avro_oneof)
+
+        if dependencies:
+            avro_record['dependencies'] = dependencies
+        avro_schema.append(avro_record)
+        for _, nested_message in proto_message_type.messages.items():
+            nested_namespace = avro_namespace + '.' + proto_message_type.name + '_types'
+            self.handle_message(nested_message, avro_schema, nested_namespace)
+        # Convert enum fields
+        for _, enum_type in proto_message_type.enums.items():
+            nested_namespace = avro_namespace + '.' + proto_message_type.name + '_types'
+            self.handle_enum(enum_type, avro_schema, nested_namespace)
+        self.generated_types[avro_record['namespace']+'.'+avro_record['name']] = "record"
+        return avro_record
+
+    def get_avro_type_for_field(self, proto_message_type: proto2parser.Message | proto3parser.Message, avro_namespace: str, avro_schema: AvroSchema, dependencies: List[str], proto_field: proto2parser.Field | proto3parser.Field):
+        """
+        Get Avro type for a Protobuf field.
+
+        Args:
+            m: The message type from the parsed proto file.
+            namespace (str): The namespace for the message.
+            dependencies (list): The list to append the dependencies.
+            f: The field from the parsed proto file.
+
+        Returns:
+            str or dict: Corresponding Avro type.
+        """
+        avro_field_type: AvroSchema = None
+        proto_field_type = proto_field.val_type if proto_field.label == 'repeated' or proto_field.type == 'map' else proto_field.type
+        is_primitive, avro_field_type = self.proto_type_to_avro_primitive(proto_field_type)
+
+        if not is_primitive:
+            if proto_field.type in self.imported_types:
+                avro_field_type = self.imported_types[proto_field.type]
+            else:
+                avro_field_type = avro_namespace + '.' + avro_field_type
+                found_in_nested_definitions = False
+                for k, nested_proto_message_type in proto_message_type.messages.items():
+                    nested_namespace = avro_namespace + '.' + proto_message_type.name + '_types'
+                    if nested_proto_message_type.name == proto_field_type:
+                        avro_field_type = self.handle_message(nested_proto_message_type, avro_schema, nested_namespace)
+                        del proto_message_type.messages[k]
+                        if 'dependencies' in avro_field_type:
+                            dependencies.extend(avro_field_type['dependencies'])
+                            del avro_field_type['dependencies']
+                        found_in_nested_definitions = True
+                        break
+                if not found_in_nested_definitions:
+                    for k, nested_proto_enum_type in proto_message_type.enums.items():
+                        nested_namespace = avro_namespace + '.' + proto_message_type.name + '_types'
+                        if nested_proto_enum_type.name == proto_field_type:
+                            avro_field_type = self.handle_enum(nested_proto_enum_type, avro_schema, nested_namespace)
+                            del proto_message_type.enums[k]
+                            found_in_nested_definitions = True
+                            break
+                if not found_in_nested_definitions:
+                    dependency_avro_field_type = avro_field_type
+                    while '.' in dependency_avro_field_type:
+                        if dependency_avro_field_type in self.forward_references:
+                            dependencies.append(dependency_avro_field_type)
+                            break
+                        n = dependency_avro_field_type.split('.')
+                        dependency_avro_field_type = '.'.join(n[:-2]+[n[-1]])
+
+        if proto_field.label == 'optional':
+            avro_field_type = ["null", avro_field_type]
+        if proto_field.label == 'repeated':
+            avro_type: AvroSchema = {
+                "type": "array",
+                "items": avro_field_type
+            }
+        elif proto_field.type == 'map':
+            avro_type: AvroSchema = {
+                "type": "map",
+                "values": avro_field_type,
+            }
+        else:
+            avro_type = avro_field_type
+        return avro_type
+
+
+def convert_proto_to_avro(proto_file_path: str, avro_schema_path: str, namespace: str = None, message_type: str = None, proto_root: str = None):
+    """
+    Convert Protobuf .proto file to Avro schema.
+
+    Args:
+        proto_file_path (str): Path to the Protobuf .proto file.
+        avro_schema_path (str): Path to save the Avro schema .avsc file.
+        namespace (str): Optional namespace for the Avro schema.
+        message_type (str): Optional specific message type to extract.
+        proto_root (str): Optional root directory for resolving proto imports.
+                         When provided, imports are resolved relative to this directory.
+
+    Raises:
+        FileNotFoundError: If the proto file does not exist.
+        ValueError: If the file extensions are incorrect.
+    """
+    if not os.path.exists(proto_file_path):
+        raise FileNotFoundError(f'Proto file {proto_file_path} does not exist.')
+
+    converter = ProtoToAvroConverter(proto_root=proto_root)
+    if not namespace:
+        namespace = pascal(os.path.basename(proto_file_path).replace('.proto', ''))
+    avro_schema = converter.convert_proto_to_avro_schema(proto_file_path, namespace, message_type)
+
+    # Convert the Avro schema to JSON and write it to the file
+    with open(avro_schema_path, 'w', encoding='utf-8') as avro_file:
+        avro_file.write(json.dumps(avro_schema, indent=2))
+