Add more type checking to Ramble

Signed-off-by: Lin Guo <linsword13@gmail.com>

Author: linsword13

lib/ramble/ramble/analysis/backwards.py

@@ -9,6 +9,7 @@
 """Define the backwards-reading analysis strategy"""
 
 import os
+from typing import Any, Dict, FrozenSet, Tuple
 
 import ramble.success_criteria
 import ramble.util.lock as lk
@@ -121,9 +122,9 @@ def __call__(self, workspace):
 
         exp_lock = app.experiment_lock
 
-        fom_values = {}
-        context_metadata = {}
-        null_key = (_NULL_CONTEXT, _NULL_CONTEXT, frozenset())
+        fom_values: Dict[Tuple[str, str, FrozenSet[Any]], Dict[str, Any]] = {}
+        context_metadata: Dict[Tuple[str, str, FrozenSet[Any]], Dict[str, Any]] = {}
+        null_key: Tuple[str, str, FrozenSet[Any]] = (_NULL_CONTEXT, _NULL_CONTEXT, frozenset())
         context_metadata[null_key] = {
             "name": _NULL_CONTEXT,
             "def_name": _NULL_CONTEXT,

lib/ramble/ramble/analysis/forward.py

@@ -10,6 +10,7 @@
 
 import os
 import string
+from typing import Any, Dict, FrozenSet, Tuple
 
 import ramble.success_criteria
 import ramble.util.lock as lk
@@ -61,9 +62,9 @@ def format_context(context_match, context_format):
 
         exp_lock = app.experiment_lock
 
-        fom_values = {}
-        context_metadata = {}
-        null_key = (_NULL_CONTEXT, _NULL_CONTEXT, frozenset())
+        fom_values: Dict[Tuple[str, str, FrozenSet[Any]], Dict[str, Any]] = {}
+        context_metadata: Dict[Tuple[str, str, FrozenSet[Any]], Dict[str, Any]] = {}
+        null_key: Tuple[str, str, FrozenSet[Any]] = (_NULL_CONTEXT, _NULL_CONTEXT, frozenset())
         context_metadata[null_key] = {
             "name": _NULL_CONTEXT,
             "def_name": _NULL_CONTEXT,
@@ -75,7 +76,7 @@ def format_context(context_match, context_format):
             for file, file_conf in files.items():
 
                 # Start with no active contexts in a file.
-                active_contexts = {}
+                active_contexts: Dict[str, Any] = {}
                 logger.debug(f"Reading log file: {file}")
 
                 if not os.path.exists(file):

lib/ramble/ramble/cmd/__init__.py

@@ -9,6 +9,7 @@
 
 import os
 import re
+from typing import List, Optional
 
 from llnl.util.lang import attr_setdefault
 
@@ -50,7 +51,7 @@ def require_cmd_name(cname):
 
 
 #: global, cached list of all commands -- access through all_commands()
-_all_commands = None
+_all_commands: Optional[List[str]] = None
 
 
 def all_commands():

lib/ramble/ramble/cmd/common/arguments.py

@@ -13,7 +13,17 @@
 
 from ramble.util.logger import logger
 
-from spack.util.pattern import Args
+
+class Args:
+    """Class to hold positional flags and keyword arguments for parser.add_argument."""
+
+    flags: tuple
+    kwargs: dict
+
+    def __init__(self, *flags, **kwargs):
+        self.flags = tuple(flags)
+        self.kwargs = kwargs
+
 
 __all__ = [
     "add_common_arguments",

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

@@ -9,6 +9,7 @@
 import enum
 import fnmatch
 import textwrap
+from collections.abc import Iterable
 
 from llnl.util.tty.colify import colified
 
@@ -376,7 +377,7 @@ def print_single_attribute(obj, attr, verbose=False, pattern="*", format=support
         # Otherwise, we print the attribute's value directly.
         if isinstance(internal_attr, dict):
             to_print = list(internal_attr.keys())
-        elif hasattr(internal_attr, "default_variants"):
+        elif internal_attr is not None and hasattr(internal_attr, "default_variants"):
             to_print = []
             for variant in internal_attr.default_variants.values():
                 to_print.append(variant)
@@ -387,7 +388,7 @@ def print_single_attribute(obj, attr, verbose=False, pattern="*", format=support
                 to_print.append(variant)
             for variant in internal_attr.version_variants.values():
                 to_print.append(variant)
-        elif hasattr(internal_attr, "family_type"):
+        elif internal_attr is not None and hasattr(internal_attr, "family_type"):
             to_print = [f"{internal_attr.family_type}={family}" for family in internal_attr]
         else:
             to_print = internal_attr
@@ -396,9 +397,7 @@ def print_single_attribute(obj, attr, verbose=False, pattern="*", format=support
         #     if it's a list of dicts, convert the keys like above and print
         #     otherwise filter it and print using the format specification
         # Otherwise, print it as a raw string.
-        if isinstance(to_print, (list, set, tuple)) or (
-            hasattr(to_print, "__iter__") and not isinstance(to_print, str)
-        ):
+        if isinstance(to_print, Iterable) and not isinstance(to_print, str):
             to_print = list(to_print)
             if (
                 internal_attr
@@ -415,9 +414,10 @@ def print_single_attribute(obj, attr, verbose=False, pattern="*", format=support
         if isinstance(internal_attr, dict):
             _print_verbose_dict_attr(internal_attr, pattern=pattern, indentation=indentation)
         elif (
-            isinstance(internal_attr, (list, set, tuple))
-            or (hasattr(internal_attr, "__iter__") and not isinstance(internal_attr, str))
-        ) and not hasattr(internal_attr, "as_str"):
+            isinstance(internal_attr, Iterable)
+            and not isinstance(internal_attr, str)
+            and not hasattr(internal_attr, "as_str")
+        ):
             internal_list = list(internal_attr)
             # If it's a list of dicts, print each
             if internal_list and isinstance(internal_list[0], dict):
@@ -440,7 +440,7 @@ def print_single_attribute(obj, attr, verbose=False, pattern="*", format=support
                         color.cprint(f"{colified(to_print, tty=True, indent=4)}")
                     color.cprint("")
         else:
-            if hasattr(internal_attr, "as_str"):
+            if internal_attr is not None and hasattr(internal_attr, "as_str"):
                 color.cprint(internal_attr.as_str(verbose=True))
             else:
                 color.cprint(f"{indentation}{internal_attr}\n")

lib/ramble/ramble/cmd/common/list.py

@@ -169,8 +169,8 @@ def perform_list(args):
     # Filter by tags
     if args.tags:
         objects_with_tags = set(ramble.repository.paths[object_type].objects_with_tags(*args.tags))
-        sorted_objects = set(sorted_objects) & objects_with_tags
-        sorted_objects = sorted(sorted_objects)
+        matching_tags = set(sorted_objects) & objects_with_tags
+        sorted_objects = sorted(matching_tags)
 
     if not sorted_objects:
         filter_strs = []

lib/ramble/ramble/cmd/config.py

@@ -26,6 +26,8 @@
 section = "config"
 level = "long"
 
+_add_parser = None
+
 
 def setup_parser(subparser):
     scopes_metavar = ramble.config.scopes_metavar
@@ -94,7 +96,8 @@ def setup_parser(subparser):
     )
 
     # Make the add parser available later
-    setup_parser.add_parser = add_parser
+    global _add_parser
+    _add_parser = add_parser
 
     update = sp.add_parser("update", help="update configuration files to the latest format")
     ramble.cmd.common.arguments.add_common_arguments(update, ["yes_to_all"])
@@ -207,7 +210,8 @@ def config_add(args):
     This is a stateful operation that edits the config files."""
     if not (args.file or args.path):
         logger.error("No changes requested. Specify a file or value.")
-        setup_parser.add_parser.print_help()
+        if _add_parser:
+            _add_parser.print_help()
         exit(1)
 
     scope, _ = _get_scope_and_section(args)

lib/ramble/ramble/cmd/data.py

@@ -42,7 +42,7 @@ def data_create_db(args):
         raise ConfigError(f"Upload type {uploader_type_str} is not valid.")
 
     uploader_type = getattr(ramble.uploader.uploader_types, uploader_type_str)
-
+    uploader: ramble.uploader.Uploader
     if uploader_type == ramble.uploader.uploader_types.BigQuery:
         uploader = ramble.uploader.BigQueryUploader()
     elif uploader_type == ramble.uploader.uploader_types.SQLite:

lib/ramble/ramble/cmd/edit.py

@@ -8,6 +8,7 @@
 
 import glob
 import os
+from typing import Any
 
 import ramble.cmd
 import ramble.paths
@@ -238,6 +239,7 @@ def edit(parser, args):
             # It's an object type. Let's find what path it would have been at.
             try:
                 obj_type = ramble.repository.ObjectTypes[type_name]
+                repo: Any
                 if args.repo:
                     repo = ramble.repository.Repo(args.repo, object_type=obj_type)
                 elif effective_namespace:

lib/ramble/ramble/cmd/filter_groups.py

@@ -7,6 +7,7 @@
 # except according to those terms.
 
 import copy
+from typing import Any, Dict, List
 
 from llnl.util.tty.colify import colify
 
@@ -203,7 +204,7 @@ def print_filter_groups(resolved_scope_name=None, original_scope_name=None, verb
                 lines.extend(f"        - {ew}" for ew in definition["exclude_where"])
             color.cprint("\n".join(lines))
     else:
-        scope_groups = {}
+        scope_groups: Dict[str, List[str]] = {}
         for item in groups_to_print:
             scope = item["scope"]
             name = item["name"]
@@ -212,7 +213,7 @@ def print_filter_groups(resolved_scope_name=None, original_scope_name=None, verb
             scope_groups[scope].append(name)
 
         out_stream = logger.active_stream()
-        colify_opts = {"indent": 4, "padding": 2}
+        colify_opts: Dict[str, Any] = {"indent": 4, "padding": 2}
         if out_stream:
             colify_opts["output"] = out_stream