Change Parallelizable and Collect to structural subtypes

hamilton/htypes.py

@@ -1,8 +1,7 @@
 import inspect
 import sys
 import typing
-from abc import ABC
-from typing import Any, Generator, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Iterable, Optional, Protocol, Tuple, Type, TypeVar, Union
 
 import typing_inspect
 
@@ -137,7 +136,11 @@ def types_match(param_type: Type[Type], required_node_type: Any) -> bool:
 
 
 _sys_version_info = sys.version_info
-_version_tuple = (_sys_version_info.major, _sys_version_info.minor, _sys_version_info.micro)
+_version_tuple = (
+    _sys_version_info.major,
+    _sys_version_info.minor,
+    _sys_version_info.micro,
+)
 
 """
 The following is purely for backwards compatibility
@@ -269,26 +272,43 @@ def get_type_information(some_type: Any) -> Tuple[Type[Type], list]:
 
 
 # Type variables for annotations below
-T = TypeVar("T")
-U = TypeVar("U")
-V = TypeVar("V")
+SequentialElement = TypeVar("SequentialElement", covariant=True)
+ParallelizableElement = TypeVar("ParallelizableElement", covariant=True)
+CollectElement = TypeVar("CollectElement", covariant=True)
 
 
 # TODO -- support sequential operation
-# class Sequential(Generator[T, None, None], ABC):
+# class Sequential(Iterable[SequentialElement], Protocol[SequentialElement]):
 #     pass
 
 
-class Parallelizable(Generator[U, None, None], ABC):
+class Parallelizable(Iterable[ParallelizableElement], Protocol[ParallelizableElement]):
+    """Marks the output of a function node as parallelizable.
+
+    Parallelizable outputs are expected to be iterable, where each element dynamically
+    generates a node. When using dynamic execution, each of these dynamic nodes can be
+    executed in parallel.
+
+    Because this uses dynamic execution, the builder method `enable_dynamic_execution`
+    must be called with `allow_experimental_mode=True`.
+    """
+
     pass
 
 
 def is_parallelizable_type(type_: Type) -> bool:
-    return issubclass(type_, Parallelizable)
+    return _get_origin(type_) == Parallelizable
 
 
-class Collect(Generator[V, None, None], ABC):
-    pass
+class Collect(Iterable[CollectElement], Protocol[CollectElement]):
+    """Marks a function node parameter as collectable.
+
+    Collectable inputs are expected to be iterable, where each element is populated with
+    the results of dynamic nodes derived from parallelizable outputs.
+
+    Because this uses dynamic execution, the builder method `enable_dynamic_execution`
+    must be called with `allow_experimental_mode=True`.
+    """
 
 
 def check_input_type(node_type: Type, input_value: Any) -> bool: