Typechecking for assemble_workflow (#1230)

* Add typechecking for assemble

Signed-off-by: Peter Nied <[email protected]>

.pre-commit-config.yaml

@@ -18,7 +18,7 @@ repos:
   - id: mypy
     stages: [commit]
     name: mypy
-    entry: bash -c 'pipenv run mypy src/manifests tests/tests_manifests'
+    entry: bash -c 'pipenv run mypy src/run_assemble.py tests/test_run_assemble.py src/assemble_workflow tests/tests_assemble_workflow src/manifests tests/tests_manifests'
     language: system
   - id: pytest
     stages: [commit]

src/assemble_workflow/assemble_args.py

@@ -6,13 +6,14 @@
 
 import argparse
 import logging
+from typing import IO
 
 
 class AssembleArgs:
-    manifest: str
+    manifest: IO
     keep: bool
 
-    def __init__(self):
+    def __init__(self) -> None:
         parser = argparse.ArgumentParser(description="Assemble an OpenSearch Distribution")
         parser.add_argument("manifest", type=argparse.FileType("r"), help="Manifest file.")
         parser.add_argument("-b", "--base-url", dest="base_url", help="The base url to download the artifacts.")

src/assemble_workflow/bundle.py

@@ -10,8 +10,11 @@
 import shutil
 import subprocess
 from abc import ABC, abstractmethod
+from typing import Any, List
 
+from assemble_workflow.bundle_recorder import BundleRecorder
 from assemble_workflow.dist import Dist
+from manifests.build_manifest import BuildComponent, BuildComponents, BuildManifest
 from paths.script_finder import ScriptFinder
 from system.temporary_directory import TemporaryDirectory
 
@@ -23,28 +26,28 @@
 
 
 class Bundle(ABC):
-    def __enter__(self):
+    def __enter__(self) -> 'Bundle':
         return self
 
-    def __exit__(self, exc_type, exc_value, exc_traceback):
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_traceback: Any) -> None:
         self.tmp_dir.__exit__(exc_type, exc_value, exc_traceback)
 
-    def __init__(self, build_manifest, artifacts_dir, bundle_recorder, keep=False):
+    def __init__(self, build_manifest: BuildManifest, artifacts_dir: str, bundle_recorder: BundleRecorder, keep: bool = False) -> None:
         """
         Construct a new Bundle instance.
         :param build_manifest: A BuildManifest created from the build workflow.
         :param artifacts_dir: Dir location where build artifacts can be found locally
         :param bundle_recorder: The bundle recorder that will capture and build a BundleManifest
         """
-        self.plugins = self.__get_plugins(build_manifest.components.values())
+        self.plugins = self.__get_plugins(build_manifest.components)
         self.artifacts_dir = artifacts_dir
         self.bundle_recorder = bundle_recorder
         self.tmp_dir = TemporaryDirectory(keep=keep)
-        self.min_dist = self.__get_min_dist(build_manifest.components.values())
-        self.installed_plugins = []
+        self.min_dist = self.__get_min_dist(build_manifest.components)
+        self.installed_plugins: List[str] = []
         self.build = build_manifest.build
 
-    def install_min(self):
+    def install_min(self) -> None:
         install_script = ScriptFinder.find_install_script(self.min_dist.name)
         install_command = " ".join(
             [
@@ -64,7 +67,7 @@ def install_min(self):
         )
         self._execute(install_command)
 
-    def install_plugins(self):
+    def install_plugins(self) -> None:
         for plugin in self.plugins:
             logging.info(f"Installing {plugin.name}")
             self.install_plugin(plugin)
@@ -73,7 +76,7 @@ def install_plugins(self):
             self.installed_plugins = os.listdir(plugins_path)
 
     @abstractmethod
-    def install_plugin(self, plugin):
+    def install_plugin(self, plugin: BuildComponent) -> None:
         install_script = ScriptFinder.find_install_script(plugin.name)
         install_command = " ".join(
             [
@@ -93,23 +96,23 @@ def install_plugin(self, plugin):
         )
         self._execute(install_command)
 
-    def package(self, dest):
+    def package(self, dest: str) -> None:
         self.min_dist.build(self.bundle_recorder.package_name, dest)
 
-    def _execute(self, command):
+    def _execute(self, command: str) -> None:
         logging.info(f'Executing "{command}" in {self.min_dist.archive_path}')
         subprocess.check_call(command, cwd=self.min_dist.archive_path, shell=True)
 
-    def _copy_component(self, component, component_type):
+    def _copy_component(self, component: BuildComponent, component_type: str) -> str:
         rel_path = self.__get_rel_path(component, component_type)
         tmp_path = self.__copy_component_files(rel_path, self.tmp_dir.name)
         self.bundle_recorder.record_component(component, rel_path)
         return tmp_path
 
-    def __get_rel_path(self, component, component_type):
+    def __get_rel_path(self, component: BuildComponent, component_type: str) -> str:
         return next(iter(component.artifacts.get(component_type, [])), None)
 
-    def __copy_component_files(self, rel_path, dest):
+    def __copy_component_files(self, rel_path: str, dest: str) -> str:
         local_path = os.path.join(self.artifacts_dir, rel_path)
         dest_path = os.path.join(dest, os.path.basename(local_path))
         if os.path.isfile(local_path):
@@ -119,11 +122,11 @@ def __copy_component_files(self, rel_path, dest):
         else:
             raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), local_path)
 
-    def __get_plugins(self, build_components):
-        return [c for c in build_components if "plugins" in c.artifacts]
+    def __get_plugins(self, build_components: BuildComponents) -> List[BuildComponent]:
+        return [c for c in build_components.values() if "plugins" in c.artifacts]
 
-    def __get_min_dist(self, build_components):
-        min_bundle = next(iter([c for c in build_components if "dist" in c.artifacts]), None)
+    def __get_min_dist(self, build_components: BuildComponents) -> Dist:
+        min_bundle = next(iter([c for c in build_components.values() if "dist" in c.artifacts]), None)
         if min_bundle is None:
             raise ValueError('Missing min "dist" in input artifacts.')
         min_dist_path = self._copy_component(min_bundle, "dist")

src/assemble_workflow/bundle_opensearch.py

@@ -7,15 +7,16 @@
 import os
 
 from assemble_workflow.bundle import Bundle
+from manifests.build_manifest import BuildComponent
 from system.os import current_platform
 
 
 class BundleOpenSearch(Bundle):
     @property
-    def install_plugin_script(self):
+    def install_plugin_script(self) -> str:
         return "opensearch-plugin.bat" if current_platform() == "windows" else "opensearch-plugin"
 
-    def install_plugin(self, plugin):
+    def install_plugin(self, plugin: BuildComponent) -> None:
         tmp_path = self._copy_component(plugin, "plugins")
         cli_path = os.path.join(self.min_dist.archive_path, "bin", self.install_plugin_script)
         self._execute(f"{cli_path} install --batch file:{tmp_path}")

src/assemble_workflow/bundle_opensearch_dashboards.py

@@ -7,10 +7,11 @@
 import os
 
 from assemble_workflow.bundle import Bundle
+from manifests.build_manifest import BuildComponent
 
 
 class BundleOpenSearchDashboards(Bundle):
-    def install_plugin(self, plugin):
+    def install_plugin(self, plugin: BuildComponent) -> None:
         tmp_path = self._copy_component(plugin, "plugins")
         cli_path = os.path.join(self.min_dist.archive_path, "bin", "opensearch-dashboards-plugin")
         self._execute(f"{cli_path} --allow-root install file:{tmp_path}")

src/assemble_workflow/bundle_recorder.py

@@ -5,13 +5,15 @@
 # compatible open source license.
 
 import os
+from typing import Any, Dict
 from urllib.parse import urljoin
 
+from manifests.build_manifest import BuildComponent, BuildManifest
 from manifests.bundle_manifest import BundleManifest
 
 
 class BundleRecorder:
-    def __init__(self, build, output_dir, artifacts_dir, base_url):
+    def __init__(self, build: BuildManifest.Build, output_dir: str, artifacts_dir: str, base_url: str) -> None:
         self.output_dir = output_dir
         self.build_id = build.id
         self.base_url = base_url
@@ -28,7 +30,7 @@ def __init__(self, build, output_dir, artifacts_dir, base_url):
             self.__get_package_location(),
         )
 
