Make ContainerSet full Hera object

docs/examples/workflows/container_set_with_env.md

@@ -0,0 +1,86 @@
+# Container Set With Env
+
+
+
+
+
+
+=== "Hera"
+
+    ```python linenums="1"
+    from hera.workflows import (
+        ConfigMapEnv,
+        ConfigMapEnvFrom,
+        ContainerNode,
+        ContainerSet,
+        Env,
+        ResourceEnv,
+        SecretEnv,
+        SecretEnvFrom,
+        Workflow,
+    )
+
+    with Workflow(generate_name="secret-env-from-", entrypoint="whalesay") as w:
+        with ContainerSet(name="whalesay"):
+            ContainerNode(
+                name="node",
+                image="docker/whalesay:latest",
+                command=["cowsay"],
+                env_from=[
+                    SecretEnvFrom(prefix="abc", name="secret", optional=False),
+                    ConfigMapEnvFrom(prefix="abc", name="configmap", optional=False),
+                ],
+                env=[
+                    Env(name="test", value="1"),
+                    SecretEnv(name="s1", secret_key="s1", secret_name="abc"),
+                    ResourceEnv(name="r1", resource="abc"),
+                    ConfigMapEnv(name="c1", config_map_key="c1", config_map_name="abc"),
+                ],
+            )
+    ```
+
+=== "YAML"
+
+    ```yaml linenums="1"
+    apiVersion: argoproj.io/v1alpha1
+    kind: Workflow
+    metadata:
+      generateName: secret-env-from-
+    spec:
+      entrypoint: whalesay
+      templates:
+      - containerSet:
+          containers:
+          - command:
+            - cowsay
+            env:
+            - name: test
+              value: '1'
+            - name: s1
+              valueFrom:
+                secretKeyRef:
+                  key: s1
+                  name: abc
+            - name: r1
+              valueFrom:
+                resourceFieldRef:
+                  resource: abc
+            - name: c1
+              valueFrom:
+                configMapKeyRef:
+                  key: c1
+                  name: abc
+            envFrom:
+            - prefix: abc
+              secretRef:
+                name: secret
+                optional: false
+            - configMapRef:
+                name: configmap
+                optional: false
+              prefix: abc
+            image: docker/whalesay:latest
+            name: node
+        name: whalesay
+    ```
+

examples/workflows/container-set-with-env.yaml

@@ -0,0 +1,40 @@
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+metadata:
+  generateName: secret-env-from-
+spec:
+  entrypoint: whalesay
+  templates:
+  - containerSet:
+      containers:
+      - command:
+        - cowsay
+        env:
+        - name: test
+          value: '1'
+        - name: s1
+          valueFrom:
+            secretKeyRef:
+              key: s1
+              name: abc
+        - name: r1
+          valueFrom:
+            resourceFieldRef:
+              resource: abc
+        - name: c1
+          valueFrom:
+            configMapKeyRef:
+              key: c1
+              name: abc
+        envFrom:
+        - prefix: abc
+          secretRef:
+            name: secret
+            optional: false
+        - configMapRef:
+            name: configmap
+            optional: false
+          prefix: abc
+        image: docker/whalesay:latest
+        name: node
+    name: whalesay

examples/workflows/container_set_with_env.py

@@ -0,0 +1,29 @@
+from hera.workflows import (
+    ConfigMapEnv,
+    ConfigMapEnvFrom,
+    ContainerNode,
+    ContainerSet,
+    Env,
+    ResourceEnv,
+    SecretEnv,
+    SecretEnvFrom,
+    Workflow,
+)
+
+with Workflow(generate_name="secret-env-from-", entrypoint="whalesay") as w:
+    with ContainerSet(name="whalesay"):
+        ContainerNode(
+            name="node",
+            image="docker/whalesay:latest",
+            command=["cowsay"],
+            env_from=[
+                SecretEnvFrom(prefix="abc", name="secret", optional=False),
+                ConfigMapEnvFrom(prefix="abc", name="configmap", optional=False),
+            ],
+            env=[
+                Env(name="test", value="1"),
+                SecretEnv(name="s1", secret_key="s1", secret_name="abc"),
+                ResourceEnv(name="r1", resource="abc"),
+                ConfigMapEnv(name="c1", config_map_key="c1", config_map_name="abc"),
+            ],
+        )

src/hera/workflows/_mixins.py

@@ -230,7 +230,7 @@ def _build_env(self) -> Optional[List[EnvVar]]:
         for e in env:
             if isinstance(e, EnvVar):
                 result.append(e)
-            elif isinstance(e, _BaseEnv):
+            elif issubclass(e.__class__, _BaseEnv):
                 result.append(e.build())
             elif isinstance(e, dict):
                 for k, v in e.items():
@@ -246,7 +246,7 @@ def _build_env_from(self) -> Optional[List[EnvFromSource]]:
         for e in env_from:
             if isinstance(e, EnvFromSource):
                 result.append(e)
-            elif isinstance(e, _BaseEnvFrom):
+            elif issubclass(e.__class__, _BaseEnvFrom):
                 result.append(e.build())
         return result
 

src/hera/workflows/container_set.py

@@ -7,6 +7,7 @@
     ContainerMixin,
     ContextMixin,
     EnvIOMixin,
+    EnvMixin,
     ResourceMixin,
     SubNodeMixin,
     TemplateMixin,
@@ -17,12 +18,40 @@
     ContainerNode as _ModelContainerNode,
     ContainerSetRetryStrategy,
     ContainerSetTemplate as _ModelContainerSetTemplate,
+    Lifecycle,
+    SecurityContext,
     Template as _ModelTemplate,
 )
 
 
-class ContainerNode(_ModelContainerNode, SubNodeMixin):
+class ContainerNode(ContainerMixin, VolumeMountMixin, ResourceMixin, EnvMixin, SubNodeMixin):
+    """A regular container that can be used as part of a `hera.workflows.ContainerSet`.
+
+    See Also
+    --------
+    https://argoproj.github.io/argo-workflows/container-set-template/
+    """
+
+    name: str
+    args: Optional[List[str]] = None
+    command: Optional[List[str]] = None
+    dependencies: Optional[List[str]] = None
+    lifecycle: Optional[Lifecycle] = None
+    security_context: Optional[SecurityContext] = None
+    working_dir: Optional[str] = None
+
     def next(self, other: ContainerNode) -> ContainerNode:
+        """Sets the given container as a dependency of this container and returns the given container.
+
+        Examples
+        --------
+        >>> from hera.workflows import ContainerNode
+        >>> # normally, you use the following within a `hera.workflows.ContainerSet` context.
+        >>> a, b = ContainerNode(name="a"), ContainerNode(name="b")
+        >>> a.next(b)
+        >>> b.dependencies
+        ['a']
+        """
         assert issubclass(other.__class__, ContainerNode)
         if other.dependencies is None:
             other.dependencies = [self.name]
@@ -32,6 +61,19 @@ def next(self, other: ContainerNode) -> ContainerNode:
         return other
 
     def __rrshift__(self, other: List[ContainerNode]) -> ContainerNode:
+        """Sets `self` as a dependent of the given list of other `hera.workflows.ContainerNode`.
+
+        Practically, the `__rrshift__` allows us to express statements such as `[a, b, c] >> d`, where `d` is `self.`
+
+        Examples
+        --------
+        >>> from hera.workflows import ContainerNode
+        >>> # normally, you use the following within a `hera.workflows.ContainerSet` context.
+        >>> a, b, c = ContainerNode(name="a"), ContainerNode(name="b"), ContainerNode(name="c")
+        >>> [a, b]
+        >>> c.dependencies
+        ['a', 'b']
+        """
         assert isinstance(other, list), f"Unknown type {type(other)} specified using reverse right bitshift operator"
         for o in other:
             o.next(self)
@@ -40,6 +82,17 @@ def __rrshift__(self, other: List[ContainerNode]) -> ContainerNode:
     def __rshift__(
         self, other: Union[ContainerNode, List[ContainerNode]]
     ) -> Union[ContainerNode, List[ContainerNode]]:
+        """Sets the given container as a dependency of this container and returns the given container.
+
+        Examples
+        --------
+        >>> from hera.workflows import ContainerNode
+        >>> # normally, you use the following within a `hera.workflows.ContainerSet` context.
+        >>> a, b = ContainerNode(name="a"), ContainerNode(name="b")
+        >>> a >> b
+        >>> b.dependencies
+        ['a']
+        """
         if isinstance(other, ContainerNode):
             return self.next(other)
         elif isinstance(other, list):
@@ -51,6 +104,33 @@ def __rshift__(
             return other
         raise ValueError(f"Unknown type {type(other)} provided to `__rshift__`")
 
+    def _build_container_node(self) -> _ModelContainerNode:
+        return _ModelContainerNode(
+            args=self.args,
+            command=self.command,
+            dependencies=self.dependencies,
+            env=self._build_env(),
+            env_from=self._build_env_from(),
+            image=self.image,
+            image_pull_policy=self._build_image_pull_policy(),
+            lifecycle=self.lifecycle,
+            liveness_probe=self.liveness_probe,
+            name=self.name,
+            ports=self.ports,
+            readiness_probe=self.readiness_probe,
+            resources=self._build_resources(),
+            security_context=self.security_context,
+            startup_probe=self.startup_probe,
+            stdin=self.stdin,
+            stdin_once=self.stdin_once,
+            termination_message_path=self.termination_message_path,
+            termination_message_policy=self.termination_message_policy,
+            tty=self.tty,
+            volume_devices=self.volume_devices,
+            volume_mounts=self._build_volume_mounts(),
+            working_dir=self.working_dir,
+        )
+
 
 class ContainerSet(
     EnvIOMixin,
@@ -61,7 +141,7 @@ class ContainerSet(
     VolumeMountMixin,
     ContextMixin,
 ):
-    containers: List[ContainerNode] = []
+    containers: List[Union[ContainerNode, _ModelContainerNode]] = []
     container_set_retry_strategy: Optional[ContainerSetRetryStrategy] = None
 
     def _add_sub(self, node: Any):
@@ -71,8 +151,9 @@ def _add_sub(self, node: Any):
         self.containers.append(node)
 
     def _build_container_set(self) -> _ModelContainerSetTemplate:
+        containers = [c._build_container_node() if isinstance(c, ContainerNode) else c for c in self.containers]
         return _ModelContainerSetTemplate(
-            containers=self.containers,
+            containers=containers,
             retry_strategy=self.container_set_retry_strategy,
             volume_mounts=self.volume_mounts,
         )