From 2469570baf02397d28f17e40ddb44a8f3329d56a Mon Sep 17 00:00:00 2001 From: Pierre Trespeuch Date: Fri, 12 Aug 2022 17:23:24 +0200 Subject: [PATCH] Enable arbitrary numbers of construct layers It was not possible until now to have more than 2 layers of constructs. --- src/e3/aws/troposphere/__init__.py | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/src/e3/aws/troposphere/__init__.py b/src/e3/aws/troposphere/__init__.py index b3a224c0..c5492e1e 100644 --- a/src/e3/aws/troposphere/__init__.py +++ b/src/e3/aws/troposphere/__init__.py @@ -1,6 +1,8 @@ from __future__ import annotations from abc import ABC, abstractmethod +from itertools import chain from troposphere import AWSObject, Template + from e3.aws import cfn, name_to_id, Session from e3.aws.cfn.main import CFNMain from e3.aws.troposphere.iam.policy_document import PolicyDocument @@ -85,6 +87,23 @@ def __init__( self.dry_run = dry_run self.template = Template() + def construct_to_objects(self, construct: Construct | AWSObject) -> list[AWSObject]: + """Return list of AWS objects resources from a construct. + + :param construct: construct to list resources from + """ + if isinstance(construct, AWSObject): + return [construct] + else: + return list( + chain.from_iterable( + [ + self.construct_to_objects(el) + for el in construct.resources(stack=self) + ] + ) + ) + def add(self, element: Union[AWSObject, Construct, Stack]) -> Stack: """Add a Construct or AWSObject to the stack. @@ -103,14 +122,7 @@ def add(self, element: Union[AWSObject, Construct, Stack]) -> Stack: # Update the template resources = [] for construct in constructs: - if isinstance(construct, Construct): - for el in construct.resources(stack=self): - if isinstance(el, Construct): - resources.extend(el.resources(stack=self)) - else: - resources.append(el) - if isinstance(construct, AWSObject): - resources.append(construct) + resources.extend(self.construct_to_objects(construct)) self.template.add_resource(resources) return self