Added support for protocols containing single message types

docs/ba_module.md

@@ -1,5 +1,5 @@
 <!-- THIS FILE IS AUTO GENERATED; DO NOT EDIT BY HAND -->
-<h4><em>last updated on 2021-09-23 for Ballistica version 1.6.5 build 20393</em></h4>
+<h4><em>last updated on 2021-09-24 for Ballistica version 1.6.5 build 20393</em></h4>
 <p>This page documents the Python classes and functions in the 'ba' module,
  which are the ones most relevant to modding in Ballistica. If you come across something you feel should be included here or could be better explained, please <a href="mailto:support@froemling.net">let me know</a>. Happy modding!</p>
 <hr>

tests/test_efro/test_message.py

@@ -72,6 +72,35 @@ class _TResp3(Message):
     fval: float
 
 
+# Generated sender with a single message type:
+# SEND_SINGLE_CODE_TEST_BEGIN
+
+
+class _TestMessageSenderSingle(MessageSender):
+    """Protocol-specific sender."""
+
+    def __init__(self) -> None:
+        protocol = TEST_PROTOCOL_SINGLE
+        super().__init__(protocol)
+
+    def __get__(self,
+                obj: Any,
+                type_in: Any = None) -> _BoundTestMessageSenderSingle:
+        return _BoundTestMessageSenderSingle(obj, self)
+
+
+class _BoundTestMessageSenderSingle(BoundMessageSender):
+    """Protocol-specific bound sender."""
+
+    def send(self, message: _TMsg1) -> _TResp1:
+        """Send a message synchronously."""
+        out = self._sender.send(self._obj, message)
+        assert isinstance(out, _TResp1)
+        return out
+
+
+# SEND_SINGLE_CODE_TEST_END
+
 # Generated sender supporting both sync and async sending:
 # SEND_SYNC_CODE_TEST_BEGIN
 
@@ -205,6 +234,46 @@ class _BoundTestMessageSenderBoth(BoundMessageSender):
 
 # SEND_BOTH_CODE_TEST_END
 
+# Generated receiver with a single message type:
+# RCV_SINGLE_CODE_TEST_BEGIN
+
+
+class _TestSingleMessageReceiver(MessageReceiver):
+    """Protocol-specific synchronous receiver."""
+
+    is_async = False
+
+    def __init__(self) -> None:
+        protocol = TEST_PROTOCOL_SINGLE
+        super().__init__(protocol)
+
+    def __get__(
+        self,
+        obj: Any,
+        type_in: Any = None,
+    ) -> _BoundTestSingleMessageReceiver:
+        return _BoundTestSingleMessageReceiver(obj, self)
+
+    def handler(
+        self,
+        call: Callable[[Any, _TMsg1], _TResp1],
+    ) -> Callable[[Any, _TMsg1], _TResp1]:
+        """Decorator to register message handlers."""
+        from typing import cast, Callable, Any
+        self.register_handler(cast(Callable[[Any, Message], Response], call))
+        return call
+
+
+class _BoundTestSingleMessageReceiver(BoundMessageReceiver):
+    """Protocol-specific bound receiver."""
+
+    def handle_raw_message(self, message: str) -> str:
+        """Synchronously handle a raw incoming message."""
+        return self._receiver.handle_raw_message(self._obj, message)
+
+
+# RCV_SINGLE_CODE_TEST_END
+
 # Generated receiver supporting sync handling:
 # RCV_SYNC_CODE_TEST_BEGIN
 
@@ -334,6 +403,17 @@ TEST_PROTOCOL = MessageProtocol(
     log_remote_exceptions=False,
 )
 
+TEST_PROTOCOL_SINGLE = MessageProtocol(
+    message_types={
+        0: _TMsg1,
+    },
+    response_types={
+        0: _TResp1,
+    },
+    trusted_sender=True,
+    log_remote_exceptions=False,
+)
+
 
 def test_protocol_creation() -> None:
     """Test protocol creation."""
@@ -349,6 +429,39 @@ def test_protocol_creation() -> None:
                                 response_types={0: _TResp1})
 
 
+def test_sender_module_single_embedded() -> None:
+    """Test generation of protocol-specific sender modules for typing/etc."""
+    # NOTE: Ideally we should be testing efro.message.create_sender_module()
+    # here, but it requires us to pass code which imports this test module
+    # to get at the protocol, and that currently fails in our static mypy
+    # tests.
+    smod = TEST_PROTOCOL_SINGLE.do_create_sender_module(
+        'TestMessageSenderSingle',
+        protocol_create_code='protocol = TEST_PROTOCOL_SINGLE',
+        enable_sync_sends=True,
+        enable_async_sends=False,
+        private=True,
+    )
+
+    # Clip everything up to our first class declaration.
+    lines = smod.splitlines()
+    classline = lines.index('class _TestMessageSenderSingle(MessageSender):')
+    clipped = '\n'.join(lines[classline:])
+
+    # This snippet should match what we've got embedded above;
+    # If not then we need to update our embedded version.
+    with open(__file__, encoding='utf-8') as infile:
+        ourcode = infile.read()
+
+    emb = (f'# SEND_SINGLE_CODE_TEST_BEGIN'
+           f'\n\n\n{clipped}\n\n\n# SEND_SINGLE_CODE_TEST_END\n')
+    if emb not in ourcode:
+        print(f'EXPECTED EMBEDDED CODE:\n{emb}')
+        raise RuntimeError('Generated sender module does not match embedded;'
+                           ' test code needs to be updated.'
+                           ' See test stdout for new code.')
+
+
 def test_sender_module_sync_embedded() -> None:
     """Test generation of protocol-specific sender modules for typing/etc."""
     # NOTE: Ideally we should be testing efro.message.create_sender_module()
@@ -448,6 +561,40 @@ def test_sender_module_both_embedded() -> None:
                            ' See test stdout for new code.')
 
 
+def test_receiver_module_single_embedded() -> None:
+    """Test generation of protocol-specific sender modules for typing/etc."""
+    # NOTE: Ideally we should be testing efro.message.create_receiver_module()
+    # here, but it requires us to pass code which imports this test module
+    # to get at the protocol, and that currently fails in our static mypy
+    # tests.
+    smod = TEST_PROTOCOL_SINGLE.do_create_receiver_module(
+        'TestSingleMessageReceiver',
+        'protocol = TEST_PROTOCOL_SINGLE',
+        is_async=False,
+        private=True,
+    )
+
+    # Clip everything up to our first class declaration.
+    lines = smod.splitlines()
+    classline = lines.index(
+        'class _TestSingleMessageReceiver(MessageReceiver):')
+    clipped = '\n'.join(lines[classline:])
+
+    # This snippet should match what we've got embedded above;
+    # If not then we need to update our embedded version.
+    with open(__file__, encoding='utf-8') as infile:
+        ourcode = infile.read()
+
+    emb = (f'# RCV_SINGLE_CODE_TEST_BEGIN'
+           f'\n\n\n{clipped}\n\n\n# RCV_SINGLE_CODE_TEST_END\n')
+    if emb not in ourcode:
+        print(f'EXPECTED SINGLE RECEIVER EMBEDDED CODE:\n{emb}')
+        raise RuntimeError(
+            'Generated single receiver module does not match embedded;'
+            ' test code needs to be updated.'
+            ' See test stdout for new code.')
+
+
 def test_receiver_module_sync_embedded() -> None:
     """Test generation of protocol-specific sender modules for typing/etc."""
     # NOTE: Ideally we should be testing efro.message.create_receiver_module()

tools/efro/message.py

@@ -350,18 +350,12 @@ class MessageProtocol:
             t for t in self.message_ids_by_type if issubclass(t, Message)
         ]
 
-        # Ew; @overload requires at least 2 different signatures so
-        # we need to simply write a single function if we have < 2.
-        if len(msgtypes) == 1:
-            raise RuntimeError('FIXME: currently we require at least 2'
-                               ' registered message types; found 1.')
-
         def _filt_tp_name(rtype: Type[Response]) -> str:
             # We accept None to equal EmptyResponse so reflect that
             # in the type annotation.
             return 'None' if rtype is EmptyResponse else rtype.__name__
 
-        if len(msgtypes) > 1:
+        if msgtypes:
             for async_pass in False, True:
                 if async_pass and not enable_async_sends:
                     continue
@@ -371,7 +365,11 @@ class MessageProtocol:
                 sfx = '_async' if async_pass else ''
                 awt = 'await ' if async_pass else ''
                 how = 'asynchronously' if async_pass else 'synchronously'
-                for msgtype in msgtypes:
+
+                if len(msgtypes) == 1:
+                    # Special case: with a single message types we don't
+                    # use overloads.
+                    msgtype = msgtypes[0]
                     msgtypevar = msgtype.__name__
                     rtypes = msgtype.get_response_types()
                     if len(rtypes) > 1:
@@ -380,17 +378,36 @@ class MessageProtocol:
                     else:
                         rtypevar = _filt_tp_name(rtypes[0])
                     out += (f'\n'
-                            f'    @overload\n'
                             f'    {pfx}def send{sfx}(self,'
                             f' message: {msgtypevar})'
                             f' -> {rtypevar}:\n'
-                            f'        ...\n')
-                out += (f'\n'
-                        f'    {pfx}def send{sfx}(self, message: Message)'
-                        f' -> Optional[Response]:\n'
-                        f'        """Send a message {how}."""\n'
-                        f'        return {awt}self._sender.'
-                        f'send{sfx}(self._obj, message)\n')
+                            f'        """Send a message {how}."""\n'
+                            f'        out = {awt}self._sender.'
+                            f'send{sfx}(self._obj, message)\n'
+                            f'        assert isinstance(out, {rtypevar})\n'
+                            f'        return out\n')
+                else:
+
+                    for msgtype in msgtypes:
+                        msgtypevar = msgtype.__name__
+                        rtypes = msgtype.get_response_types()
+                        if len(rtypes) > 1:
+                            tps = ', '.join(_filt_tp_name(t) for t in rtypes)
+                            rtypevar = f'Union[{tps}]'
+                        else:
+                            rtypevar = _filt_tp_name(rtypes[0])
+                        out += (f'\n'
+                                f'    @overload\n'
+                                f'    {pfx}def send{sfx}(self,'
+                                f' message: {msgtypevar})'
+                                f' -> {rtypevar}:\n'
+                                f'        ...\n')
+                    out += (f'\n'
+                            f'    {pfx}def send{sfx}(self, message: Message)'
+                            f' -> Optional[Response]:\n'
+                            f'        """Send a message {how}."""\n'
+                            f'        return {awt}self._sender.'
+                            f'send{sfx}(self._obj, message)\n')
 
         return out
 
@@ -429,21 +446,18 @@ class MessageProtocol:
             t for t in self.message_ids_by_type if issubclass(t, Message)
         ]
 
-        # Ew; @overload requires at least 2 different signatures so
-        # we need to simply write a single function if we have < 2.
-        if len(msgtypes) == 1:
-            raise RuntimeError('FIXME: currently require at least 2'
-                               ' registered message types; found 1.')
-
         def _filt_tp_name(rtype: Type[Response]) -> str:
             # We accept None to equal EmptyResponse so reflect that
             # in the type annotation.
             return 'None' if rtype is EmptyResponse else rtype.__name__
 
-        if len(msgtypes) > 1:
+        if msgtypes:
             cbgn = 'Awaitable[' if is_async else ''
             cend = ']' if is_async else ''
-            for msgtype in msgtypes:
+            if len(msgtypes) == 1:
+                # Special case: when we have a single message type we don't
+                # use overloads.
+                msgtype = msgtypes[0]
                 msgtypevar = msgtype.__name__
                 rtypes = msgtype.get_response_types()
                 if len(rtypes) > 1:
@@ -454,14 +468,38 @@ class MessageProtocol:
                 rtypevar = f'{cbgn}{rtypevar}{cend}'
                 out += (
                     f'\n'
-                    f'    @overload\n'
                     f'    def handler(\n'
                     f'        self,\n'
                     f'        call: Callable[[Any, {msgtypevar}], '
                     f'{rtypevar}],\n'
-                    f'    ) -> Callable[[Any, {msgtypevar}], {rtypevar}]:\n'
-                    f'        ...\n')
-            out += ('\n'
+                    f'    )'
+                    f' -> Callable[[Any, {msgtypevar}], {rtypevar}]:\n'
+                    f'        """Decorator to register message handlers."""\n'
+                    f'        from typing import cast, Callable, Any\n'
+                    f'        self.register_handler(cast(Callable'
+                    f'[[Any, Message], Response], call))\n'
+                    f'        return call\n')
+            else:
+                for msgtype in msgtypes:
+                    msgtypevar = msgtype.__name__
+                    rtypes = msgtype.get_response_types()
+                    if len(rtypes) > 1:
+                        tps = ', '.join(_filt_tp_name(t) for t in rtypes)
+                        rtypevar = f'Union[{tps}]'
+                    else:
+                        rtypevar = _filt_tp_name(rtypes[0])
+                    rtypevar = f'{cbgn}{rtypevar}{cend}'
+                    out += (f'\n'
+                            f'    @overload\n'
+                            f'    def handler(\n'
+                            f'        self,\n'
+                            f'        call: Callable[[Any, {msgtypevar}], '
+                            f'{rtypevar}],\n'
+                            f'    )'
+                            f' -> Callable[[Any, {msgtypevar}], {rtypevar}]:\n'
+                            f'        ...\n')
+                out += (
+                    '\n'
                     '    def handler(self, call: Callable) -> Callable:\n'
                     '        """Decorator to register message handlers."""\n'
                     '        self.register_handler(call)\n'

tools/efrotools/message.py

@@ -0,0 +1,78 @@
+# Released under the MIT License. See LICENSE for details.
+#
+"""Message related tools functionality."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from efrotools.code import format_yapf_str
+
+if TYPE_CHECKING:
+    from typing import Set, List, Dict, Any, Union, Optional
+
+
+def standard_message_sender_gen_pcommand(
+    projroot: Path,
+    basename: str,
+    source_module: str,
+    enable_sync_sends: bool,
+    enable_async_sends: bool,
+) -> None:
+    """Used by pcommands taking a single filename argument."""
+
+    import efro.message
+    from efro.terminal import Clr
+    from efro.error import CleanError
+
+    if len(sys.argv) != 3:
+        raise CleanError('Expected 1 arg: out-path.')
+
+    dst = sys.argv[2]
+    out = format_yapf_str(
+        projroot,
+        efro.message.create_sender_module(
+            basename,
+            protocol_create_code=(f'from {source_module} import get_protocol\n'
+                                  f'protocol = get_protocol()'),
+            enable_sync_sends=enable_sync_sends,
+            enable_async_sends=enable_async_sends,
+        ))
+
+    print(f'Meta-building {Clr.BLD}{dst}{Clr.RST}')
+    Path(dst).parent.mkdir(parents=True, exist_ok=True)
+    with open(dst, 'w', encoding='utf-8') as outfile:
+        outfile.write(out)
+
+
+def standard_message_receiver_gen_pcommand(
+    projroot: Path,
+    basename: str,
+    source_module: str,
+    is_async: bool,
+) -> None:
+    """Used by pcommands generating efro.message receiver modules."""
+
+    import efro.message
+    from efro.terminal import Clr
+    from efro.error import CleanError
+
+    if len(sys.argv) != 3:
+        raise CleanError('Expected 1 arg: out-path.')
+
+    dst = sys.argv[2]
+    out = format_yapf_str(
+        projroot,
+        efro.message.create_receiver_module(
+            basename,
+            protocol_create_code=(f'from {source_module} import get_protocol\n'
+                                  f'protocol = get_protocol()'),
+            is_async=is_async,
+        ))
+
+    print(f'Meta-building {Clr.BLD}{dst}{Clr.RST}')
+    Path(dst).parent.mkdir(parents=True, exist_ok=True)
+    with open(dst, 'w', encoding='utf-8') as outfile:
+        outfile.write(out)