Add annotations for returned data

pyproject.toml

@@ -236,8 +236,6 @@ max-returns = 11
     "ANN002",  # Missing type annotation for `*args`
     "ANN003",  # Missing type annotation for `**kwargs`
     "ANN201",  # Missing return type annotation for public function
-    "ANN202",  # Missing return type annotation for private function
-    "ANN206",  # Missing return type annotation for classmethod
     "ARG001",  # Unused function argument
     "B007",    # Loop control variable `host` not used within loop body
     "C414",    # Unnecessary `list` call within `sorted()`

tests/conftest.py

@@ -1,5 +1,5 @@
 import os
-from typing import List
+from typing import Any, Dict, List, Type, TypeVar, Union
 
 import pytest
 import ruamel.yaml
@@ -18,14 +18,16 @@
 from nornir.core.state import GlobalState
 from nornir.core.task import AggregatedResult, Task
 
+ElementType = TypeVar("ElementType", bound=Union[Group, Host])
+
 global_data = GlobalState(dry_run=True)
 
 
 def inventory_from_yaml():
     dir_path = os.path.dirname(os.path.realpath(__file__))
     yml = ruamel.yaml.YAML(typ="safe")
 
-    def get_connection_options(data):
+    def get_connection_options(data) -> Dict[str, ConnectionOptions]:
         cp = {}
         for cn, c in data.items():
             cp[cn] = ConnectionOptions(
@@ -38,7 +40,7 @@ def get_connection_options(data):
             )
         return cp
 
-    def get_defaults():
+    def get_defaults() -> Defaults:
         defaults_file = f"{dir_path}/inventory_data/defaults.yaml"
         with open(defaults_file, "r") as f:
             defaults_dict = yml.load(f)
@@ -55,7 +57,9 @@ def get_defaults():
                 ),
             )
 
-    def get_inventory_element(typ, data, name, defaults):
+    def get_inventory_element(
+        typ: Type[ElementType], data: Dict[str, Any], name: str, defaults: Union[Defaults, None]
+    ) -> ElementType:
         return typ(
             name=name,
             hostname=data.get("hostname"),

tests/core/test_InitNornir.py

@@ -167,7 +167,7 @@ def test_InitNornir_different_transform_function_by_string_with_bad_options(self
 
 class TestLogging:
     @classmethod
-    def cleanup(cls):
+    def cleanup(cls) -> None:
         # this does not work as setup_method, because pytest injects
         # _pytest.logging.LogCaptureHandler handler to the root logger
         # and StreamHandler to _pytest.capture.EncodedFile to other loggers
@@ -183,7 +183,7 @@ def cleanup(cls):
             logger_.setLevel(logging.NOTSET)
 
     @classmethod
-    def teardown_class(cls):
+    def teardown_class(cls) -> None:
         cls.cleanup()
 
     def test_InitNornir_logging_defaults(self):

tests/core/test_connections.py

@@ -106,7 +106,7 @@ def validate_params(task, conn, params, nornir_config):
 
 class Test:
     @classmethod
-    def setup_class(cls):
+    def setup_class(cls) -> None:
         ConnectionPluginRegister.deregister_all()
         ConnectionPluginRegister.register("dummy", DummyConnectionPlugin)
         ConnectionPluginRegister.register("dummy2", DummyConnectionPlugin)

tests/core/test_inventory.py

@@ -397,7 +397,7 @@ def test_filtering_func(self, inv):
         )
         assert long_names == ["dev1.group_1", "dev4.group_2", "dev6.group_3"]
 
-        def longer_than(dev, length):
+        def longer_than(dev, length) -> bool:
             return len(dev["my_var"]) > length
 
         long_names = sorted(list(inv.filter(filter_func=longer_than, length=20).hosts.keys()))

tests/wrapper.py

@@ -15,7 +15,7 @@ def wrap_cli_test(output, save_output=False):
     """
 
     @decorator
-    def run_test(func, *args, **kwargs):
+    def run_test(func, *args, **kwargs) -> None:
         stdout = StringIO()
         backup_stdout = sys.stdout
         sys.stdout = stdout