Repository: getmoto/moto
Pull request #5854 (2023-01-18 21:49:16)

Techdebt: MyPy EBS


Patch:
diff --git a/moto/dynamodb/parsing/key_condition_expression.py b/moto/dynamodb/parsing/key_condition_expression.py
index cb9953450..3bdbaa7e1 100644
--- a/moto/dynamodb/parsing/key_condition_expression.py
+++ b/moto/dynamodb/parsing/key_condition_expression.py
@@ -1,70 +1,6 @@
 from enum import Enum
 from moto.dynamodb.exceptions import MockValidationException
-
-
-class KeyConditionExpressionTokenizer:
-    """
-    Tokenizer for a KeyConditionExpression. Should be used as an iterator.
-    The final character to be returned will be an empty string, to notify the caller that we've reached the end.
-    """
-
-    def __init__(self, expression):
-        self.expression = expression
-        self.token_pos = 0
-
-    def __iter__(self):
-        self.token_pos = 0
-        return self
-
-    def is_eof(self):
-        return self.peek() == ""
-
-    def peek(self):
-        """
-        Peek the next character without changing the position
-        """
-        try:
-            return self.expression[self.token_pos]
-        except IndexError:
-            return ""
-
-    def __next__(self):
-        """
-        Returns the next character, or an empty string if we've reached the end of the string.
-        Calling this method again will result in a StopIterator
-        """
-        try:
-            result = self.expression[self.token_pos]
-            self.token_pos += 1
-            return result
-        except IndexError:
-            if self.token_pos == len(self.expression):
-                self.token_pos += 1
-                return ""
-            raise StopIteration
-
-    def skip_characters(self, phrase, case_sensitive=False) -> None:
-        """
-        Skip the characters in the supplied phrase.
-        If any other character is encountered instead, this will fail.
-        If we've already reached the end of the iterator, this will fail.
-        """
-        for ch in phrase:
-            if case_sensitive:
-                assert self.expression[self.token_pos] == ch
-            else:
-                assert self.expression[self.token_pos] in [ch.lower(), ch.upper()]
-            self.token_pos += 1
-
-    def skip_white_space(self):
-        """
-        Skip the any whitespace characters that are coming up
-        """
-        try:
-            while self.peek() == " ":
-                self.token_pos += 1
-        except IndexError:
-            pass
+from moto.utilities.tokenizer import GenericTokenizer
 
 
 class EXPRESSION_STAGES(Enum):
@@ -100,7 +36,7 @@ def parse_expression(
     key_name = comparison = None
     key_values = []
     results = []
-    tokenizer = KeyConditionExpressionTokenizer(key_condition_expression)
+    tokenizer = GenericTokenizer(key_condition_expression)
     for crnt_char in tokenizer:
         if crnt_char == " ":
             if current_stage == EXPRESSION_STAGES.INITIAL_STAGE:
diff --git a/moto/ebs/models.py b/moto/ebs/models.py
index 58f05e313..e8ac96437 100644
--- a/moto/ebs/models.py
+++ b/moto/ebs/models.py
@@ -2,13 +2,17 @@
 
 from moto.core import BaseBackend, BackendDict, BaseModel
 from moto.core.utils import unix_time
-from moto.ec2 import ec2_backends
+from moto.ec2.models import ec2_backends, EC2Backend
 from moto.ec2.models.elastic_block_store import Snapshot
 from moto.moto_api._internal import mock_random
 
+from typing import Any, Dict, List, Optional, Tuple
+
 
 class Block(BaseModel):
-    def __init__(self, block_data, checksum, checksum_algorithm, data_length):
+    def __init__(
+        self, block_data: str, checksum: str, checksum_algorithm: str, data_length: str
+    ):
         self.block_data = block_data
         self.checksum = checksum
         self.checksum_algorithm = checksum_algorithm
@@ -17,7 +21,7 @@ class Block(BaseModel):
 
 
 class EBSSnapshot(BaseModel):
-    def __init__(self, account_id, snapshot: Snapshot):
+    def __init__(self, account_id: str, snapshot: Snapshot):
         self.account_id = account_id
         self.snapshot_id = snapshot.id
         self.status = "pending"
@@ -29,15 +33,20 @@ class EBSSnapshot(BaseModel):
         ]
         self.description = snapshot.description
 
