diff --git a/nominal_refactor_advisor/refactors/return_shape.py b/nominal_refactor_advisor/refactors/return_shape.py index 12ca83b6..2cb04e15 100644 --- a/nominal_refactor_advisor/refactors/return_shape.py +++ b/nominal_refactor_advisor/refactors/return_shape.py @@ -30,6 +30,8 @@ from nominal_refactor_advisor.refactors.answered_method import AnsweredMethod, AnsweredUse from nominal_refactor_advisor.type_checker_session import SourceTexts, TypeCheckerSession +_shared_type_checker_session: TypeCheckerSession | None = None + def _tuple_return_kind(node: ast.Return) -> str | None: """Classify a return statement as zero-or-one tuple, or None if other.""" @@ -179,11 +181,16 @@ def add(file_path: str, start: int, end: int, text: str) -> None: add(resolved_class.file_path, start, end, new_annotation) answered = AnsweredMethod(snapshot, base, function) - session = TypeCheckerSession.started(SourceTexts(snapshot.parsed_modules)) + shared = _shared_type_checker_session + owned_session = shared is None + session = shared or TypeCheckerSession.started( + SourceTexts(snapshot.parsed_modules) + ) try: calls = answered.answering_calls(session) finally: - session.stop() + if owned_session: + session.stop() for use in calls: geo = geometry(use.file_path)