IWF-684 allow None as internal channel prefix type (#93)

* IWF-684 allow `None` as internal channel prefix type

---------

Co-authored-by: Nathan Givens <[email protected]>

Author: N-give

iwf/communication_schema.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Union
 
 
 class CommunicationMethodType(Enum):
@@ -22,15 +22,23 @@ def signal_channel_def(cls, name: str, value_type: type):
         )
 
     @classmethod
-    def internal_channel_def(cls, name: str, value_type: type):
+    def internal_channel_def(cls, name: str, value_type: Union[type, None]):
         return CommunicationMethod(
-            name, CommunicationMethodType.InternalChannel, value_type, False
+            name,
+            CommunicationMethodType.InternalChannel,
+            value_type if value_type is not None else type(None),
+            False,
         )
 
     @classmethod
-    def internal_channel_def_by_prefix(cls, name_prefix: str, value_type: type):
+    def internal_channel_def_by_prefix(
+        cls, name_prefix: str, value_type: Union[type, None]
+    ):
         return CommunicationMethod(
-            name_prefix, CommunicationMethodType.InternalChannel, value_type, True
+            name_prefix,
+            CommunicationMethodType.InternalChannel,
+            value_type if value_type is not None else type(None),
+            True,
         )
 
 

iwf/tests/test_internal_channel.py

@@ -120,7 +120,7 @@ def get_workflow_states(self) -> StateSchema:
     def get_communication_schema(self) -> CommunicationSchema:
         return CommunicationSchema.create(
             CommunicationMethod.internal_channel_def(test_channel_name1, int),
-            CommunicationMethod.internal_channel_def(test_channel_name2, type(None)),
+            CommunicationMethod.internal_channel_def(test_channel_name2, None),
             CommunicationMethod.internal_channel_def(test_channel_name3, int),
             CommunicationMethod.internal_channel_def(test_channel_name4, str),
             CommunicationMethod.internal_channel_def_by_prefix(

iwf/tests/test_internal_channel_with_no_prefix_channel.py

@@ -94,7 +94,7 @@ def get_workflow_states(self) -> StateSchema:
 
     def get_communication_schema(self) -> CommunicationSchema:
         return CommunicationSchema.create(
-            CommunicationMethod.internal_channel_def(internal_channel_name, type(None)),
+            CommunicationMethod.internal_channel_def(internal_channel_name, None),
             # Defining a standard channel (non-prefix) to make sure messages to the channel with a suffix added will not be accepted
             CommunicationMethod.internal_channel_def(test_non_prefix_channel_name, str),
         )

iwf/type_store.py

@@ -23,12 +23,15 @@ def __init__(self, class_type: Type):
         self._prefix_to_type_store = dict()
 
     def is_valid_name_or_prefix(self, name: str) -> bool:
-        t = self._do_get_type(name)
-        return t is not None
+        return self._validate_name(name)
 
     def get_type(self, name: str) -> type:
-        t = self._do_get_type(name)
+        is_registered = self._validate_name(name)
 
+        if not is_registered:
+            raise NotRegisteredError(f"{self._class_type} not registered: {name}")
+
+        t = self._do_get_type(name)
         if t is None:
             raise NotRegisteredError(f"{self._class_type} not registered: {name}")
 
@@ -41,18 +44,26 @@ def add_internal_channel_def(self, obj: CommunicationMethod):
             )
         self._do_add_to_store(obj.is_prefix, obj.name, obj.value_type)
 
-    def _do_get_type(self, name: str) -> Optional[type]:
+    def _validate_name(self, name: str) -> bool:
         if name in self._name_to_type_store:
-            return self._name_to_type_store[name]
+            return True
+
+        for prefix in self._prefix_to_type_store.keys():
+            if name.startswith(prefix):
+                return True
 
-        prefixes = self._prefix_to_type_store.keys()
+        return False
 
-        first = next((prefix for prefix in prefixes if name.startswith(prefix)), None)
+    def _do_get_type(self, name: str) -> Optional[type]:
+        if name in self._name_to_type_store:
+            t = self._name_to_type_store[name]
+            return t if t is not None else type(None)
 
-        if first is None:
-            return None
+        for prefix, t in self._prefix_to_type_store.items():
+            if name.startswith(prefix):
+                return t if t is not None else type(None)
 
-        return self._prefix_to_type_store.get(first, None)
+        return None
 
     def _do_add_to_store(self, is_prefix: bool, name: str, t: Optional[type]):
         if is_prefix: