Merge pull request #1286 from linsword13/more-typing

Extend type-checking to all lib/ramble/ramble/

Author: douglasjacobsen

lib/ramble/ramble/cmd/common/info.py

@@ -18,7 +18,7 @@
 from ramble.definitions.variables import Variable
 from ramble.util.logger import logger
 
-supported_formats = enum.Enum("formats", ["text", "lists"])
+supported_formats = enum.Enum("supported_formats", ["text", "lists"])
 
 obj_attribute_map = {
     "maintainers": None,

lib/ramble/ramble/cmd/software_definitions.py

@@ -7,6 +7,7 @@
 # except according to those terms.
 
 import sys
+from typing import Dict, List
 
 import llnl.util.tty.color as color
 from llnl.util.tty.colify import colify
@@ -22,7 +23,7 @@
 definitions = {}
 conflicts = {}
 used_by = {}
-specs = {"pkg_spec": {}, "compiler_spec": {}}
+specs: Dict[str, Dict[str, List[str]]] = {"pkg_spec": {}, "compiler_spec": {}}
 spec_headers = {
     "pkg_spec": "Software Packages",
     "compiler_spec": "Compiler Definitions",

lib/ramble/ramble/cmd/workspace.py

@@ -11,6 +11,7 @@
 import os
 import sys
 import tempfile
+from typing import Callable, Dict
 
 import llnl.util.tty as tty
 import llnl.util.tty.color as color
@@ -1770,7 +1771,7 @@ def workspace(parser, args, unknown_args):
         action(args)
 
 
-manage_subcommand_functions = {}
+manage_subcommand_functions: Dict[str, Callable] = {}
 
 
 def workspace_manage(args):

lib/ramble/ramble/expander.py

@@ -16,7 +16,7 @@
 import sys
 import warnings
 from contextlib import contextmanager
-from typing import Dict, FrozenSet, List, Union
+from typing import Dict, FrozenSet, List, Optional, Union
 
 import ramble.error
 import ramble.keywords
@@ -598,11 +598,11 @@ def expand_lists(self, var):
     def expand_var_name(
         self,
         var_name: str,
-        extra_vars: Dict = None,
+        extra_vars: Optional[Dict] = None,
         allow_passthrough: bool = True,
         typed: bool = False,
         merge_used_stage: bool = True,
-        replace_escaped_braces: bool = None,
+        replace_escaped_braces: Optional[bool] = None,
     ):
         """Convert a variable name to an expansion string, and expand it
 
@@ -633,11 +633,11 @@ def expand_var_name(
     def expand_var(
         self,
         var: str,
-        extra_vars: Dict = None,
+        extra_vars: Optional[Dict] = None,
         allow_passthrough: bool = True,
         typed: bool = False,
         merge_used_stage: bool = True,
-        replace_escaped_braces=None,
+        replace_escaped_braces: Optional[bool] = None,
     ):
         """Perform expansion of a string
 

lib/ramble/ramble/experiment_set.py

@@ -37,7 +37,7 @@ class ExperimentSet:
 
     # In order of lowest to highest precedence
     _contexts = Enum(
-        "contexts",
+        "_contexts",
         ["global_conf", "base", "workspace", "application", "workload", "experiment", "required"],
     )
 

lib/ramble/ramble/fetch_strategy.py

@@ -33,6 +33,7 @@
 import shutil
 import sys
 import urllib.parse
+from typing import List, Optional
 
 import llnl.util.tty as tty
 from llnl.util.filesystem import (
@@ -105,12 +106,13 @@ class FetchStrategy:
     #: The URL attribute must be specified either at the package class
     #: level, or as a keyword argument to ``version()``.  It is used to
     #: distinguish fetchers for different versions in the package DSL.
-    url_attr = None
+    url_attr: Optional[str] = None
 
     #: Optional attributes can be used to distinguish fetchers when :
     #: classes have multiple ``url_attrs`` at the top-level.
     # optional attributes in version() args.
-    optional_attrs = []
+    optional_attrs: List[str] = []
+    url: Optional[str] = None
 
     def __init__(self, **kwargs):
         # The stage is initialized late, so that fetch strategies can be
@@ -681,6 +683,11 @@ class VCSFetchStrategy(FetchStrategy):
 
     """
 
+    branch: Optional[str] = None
+    tag: Optional[str] = None
+    commit: Optional[str] = None
+    revision: Optional[str] = None
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
@@ -836,6 +843,10 @@ class GitFetchStrategy(VCSFetchStrategy):
 
     git_version_re = r"git version (\S+)"
 
+    submodules: bool = False
+    submodules_delete: bool = False
+    get_full_repo: bool = False
+
     def __init__(self, **kwargs):
         # Discards the keywords in kwargs that may conflict with the next call
         # to __init__

lib/ramble/ramble/graphs.py

@@ -11,6 +11,7 @@
 import itertools
 import re
 from collections import defaultdict
+from typing import DefaultDict
 
 import ramble.error
 import ramble.expander
@@ -462,7 +463,7 @@ class FormattedExecutableGraph(AttributeGraph):
     def __init__(self, formatted_execs: dict, obj_inst):
         """Constructs a new FormattedExecutableGraph and evaluates dependencies"""
         super().__init__(obj_inst)
-        self._formatted_executable_dependencies = defaultdict(list)
+        self._formatted_executable_dependencies: DefaultDict[str, list] = defaultdict(list)
 
         # Define all graph nodes
         for exec_name, exec_def in formatted_execs.items():

lib/ramble/ramble/keywords.py

@@ -11,8 +11,8 @@
 import ramble.error
 from ramble.util.logger import logger
 
-key_type = Enum("type", ["reserved", "optional", "required"])
-output_level = Enum("level", ["key", "variable"])
+key_type = Enum("key_type", ["reserved", "optional", "required"])
+output_level = Enum("output_level", ["key", "variable"])
 default_keys = {
     "workspace_name": {"type": key_type.reserved, "level": output_level.variable},
     "workspace": {"type": key_type.reserved, "level": output_level.variable},
@@ -96,6 +96,56 @@ class Keywords:
                   specific inputs to further configure the experiment.
     """
 
+    workspace_name: str
+    workspace: str
+    workspace_root: str
+    workspace_configs: str
+    workspace_software: str
+    workspace_logs: str
+    workspace_inputs: str
+    workspace_experiments: str
+    workspace_shared: str
+    workspace_archives: str
+    workspace_deployments: str
+    application_name: str
+    application_run_dir: str
+    application_input_dir: str
+    application_namespace: str
+    simplified_application_namespace: str
+    workload_name: str
+    workload_run_dir: str
+    workload_input_dir: str
+    workload_namespace: str
+    simplified_workload_namespace: str
+    license_input_dir: str
+    experiments_file: str
+    experiment_name: str
+    experiment_hash: str
+    experiment_run_dir: str
+    experiment_status: str
+    RAMBLE_STATUS: str
+    experiment_index: str
+    experiment_namespace: str
+    simplified_experiment_namespace: str
+    log_dir: str
+    log_file: str
+    err_file: str
+    env_path: str
+    input_name: str
+    repeat_index: str
+    spec_name: str
+    env_name: str
+    n_ranks: str
+    n_nodes: str
+    processes_per_node: str
+    n_threads: str
+    batch_submit: str
+    mpi_command: str
+    workload_template_name: str
+    experiment_template_name: str
+    unformatted_command: str
+    unformatted_command_without_logs: str
+
     def __init__(self, extra_keys=None):
         # Merge in additional Keys:
         self.keys = default_keys.copy()

lib/ramble/ramble/pipeline.py

@@ -783,14 +783,14 @@ def _upload_file(src_file, dest_file):
 pipelines = Enum(
     "pipelines",
     [
-        AnalyzePipeline.name,
-        ArchivePipeline.name,
-        MirrorPipeline.name,
-        SetupPipeline.name,
-        PushToCachePipeline.name,
-        ExecutePipeline.name,
-        PushDeploymentPipeline.name,
-        LogsPipeline.name,
+        "analyze",
+        "archive",
+        "mirror",
+        "setup",
+        "pushtocache",
+        "execute",
+        "pushdeployment",
+        "logs",
     ],
 )
 

lib/ramble/ramble/reports.py

@@ -11,6 +11,7 @@
 import os
 import re
 from enum import Enum
+from typing import List
 
 import llnl.util.filesystem as fs
 
@@ -754,7 +755,7 @@ def generate_plot_data(self, pdf_report):
 
 class MultiLinePlot(ScalingPlotGenerator):
     plot_type = "multi_line"
-    series_to_plot = []
+    series_to_plot: List[str] = []
 
     def default_better(self):
         return BetterDirection.HIGHER