Include external srcs for mypy dep analysis with pip_install

mypy.bzl

@@ -65,12 +65,21 @@ def _extract_srcs(srcs):
                 direct_src_files.append(f)
     return direct_src_files
 
-def _extract_transitive_deps(deps):
+def _extract_transitive_deps(deps, include_imports):
     transitive_deps = []
+    transitive_imports = []
+    seen_imports = {}  # No sets in Starlark, so use a dict.
     for dep in deps:
-        if MyPyStubsInfo not in dep and PyInfo in dep and not _is_external_dep(dep):
-            transitive_deps.append(dep[PyInfo].transitive_sources)
-    return transitive_deps
+        if MyPyStubsInfo not in dep and PyInfo in dep:
+            if include_imports:
+                transitive_deps.append(dep[PyInfo].transitive_sources)
+                for imp in dep[PyInfo].imports.to_list():
+                    if imp not in seen_imports:
+                        seen_imports[imp] = None
+                        transitive_imports.append(imp)
+            elif not _is_external_dep(dep):
+                transitive_deps.append(dep[PyInfo].transitive_sources)
+    return transitive_deps, transitive_imports
 
 def _extract_stub_deps(deps):
     # Need to add the .py files AND the .pyi files that are
@@ -110,22 +119,37 @@ def _mypy_rule_impl(ctx, is_aspect = False):
     transitive_srcs_depsets = []
     stub_files = []
 
+    include_imports = hasattr(base_rule.attr, "include_imports") and base_rule.attr.include_imports
+
     if hasattr(base_rule.attr, "srcs"):
         direct_src_files = _extract_srcs(base_rule.attr.srcs)
 
     if hasattr(base_rule.attr, "deps"):
-        transitive_srcs_depsets = _extract_transitive_deps(base_rule.attr.deps)
+        transitive_srcs_depsets, transitive_imports = _extract_transitive_deps(base_rule.attr.deps, include_imports)
         stub_files = _extract_stub_deps(base_rule.attr.deps)
+        if transitive_imports:
+            rel_workspace_root = ''
+            # If in a package, imports need to be made relative to the
+            # workspace root.
+            if ctx.label.package:
+                rel_workspace_root = '../' * (ctx.label.package.count('/') + 1)
+            mypypath_parts += [rel_workspace_root + x for x in transitive_imports]
 
     if hasattr(base_rule.attr, "imports"):
-        mypypath_parts = _extract_imports(base_rule.attr.imports, ctx.label)
+        mypypath_parts += _extract_imports(base_rule.attr.imports, ctx.label)
 
     final_srcs_depset = depset(transitive = transitive_srcs_depsets +
                                             [depset(direct = direct_src_files)])
-    src_files = [f for f in final_srcs_depset.to_list() if not _is_external_src(f)]
-    if not src_files:
+    input_src_files = final_srcs_depset.to_list()
+    target_src_files = [f for f in input_src_files if not _is_external_src(f)]
+    if not target_src_files:
         return None
 
+    # If imports aren't being included, the input src files are restricted to
+    # only the direct targets.
+    if not include_imports:
+        input_src_files = target_src_files
+
     mypypath_parts += [src_f.dirname for src_f in stub_files]
     mypypath = ":".join(mypypath_parts)
 
@@ -149,34 +173,43 @@ def _mypy_rule_impl(ctx, is_aspect = False):
     # Compose a list of the files needed for use. Note that aspect rules can use
     # the project version of mypy however, other rules should fall back on their
     # relative runfiles.
-    runfiles = ctx.runfiles(files = src_files + stub_files + [mypy_config_file])
+    runfiles = ctx.runfiles(files = input_src_files + stub_files + [mypy_config_file])
     if not is_aspect:
         runfiles = runfiles.merge(ctx.attr._mypy_cli.default_runfiles)
 
     src_root_paths = sets.to_list(
-        sets.make([f.root.path for f in src_files]),
+        sets.make([f.root.path for f in input_src_files]),
     )
 
+    follow_imports = ""
+    if include_imports:
+        # --follow-imports=silent is passed in order to suppress errors on
+        # non-target (imported) libraries.
+        # 0.810 has a --exclude flag which may work better:
+        # https://github.com/python/mypy/pull/9992
+        follow_imports = "--follow-imports=silent"
+
     ctx.actions.expand_template(
         template = ctx.file._template,
         output = exe,
         substitutions = {
             "{MYPY_EXE}": ctx.executable._mypy_cli.path,
             "{MYPY_ROOT}": ctx.executable._mypy_cli.root.path,
-            "{CACHE_MAP_TRIPLES}": " ".join(_sources_to_cache_map_triples(src_files, is_aspect)),
+            "{CACHE_MAP_TRIPLES}": " ".join(_sources_to_cache_map_triples(input_src_files, is_aspect)),
             "{PACKAGE_ROOTS}": " ".join([
                 "--package-root " + shell.quote(path or ".")
                 for path in src_root_paths
             ]),
             "{SRCS}": " ".join([
                 shell.quote(f.path) if is_aspect else shell.quote(f.short_path)
-                for f in src_files
+                for f in target_src_files
             ]),
             "{VERBOSE_OPT}": "--verbose" if DEBUG else "",
             "{VERBOSE_BASH}": "set -x" if DEBUG else "",
             "{OUTPUT}": out.path if out else "",
             "{MYPYPATH_PATH}": mypypath if mypypath else "",
             "{MYPY_INI_PATH}": mypy_config_file.path,
+            "{FOLLOW_IMPORTS}": follow_imports,
         },
         is_executable = True,
     )
@@ -234,5 +267,9 @@ mypy_test = rule(
     implementation = _mypy_test_impl,
     test = True,
     attrs = dict(DEFAULT_ATTRS.items() +
-                 [("deps", attr.label_list(aspects = [mypy_aspect]))]),
+                 [("deps", attr.label_list(aspects = [mypy_aspect])),
+                  ("include_imports",
+                   attr.bool(doc = "Set to true to include imported Python files for mypy. This is required for use with pip `requirement()` rules.")),
+                 ]
+    ),
 )

templates/mypy.sh.tpl

@@ -27,7 +27,7 @@ main() {
   fi
 
   set +o errexit
-  output=$($mypy {VERBOSE_OPT} --bazel {PACKAGE_ROOTS} --config-file {MYPY_INI_PATH} --cache-map {CACHE_MAP_TRIPLES} -- {SRCS} 2>&1)
+  output=$($mypy {VERBOSE_OPT} --bazel {PACKAGE_ROOTS} --config-file {MYPY_INI_PATH} --cache-map {CACHE_MAP_TRIPLES} {FOLLOW_IMPORTS} -- {SRCS} 2>&1)
   status=$?
   set -o errexit