Merge pull request #1285 from linsword13/more-typing

Expand type checks

Author: douglasjacobsen

lib/ramble/ramble/definitions/variables.py

@@ -7,7 +7,7 @@
 # except according to those terms.
 
 import copy
-from typing import List, Optional
+from typing import Any, Dict, List, Optional, Set
 
 import ramble.util.colors as rucolor
 
@@ -35,7 +35,7 @@ def __init__(
         self,
         name: str,
         default=None,
-        description: str = None,
+        description: Optional[str] = None,
         values=None,
         expandable: bool = True,
         track_used: bool = True,
@@ -182,7 +182,7 @@ def __init__(
         self,
         name: str,
         value=None,
-        description: str = None,
+        description: Optional[str] = None,
         method: str = "set",
         append_separator: str = ",",
         when=None,
@@ -278,10 +278,10 @@ def __init__(
         """
         self.name = name
         self.when = when.copy() if when else []
-        self.set = {}
-        self.unset = {}
-        self.prepend = []
-        self.append = []
+        self.set: Dict[str, str] = {}
+        self.unset: Set[str] = set()
+        self.prepend: List[Dict[str, Dict[str, str]]] = []
+        self.append: List[Dict[str, Any]] = []
 
         self.add_modification(
             modification=modification,

lib/ramble/ramble/language/language_base.py

@@ -14,7 +14,7 @@
 import functools
 import inspect
 from collections.abc import Sequence  # novm
-from typing import Any, Dict, List
+from typing import Any, Callable, Dict, List, Set
 
 import llnl.util.lang
 
@@ -25,7 +25,7 @@
 
 #: These are variant names used by ramble internally; applications can't use
 #: them
-reserved_names = []
+reserved_names: List[str] = []
 
 namespaces = [
     "ramble.app",
@@ -61,12 +61,12 @@ class DirectiveMeta(type):
     """
 
     # Set of all known directives
-    _directive_names = set()
-    _directive_init_values = {}
-    _directives_to_be_executed = []
-    _directive_functions = {}
-    _directive_classes = {}
-    _when_constraints_from_context = []
+    _directive_names: Set[str] = set()
+    _directive_init_values: Dict[str, Any] = {}
+    _directives_to_be_executed: List[Callable[..., Any]] = []
+    _directive_functions: Dict[str, Callable[..., Any]] = {}
+    _directive_classes: Dict[str, type] = {}
+    _when_constraints_from_context: List[str] = []
     _default_args: List[dict] = []
 
     push_to_context = _push_to_context

lib/ramble/ramble/language/language_helpers.py

@@ -8,7 +8,7 @@
 
 import fnmatch
 from collections import OrderedDict
-from typing import Any, List, Union
+from typing import Any, List, Optional, Union
 
 from ramble.language.language_base import DirectiveError
 
@@ -131,7 +131,11 @@ def require_definition(
 
 
 def require_condition(
-    obj, directive_name: str, single_arg_name: str = None, multiple_arg_name: str = None, **kwargs
+    obj,
+    directive_name: str,
+    single_arg_name: Optional[str] = None,
+    multiple_arg_name: Optional[str] = None,
+    **kwargs,
 ):
     """Require at least one condition for a type in a directive, and converts all conditions to
     when conditions
@@ -155,25 +159,26 @@ def require_condition(
     Returns:
         List of all when conditions
     """
+    single_arg_val = kwargs.get(single_arg_name) if single_arg_name else None
+    multiple_arg_val = kwargs.get(multiple_arg_name) if multiple_arg_name else None
+    when_arg_val = kwargs.get("when")
 
-    if not (kwargs[single_arg_name] or kwargs[multiple_arg_name] or kwargs["when"]):
+    if not (single_arg_val or multiple_arg_val or when_arg_val):
         raise DirectiveError(
             f"Directive {directive_name} requires at least one of "
             f"{single_arg_name} or {multiple_arg_name} or when to be defined."
         )
 
-    when_list = []
-    if "when" in kwargs:
+    if when_arg_val is not None:
         when_list = build_when_list(kwargs["when"], obj, obj.name, directive_name)
-
-    single_arg_val = kwargs[single_arg_name] if single_arg_name in kwargs else ""
-    multiple_arg_vals = kwargs[multiple_arg_name] if multiple_arg_name in kwargs else []
+    else:
+        when_list = []
 
     # If args are modifier modes, convert to when conditions
     if single_arg_name == "mode" or multiple_arg_name == "modes":
         all_modes = merge_definitions(
             single_arg_val,
-            multiple_arg_vals,
+            multiple_arg_val,
             obj.modes,
             "mode",
             "modes",

lib/ramble/ramble/language/modifier_language.py

@@ -325,8 +325,8 @@ def _define_modifier_variable(mod):
 def package_manager_requirement(
     command: str,
     validation_type: str,
-    mode: str = None,
-    modes: list = None,
+    mode: Optional[str] = None,
+    modes: Optional[list] = None,
     regex=None,
     package_manager: str = "*",
     when=None,

lib/ramble/ramble/variants.py

@@ -196,9 +196,9 @@ def multi_value_variant(self, name: str, value: Any):
     def _define_variant(
         self,
         name: str,
-        variant_type: int,
+        variant_type: variant_types,
         default: Optional[Any] = None,
-        description: str = "",
+        description: Optional[str] = "",
         values: Optional[Union[Sequence, Callable[[Any], bool]]] = None,
     ):
         """Define a variant within this set.
@@ -299,7 +299,7 @@ def __init__(
         self,
         name: str,
         default: Optional[Any] = None,
-        description: str = "",
+        description: Optional[str] = "",
         values: Optional[Union[Sequence, Callable[[Any], bool]]] = None,
     ):
         self.name = name

lib/ramble/ramble/workload.py

@@ -6,7 +6,7 @@
 # option. This file may not be copied, modified, or distributed
 # except according to those terms.
 
-from typing import List, Optional
+from typing import Dict, FrozenSet, List, Optional
 
 import ramble.util.colors as rucolor
 from ramble.definitions.variables import EnvironmentVariable, Variable
@@ -15,6 +15,10 @@
 class Workload:
     """Class representing a single workload"""
 
+    executables: List[str]
+    inputs: List[str]
+    tags: List[str]
+
     def __init__(
         self,
         name: str,
@@ -36,8 +40,8 @@ def __init__(
             tags = []
 
         self.name = name
-        self.variables = {}
-        self.environment_variables = {}
+        self.variables: Dict[FrozenSet[str], List[Variable]] = {}
+        self.environment_variables: Dict[FrozenSet[str], List[EnvironmentVariable]] = {}
 
         attr_names = ["executables", "inputs", "tags"]
         attr_vals = [executables, inputs, tags]

pyproject.toml

@@ -108,7 +108,12 @@ directory = "htmlcov"
 [tool.mypy]
 # TODO: Add in more files for type-checking
 # For now, only type-check the util directory
-files = ["lib/ramble/ramble/util/**/*.py"]
+files = [
+  "lib/ramble/ramble/util/**/*.py",
+  "lib/ramble/ramble/schema/**/*.py",
+  "lib/ramble/ramble/language/**/*.py",
+  "lib/ramble/ramble/definitions/**/*.py",
+]
 ignore_errors = true
 ignore_missing_imports = true
 allow_redefinition = true