Merge pull request #572 from valory-xyz/feature/protocol_union_fixes

protocol generator union support backport

aea/protocols/generator/base.py

@@ -20,6 +20,7 @@
 """This module contains the protocol generator."""
 import itertools
 import os
+import re
 import shutil
 
 # pylint: skip-file
@@ -50,6 +51,7 @@
     _create_protocol_file,
     _get_sub_types_of_compositional_types,
     _includes_custom_type,
+    _is_compositional_type,
     _python_pt_or_ct_type_to_proto_type,
     _to_camel_case,
     _union_sub_type_to_protobuf_variable_name,
@@ -317,8 +319,10 @@ def _to_custom_custom(self, content_type: str) -> str:
         new_content_type = content_type
         if _includes_custom_type(content_type):
             for custom_type in self.spec.all_custom_types:
-                new_content_type = new_content_type.replace(
-                    custom_type, self.spec.custom_custom_types[custom_type]
+                new_content_type = re.sub(
+                    rf"(^|[ \[\,])({custom_type})($|[, \]])",
+                    rf"\g<1>{self.spec.custom_custom_types[custom_type]}\g<3>",
+                    new_content_type,
                 )
         return new_content_type
 
@@ -1360,58 +1364,109 @@ def _custom_types_module_str(self) -> str:
             self._change_indent(-2)
         return cls_str
 
+    def _to_python_type(self, content_type: str) -> str:
+        """
+        Return python type.
+
+        :param  content_type: str
+        :return: str
+        """
+        if not _is_compositional_type(content_type):
+            return content_type
+        m = re.search(r"^(pt:)?([a-zA-Z0-9_]+)(\[.*)?", content_type)
+        if not m:
+            return content_type
+        type_ = m.groups()[1].lower()
+        if type_ == "frozenset":
+            return "(set, frozenset)"
+        if type_ == "tuple":
+            return "(list, tuple)"
+
+        return type_
+
     def _encoding_message_content_from_python_to_protobuf(
         self,
         content_name: str,
         content_type: str,
+        performative_name: Optional[str] = None,
     ) -> str:
         """
         Produce the encoding of message contents for the serialisation class.
 
         :param content_name: the name of the content to be encoded
         :param content_type: the type of the content to be encoded
+        :param performative_name: optional performative name of the content to be encoded
 
         :return: the encoding string
         """
+        performative_name = performative_name or content_name
         encoding_str = ""
         if content_type in PYTHON_TYPE_TO_PROTO_TYPE.keys():
             encoding_str += self.indent + "{} = msg.{}\n".format(
-                content_name, content_name
+                performative_name, content_name
             )
             encoding_str += self.indent + "performative.{} = {}\n".format(
-                content_name, content_name
+                performative_name, performative_name
             )
         elif content_type.startswith("FrozenSet") or content_type.startswith("Tuple"):
             encoding_str += self.indent + "{} = msg.{}\n".format(
                 content_name, content_name
             )
             encoding_str += self.indent + "performative.{}.extend({})\n".format(
-                content_name, content_name
+                performative_name, content_name
             )
         elif content_type.startswith("Dict"):
             encoding_str += self.indent + "{} = msg.{}\n".format(
                 content_name, content_name
             )
             encoding_str += self.indent + "performative.{}.update({})\n".format(
-                content_name, content_name
+                performative_name, content_name
             )
+
         elif content_type.startswith("Union"):
             sub_types = _get_sub_types_of_compositional_types(content_type)
+            encoding_str += self.indent + f'if msg.is_set("{content_name}"):\n'
+            self._change_indent(1)
+            elif_add = ""
             for sub_type in sub_types:
                 sub_type_name_in_protobuf = _union_sub_type_to_protobuf_variable_name(
                     content_name, sub_type
                 )
-                encoding_str += self.indent + 'if msg.is_set("{}"):\n'.format(
-                    sub_type_name_in_protobuf
+                extra = ""
+                if _is_compositional_type(sub_type):
+                    subt = _get_sub_types_of_compositional_types(sub_type)
+                    if "dict" in sub_type.lower():
+                        extra = f" and all(map(lambda x: isinstance(x[0], {subt[0]}) and isinstance(x[1], {subt[1]}), msg.{content_name}.items()))"
+                    else:
+                        extra = f" and all(map(lambda x: isinstance(x, {subt[0]}), msg.{content_name}))"
+                encoding_str += (
+                    self.indent
+                    + f"{elif_add}if isinstance(msg.{content_name}, {self._to_python_type(sub_type)}){extra}:\n"
                 )
                 self._change_indent(1)
                 encoding_str += self.indent + "performative.{}_is_set = True\n".format(
                     sub_type_name_in_protobuf
                 )
                 encoding_str += self._encoding_message_content_from_python_to_protobuf(
-                    sub_type_name_in_protobuf, sub_type
+                    content_name, sub_type, performative_name=sub_type_name_in_protobuf
                 )
                 self._change_indent(-1)
+                elif_add = "el"
+
+            encoding_str += self.indent + f"elif msg.{content_name} is None:\n"
+            self._change_indent(1)
+            encoding_str += self.indent + "pass\n"
+            self._change_indent(-1)
+
+            encoding_str += self.indent + "else:\n"
+            self._change_indent(1)
+            encoding_str += (
+                self.indent
+                + f"raise ValueError(f'Bad value set to `{content_name}` {{msg.{content_name} }}')\n"
+            )
+            self._change_indent(-1)
+
+            self._change_indent(-1)
         elif content_type.startswith("Optional"):
             sub_type = _get_sub_types_of_compositional_types(content_type)[0]
             if not sub_type.startswith("Union"):
@@ -1429,10 +1484,10 @@ def _encoding_message_content_from_python_to_protobuf(
                 self._change_indent(-1)
         else:
             encoding_str += self.indent + "{} = msg.{}\n".format(
-                content_name, content_name
+                performative_name, content_name
             )
             encoding_str += self.indent + "{}.encode(performative.{}, {})\n".format(
-                content_type, content_name, content_name
+                content_type, performative_name, performative_name
             )
         return encoding_str
 
@@ -1474,7 +1529,7 @@ def _decoding_message_content_from_protobuf_to_python(
                 content_name,
                 self.protocol_specification.name,
                 performative,
-                content_name,
+                variable_name,
             )
             decoding_str += self.indent + "{}_frozenset = frozenset({})\n".format(
                 content_name, content_name
@@ -1490,7 +1545,7 @@ def _decoding_message_content_from_protobuf_to_python(
                 content_name,
                 self.protocol_specification.name,
                 performative,
-                content_name,
+                variable_name,
             )
             decoding_str += self.indent + "{}_tuple = tuple({})\n".format(
                 content_name, content_name
@@ -1506,7 +1561,7 @@ def _decoding_message_content_from_protobuf_to_python(
                 content_name,
                 self.protocol_specification.name,
                 performative,
-                content_name,
+                variable_name,
             )
             decoding_str += self.indent + "{}_dict = dict({})\n".format(
                 content_name, content_name
@@ -1836,7 +1891,7 @@ def _content_to_proto_field_str(
                     content_name, sub_type
                 )
                 content_to_proto_field_str, tag_no = self._content_to_proto_field_str(
-                    sub_type_name, sub_type, tag_no
+                    sub_type_name, f"Optional[{sub_type}]", tag_no
                 )
                 entry += content_to_proto_field_str
         elif content_type.startswith("Optional"):  # it is an <O>

aea/protocols/generator/common.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 # ------------------------------------------------------------------------------
 #
-#   Copyright 2022 Valory AG
+#   Copyright 2022-2023 Valory AG
 #   Copyright 2018-2021 Fetch.AI Limited
 #
 #   Licensed under the Apache License, Version 2.0 (the "License");
@@ -288,6 +288,7 @@ def _includes_custom_type(content_type: str) -> bool:
     :param content_type: the content type
     :return: Boolean result
     """
+
     if content_type.startswith("Optional"):
         sub_type = _get_sub_types_of_compositional_types(content_type)[0]
         result = _includes_custom_type(sub_type)
@@ -575,3 +576,17 @@ def apply_protolint(path_to_proto_file: str, name: str) -> Tuple[bool, str]:
                 lines_to_show.append(line)
         error_message = "\n".join(lines_to_show)
         return False, error_message
+
+
+def _is_compositional_type(content_type: str) -> bool:
+    """Checks if content_type is compositional.
+
+    :param content_type: the type string.
+    :return: bool.
+    """
+    for valid_compositional_type in (
+        SPECIFICATION_COMPOSITIONAL_TYPES + PYTHON_COMPOSITIONAL_TYPES
+    ):
+        if content_type.startswith(valid_compositional_type):
+            return True
+    return False

aea/protocols/generator/extract_specification.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 # ------------------------------------------------------------------------------
 #
-#   Copyright 2022 Valory AG
+#   Copyright 2022-2023 Valory AG
 #   Copyright 2018-2021 Fetch.AI Limited
 #
 #   Licensed under the Apache License, Version 2.0 (the "License");
@@ -29,6 +29,7 @@
 from aea.protocols.generator.common import (
     SPECIFICATION_PRIMITIVE_TYPES,
     _get_sub_types_of_compositional_types,
+    _is_compositional_type,
 )
 
 
@@ -213,6 +214,15 @@ def extract(
             if content_type.startswith("ct:"):
                 all_custom_types_set.add(pythonic_content_type)
 
+            for sub_type in (
+                list(_get_sub_types_of_compositional_types(content_type))
+                if _is_compositional_type(content_type)
+                else []
+            ):
+                if sub_type.startswith("ct:"):
+                    pythonic_content_type = _specification_type_to_python_type(sub_type)
+                    all_custom_types_set.add(pythonic_content_type)
+
     # sort the sets
     spec.all_performatives = sorted(all_performatives_set)
     spec.all_custom_types = sorted(all_custom_types_set)

aea/protocols/generator/validate.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 # ------------------------------------------------------------------------------
 #
-#   Copyright 2022 Valory AG
+#   Copyright 2022-2023 Valory AG
 #   Copyright 2018-2021 Fetch.AI Limited
 #
 #   Licensed under the Apache License, Version 2.0 (the "License");
@@ -28,6 +28,7 @@
     SPECIFICATION_PRIMITIVE_TYPES,
     _get_sub_types_of_compositional_types,
     _has_matched_brackets,
+    _is_compositional_type,
 )
 
 
@@ -39,7 +40,10 @@
 PERFORMATIVE_REGEX_PATTERN = "^[a-zA-Z0-9]+$|^[a-zA-Z0-9]+(_?[a-zA-Z0-9]+)+$"
 CONTENT_NAME_REGEX_PATTERN = "^[a-zA-Z0-9]+$|^[a-zA-Z0-9]+(_?[a-zA-Z0-9]+)+$"
 
-CT_CONTENT_TYPE_REGEX_PATTERN = "^ct:([A-Z]+[a-z]*)+$"  # or maybe "ct:(?:[A-Z][a-z]+)+" or # "^ct:[A-Z][a-zA-Z0-9]*$"
+CT_NAME_RE = (
+    "[A-Z][a-zA-Z0-9]*"  # or maybe "ct:(?:[A-Z][a-z]+)+" or # "^ct:([A-Z]+[a-z]*)+$"
+)
+CT_CONTENT_TYPE_REGEX_PATTERN = f"^ct:{CT_NAME_RE}$"
 
 ROLE_REGEX_PATTERN = "^[a-zA-Z0-9]+$|^[a-zA-Z0-9]+(_?[a-zA-Z0-9]+)+$"
 END_STATE_REGEX_PATTERN = "^[a-zA-Z0-9]+$|^[a-zA-Z0-9]+(_?[a-zA-Z0-9]+)+$"
@@ -226,9 +230,9 @@ def _is_valid_union(content_type: str) -> bool:
         if not (
             _is_valid_ct(sub_type)
             or _is_valid_pt(sub_type)
-            or _is_valid_set(sub_type)
-            or _is_valid_list(sub_type)
             or _is_valid_dict(sub_type)
+            or _is_valid_list(sub_type)
+            or _is_valid_set(sub_type)
         ):
             return False
 
@@ -481,8 +485,13 @@ def _validate_speech_acts_section(
 
             content_names_types[content_name] = (performative, content_type)
 
-            if _is_valid_ct(content_type):
-                custom_types_set.add(content_type.strip())
+            for sub_type in (
+                list(_get_sub_types_of_compositional_types(content_type))
+                if _is_compositional_type(content_type)
+                else []
+            ) + [content_type]:
+                if _is_valid_ct(sub_type):
+                    custom_types_set.add(sub_type.strip())
 
     return True, "Speech-acts are valid.", performatives_set, custom_types_set