Fix some ignored type hinting

pretf.aws/pretf/aws.py

@@ -5,13 +5,17 @@
 from time import sleep
 from typing import Any, Callable, Optional
 
+from boto3 import Session
+
 from pretf.api import block, log
 from pretf.blocks import Block
 
 try:
-    from boto_source_profile_mfa import get_session as Session  # type: ignore
+    import boto_source_profile_mfa
+
+    use_boto_source_profile_mfa = True
 except ImportError:
-    from boto3 import Session  # type: ignore
+    use_boto_source_profile_mfa = False
 
 
 # Use this lock on anything that might trigger an MFA prompt,
@@ -239,7 +243,10 @@ def get_frozen_credentials(
 
 @lru_cache()
 def get_session(**kwargs: Any) -> Session:
-    return Session(**kwargs)
+    if use_boto_source_profile_mfa:
+        return boto_source_profile_mfa.get_session(**kwargs)
+    else:
+        return Session(**kwargs)
 
 
 @locked

pretf.aws/setup.py

@@ -18,6 +18,6 @@ def get_version():
     author_email="randomy@gmail.com",
     license="MIT License",
     packages=["pretf"],
-    install_requires=["boto3", "pretf"],
+    install_requires=["boto3", "boto3-stubs", "pretf"],
     zip_safe=False,
 )

pretf/pretf/test.py

@@ -7,7 +7,8 @@
 from json import loads as json_loads
 from pathlib import Path
 from subprocess import CompletedProcess
-from typing import Any, Callable, Dict, Generator, List, Optional, Union
+from types import TracebackType
+from typing import Any, Callable, Dict, Generator, List, Optional, Type, Union
 
 import pytest
 
@@ -35,7 +36,12 @@ def __call__(self, cwd: Union[Path, str] = "") -> "TerraformProxy":
     def __enter__(self) -> "TerraformProxy":
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
         return None
 
     # Terraform command.
@@ -134,7 +140,12 @@ def execute(self, *args: str) -> CompletedProcess:
 
 
 class SimpleTestMeta(type):
-    def __new__(cls, name, bases, dct):  # type: ignore
+    def __new__(
+        cls: Type["SimpleTestMeta"],
+        name: str,
+        bases: tuple,
+        dct: dict,
+    ) -> "SimpleTestMeta":
         """
         Wraps all test methods with the pretf_test_function() decorator.
 
@@ -146,7 +157,7 @@ def __new__(cls, name, bases, dct):  # type: ignore
 
         return super().__new__(cls, name, bases, dct)
 
-    def __init__(self, name, bases, dct):  # type: ignore
+    def __init__(self, name: str, bases: tuple, dct: dict) -> None:
         """
         Adds any test method using the @always decorator to cls._always
         so the pretf_test_function() can run it even when previous tests

pretf/pretf/variables.py

@@ -94,7 +94,7 @@ def proxy(self, consumer: Any) -> VariableProxy:
 
 class TerraformVariableStore(VariableStore):
     def __init__(self, files_to_create: dict) -> None:
-        super().__init__()  # type: ignore
+        super().__init__()
         self._files_to_create = files_to_create
         self._files_done: Set[Path] = set()
         self._tfvars_waiting: Set[Path] = set()

pretf/pretf/workflow.py

@@ -113,12 +113,14 @@ def custom(
 
     with import_file(path) as module:
 
-        if not hasattr(module, "pretf_workflow"):
+        try:
+            func = getattr(module, "pretf_workflow")
+        except AttributeError:
             raise log.bad(f"workflow: {path} does not have a 'pretf_workflow' function")
 
         # Call the pretf_workflow() function,
         # passing in "path" and "terraform" if required.
-        result = call_pretf_function(func=module.pretf_workflow, context=context)  # type: ignore
+        result = call_pretf_function(func=func, context=context)
 
     if isinstance(result, int):
         result = CompletedProcess(args=[str(path)], returncode=result)

stubs/boto_source_profile_mfa.pyi

@@ -0,0 +1,5 @@
+from typing import Any
+
+import boto3
+
+def get_session(**kwargs: Any) -> boto3.Session: ...