diff --git a/mypy/stubgen.py b/mypy/stubgen.py index fe90c3d256b7..e312e7209e26 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -48,6 +48,7 @@ import sys import traceback from collections.abc import Iterable, Iterator +from functools import cache from typing import Final import mypy.build @@ -126,7 +127,11 @@ from mypy.semanal_shared import find_dataclass_transform_spec from mypy.sharedparse import MAGIC_METHODS_POS_ARGS_ONLY from mypy.stubdoc import ArgSig, FunctionSig -from mypy.stubgenc import InspectionStubGenerator, generate_stub_for_c_module +from mypy.stubgenc import ( + ExternalSignatureGenerator, + InspectionStubGenerator, + generate_stub_for_c_module, +) from mypy.stubutil import ( TYPING_BUILTIN_REPLACEMENTS, BaseStubGenerator, @@ -1872,7 +1877,8 @@ def generate_stubs(options: Options) -> None: all_modules=all_module_names, ) - # Separately analyse C modules using different logic. + # Reuse the external signature index for C modules in this invocation. + load_doc_sigs = cache(ExternalSignatureGenerator.from_doc_dir) for mod in c_modules: if any(py_mod.module.startswith(mod.module + ".") for py_mod in all_modules): target = mod.module.replace(".", "/") + "/__init__.pyi" @@ -1886,6 +1892,7 @@ def generate_stubs(options: Options) -> None: target, known_modules=all_module_names, doc_dir=options.doc_dir, + load_doc_sigs=load_doc_sigs, include_private=options.include_private, export_less=options.export_less, include_docstrings=options.include_docstrings, diff --git a/mypy/stubgenc.py b/mypy/stubgenc.py index ba312481c51c..e2ae8aaab2ff 100755 --- a/mypy/stubgenc.py +++ b/mypy/stubgenc.py @@ -150,6 +150,7 @@ def generate_stub_for_c_module( known_modules: list[str], doc_dir: str = "", *, + load_doc_sigs: Callable[[str], ExternalSignatureGenerator] | None = None, include_private: bool = False, export_less: bool = False, include_docstrings: bool = False, @@ -174,6 +175,7 @@ def generate_stub_for_c_module( module_name, known_modules, doc_dir, + load_doc_sigs=load_doc_sigs, include_private=include_private, export_less=export_less, include_docstrings=include_docstrings, @@ -243,8 +245,11 @@ def __init__( export_less: bool = False, include_docstrings: bool = False, module: ModuleType | None = None, + *, + load_doc_sigs: Callable[[str], ExternalSignatureGenerator] | None = None, ) -> None: self.doc_dir = doc_dir + self._load_doc_sigs = load_doc_sigs or ExternalSignatureGenerator.from_doc_dir if module is None: self.module = importlib.import_module(module_name) else: @@ -397,7 +402,7 @@ def get_sig_generators(self) -> list[SignatureGenerator]: sig_generators: list[SignatureGenerator] = [DocstringSignatureGenerator()] if self.doc_dir: # Collect info from docs (if given). Always check these first. - sig_generators.insert(0, ExternalSignatureGenerator.from_doc_dir(self.doc_dir)) + sig_generators.insert(0, self._load_doc_sigs(self.doc_dir)) return sig_generators def strip_or_import(self, type_name: str) -> str: diff --git a/mypy/test/teststubgen.py b/mypy/test/teststubgen.py index b5cdbf479f1c..c3f5552e3be4 100644 --- a/mypy/test/teststubgen.py +++ b/mypy/test/teststubgen.py @@ -7,8 +7,10 @@ import sys import tempfile import unittest +from pathlib import Path from types import ModuleType from typing import Any +from unittest.mock import patch import pytest @@ -35,7 +37,7 @@ mypy_options, parse_options, ) -from mypy.stubgenc import InspectionStubGenerator, infer_c_method_args +from mypy.stubgenc import ExternalSignatureGenerator, InspectionStubGenerator, infer_c_method_args from mypy.stubutil import ( ClassInfo, FunctionContext, @@ -120,6 +122,25 @@ def test_module_not_found(self) -> None: sys.stdout = sys.__stdout__ os.chdir(current) + def test_doc_dir_reused_within_one_invocation(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + doc = Path(tmp, "functions.rst") + output = Path(tmp, "out") + options = parse_options( + ["--doc-dir", tmp, "-m", "math", "-m", "cmath", "-o", str(output)] + ) + with patch.object( + ExternalSignatureGenerator, + "from_doc_dir", + wraps=ExternalSignatureGenerator.from_doc_dir, + ) as load: + for arg in ["first_argument", "updated_argument"]: + doc.write_text(f".. function:: sqrt({arg})\n") + generate_stubs(options) + for module in ["math", "cmath"]: + assert f"def sqrt({arg}" in (output / f"{module}.pyi").read_text() + assert load.call_count == 2 + def make_file(self, *path: str, content: str = "") -> None: file = os.path.join(*path) with open(file, "w") as f: