Improved type annotations as suggested by happz

newa/__init__.py

@@ -6,23 +6,37 @@
 import os
 import re
 import time
+from collections.abc import Iterator
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Literal,
+    Optional,
+    TypedDict,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+    )
 
 import attrs
 import jinja2
 import requests
 import ruamel.yaml
 import ruamel.yaml.nodes
 import ruamel.yaml.representer
+import urllib3.response
 from attrs import define, field, frozen, validators
 from requests_kerberos import HTTPKerberosAuth
 
 if TYPE_CHECKING:
     from typing_extensions import Self, TypeAlias
 
     ErratumId: TypeAlias = str
+    JSON: TypeAlias = Any
 
 
 T = TypeVar('T')
@@ -98,6 +112,50 @@ class ResponseContentType(Enum):
     BINARY = 'binary'
 
 
+@overload
+def get_request(
+        *,
+        url: str,
+        krb: bool = False,
+        attempts: int = 5,
+        delay: int = 5,
+        response_content: Literal[ResponseContentType.TEXT]) -> str:
+    pass
+
+
+@overload
+def get_request(
+        *,
+        url: str,
+        krb: bool = False,
+        attempts: int = 5,
+        delay: int = 5,
+        response_content: Literal[ResponseContentType.BINARY]) -> bytes:
+    pass
+
+
+@overload
+def get_request(
+        *,
+        url: str,
+        krb: bool = False,
+        attempts: int = 5,
+        delay: int = 5,
+        response_content: Literal[ResponseContentType.JSON]) -> JSON:
+    pass
+
+
+@overload
+def get_request(
+        *,
+        url: str,
+        krb: bool = False,
+        attempts: int = 5,
+        delay: int = 5,
+        response_content: Literal[ResponseContentType.RAW]) -> urllib3.response.HTTPResponse:
+    pass
+
+
 def get_request(
         url: str,
         krb: bool = False,
@@ -202,7 +260,7 @@ def from_yaml_file(cls: type[SerializableT], filepath: Path) -> SerializableT:
 
     @classmethod
     def from_yaml_url(cls: type[SerializableT], url: str) -> SerializableT:
-        r = get_request(url, response_content=ResponseContentType.TEXT)
+        r = get_request(url=url, response_content=ResponseContentType.TEXT)
         return cls.from_yaml(r)
 
 
@@ -233,15 +291,15 @@ def _url_factory(self) -> str:
         raise Exception("NEWA_ET_URL envvar is required.")
 
     # TODO: Not used at this point because we only consume builds now
-    def fetch_info(self, erratum_id: str) -> Any:
+    def fetch_info(self, erratum_id: str) -> JSON:
         return get_request(
-            f"{self.url}/advisory/{erratum_id}.json",
+            url=f"{self.url}/advisory/{erratum_id}.json",
             krb=True,
             response_content=ResponseContentType.JSON)
 
-    def fetch_releases(self, erratum_id: str) -> Any:
+    def fetch_releases(self, erratum_id: str) -> JSON:
         return get_request(
-            f"{self.url}/advisory/{erratum_id}/builds.json",
+            url=f"{self.url}/advisory/{erratum_id}/builds.json",
             krb=True,
             response_content=ResponseContentType.JSON)
 
@@ -338,20 +396,42 @@ class Recipe(Cloneable, Serializable):
     url: str
 
 
+# A tmt context for a recipe, dimension -> value mapping.
+RecipeContext = dict[str, str]
+
+# An environment for e recipe, name -> value mapping.
+RecipeEnvironment = dict[str, str]
+
+
+class RawRecipeConfigDimension(TypedDict, total=False):
+    context: RecipeContext
+    environment: RecipeEnvironment
+    git_url: Optional[str]
+    git_ref: Optional[str]
+
+
+_RecipeConfigDimensionKey = Literal['context', 'environment', 'git_url', 'git_ref']
+
+
+# A list of recipe config dimensions, as stored in a recipe config file.
+RawRecipeConfigDimensions = dict[str, list[RawRecipeConfigDimension]]
+
+
 @define
 class RecipeConfig(Cloneable, Serializable):
     """ A job recipe configuration """
 
-    fixtures: dict[str, dict[str, Any]] = field(factory=dict)
-    dimensions: dict[str, dict[str, Any]] = field(factory=dict)
-
-    def build_requests(self) -> list[Request]:
+    fixtures: RawRecipeConfigDimension = field(
+        factory=cast(Callable[[], RawRecipeConfigDimension], dict))
+    dimensions: RawRecipeConfigDimensions = field(
+        factory=cast(Callable[[], RawRecipeConfigDimensions], dict))
 
+    def build_requests(self) -> Iterator[Request]:
         # this is here to generate unique recipe IDs
         recipe_id_gen = itertools.count(start=1)
 
         # get all options from dimentions
-        options = []
+        options: list[list[RawRecipeConfigDimension]] = []
         for dimension in self.dimensions:
             options.append(self.dimensions[dimension])
         # generate combinations
@@ -360,32 +440,43 @@ def build_requests(self) -> list[Request]:
         for i in range(len(combinations)):
             combinations[i] = (self.fixtures,) + (combinations[i])
 
-        def merge_combination_data(combination: list[dict[str, Any]]) -> dict[str, Any]:
-            merged = {}
+        # Note: moved into its own function to avoid being indented too much;
+        # mypy needs to be silenced because we use `key` variable instead of
+        # literal keys defined in the corresponding typeddicts. And being nested
+        # too much, autopep8 was reformatting and misplacing `type: ignore`.
+        def _merge_key(
+                dest: RawRecipeConfigDimension,
+                src: RawRecipeConfigDimension,
+                key: str) -> None:
+            if key not in dest:
+                # we need to do a deep copy so we won't corrupt the original data
+                dest[key] = copy.deepcopy(src[key])  # type: ignore[literal-required]
+            elif isinstance(dest[key], dict) and isinstance(src[key], dict):  # type: ignore[literal-required]
+                dest[key].update(src[key])  # type: ignore[literal-required]
+            else:
+                raise Exception(f"Don't know how to merge record type {key}")
+
+        def merge_combination_data(
+                combination: tuple[RawRecipeConfigDimension, ...]) -> RawRecipeConfigDimension:
+            merged: RawRecipeConfigDimension = {}
             for record in combination:
                 for key in record:
-                    if key not in merged:
-                        # we need to do a deep copy so we won't corrupt the original data
-                        merged[key] = copy.deepcopy(record[key])
-                    elif isinstance(merged[key], dict) and isinstance(record[key], dict):
-                        merged[key].update(record[key])
-                    else:
-                        raise Exception(f"Don't know how to merge record type {key}")
+                    _merge_key(merged, record, key)
             return merged
 
         # now for each combination merge data from individual dimensions
         merged_combinations = list(map(merge_combination_data, combinations))
-        return [Request(id=f'REQ-{next(recipe_id_gen)}', **combination)
-                for combination in merged_combinations]
+        for combination in merged_combinations:
+            yield Request(id=f'REQ-{next(recipe_id_gen)}', **combination)
 
 
 @define
 class Request(Cloneable, Serializable):
     """ A test job request configuration """
 
     id: str
-    context: dict[str, str] = field(factory=dict)
-    environment: dict[str, str] = field(factory=dict)
+    context: RecipeContext = field(factory=dict)
+    environment: RecipeEnvironment = field(factory=dict)
     git_url: Optional[str] = None
     git_ref: Optional[str] = None
     tmt_path: Optional[str] = None

newa/cli.py

@@ -232,10 +232,9 @@ def cmd_schedule(ctx: CLIContext) -> None:
         # prepare a list of Request objects
 
         config = RecipeConfig.from_yaml_url(jira_job.recipe.url)
-        requests = config.build_requests()
 
         # create few fake Issue objects for now
-        for request in requests:
+        for request in config.build_requests():
             request_job = RequestJob(
                 event=jira_job.event,
                 erratum=jira_job.erratum,