From eb816b05ce70f8bf3c86ecf85a6ac64c10a17aeb Mon Sep 17 00:00:00 2001
From: Shantanu <12621235+hauntsaninja@users.noreply.github.com>
Date: Tue, 15 Oct 2024 07:58:42 -0700
Subject: [PATCH] Use fast path in modulefinder more often (#17950)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

See https://github.com/python/mypy/issues/17948

This is about 1.06x faster on `mypy -c 'import torch'` (in both the
clean and openai environments)
- 19.094 -> 17.896
- 34.161 -> 32.214

```
λ hyperfine -w 1 -M 3 '/tmp/mypy_primer/timer_mypy_36738b392/venv/bin/mypy  -c "import torch" --no-incremental --python-executable clean/bin/python'
Benchmark 1: /tmp/mypy_primer/timer_mypy_36738b392/venv/bin/mypy  -c "import torch" --no-incremental --python-executable clean/bin/python
  Time (mean ± σ):     17.896 s ±  0.130 s    [User: 16.472 s, System: 1.408 s]
  Range (min … max):   17.757 s … 18.014 s    3 runs

 λ hyperfine -w 1 -M 3 '/tmp/mypy_primer/timer_mypy_36738b392/venv/bin/mypy  -c "import torch" --no-incremental --python-executable /opt/oai/bin/python'
Benchmark 1: /tmp/mypy_primer/timer_mypy_36738b392/venv/bin/mypy  -c "import torch" --no-incremental --python-executable /opt/oai/bin/python
  Time (mean ± σ):     32.214 s ±  0.106 s    [User: 29.468 s, System: 2.722 s]
  Range (min … max):   32.098 s … 32.305 s    3 runs
```
---
 mypy/build.py        | 8 +++++---
 mypy/modulefinder.py | 2 +-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/mypy/build.py b/mypy/build.py
index 964da5aac8b0..043b52f0a241 100644
--- a/mypy/build.py
+++ b/mypy/build.py
@@ -664,7 +664,7 @@ def __init__(
         for module in CORE_BUILTIN_MODULES:
             if options.use_builtins_fixtures:
                 continue
-            path = self.find_module_cache.find_module(module)
+            path = self.find_module_cache.find_module(module, fast_path=True)
             if not isinstance(path, str):
                 raise CompileError(
                     [f"Failed to find builtin module {module}, perhaps typeshed is broken?"]
@@ -2725,7 +2725,9 @@ def exist_added_packages(suppressed: list[str], manager: BuildManager, options:
 
 def find_module_simple(id: str, manager: BuildManager) -> str | None:
     """Find a filesystem path for module `id` or `None` if not found."""
-    x = find_module_with_reason(id, manager)
+    t0 = time.time()
+    x = manager.find_module_cache.find_module(id, fast_path=True)
+    manager.add_stats(find_module_time=time.time() - t0, find_module_calls=1)
     if isinstance(x, ModuleNotFoundReason):
         return None
     return x
@@ -2734,7 +2736,7 @@ def find_module_simple(id: str, manager: BuildManager) -> str | None:
 def find_module_with_reason(id: str, manager: BuildManager) -> ModuleSearchResult:
     """Find a filesystem path for module `id` or the reason it can't be found."""
     t0 = time.time()
-    x = manager.find_module_cache.find_module(id)
+    x = manager.find_module_cache.find_module(id, fast_path=False)
     manager.add_stats(find_module_time=time.time() - t0, find_module_calls=1)
     return x
 
diff --git a/mypy/modulefinder.py b/mypy/modulefinder.py
index 0cfe8f3b9d2f..59a71025f71e 100644
--- a/mypy/modulefinder.py
+++ b/mypy/modulefinder.py
@@ -557,7 +557,7 @@ def _find_module(self, id: str, use_typeshed: bool) -> ModuleSearchResult:
             return ModuleNotFoundReason.NOT_FOUND
 
     def find_modules_recursive(self, module: str) -> list[BuildSource]:
-        module_path = self.find_module(module)
+        module_path = self.find_module(module, fast_path=True)
         if isinstance(module_path, ModuleNotFoundReason):
             return []
         sources = [BuildSource(module_path, module, None)]