From d5f56a077d05f5b3f09db5596dade64880aef618 Mon Sep 17 00:00:00 2001
From: Jonathon Belotti <jonathon.bel.melbourne@gmail.com>
Date: Sat, 18 Jul 2020 22:08:37 +1000
Subject: [PATCH] Fix
 https://github.com/thundergolfer/bazel-mypy-integration/pull/16\#issuecomment-657167922
 and introduce regression test

---
 mypy.bzl                |  8 ++++----
 test/foo/BUILD          | 11 +++++++++++
 test/foo/bar.py         |  4 ++++
 test/shell/test_mypy.sh |  6 ++++++
 test/types/BUILD        |  9 +++++++++
 test/types/__init__.py  | 12 ++++++++++++
 6 files changed, 46 insertions(+), 4 deletions(-)
 create mode 100644 test/foo/BUILD
 create mode 100644 test/foo/bar.py
 create mode 100644 test/types/BUILD
 create mode 100644 test/types/__init__.py

diff --git a/mypy.bzl b/mypy.bzl
index 85d4dd8..7c43636 100644
--- a/mypy.bzl
+++ b/mypy.bzl
@@ -92,10 +92,6 @@ def _mypy_rule_impl(ctx, is_aspect = False, exe = None, out_path = None):
     transitive_srcs_depsets = []
     stub_files = []
 
-    direct_src_root_paths = sets.to_list(
-        sets.make([f.root.path for f in direct_src_files]),
-    )
-
     if hasattr(base_rule.attr, "srcs"):
         direct_src_files = _extract_srcs(base_rule.attr.srcs)
 
@@ -133,6 +129,10 @@ def _mypy_rule_impl(ctx, is_aspect = False, exe = None, out_path = None):
     if not is_aspect:
         runfiles = runfiles.merge(ctx.attr._mypy_cli.default_runfiles)
 
+    direct_src_root_paths = sets.to_list(
+        sets.make([f.root.path for f in direct_src_files]),
+    )
+
     ctx.actions.expand_template(
         template = ctx.file._template,
         output = exe,
diff --git a/test/foo/BUILD b/test/foo/BUILD
new file mode 100644
index 0000000..667460a
--- /dev/null
+++ b/test/foo/BUILD
@@ -0,0 +1,11 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+py_library(
+    name = "foo",
+    srcs = ["bar.py"],
+    deps = [
+        "//test/types",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//test:__subpackages__"],
+)
diff --git a/test/foo/bar.py b/test/foo/bar.py
new file mode 100644
index 0000000..6db950d
--- /dev/null
+++ b/test/foo/bar.py
@@ -0,0 +1,4 @@
+from test.types import FooEnumType, SeqOfMaybeInts
+
+blah: SeqOfMaybeInts = [1, None, 2, None, 3, 4, 5]
+blurgh: FooEnumType = FooEnumType.BAR
diff --git a/test/shell/test_mypy.sh b/test/shell/test_mypy.sh
index 07cd19b..0f53c8a 100644
--- a/test/shell/test_mypy.sh
+++ b/test/shell/test_mypy.sh
@@ -8,6 +8,11 @@ test_ok_on_valid_mypy_typings() {
   action_should_succeed build --aspects //:mypy.bzl%mypy_aspect --output_groups=mypy //test:correct_mypy_typings
 }
 
+# Test for regression originally introduced in https://github.com/thundergolfer/bazel-mypy-integration/pull/16/files
+test_ok_for_package_roots_regression() {
+  action_should_succeed build --aspects //:mypy.bzl%mypy_aspect --output_groups=mypy //test/foo:foo
+}
+
 test_fails_on_broken_mypy_typings() {
   action_should_fail build --aspects //:mypy.bzl%mypy_aspect --output_groups=mypy //test:broken_mypy_typings
 }
@@ -25,6 +30,7 @@ test_fails_on_empty_mypy_test() {
 }
 
 $runner test_ok_on_valid_mypy_typings
+$runner test_ok_for_package_roots_regression
 $runner test_fails_on_broken_mypy_typings
 $runner test_ok_on_valid_mypy_test
 $runner test_fails_on_broken_mypy_test
diff --git a/test/types/BUILD b/test/types/BUILD
new file mode 100644
index 0000000..77fd688
--- /dev/null
+++ b/test/types/BUILD
@@ -0,0 +1,9 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+py_library(
+    name = "types",
+    srcs = ["__init__.py"],
+    deps = [],
+    srcs_version = "PY3",
+    visibility = ["//test:__subpackages__"],
+)
diff --git a/test/types/__init__.py b/test/types/__init__.py
new file mode 100644
index 0000000..8cf324d
--- /dev/null
+++ b/test/types/__init__.py
@@ -0,0 +1,12 @@
+import enum
+from typing import List, Optional
+
+
+class FooEnumType(enum.IntEnum):
+    BAR = 1
+    BAZ = 2
+    BEE = 3
+    BOO = 4
+
+
+SeqOfMaybeInts = List[Optional[int]]