From bd6204c31da44db67e23672a4b7a332c52b3f4b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sondre=20Lilleb=C3=B8=20Gundersen?= <sondrelg@live.no>
Date: Mon, 31 May 2021 08:58:24 +0200
Subject: [PATCH] Add mypy and flake8-type-checking config

---
 .pre-commit-config.yaml    | 28 ++++++++++++++++------------
 setup.cfg                  |  3 +--
 setup.py                   |  3 +--
 src/pytest_split/plugin.py | 23 ++++++++++++-----------
 4 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 89235a0..b90185e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -4,18 +4,6 @@ repos:
     hooks:
       - id: black
         args: ['--quiet']
-  - repo: https://gitlab.com/pycqa/flake8
-    rev: 3.8.4
-    hooks:
-      - id: flake8
-        additional_dependencies: [
-          'flake8-bugbear',
-          'flake8-comprehensions',
-          'flake8-deprecated',
-          'flake8-print',
-          'flake8-annotations',
-          'flake8-type-checking',
-        ]
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v3.4.0
     hooks:
@@ -31,3 +19,19 @@ repos:
       - id: trailing-whitespace
       - id: mixed-line-ending
       - id: trailing-whitespace
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: 3.8.4
+    hooks:
+      - id: flake8
+        additional_dependencies: [
+            'flake8-bugbear',
+            'flake8-comprehensions',
+            'flake8-deprecated',
+            'flake8-print',
+            'flake8-type-checking',
+        ]
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v0.812'
+    hooks:
+      - id: mypy
+        additional_dependencies: [django, djangorestframework, inflection, openapi-spec-validator, prance, pyYAML, django-stubs, djangorestframework-stubs, drf_yasg, drf-spectacular]
diff --git a/setup.cfg b/setup.cfg
index 1d4803a..b4fc7a0 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,4 +1,3 @@
 [flake8]
 max-line-length = 88
-exclude =
-    tests/*
+enable-extensions=TC,TC1
diff --git a/setup.py b/setup.py
index efeb630..caec675 100644
--- a/setup.py
+++ b/setup.py
@@ -22,8 +22,7 @@
     install_requires=["pytest"],
     extras_require={"testing": tests_require},
     classifiers=[
-        "Development Status :: 4 - Beta"
-        "Intended Audience :: Developers",
+        "Development Status :: 4 - Beta" "Intended Audience :: Developers",
         "Operating System :: OS Independent",
         "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.6",
diff --git a/src/pytest_split/plugin.py b/src/pytest_split/plugin.py
index fcf38da..a764098 100644
--- a/src/pytest_split/plugin.py
+++ b/src/pytest_split/plugin.py
@@ -1,14 +1,14 @@
 import json
 import os
 from collections import defaultdict, OrderedDict
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING
 
 from _pytest.config import create_terminal_writer
-from _pytest.main import Session
 
 if TYPE_CHECKING:
-    from typing import Generator, List
+    from typing import List, Tuple
     from _pytest.config.argparsing import Parser
+    from _pytest.main import Session
 
     from _pytest import nodes
     from _pytest.config import Config
@@ -52,9 +52,7 @@ def pytest_addoption(parser: "Parser") -> None:
     )
 
 
-def pytest_collection_modifyitems(
-    config: "Config", items: "List[nodes.Item]"
-) -> "Generator[None, None, None]":
+def pytest_collection_modifyitems(config: "Config", items: "List[nodes.Item]") -> None:
     splits = config.option.splits
     group = config.option.group
     store_durations = config.option.store_durations
@@ -62,12 +60,12 @@ def pytest_collection_modifyitems(
 
     if any((splits, group)):
         if not all((splits, group)):
-            return
+            return None
         if not os.path.isfile(durations_report_path):
-            return
+            return None
         if store_durations:
             # Don't split if we are storing durations
-            return
+            return None
     total_tests_count = len(items)
     if splits and group:
         with open(durations_report_path) as f:
@@ -86,13 +84,14 @@ def pytest_collection_modifyitems(
             )
         )
         terminal_reporter.write(message)
+    return None
 
 
 def pytest_sessionfinish(session: "Session") -> None:
     if session.config.option.store_durations:
         report_path = session.config.option.durations_path
         terminal_reporter = session.config.pluginmanager.get_plugin("terminalreporter")
-        durations = defaultdict(float)
+        durations: dict = defaultdict(float)
         for test_reports in terminal_reporter.stats.values():
             for test_report in test_reports:
                 if hasattr(test_report, "duration"):
@@ -124,7 +123,9 @@ def _calculate_suite_start_and_end_idx(
     splits: int, group: int, items: "List[nodes.Item]", stored_durations: OrderedDict
 ) -> Tuple[int, int]:
     item_node_ids = [item.nodeid for item in items]
-    stored_durations = {k: v for k, v in stored_durations.items() if k in item_node_ids}
+    stored_durations = OrderedDict(
+        {k: v for k, v in stored_durations.items() if k in item_node_ids}
+    )
     avg_duration_per_test = sum(stored_durations.values()) / len(stored_durations)
 
     durations = OrderedDict()