-        self.blocks = dict()
+        self.blocks: Dict[str, Block] = dict()
 
     def put_block(
-        self, block_idx, block_data, checksum, checksum_algorithm, data_length
-    ):
+        self,
+        block_idx: str,
+        block_data: str,
+        checksum: str,
+        checksum_algorithm: str,
+        data_length: str,
+    ) -> None:
         block = Block(block_data, checksum, checksum_algorithm, data_length)
         self.blocks[block_idx] = block
 
-    def to_json(self):
+    def to_json(self) -> Dict[str, Any]:
         return {
             "SnapshotId": self.snapshot_id,
             "OwnerId": self.account_id,
@@ -53,60 +62,65 @@ class EBSSnapshot(BaseModel):
 class EBSBackend(BaseBackend):
     """Implementation of EBS APIs."""
 
-    def __init__(self, region_name, account_id):
+    def __init__(self, region_name: str, account_id: str):
         super().__init__(region_name, account_id)
-        self.snapshots = dict()
+        self.snapshots: Dict[str, EBSSnapshot] = dict()
 
     @property
-    def ec2_backend(self):
+    def ec2_backend(self) -> EC2Backend:
         return ec2_backends[self.account_id][self.region_name]
 
-    def start_snapshot(self, volume_size, tags, description):
+    def start_snapshot(
+        self, volume_size: str, tags: Optional[List[Dict[str, str]]], description: str
+    ) -> EBSSnapshot:
         zone_name = f"{self.region_name}a"
         vol = self.ec2_backend.create_volume(size=volume_size, zone_name=zone_name)
         snapshot = self.ec2_backend.create_snapshot(
             volume_id=vol.id, description=description
         )
         if tags:
-            tags = {tag["Key"]: tag["Value"] for tag in tags}
-            snapshot.add_tags(tags)
+            snapshot.add_tags({tag["Key"]: tag["Value"] for tag in tags})
         ebs_snapshot = EBSSnapshot(account_id=self.account_id, snapshot=snapshot)
         self.snapshots[ebs_snapshot.snapshot_id] = ebs_snapshot
         return ebs_snapshot
 
-    def complete_snapshot(self, snapshot_id):
+    def complete_snapshot(self, snapshot_id: str) -> Dict[str, str]:
         self.snapshots[snapshot_id].status = "completed"
         return {"Status": "completed"}
 
     def put_snapshot_block(
         self,
-        snapshot_id,
-        block_index,
-        block_data,
-        checksum,
-        checksum_algorithm,
-        data_length,
-    ):
+        snapshot_id: str,
+        block_index: str,
+        block_data: str,
+        checksum: str,
+        checksum_algorithm: str,
+        data_length: str,
+    ) -> Tuple[str, str]:
         snapshot = self.snapshots[snapshot_id]
         snapshot.put_block(
             block_index, block_data, checksum, checksum_algorithm, data_length
         )
         return checksum, checksum_algorithm
 
-    def get_snapshot_block(self, snapshot_id, block_index):
+    def get_snapshot_block(self, snapshot_id: str, block_index: str) -> Block:
         """
         The BlockToken-parameter is not yet implemented
         """
         snapshot = self.snapshots[snapshot_id]
         return snapshot.blocks[block_index]
 
-    def list_changed_blocks(self, first_snapshot_id, second_snapshot_id):
+    def list_changed_blocks(
+        self, first_snapshot_id: str, second_snapshot_id: str
+    ) -> Tuple[Dict[str, Tuple[str, Optional[str]]], EBSSnapshot]:
         """
         The following parameters are not yet implemented: NextToken, MaxResults, StartingBlockIndex
         """
         snapshot1 = self.snapshots[first_snapshot_id]
         snapshot2 = self.snapshots[second_snapshot_id]
-        changed_blocks = dict()  # {idx: (token1, token2), ..}
+        changed_blocks: Dict[
+            str, Tuple[str, Optional[str]]
+        ] = dict()  # {idx: (token1, token2), ..}
         for idx in snapshot1.blocks:
             block1 = snapshot1.blocks[idx]
             if idx in snapshot2.blocks:
@@ -118,7 +132,7 @@ class EBSBackend(BaseBackend):
 
         return changed_blocks, snapshot1
 
-    def list_snapshot_blocks(self, snapshot_id):
+    def list_snapshot_blocks(self, snapshot_id: str) -> EBSSnapshot:
         return self.snapshots[snapshot_id]
 
 
diff --git a/moto/ebs/responses.py b/moto/ebs/responses.py
index 4f301a415..75d1549f9 100644
--- a/moto/ebs/responses.py
+++ b/moto/ebs/responses.py
@@ -1,39 +1,41 @@
 """Handles incoming ebs requests, invokes methods, returns responses."""
 import json
+from typing import Any
 
 from moto.core.responses import BaseResponse
-from .models import ebs_backends
+from moto.core.common_types import TYPE_RESPONSE
+from .models import ebs_backends, EBSBackend
 
 
 class EBSResponse(BaseResponse):
     """Handler for EBS requests and responses."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(service_name="ebs")
 
     @property
-    def ebs_backend(self):
+    def ebs_backend(self) -> EBSBackend:
         """Return backend instance specific for this region."""
         return ebs_backends[self.current_account][self.region]
 
-    def snapshots(self, request, full_url, headers):
+    def snapshots(self, request: Any, full_url: str, headers: Any) -> TYPE_RESPONSE:  # type: ignore[return]
         self.setup_class(request, full_url, headers)
         if request.method == "POST":
             return self.start_snapshot()
 
-    def snapshot_block(self, request, full_url, headers):
+    def snapshot_block(self, request: Any, full_url: str, headers: Any) -> TYPE_RESPONSE:  # type: ignore[return]
         self.setup_class(request, full_url, headers)
         if request.method == "PUT":
             return self.put_snapshot_block(full_url, headers)
         if request.method == "GET":
             return self.get_snapshot_block()
 
-    def snapshot_blocks(self, request, full_url, headers):
+    def snapshot_blocks(self, request: Any, full_url: str, headers: Any) -> TYPE_RESPONSE:  # type: ignore[return]
         self.setup_class(request, full_url, headers)
         if request.method == "GET":
             return self.list_snapshot_blocks()
 
-    def start_snapshot(self):
+    def start_snapshot(self) -> TYPE_RESPONSE:
         """
         The following parameters are not yet implemented: ParentSnapshotId, ClientToken, Encrypted, KmsKeyArn, Timeout
         """
@@ -48,7 +50,9 @@ class EBSResponse(BaseResponse):
         )
         return 200, {}, json.dumps(snapshot.to_json())
 
-    def complete_snapshot(self, request, full_url, headers):
+    def complete_snapshot(
+        self, request: Any, full_url: str, headers: Any
+    ) -> TYPE_RESPONSE:
         """
         The following parameters are not yet supported: ChangedBlocksCount, Checksum, ChecksumAlgorithm, ChecksumAggregationMethod
         """
@@ -57,7 +61,7 @@ class EBSResponse(BaseResponse):
         status = self.ebs_backend.complete_snapshot(snapshot_id=snapshot_id)
         return 200, {}, json.dumps(status)
 
-    def put_snapshot_block(self, full_url, headers):
+    def put_snapshot_block(self, full_url: str, headers: Any) -> TYPE_RESPONSE:
         """
         The following parameters are currently not taken into account: DataLength, Progress.
         The Checksum and ChecksumAlgorithm are taken at face-value, but no validation takes place.
@@ -86,7 +90,7 @@ class EBSResponse(BaseResponse):
             "{}",
         )
 
-    def get_snapshot_block(self):
+    def get_snapshot_block(self) -> TYPE_RESPONSE:
         snapshot_id = self.path.split("/")[-3]
         block_index = self.path.split("/")[-1]
         block = self.ebs_backend.get_snapshot_block(
@@ -100,12 +104,14 @@ class EBSResponse(BaseResponse):
         }
         return 200, headers, block.block_data
 
-    def snapshot_changed_blocks(self, request, full_url, headers):
+    def snapshot_changed_blocks(
+        self, request: Any, full_url: str, headers: Any
+    ) -> TYPE_RESPONSE:
         self.setup_class(request, full_url, headers)
         first_snapshot_id = self._get_params().get("firstSnapshotId")
         second_snapshot_id = self.path.split("/")[-2]
         changed_blocks, snapshot = self.ebs_backend.list_changed_blocks(
-            first_snapshot_id=first_snapshot_id,
+            first_snapshot_id=first_snapshot_id,  # type: ignore[arg-type]
             second_snapshot_id=second_snapshot_id,
         )
         blocks = [
@@ -124,7 +130,7 @@ class EBSResponse(BaseResponse):
             ),
         )
 
-    def list_snapshot_blocks(self):
+    def list_snapshot_blocks(self) -> TYPE_RESPONSE:
         """
         The following parameters are not yet implemented: NextToken, MaxResults, StartingBlockIndex
         """
diff --git a/moto/ec2/models/core.py b/moto/ec2/models/core.py
index 6712b84b3..5a7c6c9f2 100644
--- a/moto/ec2/models/core.py
+++ b/moto/ec2/models/core.py
@@ -1,10 +1,11 @@
+from typing import Dict, List
 from moto.core import BaseModel
 
 from ..exceptions import FilterNotImplementedError
 
 
 class TaggedEC2Resource(BaseModel):
-    def get_tags(self):
+    def get_tags(self) -> List[Dict[str, str]]:
         tags = []
         if self.id:
             tags = self.ec2_backend.describe_tags(filters={"resource-id": [self.id]})
@@ -13,7 +14,7 @@ class TaggedEC2Resource(BaseModel):
     def add_tag(self, key, value):
         self.ec2_backend.create_tags([self.id], {key: value})
 
-    def add_tags(self, tag_map):
+    def add_tags(self, tag_map: Dict[str, str]):
         for key, value in tag_map.items():
             self.ec2_backend.create_tags([self.id], {key: value})
 
diff --git a/moto/ec2/models/elastic_block_store.py b/moto/ec2/models/elastic_block_store.py
index 33148d4bd..c2cff2478 100644
--- a/moto/ec2/models/elastic_block_store.py
+++ b/moto/ec2/models/elastic_block_store.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 from moto.core import CloudFormationModel
 from moto.packages.boto.ec2.blockdevicemapping import BlockDeviceType
 from ..exceptions import (
@@ -240,14 +242,14 @@ class EBSBackend:
 
     def create_volume(
         self,
-        size,
-        zone_name,
-        snapshot_id=None,
-        encrypted=False,
-        kms_key_id=None,
-        volume_type=None,
-        iops=None,
-    ):
+        size: str,
+        zone_name: str,
+        snapshot_id: Optional[str] = None,
+        encrypted: bool = False,
+        kms_key_id: Optional[str] = None,
+        volume_type: Optional[str] = None,
+        iops: Optional[str] = None,
+    ) -> Volume:
         if kms_key_id and not encrypted:
             raise InvalidParameterDependency("KmsKeyId", "Encrypted")
         if encrypted and not kms_key_id:
@@ -360,7 +362,13 @@ class EBSBackend:
         volume.attachment = None
         return old_attachment
 
-    def create_snapshot(self, volume_id, description, owner_id=None, from_ami=None):
+    def create_snapshot(
+        self,
+        volume_id: str,
+        description: str,
+        owner_id: Optional[str] = None,
+        from_ami: Optional[str] = None,
+    ) -> Snapshot:
         snapshot_id = random_snapshot_id()
         volume = self.get_volume(volume_id)
         params = [self, snapshot_id, volume, description, volume.encrypted]
diff --git a/moto/ses/models.py b/moto/ses/models.py
index 7227abe12..27e1e5fef 100644
--- a/moto/ses/models.py
+++ b/moto/ses/models.py
@@ -1,4 +1,3 @@
-import re
 import json
 import email
 import datetime
@@ -23,9 +22,9 @@ from .exceptions import (
     RuleSetNameAlreadyExists,
     RuleSetDoesNotExist,
     RuleAlreadyExists,
-    MissingRenderingAttributeException,
     ConfigurationSetAlreadyExists,
 )
+from .template import parse_template
 from .utils import get_random_message_id, is_valid_address
 from .feedback import COMMON_MAIL, BOUNCE, COMPLAINT, DELIVERY
 
@@ -106,17 +105,6 @@ class SESQuota(BaseModel):
         return self.sent
 
 
-def are_all_variables_present(template, template_data):
-    subject_part = template["subject_part"]
-    text_part = template["text_part"]
-    html_part = template["html_part"]
-
-    for var in re.findall("{{(.+?)}}", subject_part + text_part + html_part):
-        if not template_data.get(var):
-            return var, False
-    return None, True
-
-
 class SESBackend(BaseBackend):
     """
     Responsible for mocking calls to SES.
@@ -469,18 +457,13 @@ class SESBackend(BaseBackend):
                 "Template rendering data is invalid"
             )
 
-        var, are_variables_present = are_all_variables_present(template, template_data)
-        if not are_variables_present:
-            raise MissingRenderingAttributeException(var)
-
         subject_part = template["subject_part"]
         text_part = template["text_part"]
         html_part = template["html_part"]
 
-        for key, value in template_data.items():
-            subject_part = str.replace(str(subject_part), "{{" + key + "}}", value)
-            text_part = str.replace(str(text_part), "{{" + key + "}}", value)
-            html_part = str.replace(str(html_part), "{{" + key + "}}", value)
+        subject_part = parse_template(str(subject_part), template_data)
+        text_part = parse_template(str(text_part), template_data)
+        html_part = parse_template(str(html_part), template_data)
 
         email_obj = MIMEMultipart("alternative")
 
diff --git a/moto/ses/template.py b/moto/ses/template.py
new file mode 100644
index 000000000..ec34cd19c
--- /dev/null
+++ b/moto/ses/template.py
@@ -0,0 +1,45 @@
+from .exceptions import MissingRenderingAttributeException
+from moto.utilities.tokenizer import GenericTokenizer
+
+
+def parse_template(template, template_data, tokenizer=None, until=None):
+    tokenizer = tokenizer or GenericTokenizer(template)
+    state = ""
+    parsed = ""
+    var_name = ""
+    for char in tokenizer:
+        if until is not None and (char + tokenizer.peek(len(until) - 1)) == until:
+            return parsed
+        if char == "{":
+            tokenizer.skip_characters("{")
+            tokenizer.skip_white_space()
+            if tokenizer.peek() == "#":
+                state = "LOOP"
+                tokenizer.skip_characters("#each")
+                tokenizer.skip_white_space()
+            else:
+                state = "VAR"
+            continue
+        if char == "}":
+            if state in ["LOOP", "VAR"]:
+                tokenizer.skip_characters("}")
+                if state == "VAR":
+                    if template_data.get(var_name) is None:
+                        raise MissingRenderingAttributeException(var_name)
+                    parsed += template_data.get(var_name)
+                else:
+                    current_position = tokenizer.token_pos
+                    for item in template_data.get(var_name, []):
+                        tokenizer.token_pos = current_position
+                        parsed += parse_template(
+                            template, item, tokenizer, until="{{/each}}"
+                        )
+                        tokenizer.skip_characters("{/each}}")
+                var_name = ""
+                state = ""
+                continue
+        if state in ["LOOP", "VAR"]:
+            var_name += char
+            continue
+        parsed += char
+    return parsed
diff --git a/moto/utilities/tokenizer.py b/moto/utilities/tokenizer.py
new file mode 100644
index 000000000..a6e76dc7f
--- /dev/null
+++ b/moto/utilities/tokenizer.py
@@ -0,0 +1,62 @@
+class GenericTokenizer:
+    """
+    Tokenizer for a KeyConditionExpression. Should be used as an iterator.
+    The final character to be returned will be an empty string, to notify the caller that we've reached the end.
+    """
+
+    def __init__(self, expression):
+        self.expression = expression
+        self.token_pos = 0
+
+    def __iter__(self):
+        return self
+
+    def is_eof(self):
+        return self.peek() == ""
+
+    def peek(self, length=1):
+        """
+        Peek the next character without changing the position
+        """
+        try:
+            return self.expression[self.token_pos : self.token_pos + length]
+        except IndexError:
+            return ""
+
+    def __next__(self):
+        """
+        Returns the next character, or an empty string if we've reached the end of the string.
+        Calling this method again will result in a StopIterator
+        """
+        try:
+            result = self.expression[self.token_pos]
+            self.token_pos += 1
+            return result
+        except IndexError:
+            if self.token_pos == len(self.expression):
+                self.token_pos += 1
+                return ""
+            raise StopIteration
+
+    def skip_characters(self, phrase, case_sensitive=False) -> None:
+        """
+        Skip the characters in the supplied phrase.
+        If any other character is encountered instead, this will fail.
+        If we've already reached the end of the iterator, this will fail.
+        """
+        for ch in phrase:
+            if case_sensitive:
+                assert self.expression[self.token_pos] == ch
+            else:
+                assert self.expression[self.token_pos] in [ch.lower(), ch.upper()]
+            self.token_pos += 1
+
+    def skip_white_space(self):
+        """
+        Skip any whitespace characters that are coming up
+        """
+        try:
+            while self.peek() == " ":
+                self.token_pos += 1
+        except IndexError:
+            pass


Test patch:
diff --git a/tests/test_ses/test_ses_boto3.py b/tests/test_ses/test_ses_boto3.py
index c2c27cdb9..0ca55c7f8 100644
--- a/tests/test_ses/test_ses_boto3.py
+++ b/tests/test_ses/test_ses_boto3.py
@@ -1268,6 +1268,35 @@ def test_render_template():
     )
 
 
+@mock_ses
+def test_render_template__with_foreach():
+    conn = boto3.client("ses", region_name="us-east-1")
+
+    kwargs = dict(
+        TemplateName="MTT",
+        TemplateData=json.dumps(
+            {
+                "items": [
+                    {"type": "dog", "name": "bobby"},
+                    {"type": "cat", "name": "pedro"},
+                ]
+            }
+        ),
+    )
+
+    conn.create_template(
+        Template={
+            "TemplateName": "MTT",
+            "SubjectPart": "..",
+            "TextPart": "..",
+            "HtmlPart": "{{#each items}} {{name}} is a {{type}}, {{/each}}",
+        }
+    )
+    result = conn.test_render_template(**kwargs)
+    result["RenderedTemplate"].should.contain("bobby is a dog")
+    result["RenderedTemplate"].should.contain("pedro is a cat")
+
+
 @mock_ses
 def test_update_ses_template():
     conn = boto3.client("ses", region_name="us-east-1")
diff --git a/tests/test_ses/test_templating.py b/tests/test_ses/test_templating.py
new file mode 100644
index 000000000..0883013b0
--- /dev/null
+++ b/tests/test_ses/test_templating.py
@@ -0,0 +1,26 @@
+from moto.ses.template import parse_template
+import sure  # noqa # pylint: disable=unused-import
+
+
+def test_template_without_args():
+    parse_template("some template", template_data={}).should.equal("some template")
+
+
+def test_template_with_simple_arg():
+    t = "Dear {{name}},"
+    parse_template(t, template_data={"name": "John"}).should.equal("Dear John,")
+
+
+def test_template_with_foreach():
+    t = "List:{{#each l}} - {{obj}}{{/each}} and other things"
+    kwargs = {"l": [{"obj": "it1"}, {"obj": "it2"}]}
+    parse_template(t, kwargs).should.equal("List: - it1 - it2 and other things")
+
+
+def test_template_with_multiple_foreach():
+    t = "{{#each l}} - {{obj}}{{/each}} and list 2 {{#each l2}} {{o1}} {{o2}} {{/each}}"
+    kwargs = {
+        "l": [{"obj": "it1"}, {"obj": "it2"}],
+        "l2": [{"o1": "ob1", "o2": "ob2"}, {"o1": "oc1", "o2": "oc2"}],
+    }
+    parse_template(t, kwargs).should.equal(" - it1 - it2 and list 2  ob1 ob2  oc1 oc2 ")