-    def __get_package_name(self, build):
+    def __get_package_name(self, build: BuildManifest.Build) -> str:
         parts = [
             build.name.lower().replace(" ", "-"),
             build.version,
@@ -37,27 +39,27 @@ def __get_package_name(self, build):
         ]
         return "-".join(parts) + (".zip" if build.platform == "windows" else ".tar.gz")
 
-    def __get_public_url_path(self, folder, rel_path):
+    def __get_public_url_path(self, folder: str, rel_path: str) -> str:
         path = "/".join((folder, rel_path))
         return urljoin(self.base_url + "/", path)
 
-    def __get_location(self, folder_name, rel_path, abs_path):
+    def __get_location(self, folder_name: str, rel_path: str, abs_path: str) -> str:
         if self.base_url:
             return self.__get_public_url_path(folder_name, rel_path)
         return abs_path
 
     # Assembled output are expected to be served from a separate "dist" folder
     # Example: https://ci.opensearch.org/ci/dbc/bundle-build/1.2.0/build-id/linux/x64/dist/
-    def __get_package_location(self):
+    def __get_package_location(self) -> str:
         return self.__get_location("dist", self.package_name, os.path.join(self.output_dir, self.package_name))
 
     # Build artifacts are expected to be served from a "builds" folder
     # Example: https://ci.opensearch.org/ci/dbc/bundle-build/1.2.0/build-id/linux/x64/builds/
-    def __get_component_location(self, component_rel_path):
+    def __get_component_location(self, component_rel_path: str) -> str:
         abs_path = os.path.join(self.artifacts_dir, component_rel_path)
         return self.__get_location("builds", component_rel_path, abs_path)
 
-    def record_component(self, component, rel_path):
+    def record_component(self, component: BuildComponent, rel_path: str) -> None:
         self.bundle_manifest.append_component(
             component.name,
             component.repository,
@@ -66,16 +68,16 @@ def record_component(self, component, rel_path):
             self.__get_component_location(rel_path),
         )
 
-    def get_manifest(self):
+    def get_manifest(self) -> BundleManifest:
         return self.bundle_manifest.to_manifest()
 
-    def write_manifest(self, folder):
+    def write_manifest(self, folder: str) -> None:
         manifest_path = os.path.join(folder, "manifest.yml")
         self.get_manifest().to_file(manifest_path)
 
     class BundleManifestBuilder:
-        def __init__(self, build_id, name, version, platform, architecture, location):
-            self.data = {}
+        def __init__(self, build_id: str, name: str, version: str, platform: str, architecture: str, location: str) -> None:
+            self.data: Dict[str, Any] = {}
             self.data["build"] = {}
             self.data["build"]["id"] = build_id
             self.data["build"]["name"] = name
@@ -88,7 +90,7 @@ def __init__(self, build_id, name, version, platform, architecture, location):
             # When we convert to a BundleManifest this will get converted back into a list
             self.data["components"] = []
 
-        def append_component(self, name, repository_url, ref, commit_id, location):
+        def append_component(self, name: str, repository_url: str, ref: str, commit_id: str, location: str) -> None:
             component = {
                 "name": name,
                 "repository": repository_url,
@@ -98,5 +100,5 @@ def append_component(self, name, repository_url, ref, commit_id, location):
             }
             self.data["components"].append(component)
 
-        def to_manifest(self):
+        def to_manifest(self) -> BundleManifest:
             return BundleManifest(self.data)

src/assemble_workflow/bundles.py

@@ -4,8 +4,12 @@
 # this file be licensed under the Apache-2.0 license or a
 # compatible open source license.
 
+
+from assemble_workflow.bundle import Bundle
 from assemble_workflow.bundle_opensearch import BundleOpenSearch
 from assemble_workflow.bundle_opensearch_dashboards import BundleOpenSearchDashboards
+from assemble_workflow.bundle_recorder import BundleRecorder
+from manifests.build_manifest import BuildManifest
 
 
 class Bundles:
@@ -15,13 +19,13 @@ class Bundles:
     }
 
     @classmethod
-    def from_name(cls, name):
+    def from_name(cls, name: str) -> Bundle:
         klass = cls.TYPES.get(name, None)
         if not klass:
             raise ValueError(f"Unsupported bundle: {name}")
-        return klass
+        return klass  # type: ignore[return-value]
 
     @classmethod
-    def create(cls, build_manifest, artifacts_dir, bundle_recorder, keep):
+    def create(cls, build_manifest: BuildManifest, artifacts_dir: str, bundle_recorder: BundleRecorder, keep: bool) -> Bundle:
         klass = cls.from_name(build_manifest.build.name)
-        return klass(build_manifest, artifacts_dir, bundle_recorder, keep)
+        return klass(build_manifest, artifacts_dir, bundle_recorder, keep)  # type: ignore[no-any-return, operator]

src/assemble_workflow/dist.py

@@ -16,15 +16,19 @@
 
 
 class Dist(ABC):
-    def __init__(self, name, path):
+    def __init__(self, name: str, path: str) -> None:
         self.name = name
         self.path = path
 
     @abstractmethod
-    def __extract__(self, dest):
+    def __extract__(self, dest: str) -> None:
         pass
 
-    def extract(self, dest):
+    @abstractmethod
+    def __build__(self, name: str, dest: str) -> None:
+        pass
+
+    def extract(self, dest: str) -> str:
         self.__extract__(dest)
 
         # OpenSearch & Dashboard tars will include only a single folder at the top level of the tar.
@@ -36,14 +40,14 @@ def extract(self, dest):
 
         raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), os.path.join(dest, "*"))
 
-    def build(self, name, dest):
+    def build(self, name: str, dest: str) -> None:
         self.__build__(name, dest)
         path = os.path.join(dest, name)
         shutil.copyfile(name, path)
         logging.info(f"Published {path}.")
 
     @classmethod
-    def from_path(cls, name, path):
+    def from_path(cls, name: str, path: str) -> 'Dist':
         ext = os.path.splitext(path)[1]
         if ext == ".gz":
             return DistTar(name, path)
@@ -54,24 +58,24 @@ def from_path(cls, name, path):
 
 
 class DistZip(Dist):
-    def __extract__(self, dest):
+    def __extract__(self, dest: str) -> None:
         with ZipFile(self.path, "r") as zip:
             zip.extractall(dest)
 
-    def __build__(self, name, dest):
+    def __build__(self, name: str, dest: str) -> None:
         with ZipFile(name, "w", zipfile.ZIP_DEFLATED) as zip:
             rootlen = len(self.archive_path) + 1
-            for base, dirs, files in os.walk(self.archive_path):
+            for base, _, files in os.walk(self.archive_path):
                 for file in files:
                     fn = os.path.join(base, file)
                     zip.write(fn, fn[rootlen:])
 
 
 class DistTar(Dist):
-    def __extract__(self, dest):
+    def __extract__(self, dest: str) -> None:
         with tarfile.open(self.path, "r:gz") as tar:
             tar.extractall(dest)
 
-    def __build__(self, name, dest):
+    def __build__(self, name: str, dest: str) -> None:
         with tarfile.open(name, "w:gz") as tar:
             tar.add(self.archive_path, arcname=os.path.basename(self.archive_path))

src/manifests/build_manifest.py

@@ -100,11 +100,11 @@ def __to_dict__(self) -> dict:
 
     class Build:
         def __init__(self, data: Any):
-            self.name = data["name"]
-            self.version = data["version"]
-            self.platform = data["platform"]
-            self.architecture = data["architecture"]
-            self.id = data["id"]
+            self.name: str = data["name"]
+            self.version: str = data["version"]
+            self.platform: str = data["platform"]
+            self.architecture: str = data["architecture"]
+            self.id: str = data["id"]
 
         def __to_dict__(self) -> dict:
             return {

src/paths/assemble_output_dir.py

@@ -8,5 +8,5 @@
 
 
 class AssembleOutputDir(OutputDir):
-    def __init__(cls, name, cwd=None, makedirs=True):
+    def __init__(cls, name: str, cwd: str = None, makedirs: bool = True) -> None:
         super().__init__("dist", name, cwd, makedirs)

src/paths/output_dir.py

@@ -9,7 +9,7 @@
 
 
 class OutputDir(abc.ABC):
-    def __init__(cls, parent_dir, name, cwd=None, makedirs=True):
+    def __init__(cls, parent_dir: str, name: str, cwd: str = None, makedirs: bool = True) -> None:
         cls.dir = os.path.join(
             cwd or os.getcwd(),
             parent_dir,