chore: Use dataclasses to make type safer

aahung · Nov 18, 2020 · b5e2b1b · b5e2b1b
1 parent 8f46e04
commit b5e2b1b
Show file tree

Hide file tree

Showing 6 changed files with 182 additions and 148 deletions.
diff --git a/requirements/base.txt b/requirements/base.txt
@@ -15,3 +15,4 @@ requests==2.23.0
 serverlessrepo==0.1.10
 aws_lambda_builders==1.1.0
 tomlkit==0.7.0
+dataclasses~=0.8; python_version<"3.7"
diff --git a/samcli/commands/init/init_templates.py b/samcli/commands/init/init_templates.py
@@ -9,8 +9,10 @@
 import platform
 import shutil
 import subprocess
+from dataclasses import dataclass
 
 from pathlib import Path  # must come after Py2.7 deprecation
+from typing import List, Optional
 
 import click
 
@@ -23,6 +25,16 @@
 LOG = logging.getLogger(__name__)
 
 
+@dataclass
+class TemplateOption:
+    directory: Optional[str]
+    display_name: Optional[str]
+    dependency_manager: str
+    app_template: str
+    init_location: Optional[str] = None
+    is_dynamic_template: Optional[str] = None
+
+
 class InvalidInitTemplateError(UserException):
     pass
 
@@ -46,8 +58,8 @@ def prompt_for_location(self, runtime, dependency_manager):
             choice_num = 1
             click.echo("\nAWS quick start application templates:")
             for o in options:
-                if o.get("displayName") is not None:
-                    msg = "\t" + str(choice_num) + " - " + o.get("displayName")
+                if o.display_name is not None:
+                    msg = "\t" + str(choice_num) + " - " + o.display_name
                     click.echo(msg)
                 else:
                     msg = (
@@ -62,55 +74,66 @@ def prompt_for_location(self, runtime, dependency_manager):
                 choice_num = choice_num + 1
             choice = click.prompt("Template selection", type=click.Choice(choices), show_choices=False)
             template_md = options[int(choice) - 1]  # zero index
-        if template_md.get("init_location") is not None:
-            return (template_md["init_location"], "hello-world")
-        if template_md.get("directory") is not None:
-            return (os.path.join(self.repo_path, template_md["directory"]), template_md["appTemplate"])
+        if template_md.init_location is not None:
+            return template_md.init_location, "hello-world"
+        if template_md.directory is not None:
+            return os.path.join(self.repo_path, template_md.directory), template_md.app_template
         raise InvalidInitTemplateError("Invalid template. This should not be possible, please raise an issue.")
 
     def location_from_app_template(self, runtime, dependency_manager, app_template):
         options = self.init_options(runtime, dependency_manager)
         try:
             template = next(item for item in options if self._check_app_template(item, app_template))
-            if template.get("init_location") is not None:
-                return template["init_location"]
-            if template.get("directory") is not None:
-                return os.path.join(self.repo_path, template["directory"])
+            if template.init_location is not None:
+                return template.init_location
+            if template.directory is not None:
+                return os.path.join(self.repo_path, template.directory)
             raise InvalidInitTemplateError("Invalid template. This should not be possible, please raise an issue.")
         except StopIteration as ex:
             msg = "Can't find application template " + app_template + " - check valid values in interactive init."
             raise InvalidInitTemplateError(msg) from ex
 
-    def _check_app_template(self, entry, app_template):
-        return entry["appTemplate"] == app_template
+    def _check_app_template(self, entry: TemplateOption, app_template):
+        return entry.app_template == app_template
 
-    def init_options(self, runtime, dependency_manager):
+    def init_options(self, runtime, dependency_manager) -> List[TemplateOption]:
         if self.clone_attempted is False:  # pylint: disable=compare-to-zero
             self._clone_repo()
         if self.repo_path is None:
             return self._init_options_from_bundle(runtime, dependency_manager)
         return self._init_options_from_manifest(runtime, dependency_manager)
 
-    def _init_options_from_manifest(self, runtime, dependency_manager):
+    def _init_options_from_manifest(self, runtime, dependency_manager) -> List[TemplateOption]:
         manifest_path = os.path.join(self.repo_path, "manifest.json")
         with open(str(manifest_path)) as fp:
             body = fp.read()
             manifest_body = json.loads(body)
-            templates = manifest_body.get(runtime)
-            if templates is None:
+            raw_templates = manifest_body.get(runtime)
+            if raw_templates is None:
                 # Fallback to bundled templates
                 return self._init_options_from_bundle(runtime, dependency_manager)
+            templates = [
+                TemplateOption(
+                    obj.get("directory", None),
+                    obj.get("displayName", None),
+                    obj.get("dependencyManager", None),
+                    obj.get("appTemplate", None),
+                    obj.get("initLocation", None),
+                    obj.get("isDynamicTemplate", None),
+                )
+                for obj in raw_templates
+            ]
             if dependency_manager is not None:
-                templates_by_dep = filter(lambda x: x["dependencyManager"] == dependency_manager, templates)
-                return list(templates_by_dep)
+                return [template for template in templates if template.dependency_manager == dependency_manager]
+
             return templates
 
-    def _init_options_from_bundle(self, runtime, dependency_manager):
-        for mapping in list(itertools.chain(*(RUNTIME_DEP_TEMPLATE_MAPPING.values()))):
-            if runtime in mapping["runtimes"] or any([r.startswith(runtime) for r in mapping["runtimes"]]):
-                if not dependency_manager or dependency_manager == mapping["dependency_manager"]:
-                    mapping["appTemplate"] = "hello-world"  # when bundled, use this default template name
-                    return [mapping]
+    def _init_options_from_bundle(self, runtime, dependency_manager) -> List[TemplateOption]:
+        for mapping in list(itertools.chain.from_iterable(RUNTIME_DEP_TEMPLATE_MAPPING.values())):
+            if runtime in mapping.runtimes or any([r.startswith(runtime) for r in mapping.runtimes]):
+                if not dependency_manager or dependency_manager == mapping.dependency_manager:
+                    app_template = "hello-world"  # when bundled, use this default template name
+                    return [TemplateOption(None, None, mapping.dependency_manager, app_template, mapping.init_location)]
         msg = "Lambda Runtime {} and dependency manager {} does not have an available initialization template.".format(
             runtime, dependency_manager
         )
@@ -224,6 +247,6 @@ def is_dynamic_schemas_template(self, app_template, runtime, dependency_manager)
         """
         options = self.init_options(runtime, dependency_manager)
         for option in options:
-            if option.get("appTemplate") == app_template:
-                return option.get("isDynamicTemplate", False)
+            if option.app_template == app_template:
+                return option.is_dynamic_template
         return False
diff --git a/samcli/lib/init/__init__.py b/samcli/lib/init/__init__.py
@@ -52,10 +52,10 @@ def generate_project(
     template = None
 
     if runtime:
-        for mapping in list(itertools.chain(*(RUNTIME_DEP_TEMPLATE_MAPPING.values()))):
-            if runtime in mapping["runtimes"] or any([r.startswith(runtime) for r in mapping["runtimes"]]):
-                if not dependency_manager or dependency_manager == mapping["dependency_manager"]:
-                    template = mapping["init_location"]
+        for mapping in list(itertools.chain.from_iterable(RUNTIME_DEP_TEMPLATE_MAPPING.values())):
+            if runtime in mapping.runtimes or any([r.startswith(runtime) for r in mapping.runtimes]):
+                if not dependency_manager or dependency_manager == mapping.dependency_manager:
+                    template = mapping.init_location
                     break
 
         if not template:

diff --git a/samcli/local/common/runtime_template.py b/samcli/local/common/runtime_template.py
@@ -5,67 +5,76 @@
 import itertools
 import os
 import pathlib
-from typing import Set
+from dataclasses import dataclass
+from typing import List, Dict
 
 _init_path = str(pathlib.Path(os.path.dirname(__file__)).parent)
 _templates = os.path.join(_init_path, "init", "templates")
 
 
+@dataclass
+class RuntimeDepInfo:
+    runtimes: List[str]
+    dependency_manager: str
+    init_location: str
+    build: bool
+
+
 # Note(TheSriram): The ordering of the runtimes list per language is based on the latest to oldest.
-RUNTIME_DEP_TEMPLATE_MAPPING = {
+RUNTIME_DEP_TEMPLATE_MAPPING: Dict[str, List[RuntimeDepInfo]] = {
     "python": [
-        {
-            "runtimes": ["python3.8", "python3.7", "python3.6", "python2.7"],
-            "dependency_manager": "pip",
-            "init_location": os.path.join(_templates, "cookiecutter-aws-sam-hello-python"),
-            "build": True,
-        }
+        RuntimeDepInfo(
+            ["python3.8", "python3.7", "python3.6", "python2.7"],
+            "pip",
+            os.path.join(_templates, "cookiecutter-aws-sam-hello-python"),
+            True,
+        )
     ],
     "ruby": [
-        {
-            "runtimes": ["ruby2.5", "ruby2.7"],
-            "dependency_manager": "bundler",
-            "init_location": os.path.join(_templates, "cookiecutter-aws-sam-hello-ruby"),
-            "build": True,
-        }
+        RuntimeDepInfo(
+            ["ruby2.5", "ruby2.7"],
+            "bundler",
+            os.path.join(_templates, "cookiecutter-aws-sam-hello-ruby"),
+            True,
+        )
     ],
     "nodejs": [
-        {
-            "runtimes": ["nodejs12.x", "nodejs10.x"],
-            "dependency_manager": "npm",
-            "init_location": os.path.join(_templates, "cookiecutter-aws-sam-hello-nodejs"),
-            "build": True,
-        }
+        RuntimeDepInfo(
+            ["nodejs12.x", "nodejs10.x"],
+            "npm",
+            os.path.join(_templates, "cookiecutter-aws-sam-hello-nodejs"),
+            True,
+        )
     ],
     "dotnet": [
-        {
-            "runtimes": ["dotnetcore3.1", "dotnetcore2.1"],
-            "dependency_manager": "cli-package",
-            "init_location": os.path.join(_templates, "cookiecutter-aws-sam-hello-dotnet"),
-            "build": True,
-        }
+        RuntimeDepInfo(
+            ["dotnetcore3.1", "dotnetcore2.1"],
+            "cli-package",
+            os.path.join(_templates, "cookiecutter-aws-sam-hello-dotnet"),
+            True,
+        )
     ],
     "go": [
-        {
-            "runtimes": ["go1.x"],
-            "dependency_manager": "mod",
-            "init_location": os.path.join(_templates, "cookiecutter-aws-sam-hello-golang"),
-            "build": False,
-        }
+        RuntimeDepInfo(
+            ["go1.x"],
+            "mod",
+            os.path.join(_templates, "cookiecutter-aws-sam-hello-golang"),
+            False,
+        )
     ],
     "java": [
-        {
-            "runtimes": ["java11", "java8", "java8.al2"],
-            "dependency_manager": "maven",
-            "init_location": os.path.join(_templates, "cookiecutter-aws-sam-hello-java-maven"),
-            "build": True,
-        },
-        {
-            "runtimes": ["java11", "java8", "java8.al2"],
-            "dependency_manager": "gradle",
-            "init_location": os.path.join(_templates, "cookiecutter-aws-sam-hello-java-gradle"),
-            "build": True,
-        },
+        RuntimeDepInfo(
+            ["java11", "java8", "java8.al2"],
+            "maven",
+            os.path.join(_templates, "cookiecutter-aws-sam-hello-java-maven"),
+            True,
+        ),
+        RuntimeDepInfo(
+            ["java11", "java8", "java8.al2"],
+            "gradle",
+            os.path.join(_templates, "cookiecutter-aws-sam-hello-java-gradle"),
+            True,
+        ),
     ],
 }
 
@@ -86,15 +95,16 @@
     "java8.al2": ["maven", "gradle"],
 }
 
-SUPPORTED_DEP_MANAGERS: Set[str] = {
-    c["dependency_manager"]  # type: ignore
-    for c in list(itertools.chain(*(RUNTIME_DEP_TEMPLATE_MAPPING.values())))
-    if c["dependency_manager"]
+
+SUPPORTED_DEP_MANAGERS = {
+    c.dependency_manager
+    for c in list(itertools.chain.from_iterable(RUNTIME_DEP_TEMPLATE_MAPPING.values()))
+    if c.dependency_manager
 }
 
-RUNTIMES: Set[str] = set(
-    itertools.chain(
-        *[c["runtimes"] for c in list(itertools.chain(*(RUNTIME_DEP_TEMPLATE_MAPPING.values())))]  # type: ignore
+RUNTIMES = set(
+    itertools.chain.from_iterable(
+        [c.runtimes for c in list(itertools.chain.from_iterable(RUNTIME_DEP_TEMPLATE_MAPPING.values()))]
     )
 )